Generative Trajectory Stitching through Diffusion Composition

Yilun Du (Google Deepmind / Harvard) · Danfei Xu (Georgia Tech & NVIDIA) · Yunhao Luo (University of Michigan - Ann Arbor) · Utkarsh Mishra (Georgia Institute of Technology)
agent state dimensionbenchmark tasksbidirectional diffusion modelcompositionallyconditional relationshipsdiffusion modelsgenerative approachinformation propagationlong-horizon planningoverlapping chunksphysically consistent connectionsrobotic decision-makingtraining data qualitytrajectory distributiontrajectory stitching

Effective trajectory stitching for long-horizon planning is a significant challenge in robotic decision-making. While diffusion models have shown promise in planning, they are limited to solving tasks similar to those seen in their training data. We propose CompDiffuser, a novel generative approach that can solve new tasks by learning to compositionally stitch together shorter trajectory chunks from previously seen tasks. Our key insight is modeling the trajectory distribution by subdividing it into overlapping chunks and learning their conditional relationships through a single bidirectional diffusion model. This allows information to propagate between segments during generation, ensuring physically consistent connections. We conduct experiments on benchmark tasks of various difficulties, covering different environment sizes, agent state dimension, trajectory types, training data quality, and show that CompDiffuser significantly outperforms existing methods.