Official codebase for "Streaming Diffusion Policy: Fast Policy Synthesis with Variable Noise Diffusion Models". This repository is based on Diffusion Policy.

See the original Diffusion Policy repo for installation.

Our contributions to the repo are:

• The definition of state- and image-based SDP, in diffusion_policy/policy, in the files tedi_unet_lowdim_policy.py and tedi_unet_hybrid_image_policy.py
• The noise scheduler in diffusion_policy/policy/schedulers.py
• The modified model architecture in diffusion_policy/model/diffusion/conditional_unet1d_tedi.py

As in Diffusion Policy, we allow for training with a single seed, or multiple seeds using a Ray-cluster.

The training uses the same command as for Diffusion Policy, only changing the .yaml file. When positioned in the root of the repo, run:

python train.py --config-dir=diffusion_policy/config/ --config-name=train_tedi_unet_lowdim_workspace.yaml

In the case of training the image-based policy, use the config train_tedi_unet_hybrid_workspace.yaml instead.

We use the same yaml-files for the multi-seed training. On a SLURM-cluster, run:

ray start --head --num-gpus=3  --port=50004

python ray_train_multirun.py --config-dir=diffusion_policy/config --config-name=train_tedi_unet_lowdim_workspace.yaml --seeds=42,43,44

Se the Diffusion Policy repo for more details. In the paper, we report max/test_mean_score as "Max" and k_min_train_loss/test_mean_score as "Avg".

We obtained all results in the paper by multi-seed training. The results for Consistency Policy are obtained by running multi-seed training with the Consistency Policy repo.

For the image-based Push-T task, we trained for 1000 epochs, and for image-based tasks in Robomimic we trained for 1500 epochs for Lift, Can, and Square, and 500 epochs for Transport and Tool-hang​.

We provide pre-trained checkpoints at the HuggingFace Hub.

If you find our work useful, please consider citing our paper:

@misc{høeg2024streamingdiffusionpolicyfast,
      title={Streaming Diffusion Policy: Fast Policy Synthesis with Variable Noise Diffusion Models}, 
      author={Sigmund H. Høeg and Yilun Du and Olav Egeland},
      year={2024},
      eprint={2406.04806},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2406.04806}, 
}

We thank the authors of Diffusion Policy and Consistency Policy for sharing their codebase.