· ArXiv · Paper · Documentation ·
CleanDiffuser is an easy-to-use modularized Diffusion Model library tailored for decision-making, which comprehensively integrates different types of diffusion algorithmic branches. CleanDiffuser offers a variety of advanced diffusion models, network structures, diverse conditions, and algorithm pipelines in a simple and user-friendly manner. Inheriting the design philosophy of CleanRL and Diffusers, CleanDiffuser emphasizes usability, simplicity, and customizability. We hope that CleanDiffuser will serve as a foundational tool library, providing long-term support for Diffusion Model research in the decision-making community, facilitating the application of research for scientists and practitioners alike. The highlight features of CleanDiffuser are:
- 🚀 Amazing features specially tailored for decision-making tasks
- 🍧 Support for multiple advanced diffusion models and network architectures
- 🧩 Build decoupled modules into integrated pipelines easily like building blocks
- 📈 Wandb logging and Hydra configuration
- 🌏 Unified environmental interface and efficient dataloader
We strongly recommend reading papers and documents to learn more about CleanDiffuser and its design philosophy.
- [2024-07-03] 💫 We provided a CleanDiffuser-based replication of ACT (action chunking with transformers) in the act branch.
- [2024-06-24] 🥰 We have added Consistency Models into CleanDifuser. With one model, you can do both Consistency Distillation and Consistency Training! Check out an example in
tutorials/sp_consistency_policy.py! (Note: Our consistency training implementation refers to the improved version, see https://arxi.org/abs/2310.14189.) - [2024-06-17] 🔥 We released arxiv version of CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Making.
$ conda create -n cleandiffuser python==3.9
$ conda activate cleandiffuserInstall torch>1.0.0,<2.3.0 that is compatible with your CUDA version. For example, PyTorch 2.2.2 with CUDA 12.1:
$ conda install pytorch==2.2.2 torchvision==0.17.2 pytorch-cuda=12.1 -c pytorch -c nvidia$ git clone https://github.com/CleanDiffuserTeam/CleanDiffuser.git
$ cd CleanDiffuser
$ pip install -e .For users who need to run pipelines and reproduce the results of the paper, they will need to install RL simulators.
First, install the dependencies related to the mujoco-py environment. For more details, see https://github.com/openai/mujoco-py#install-mujoco
$ sudo apt-get install libosmesa6-dev libgl1-mesa-glx libglfw3 libglew-dev patchelf# Install D4RL from source (recommended)
$ cd <PATH_TO_D4RL_INSTALL_DIR>
$ git clone https://github.com/Farama-Foundation/D4RL.git
$ cd D4RL
$ pip install -e .
# Install Robomimic from source (recommended)
$ cd <PATH_TO_ROBOMIMIC_INSTALL_DIR>
$ git clone https://github.com/ARISE-Initiative/robomimic.git
$ cd robomimic
$ pip install -e .
$ cd <PATH_TO_ROBOSUITE_INSTALL_DIR>
$ git clone https://github.com/ARISE-Initiative/robosuite.git
$ cd robosuite
$ pip install -e .Note: The latest version of dependencies running the
robomimic imagestill has compatibility issues, and we are actively working on a fix. The temporary solution is to downgrade thegymversion to0.21.0: pip install setuptools==65.5.0 pip==21, pip install gym==0.21.0
Try it now!
# Tutorial
$ python tutorials/1_a_minimal_DBC_implementation.py
# Reinforcement Learning
$ python pipelines/diffuser_d4rl_mujoco.py
# Imitation Learning (need to download the dataset, see below)
$ python pipelines/dp_pusht.pyIf you need to reproduce Imitation Learning environments (pusht, kitchen, robomimic), you need to download the datasets additionally. We recommend downloading the corresponding compressed files from Datasets. We provide the default dataset path as dev/:
dev/
.
├── kitchen
├── pusht
├── robomimicWe will make every effort to provide detailed tutorials for beginners in the field of Diffusion Models in Decision Making, which is also beneficial for learning the core components of CleanDiffuser and expanding them into new algorithms. Our vision is not only to offer a benchmark for the community but more importantly, to enable everyone to implement and innovate diffusion algorithms more easily based on CleanDiffuser.
Note: In the
tutorials, we generally only explain and demonstrate individual mechanisms or components, rather than a complete algorithm, and therefore ignore the extra tricks and take just a few minutes of training time. This may cause performance drop, which is normal!
We have now provided the following tutorials and are continuously updating more:
# Build the DiffusionBC algorithm with minimal code
python tutorials/1_a_minimal_DBC_implementation.py
# Customize classifier-free guidance
python tutorials/2_classifier-free_guidance.py
# Customize classifier guidance
python tutorials/3_classifier_guidance.py
# Customize diffusion network backbone
python tutorials/4_customize_your_diffusion_network_backbone.py
# Special. Consistency Policies
python tutorials/sp_consistency_policy.py If you wish to reproduce the results of the paper perfectly, we recommend using the full implementation in pipelines.
The cleandiffuser folder contains the core components of the CleanDiffuser codebase, including Diffusion Models, Network Architectures, and Guided Sampling. It also provides unified Env and Dataset Interfaces.
In CleanDiffuser, we can combine independent modules to algorithms pipelines like building blocks. In the pipelines folder, we provide all the algorithms currently implemented in CleanDiffuser. By linking with the Hydra configurations in the configs folder, you can reproduce the results presented in the papers:
You can simply run each algorithm with the default environment and configuration without any additional setup, for example:
# DiffusionPolicy with Chi_UNet in lift-ph
python pipelines/dp_pusht.py
# Diffuser in halfcheetah-medium-expert-v2
python pipelines/diffuser_d4rl_mujoco.pyThanks to Hydra, CleanDiffuser also supports flexible running of algorithms through CLI or directly modifying the corresponding configuration files. We provide some examples:
# Load PushT config
python pipelines/dp_pusht.py --config-path=../configs/dp/pusht/dit --config-name=pusht
# Load PushT config and overwrite some hyperparameters
python pipelines/dp_pusht.py --config-path=../configs/dp/pusht/dit --config-name=pusht dataset_path=path/to/dataset seed=42 device=cuda:0
# Train Diffuser in hopper-medium-v2 task
python pipelines/diffuser_d4rl_mujoco.py task=hopper-medium-v2 In CleanDiffuser, we provide a mode option to switch between training (mode=train) or inference (mode=inference) of the model:
# Imitation learning environment
python pipelines/dp_pusht.py mode=inference model_path=path/to/checkpoint
# Reinforcement learning environment
python pipelines/diffuser_d4rl_mujoco.py mode=inference ckpt=latestContributions are what make the open source community such an amazing place to learn, inspire, and create. Any contributions you make are greatly appreciated.
If you have a suggestion that would make this better, please fork the repo and create a pull request. You can also simply open an issue with the tag "enhancement". Don't forget to give the project a star! Thanks again!
- Fork the Project
- Create your Feature Branch (
git checkout -b feature/AmazingFeature) - Commit your Changes (
git commit -m 'Add some AmazingFeature') - Push to the Branch (
git push origin feature/AmazingFeature) - Open a Pull Request
Distributed under the Apache License 2.0. See LICENSE.txt for more information.
For any questions, please feel free to email [email protected] and [email protected].
If you find our work useful, please consider citing:
@article{cleandiffuser,
author = {Zibin Dong and Yifu Yuan and Jianye Hao and Fei Ni and Yi Ma and Pengyi Li and Yan Zheng},
title = {CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Making},
journal = {arXiv preprint arXiv:2406.09509},
year = {2024},
url = {https://arxiv.org/abs/2406.09509},
}