Robomimic Fork | Installation | Prepare Dataset | Train️ | Project Website
We additionally provide a separate implemention of CLASS using a forked version of robomimic. This implementation natively only supports homogeneous Square task. Use the original implementation if you want to simulate dynamic camera settings. Note that the forked version uses a different Robomimic version (v0.5) from the original implementation (v0.3).
git clone https://github.com/sean1295/CLASS.git
cd CLASS
You can install the vitual conda environment using the following command:
conda create -n CLASS python=3.10
conda activate CLASS
pip install -r requirements.txt
pip install -e .
git clone https://github.com/NVlabs/mimicgen.git
cd mimicgen
pip install -e .
cd ..
Prior to generating heterogeneous datasets, you need to first download robomimic/mimicgen hdf5 datasets. See robomimic and mimicgen documentations for the instructions. It should look like this:
python /path/to/robomimic/scripts/download_datasets.py --tasks square --dataset_types ph --hdf5_types low_dim --download_dir /path/to/dataset
python /path/to/mimicgen/scripts/download_datasets.py --tasks stack_three_d0 --dataset_type core --download_dir /path/to/dataset
Make sure to modify /path/to/dataset for the commands above and dataset_dir value in the config files (e.g., CLASS/configs/square_ph_dp_abs), which is the directory where the original and transformed datasets will be stored. Once complete, run
python CLASS/scripts/generate_dataset.py --config_name square_ph_dp_abs
Run the following command to precompute pairwise DTW distances between the actions in the dataset. It utilizes multiprocessing, but it is expected to take a while (hours) for larger datasets.
python CLASS/scripts/precompute.py --config_name square_ph_dp_abs
Set up wandb from terminal with wandb login
You can pretrain the policy using the following command:
python CLASS/scripts/pretrain.py --config_name square_ph_dp_abs
During pre-training, it periodically reports non-parametric rollout results.
You can train the policy using the following command:
python CLASS/scripts/finetune.py --config_name square_ph_dp_abs
During fine-tuning, it periodically reports both non-parametric and parametric rollout results. If you want to train from scratch without pre-training, set finetune.enabled to false in the config file.