CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Making

· 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.

🔥 News and Change Log

• [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 improvedverslon, see https://arxi.org/abs/2310.14189.)
• [2024-06-20] 🛠️ We have prepared dependency requirements for Tutorial, RL, and IL environments to avoid compatibility issues.
• [2024-06-17] 🔥 We released arxiv version of CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Making.

🛠️ Getting Started

1. Create and activate conda environment

$ conda create -n cleandiffuser python==3.9
$ conda activate cleandiffuser

Download CleanDiffuser and add this folder to your PYTHONPATH. You can also add it to .bashrc for convenience:

git clone https://github.com/CleanDiffuserTeam/CleanDiffuser.git
export PYTHONPATH=$PYTHONPATH:/path/to/CleanDiffuser

2. Install PyTorch and PyTorch3d

We recommend visiting https://pytorch.org/get-started/previous-versions/ to find the appropriate PyTorch version for your CUDA installation. Essentially, we only require torch>1.0.0,<2.3.0, as we have found conflicts between PyTorch version 2.3.0 and NumPy. Below is an example using PyTorch version 2.2.2 with CUDA 11.8.

$ conda install pytorch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 pytorch-cuda=12.1 -c pytorch -c nvidia
$ conda install pytorch3d -c pytorch3d

3. Pure CleanDiffuser installation

For users who do not need to run pipelines and wish to directly use CleanDiffuser for algorithm development, they can simply run pip install 'numpy<1.23.0' einops. That's it.

4. Additional installations

For users who need to run pipelines and reproduce the results of the paper, they will need to install requirements.txt as well as the 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 requirements
$ pip install -r requirements.txt
# Install D4RL from source (recommended)
$ cd <PATH_TO_MJRL_INSTALL_DIR>
$ git clone https://github.com/aravindr93/mjrl.git
$ cd mjrl
$ pip install -e .
$ 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 git clone https://github.com/ARISE-Initiative/robosuite.git
$ cd robosuite
$ pip install -e .

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.py

If 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
├── robomimic

🍷 Tutorials

We 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.

💻 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.py

Thanks 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=latest

🎁 Implemented Components

Category	Items	Paper
SDE/ODE with Solvers
Diffusion SDE	DDPM	✅Denoising Diffusion Probabilistic Models
	DDIM	✅Denoising Diffusion Implicit Models
	DPM-Solver	✅DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps
	DPM-Solver++	✅DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models
EDM	Eular	✅Elucidating the Design Space of Diffusion-Based Generative Models
	2nd Order Heun
Recitified Flow	Euler	✅Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow
Network Architectures
	Pearce_MLP	✅Imitating Human Behaviour with Diffusion Models
	Pearce_Transformer
	Chi_UNet1d	✅Diffusion Policy: Visuomotor Policy Learning via Action Diffusion
	Chi_Transformer
	LNResnet	✅IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies
	DQL_MLP	✅Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning
	Janner_UNet1d	✅Planning with Diffusion for Flexible Behavior Synthesis
	DiT1d	✅AlignDiff: Aligning Diverse Human Preferences via Behavior-Customisable Diffusion Model
Guided Sampling Methods
	Classifier Guidance	✅Diffusion Models Beat GANs on Image Synthesis
	Classifier-free Guidance	✅Classifier-Free Diffusion Guidance
Pipelines
Planners	Diffuser	✅Planning with Diffusion for Flexible Behavior Synthesis
	Decision Diffuser	✅Is Conditional Generative Modeling all you need for Decision-Making?
	AdaptDiffuser	✅AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners
Policies	DQL	✅Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning
	EDP	✅Efficient Diffusion Policies for Offline Reinforcement Learning
	IDQL	✅IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies
	Diffusion Policy	✅Diffusion Policy: Visuomotor Policy Learning via Action Diffusion
	DiffusionBC	✅Imitating Human Behaviour with Diffusion Models
Data Synthesizers	SynthER	✅Synthetic Experience Replay

🙏 Contributing

Contributions 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!

1. Fork the Project
2. Create your Feature Branch (git checkout -b feature/AmazingFeature)
3. Commit your Changes (git commit -m 'Add some AmazingFeature')
4. Push to the Branch (git push origin feature/AmazingFeature)
5. Open a Pull Request

🏷️ License

Distributed under the Apache License 2.0. See LICENSE.txt for more information.

💓 Acknowledgement

• huggingface diffusers
• diffuser
• diffusion policy
• robomimic

✉️ Contact

For any questions, please feel free to email zibindong@outlook.com and yuanyf@tju.edu.cn.

📝 Citation

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},
}