3D Diffusion Policy (DP3)#

Code for running DP3 based on 3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations. Adapted from linden713/ManiSkill diffusion_policy_3d and the original DP3 repository.

Integration status: DP3 uses a dedicated point-cloud data pipeline (replay -> HDF5 -> Zarr -> task YAML). It is documented in the same style as other policies, but remains operationally independent.

Directory Structure#

diffusion_policy_3d/
├── train.py
├── eval.py
├── train_policy.sh
├── eval_policy.sh
├── setup.py
├── data/
├── dataset/
│   ├── convert_hdf5_to_zarr.py
│   └── load_trajectories.py
└── diffusion_policy_3d/
    ├── config/
    │   ├── dp3.yaml
    │   └── task/
    │       ├── PickCube-v1.yaml
    │       ├── PlugCharger-v1.yaml
    │       └── PushCube-v1.yaml
    ├── policy/dp3.py
    ├── model/
    ├── dataset/
    ├── env_runner/
    └── env/

1. Installation#

conda create -n mani_skill python=3.10 -y
conda activate mani_skill

pip install torch torchvision torchaudio
pip install mani_skill

git clone https://github.com/facebookresearch/pytorch3d.git
cd pytorch3d && pip install -e .
cd ..

cd core/policies/diffusion_policy_3d
pip install -e .

conda install -y libstdcxx-ng -c conda-forge
pip install "setuptools<81"

2. Replay Trajectories#

cd /home/DMP/nmi/FGManip

python utils/replay_trajectory.py \
    --traj-path demos/PlugCharger-v1/motionplanning/trajectory.h5 \
    --use-first-env-state \
    -c pd_joint_delta_pos \
    -o pointcloud \
    --save-traj \
    --num-envs 20 \
    -b physx_cpu

Parameter	Description	Common Values
`-c`	Control mode	`pd_joint_delta_pos`, `pd_ee_delta_pose`, `pd_ee_delta_pos`
`-o`	Observation mode	`pointcloud`
`--num-envs`	Parallel envs	`20`
`-b`	Simulation backend	`physx_cpu`

Important: control mode determines action dimension, and task config must match it.

Output example:

demos/PlugCharger-v1/motionplanning/trajectory.pointcloud.pd_joint_delta_pos.cpu.h5

3. Convert to Zarr Format#

cd core/policies/diffusion_policy_3d

python dataset/convert_hdf5_to_zarr.py \
    --env_name PlugCharger-v1 \
    --num_demos 100 \
    --hdf5_path /home/DMP/nmi/FGManip/demos/PlugCharger-v1/motionplanning/trajectory.pointcloud.pd_joint_delta_pos.cpu.h5 \
    --zarr_dir ./data/

Expected output shapes:

action shape: (24142, 8)
state shape: (24142, 25)
point_cloud shape: (24142, 512, 6)

Verify Zarr data:

python -c "import zarr; z=zarr.open('./data/maniskill_PlugCharger-v1_expert.zarr','r'); print(z['data']['action'].shape)"

4. Write Task Config#

Create task YAML under diffusion_policy_3d/config/task/. All dimensions must match Zarr.

name: PlugCharger-v1
task_name: ${name}

shape_meta: &shape_meta
  obs:
    point_cloud:
      shape: [512, 3]
      type: point_cloud
    agent_pos:
      shape: [25]
      type: low_dim
  action:
    shape: [8]

env_runner:
  _target_: diffusion_policy_3d.env_runner.maniskill_runner.ManiSkillRunner
  eval_episodes: 20
  max_steps: 200
  n_obs_steps: ${n_obs_steps}
  n_action_steps: ${n_action_steps}
  control_mode: pd_joint_delta_pos
  task_name: ${task_name}
  num_eval_envs: 1
  sim_backend: cpu
  device: ${training.device}
  use_point_crop: ${policy.use_point_crop}
  use_pc_color: ${policy.use_pc_color}

dataset:
  _target_: diffusion_policy_3d.dataset.maniskill_dataset.ManiSkillDataset
  zarr_path: ./data/maniskill_PlugCharger-v1_expert.zarr
  horizon: ${horizon}
  pad_before: ${eval:'${n_obs_steps}-1'}
  pad_after: ${eval:'${n_action_steps}-1'}
  seed: 42
  val_ratio: 0.02
  max_train_episodes: 90

Control mode and action dimension reference:

Control Mode	Action Dim	Description
`pd_joint_delta_pos`	8	7 joint deltas + 1 gripper
`pd_ee_delta_pose`	7	6 EE pose deltas + 1 gripper
`pd_ee_delta_pos`	4	3 EE position deltas + 1 gripper

5. Training#

cd core/policies/diffusion_policy_3d
conda activate mani_skill

# bash train_policy.sh <alg_name> <task_name> <exp_id> <seed> <gpu_id>
bash train_policy.sh dp3 PlugCharger-v1 001 0 0

Outputs are saved to data/outputs/<task>-<alg>-<id>_seed<seed>/. Logging uses WandB.

6. Evaluation#

bash eval_policy.sh dp3 PlugCharger-v1 001 0 0

Troubleshooting#

CXXABI_1.3.15 not found

conda install -y libstdcxx-ng -c conda-forge
export LD_PRELOAD=${CONDA_PREFIX}/lib/libstdc++.so.6

ModuleNotFoundError: No module named 'pkg_resources'

pip install "setuptools<81"

Citation#

@inproceedings{ze20243d,
  title={3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations},
  author={Ze, Yanjie and Zhang, Gu and Zhang, Kangning and Hu, Chenyuan and Wang, Muhan and Xu, Huazhe},
  booktitle={ICRA 2024 Workshop on 3D Visual Representations for Robot Manipulation},
  year={2024}
}

Maintenance Checklist#

Keep replay control mode, action dimension, and task YAML action shape consistent.
Keep Zarr path naming convention stable across train/eval scripts.
Record exact train/eval commands and output directory for reproducibility.