Neural Dynamics Augmented Diffusion Policy

(ICRA 2025)

Project Page | Video

Installation

1. Clone the repository

git clone https://github.com/warshallrho/Dynamics-DP.git
cd Dynamics-DP

2. Create & activate a virtual environment

conda create -n dynamicsdp python==3.10
conda activate dynamicsdp

3. install dependencies

# install torch 
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu118
pip install -e .

Data & Checkpoints

1. Download dataset and checkpoints Download files from here and decompressed it to the project folder. The folder structure should show as below:

├── assets
├── data
│   └── insertT_sim
├── dynamicsdp
│   ├── __init__.py
│   ├── configs
│   ├── insertT
│   ├── stow
│   └── utils
├── LICENSE
├── outputs
│   └── model
│       └── insertT_sim
│           ├── dfp
│           ├── dynamics
│           └── planning
├── pyproject.toml
└── README.md

Usage

Interaction Environment

python dynamicsdp/insertT/env/insertT_env.py

You can also collection your own training dataset with following command

python dynamicsdp/insertT/teleop/human_demo_collection.py --record

Tips for human_demo_collection

• Move mouse slowly during tele-operation for better performance
• See the human_demo_data section in dynamicsdp/configs/insert_config.yaml for more option

Dynamics Model

Automaticaly generate self-play data for training dynamics model

python dynamicsdp/insertT/dynamics/data_generation.py

Train dynamics model

python dynamicsdp/insertT/dynamics/train.py

Visualize the performance of the dynamics model

python dynamicsdp/insertT/dynamics/visualization.py

It is also recommanded to test the dynamics model with interaction environment

python dynamicsdp/insertT/dynamics/interaction_demo.py

Data Augmentation

With the dynamics model, augmented demo data can be generated for training diffusion policy

python dynamicsdp/insertT/planning/mppi_data_generation.py

You can also visualize how mppi planner working during data augmentation

python dynamicsdp/insertT/planning/mppi_visualization.py

Diffusion Policy

Train diffusion policy

python dynamicsdp/insertT/dfp/train.py

Evaluate diffusion policy

python dynamicsdp/insertT/dfp/evaluation.py

Citation

@inproceedings{wu2025neural,
    title={Neural Dynamics Augmented Diffusion Policy},
    author={Wu, Ruihai and Chen, Haozhe and Zhang, Mingtong and Lu, Haoran and Li, Yitong and Li, Yunzhu},
    booktitle={IEEE International Conference on Robotics and Automation (ICRA)},
    year={2025}
}