Tao Sun

All publications

ICML 2026

Energy-based Compositional Diffusion Planning

Tao Sun, Utkarsh Aashu Mishra, Jiaxin Lu, Danfei Xu, Iro Armeni

International Conference on Machine Learning (ICML) 2026

Paper Code OpenReview BibTeX
AntMaze trajectories comparing CompDiffuser and ECD: blue paths reach the goal without hitting walls, while red paths collide.
Energy-based composition corrects the boundaries between short trajectory fragments.

Abstract

Energy-based Compositional Diffuser (ECD) composes short diffusion-generated trajectory fragments into long-horizon robot plans. It defines a global trajectory energy whose gradient corrects inconsistencies at fragment boundaries. An efficient Markov-based approximation keeps inference cost linear in the planning horizon. Experiments on OGBench stitching tasks evaluate planning success and inference efficiency.

BibTeX

@inproceedings{sun2026energy,
  title={Energy-based Compositional Diffusion Planning},
  author={Sun, Tao and Mishra, Utkarsh Aashu and Lu, Jiaxin and Xu, Danfei and Armeni, Iro},
  booktitle={International Conference on Machine Learning},
  year={2026},
  url={https://arxiv.org/abs/2606.21646}
}