Generative Trajectory Stitching through Diffusion Composition

TL;DR

This work introduces CompDiffuser, a diffusion-based planner that generates long-horizon trajectories by composing overlapping short-horizon chunks learned from short demonstrations. It models the trajectory distribution p_θ(τ|q_s,q_g) in a compositional form, with a single diffusion model conditioning on neighboring chunks and endpoints to propagate information bidirectionally and ensure dynamic consistency. Training uses a unified denoiser ε_θ with objectives 𝓛_nbr, 𝓛_start, and 𝓛_end to encode start and goal conditioning, while inference supports parallel and autoregressive sampling with exponential blending to merge chunks into a coherent plan. Extensive experiments across PointMaze and high-dimensional benchmarks demonstrate strong improvements over baselines, showing CompDiffuser can solve unseen long-horizon tasks from short-horizon data and even support test-time replanning in complex environments.

Abstract

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.

Figures (21)

• Figure 1: Compositional Trajectory Generation. CompDiffuser enables generative trajectory stitching through diffusion composition. Left: Monolithic generative planner fails to generalize to tasks of longer horizon and collapses to the maze center. Right: Our method successfully navigates the ant agent from start to goal by compositionally stitching together shorter trajectories.
• Figure 2: Illustrating the Trajectory Stitching Process. Given an unseen start (blue circle) and goal (green star), CompDiffuser generates a long-horizon plan by progressively denoising three trajectory chunks in parallel, with each chunk conditioning on its neighbors to ensure smooth transitions.
• Figure 3: Compositional Trajectory Planning: Parallel Sampling and Autoregressive Sampling. We present an illustrative example of sampling three trajectories $\tau_{1:3}$ with the proposed compositional sampling methods. Dashed lines represent cross trajectory information exchange between adjacent trajectories and black lines represent the denoising flow of each trajectory. In parallel sampling, $\tau_{1:3}$ can be denoised concurrently; while in autoregressive sampling, denoising $\tau_k$ depends on the previous trajectory $\tau_{k-1}$, e.g., the denoising of $\tau_2$ depends on $\tau_1$ (as shown in the blue horizontal dashed arrows). Additionally, start state $q_s$ and goal state $q_g$ conditioning are applied to the trajectories in the two ends, $\tau_1$ and $\tau_3$, which enables goal-conditioned planning. Trajectories $\tau_{1:3}$ will be merged to form a longer plan $\tau_{\text{comp}}$ after the full diffusion denoising process.
• Figure 4: Qualitative Comparison of DD, GSC and Comp-Diffuser on OGBench PointMaze Giant. Effective bidirection information propagation enables CompDiffuser to successfully synthesize trajectories from start (bottom left) to goal (upper right), while other methods generate o.o.d trajectories that are disconnected with the start/goal or passing through walls. See \ref{['fig:35-ncomp-qual']} for per-segment visualization.
• Figure 5: Qualitative Results of Planning in High Dimension on OGBench AntMaze Large. Original plan is sub-sampled for clearer view. Our method is able to synthesize plans of valid dynamics while reaching the goal position (bottom right). Note that our method is only trained on trajectory segments of much shorter length.