Equidiff / README.md
Lillianwei's picture
mimicgen
c1f1d32
# 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
```