Refining Compositional Diffusion for Reliable Long-Horizon Planning

[Project page] [Paper] [ArXiv]

Kyowoon Lee¹, Yunhao Luo², Anh Tong³, Jaesik Choi^1,4

¹KAIST, ²University of Michigan, ³Korea University, ⁴INEEJI

This is the official implementation for "Refining Compositional Diffusion for Reliable Long-Horizon Planning" (RCD), built on top of CompDiffuser.

RCD is a training-free guidance method for compositional diffusion planning. It combines (1) a self-reconstruction error that serves as a density proxy for the composed plan distribution and (2) an overlap consistency term that penalizes score disagreement at segment boundaries, steering compositional sampling toward high-density, globally coherent plans without any additional training.

Visual comparison with CompDiffuser on AntMaze-Giant-Stitch

Task 3	Task 5

🛠️ Installation

The following procedure should work well for a GPU machine with cuda 12.1. Our machine is installed with Ubuntu 20.04.6 LTS.

Please follow the steps below to create a conda environment to reproduce our results on OGBench environments.

1. Create a python env.

conda create -n rcd_ogbench python=3.9.20

2. Install packages in requirements.txt.

pip install -r conda_env/requirements.txt

3. Install gymnasium, gymnasium-robotics, and torch 2.5.0.

./conda_env/install_pre.sh

4. Clone the CompDiffuser-customized OGBench codebase ogbench_cpdfu_release and install it into the conda env rcd_ogbench.

## Please set the OGBench path in the script before running this line
./conda_env/install_ogb.sh

After these steps, you can use conda env rcd_ogbench to launch experiments.

🧪 Toy Experiment

A self-contained reproduction of the toy bimodal composition example (Figure 1 of the paper) is provided under toy_example/. It includes training of the three local diffusion models on length-3 segments, compositional inference at horizons $L \in {4, \ldots, 20}$, and figure assembly. See toy_example/README.md for instructions.

CompDiffuser	RCD

📊 Using Pretrained Models

Since RCD is a training-free method, one can directly reuse the pre-trained CompDiffuser planners and inverse dynamics models released by Luo et al. Please follow the CompDiffuser release for the download links and corresponding config files. After unzipping, the files should be placed under logs/ with the following structure:

└── logs
    ├── ${environment_1}
    │   ├── diffusion
    │   │   └── ${experiment_name}
    │   │       ├── model-${iter}.pt
    │   │       └── {dataset, trainer}_config.pkl
    │   └── plans
    │       └── ${experiment_name}
    │           └── ...
    │
    ├── ${environment_2}
    │   └── ...

🚀 Running RCD

We provide an evaluation entry point at diffuser/ogb_task/ogb_maze_v1/plan_ogb_stgl_sml_rcd.py, which monkey-patches the compose operators for RCD guidance and then delegates to the standard rollout loop. $GPU is the GPU index, $N_EP the number of episodes (use 100 for the full evaluation), $SEED an integer seed. For example:

AntMaze-Giant-Stitch

CUDA_VISIBLE_DEVICES=$GPU \
python diffuser/ogb_task/ogb_maze_v1/plan_ogb_stgl_sml_rcd.py \
    --config config/ogb_ant_maze/og_antM_Gi_o2d_Cd_Stgl_PadBuf_Ft64_ts512.py \
    --logbase logs --save_logbase logs_rollout_rcd \
    --pl_seeds $SEED --plan_n_ep $N_EP \
    --eval_method rcd --ev_meta_method rcd \
    --ev_density_p_ratio 0.40 \
    --ev_global_density_weight 0.25 \
    --ev_global_density_inter_rate 5 \
    --ev_global_density_n_mc 1 \
    --ev_global_density_proxy_type coupled_recon \
    --ev_global_density_proxy_overlap_weight 0.5

PointMaze-Giant-Stitch

CUDA_VISIBLE_DEVICES=$GPU \
python diffuser/ogb_task/ogb_maze_v1/plan_ogb_stgl_sml_rcd.py \
    --config config/ogb_pnt_maze/og_pntM_Gi_o2d_Cd_Stgl_PadBuf_Ft64_ts512.py \
    --logbase logs --save_logbase logs_rollout_rcd \
    --pl_seeds $SEED --plan_n_ep $N_EP \
    --eval_method rcd --ev_meta_method rcd \
    --ev_density_p_ratio 0.40 \
    --ev_global_density_weight 0.25 \
    --ev_global_density_inter_rate 5 \
    --ev_global_density_n_mc 1 \
    --ev_global_density_proxy_type coupled_recon \
    --ev_global_density_proxy_overlap_weight 0.5

Replace the --config path with any other env-specific config under config/. The --eval_method flag also accepts compdiffuser and cdgs for the baselines reported in the paper. Rollout videos and per-episode rollout JSON summaries are saved under logs_rollout_rcd/${env_name}/plans/${experiment_name}/${seed}/${experiment_time}/.

🏷️ License

This repository is released under the MIT license. See LICENSE for additional details.

🙏 Acknowledgement

• Thanks to CompDiffuser and CDGS, on which our implementation is built.
• We thank the OGBench authors for the evaluation benchmark.