CleanDiffuser is a modular Python library for decision-making tasks using diffusion models, designed for researchers and practitioners. It offers a unified interface for various diffusion models, network architectures, conditioning methods, and algorithm pipelines, inspired by CleanRL and Diffusers, emphasizing usability, simplicity, and customizability.

How It Works

CleanDiffuser integrates multiple diffusion model variants (DDPM, DDIM, DPM-Solver, EDM, Rectified Flow, Consistency Models) and network architectures (MLPs, Transformers, UNets) into flexible pipelines. It supports classifier and classifier-free guidance for enhanced control and offers pre-built pipelines for imitation learning (e.g., Diffusion Policy, DiffusionBC) and reinforcement learning (e.g., DQL, IDQL) on standard robotics benchmarks like D4RL and Robosuite.

Quick Start & Requirements

• Install: Clone the repository and install with pip install -e ..
• Prerequisites: Python 3.9+, PyTorch (e.g., 2.2.2 with CUDA 12.1). For specific pipelines, requires mujoco-py, D4RL, robomimic, and robosuite. Downgrading gym to 0.21.0 may be necessary for compatibility.
• Setup: Basic installation is quick. Setting up robotics environments and downloading datasets can take significant time.
• Docs: Documentation
• Examples: Tutorials, Pipelines

Highlighted Details

• Supports advanced diffusion techniques like Consistency Models and Rectified Flow.
• Includes implementations of recent diffusion planners and policies (DiffuserLite, SfBC, QGPO).
• Features modular design for easy customization and extension of algorithms.
• Integrates Hydra for flexible configuration management and Wandb for logging.

Maintenance & Community

The project is actively developed, with recent additions including a diffusion planner (Diffusion Veteran) and new policies. The lightning branch is recommended for improved implementation and broader support. Contact emails are provided for questions.

Licensing & Compatibility

Distributed under the Apache License 2.0, permitting commercial use and integration with closed-source projects.

Limitations & Caveats

The README notes potential compatibility issues with the latest robomimic image dependencies, recommending a specific downgrade of the gym package. Reproducing paper results may require downloading large datasets separately.