Energy-based Compositional Diffusion Planning
Abstract
Compositional diffusion planners aim to solve long-horizon robotic tasks using short training trajectories. Yet, current approaches often rely on the heuristic stitching of local predictions. We show that the resulting stitched update is generally a non-conservative field} that does not mathematically correspond to any valid global trajectory log-density function. We propose Energy-based Compositional Diffuser (ECD), a framework that formulates the global trajectory as the minimizer of the sum of local bridge potentials. This energy-based perspective defines a conservative correction field and contains a boundary reaction term that heuristic stitching omits. To enable efficient inference, we further introduce a Markov-based score approximation that computes the reaction term via a single block-tridiagonal solve, maintaining time complexity linear in the planning horizon. Empirically, ECD achieves state-of-the-art success rates on a range of OGBench stitching tasks, while nearly matching the inference speed of heuristic stitching methods. Code is available at https://github.com/GradientSpaces/ECD.
Lay Summary
Robots often need to plan long tasks, but their training data usually contains only short motion clips. Existing diffusion planners extend these clips by stitching together many short predictions, like patching small road segments into a full route. We show that this stitching can create inconsistent plans because it does not follow any single global objective. We propose Energy-based Compositional Diffuser (ECD), which instead treats the whole plan as minimizing one global energy built from local chunks. This reveals a boundary-reaction correction that passes consistency information between neighboring chunks. We also make this correction efficient with a Markov approximation, so runtime stays close to fast stitching methods. On OGBench long-horizon planning tasks, ECD produces more reliable trajectories while remaining practical to run.