Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals

Paper, Project Page, RSS 2024

Moritz Reuss¹, Ömer Erdinç Yağmurlu, Fabian Wenzel, Rudolf Lioutikov¹

¹Intuitive Robots Lab, Karlsruhe Institute of Technology

This is the official code repository for the paper Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals.

Pre-trained models are available here.

Performance

Results on the CALVIN benchmark (1000 chains):

*: 3.72±(0.05) (D) and 4.52±(0.02) (ABCD) in the paper. Performance is higher than reported given some fixes in the camera-ready code version.

Installation

To begin, clone the repository locally:

git clone --recurse-submodules [email protected]:intuitive-robots/mdt_policy.git
export MDT_ROOT=$(pwd)/mdt_policy

Install the requirements: (Note we provided a changed verison of pyhash, given numerous problems when installing it manually)

cd $MDT_ROOT
conda create -n mdt_env python=3.8
conda activate mdt_env
cd calvin_env/tacto
pip install -e .
cd ..
pip install -e .
cd ..
pip install setuptools==57.5.0
cd pyhash-0.9.3
python setup.py build
python setup.py install
cd ..

Next we can install the rest of the missing packages:

pip install -r requirements.txt

Evaluation

Step 1 - Download CALVIN Datasets

If you want to train on the CALVIN dataset, choose a split with:

cd $MDT_ROOT/dataset
sh download_data.sh D | ABCD

(Optional) Preprocessing with CALVIN

Since MDT uses action chunking, it needs to load multiple (~10) episode_{}.npz files for each inference. In combination with batching, this results in a large disk bandwidth needed for each iteration (usually ~2000MB/iteration). This has the potential of significantly reducing your GPU utilization rate during training depending on your hardware. Therefore, you can use the script extract_by_key.py to extract the data into a single file, avoiding opening too many episode files when using the CALVIN dataset.

Usage example:

python preprocess/extract_by_key.py -i /YOUR/PATH/TO/CALVIN/ \
    --in_task all

Params:

Run this command to see more detailed information:

python preprocess/extract_by_key.py -h

Important params:

• --in_root: /YOUR/PATH/TO/CALVIN/, e.g /data3/geyuan/datasets/CALVIN/
• --extract_key: A key of dict(episode_xxx.npz), default is 'rel_actions', the saved file name depends on this (i.e ep_{extract_key}.npy)

Optional params:

• --in_task: default is 'all', meaning all task folders (e.g task_ABCD_D/) of CALVIN
• --in_split: default is 'all', meaning both training/ and validation/
• --out_dir: optional, default is 'None', and will be converted to {in_root}/{in_task}/{in_split}/extracted/
• --force: whether to overwrite existing extracted data

Thanks to @ygtxr1997 for debugging the GPU utilization and providing a merge request.

Step 2 - Download Pre-trained Models

You can find all of the aforementioned models under here.

Step 3 - Run

Adjust conf/mdt_evaluate.conf according to the model you downloaded. Important keys are:

• voltron_cache: set a rw path in order to avoid downloading Voltron from huggingface each run.
• num_videos: n first rollouts of the model will be recorded as a gif.
• dataset_path and train_folder: location of the downloaded files.
• sigma_min: use the values in the table for best results with a given model.

Then run:

python mdt/evaluation/mdt_evaluate.py

Training

To train the MDT model with the maximum amount of available GPUS, run:

python mdt/training.py

For replication of the originial training results I recommend to use 4 GPUs with a batch_size of 128 and train them for 30 (20) epochs for ABCD (D only). You can find the exact configuration used for each of our pretrained models in .hydra/config.yaml under the respective model's folder.

Acknowledgements

This work is only possible because of the code from the following open-source projects and datasets:

CALVIN

Original: https://github.com/mees/calvin License: MIT

BESO

Original: https://github.com/intuitive-robots/beso License: MIT

Voltron

Original: https://github.com/siddk/voltron-robotics License: MIT

OpenAI CLIP

Original: https://github.com/openai/CLIP License: MIT

HULC

Original: https://github.com/lukashermann/hulc License: MIT

Citation

If you found the code usefull, please cite our work:

@inproceedings{
    reuss2024multimodal,
    title={Multimodal Diffusion Transformer: Learning Versatile Behavior from Multimodal Goals},
    author={Moritz Reuss and {\"O}mer Erdin{\c{c}} Ya{\u{g}}murlu and Fabian Wenzel and Rudolf Lioutikov},
    booktitle={Robotics: Science and Systems},
    year={2024}
}