# Equidiff

- folder_name: the name of the folder
- file_name: the name of your file

## Prepare data
Use mimicgen to generate data.

Use `EmbodiedBM/equidiff/combinehdf5.py` to combine data from multiple .hdf5 files if needed.

Put hdf5 data at `EmbodiedBM/equidiff/data/robomimic/datasets` with the format [folder_name]/[file_name].hdf5

## Convert data
```bash
python equi_diffpo/scripts/robomimic_dataset_conversion.py -i data/robomimic/datasets/square_d2_test/demo.hdf5 -o data/robomimic/datasets/square_d2_test/demo_abs.hdf5 -n 12
```

## Train

Use another CUDA device if 7 is currently in use.

```bash
CUDA_VISIBLE_DEVICES=5 MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa HYDRA_FULL_ERROR=1 python train.py --config-name=train_sq2_5000 folder_name=square_d2_5000 file_name=demo n_demo=5000
```

If you use another task than square_d2, you should change the task_name config by adding task_name=[task_name]

## Test
Change the `ckpt_path` to the trained policy's weight's path in `EmbodiedBM/equidiff/equi_diffpo/config/test_sq2.yaml`

If you use another task than square_d2, you should change the dataset config in test_sq2.yaml and download the corresponding dataset from [Huggingface](https://huggingface.co/datasets/amandlek/mimicgen_datasets/tree/main/core).

```bash
python test.py
```