Diffusion Policy

To implement a new task, you must create a new dataset class, an environment runner, and a configuration file. Follow these steps:

1. Implement the Dataset: Create a new dataset class (e.g., by imitating diffusion_policy/dataset/pusht_image_dataset.py).
2. Implement the Environment Runner: Create a new runner class (e.g., by imitating diffusion_policy/env_runner/pusht_image_runner.py).
3. Create a Config File: Define the task configuration in a YAML file (e.g., diffusion_policy/config/task/pusht_image.yaml).

Requirements:

• Ensure shape_meta correctly corresponds to the input and output shapes for your specific task.
• Update env_runner._target_ and dataset._target_ in your configuration to point to your newly created classes.
• When training, pass the task name as an argument to train.py using task=<your_task_name>.

The easiest way to experiment with Diffusion Policy without local installation is using the provided Google Colab notebooks. There are separate notebooks for different environment types:

• State-based environment
• Vision-based environment

To reproduce simulation benchmark results on a Linux machine with an Nvidia GPU, follow these steps:

1. Install Mujoco dependencies (on Ubuntu 20.04):

$ sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf

2. Create the Conda environment using Mambaforge (recommended) or Conda:

$ mamba env create -f conda_environment.yaml
# OR
$ conda env create -f conda_environment.yaml

Note: conda_environment_macos.yaml is available for MacOS development but does not fully support benchmarks.

Before running simulations, you must download and extract the training data and the experiment configuration file.

1. Create data directory and download training data:

[diffusion_policy]$ mkdir data && cd data
[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip
[data]$ unzip pusht.zip && rm -f pusht.zip && cd ..

2. Download the experiment config file (example for image_pusht_diffusion_policy_cnn):

[diffusion_policy]$ wget -O image_pusht_diffusion_policy_cnn.yaml https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/config.yaml

Experiment logs, including config.yaml, training/eval logs.json.txt, and checkpoints (epoch=*-test_mean_score=*.ckpt and latest.ckpt), are hosted on the project website.

Logs are organized by modality, task, and method at: https://diffusion-policy.cs.columbia.edu/data/experiments/<image|low_dim>/<task>/<method>/

To download an entire experiment subdirectory recursively, use wget with the following flags:

$ wget --recursive --no-parent --no-host-directories --relative --reject="index.html*" https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/square_ph/diffusion_policy_cnn/

To run on a real robot (Push-T setup), ensure you have the required hardware (UR5/UR5e, RealSense D415, SpaceMouse, etc.) and follow these software setup steps:

1. OS: Ubuntu 20.04.3 (tested).
2. Mujoco dependencies: sudo apt install libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
3. RealSense SDK: Install via the official Intel documentation.
4. Spacemouse dependencies: sudo apt install libspnav-dev spacenavd; sudo systemctl start spacenavd
5. Conda environment: mamba env create -f conda_environment_real.yaml

To implement a new training method (workspace/policy combination), follow these steps:

1. Implement the Policy: Create a new policy class (e.g., by imitating diffusion_policy/policy/diffusion_unet_image_policy.py).
2. Implement the Workspace: Create a new workspace class (e.g., by imitating diffusion_policy/workspace/train_diffusion_unet_image_workspace.py).
3. Create a Workspace Config: Define the workspace configuration in a YAML file (e.g., diffusion_policy/config/train_diffusion_unet_image_workspace.yaml).

Requirement:

• Ensure your workspace YAML's _target_ key points to the new workspace class you created.

To evaluate a downloaded checkpoint, use the eval.py script. This generates an evaluation log and rollout videos.

(robodiff)[diffusion_policy]$ python eval.py --checkpoint data/0550-test_mean_score=0.969.ckpt --output_dir data/pusht_eval_output --device cuda:0

Results are stored in --output_dir. The eval_log.json file contains metrics such as test/mean_score and paths to rollout videos in the media/ subdirectory.

To run multiple seeds in parallel, launch a local Ray cluster and use ray_train_multirun.py.

1. Start Ray cluster (specifying GPUs to manage):

(robodiff)[diffusion_policy]$ export CUDA_VISIBLE_DEVICES=0,1,2
(robodiff)[diffusion_policy]$ ray start --head --num-gpus=3

2. Launch multi-run training:

(robodiff)[diffusion_policy]$ python ray_train_multirun.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml --seeds=42,43,44 --monitor_key=test/mean_score -- multi_run.run_dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' multi_run.wandb_name_base='${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}'

This will log aggregated metrics to the diffusion_policy_metrics project on WandB and save individual seed results in a structured directory under data/outputs/.

To train a model for a single seed, activate your environment, log in to wandb, and run train.py.

Example command to launch training with seed 42 on GPU 0 using a specific Hydra config:

(robodiff)[diffusion_policy]$ conda activate robodiff
(robodiff)[diffusion_policy]$ wandb login
(robodiff)[diffusion_policy]$ python train.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'

Outputs (checkpoints, logs, and media) are saved to data/outputs/yyyy.mm.dd/hh.mm.ss_<method_name>_<task_name>.

An EnvRunner abstracts the differences between task environments. It inherits from a base class and provides a run method.

Key Features:

• Evaluation: The run method takes a Policy object for evaluation and returns a dictionary of logs and metrics compatible with wandb.log.
• Vectorization: To maximize speed, environments are often vectorized using a modified gym.vector.AsyncVectorEnv, which runs environments in separate processes to bypass the Python GIL.

⚠️ Warning: OpenGL and Subprocesses

When using fork on Linux (as in AsyncVectorEnv), environments that initialize an OpenGL context (like robosuite) during initialization may cause segmentation faults. Workaround: Provide a dummy_env_fn that constructs an environment without initializing OpenGL.