KDPE: A Kernel Density Estimation Strategy for Diffusion Policy Trajectory Selection

Project Page | Paper

(Official implementation of KDPE – Accepted at CoRL 2025)
Fork of Diffusion Policy with population sampling + density filtering for robust visuomotor control

This is the code for our work KDPE: A Kernel Density Estimation Strategy for Diffusion Policy Trajectory Selection, which is based on the original Diffusion Policy implementation.

Table of Contents

• Installation
• Datasets
• Running Experiments
• License
• Citing this work

Installation

To set up the environment, please use Conda and run:

conda env create --file conda_environment.yaml
conda activate kdpe

If you are running it on a remote cluster with no sudo privileges you might need to apply this fix

Datasets

The datasets should be organized in the following structure:

data/datasets/
├── tool_hang/
│   └── ph/
│       └── image_abs.hdf5
├── square/
│   ├── ph/
│   │   └── image_abs.hdf5
│   └── mh/
│       └── image_abs.hdf5
├── assembly/
│   └── image_abs.hdf5
...

The tasks comes from two different sources:

robotmimic = [lift, can, square, tool_hang]
mimicgen = [stack, coffee, assembly]

You can download the datasets from from the following links: robomimic_dataset, mimicgen_dataset.

mimicgen datasets use relative actions and need to be converted to absolute actions. You can use the following script for conversion:

python diffusion_policy/scripts/robomimic_dataset_conversion.py --input /path/to/your/dataset.hdf5 --output /path/to/your/output_dataset.hdf5

Running Experiments

First, you need to train a model to obtain a checkpoint.

python train.py \
    --config-dir diffusion_policy/config \
    --config-name train_diffusion_unet_hybrid_workspace.yaml \
    task=tool_hang

This will produce a checkpoint named step_79999.ckpt under the path data/weights/diffusion_cnn_tool_hang_ph/checkpoints.

To test our method (KDPE), run:

python eval.py --model cnn --task tool_hang_ph --Filter.name kde

To test the baseline, run:

python eval.py --model cnn --task tool_hang_ph --Filter.name passall

License

The original Diffusion Policy code is licensed under the MIT License (licenses/LICENSE_DP). The modifications introduced in this repository from commit e68659ad98540c9886b79d2deac890dada097d22 onwards are licensed under the Apache License (licenses/LICENSE_KDPE).

Citation

If you find our method or visualizer useful, consider citing it.