This repository contains the code to our paper "Diffusion Predictive Control with Constraints". We build upon the temporal U-Net implementation from Diffuser and use the Avoiding environment from D3IL.

Installation

Clone our repo and run:

pip install torch==2.2.2+cu121 torchvision==0.17.2+cu121 torchaudio==2.2.2+cu121 -f https://download.pytorch.org/whl/torch_stable.html
pip install -r requirements.txt

You also need to install D3IL for the simulation environment.

Training

To train the trajectory diffusion model, run:

python scripts/train.py

You can also visualize the training data without constraints and the novel test-time constraints:

python visualize_data_constraints/train.py

Testing

To evaluate DPCC and reproduce the results reported in the paper, run:

python scripts/eval.py
python scripts/load_results.py