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.

• The paper presents an energy-based framework (ECD) that composes short-horizon predictions into globally-consistent trajectories.
• It employs a mathematically-grounded correction field, including boundary reaction terms, to ensure conservative energy integration and mode commitment.
• Empirical results across high-dimensional tasks demonstrate enhanced success rates and efficiency compared to heuristic stitching methods.

Energy-based Compositional Diffusion Planning: A Formal Summary

Introduction

This work introduces Energy-based Compositional Diffuser (ECD), a principled framework for long-horizon trajectory planning using diffusion models when only short-horizon trajectories are available during training. Prior compositional diffusion planning methods (e.g., CompDiffuser, CD) rely on heuristic stitching of local predictions obtained from overlapping short-horizon chunks, but these heuristics yield non-conservative (non-integrable) update fields that cannot be interpreted as correct global score functions for any energy landscape. This introduces failures in global trajectory consistency—particularly evident in multimodal trajectory spaces or when task feasibility requires sustained mode commitment. ECD eliminates this pathology with a mathematically-grounded energy-based approach that composes chunk-level local energies and computes the global plan as the energy minimizer using conservative gradient-based updates. This formulation exposes the necessary boundary reaction terms omitted by heuristic stitching and enables efficient inference via a specialized Markov-structured approximation.

Background and Formal Problem Statement

Diffusion-based decision-making and planning operate by generatively sampling entire trajectories conditioned on boundary states (start/goal) and possibly task or environment constraints. In many robotic and sequential decision-making pipelines, long-horizon demonstrations are scarce; available datasets consist of fragmented short-horizon skills or behaviors. The compositional planning challenge is to synthesize globally-consistent, long-horizon trajectories by composing local short-horizon statistics, requiring robust stitching mechanisms.

Traditional methods such as CompDiffuser perform inference by overlaying local chunk predictions and aggregating conflicting assignments in overlapping regions via averaging or overwriting. These approaches are computationally tractable and offer strong local multimodal support, but, as shown by this work, the resulting field lacks global integrability—leading to poor commitment in multimodal spaces and spurious mode-switching.

Theoretical Analysis: Non-Conservativity of Heuristic Stitching

The paper provides a detailed theoretical analysis demonstrating that heuristic stitching, by ignoring the dependence of chunk predictions on boundary conditions, omits necessary chain-rule terms. Specifically, the effect of perturbing boundary nodes on the interior prediction (the "boundary reaction") is dropped. The result is a non-conservative update field with nonzero curl, meaning the trajectory update cannot be interpreted as the gradient of any global log-probability.

Figure 1: ECD achieves the best trade-off between planning success rate and inference runtime among prior diffusion-based and RL-based planners on OGBench stitch tasks.

Empirically, this failure manifests as poor or slow mode commitment, requiring aggressive and computationally expensive resampling heuristics (e.g., CDGS) to recover reliable plans. The non-conservative nature is confirmed both by direct mathematical proofs and by empirical studies on multimodal toy problems (see Figure 2).

Figure 2: ECD preserves mode-consistent sampling for long-horizon multimodal planning, while CD exhibits inconsistent crossings between modes when horizon increases.

Method: Energy-based Compositional Diffuser (ECD)

ECD reformulates compositional planning as energy minimization. Local chunk denoisers define prediction energies (squared Mahalanobis distances between current local states and their predicted denoised means, conditioned on boundary states). The global energy is a sum over all chunk energies, inducing a globally-consistent energy landscape.

The correction field for updating the full trajectory at each diffusion reverse step is then defined as the negative gradient of this energy. Critically, this derivative includes both:

• The standard "interior" correction, pulling interior chunk coordinates toward their local mode.
• The "boundary reaction" term, communicating chunk residuals back to the boundary variables, mediated by the derivative of the local denoiser with respect to its boundary input.

This exact reaction term can be computed by backpropagation through the denoiser, but is approximated efficiently via Markov structure assumptions using a block-tridiagonal linear solve, resulting in $\mathcal{O}(L)$ inference time per denoising step.

Figure 3: Method comparison schematic—CD ignores the feedback loop induced by boundary condition perturbations, while ECD’s correction is derived from global, conservative energy landscape gradients.

Figure 4: Schematic illustrating the three physical properties enforced by ECD’s energy formulation: interior update, boundary reaction, and chunk consensus.

Efficient Implementation and Approximate Inference

An efficient Markovian approximation is derived for the reaction term, leveraging the fact that, in planning settings, state dependencies are typically local (first-order Markov). This yields sparse linear systems amenable to fast inversion, and the overall computational complexity is competitive with existing stitching methods. Empirically, the approximation very closely tracks the exact reaction for most of the denoising process, with largest deviations only at the highest noise levels, where precise chunk-level feedback is less critical.

Figure 5: Cosine similarity and norm ratio statistics show that the Markov-based reaction term closely approximates the true Jacobian message, especially as diffusion approaches deterministic regime.

Empirical Evaluation

ECD is evaluated on OGBench long-horizon stitching tasks across PointMaze, AntMaze, HumanoidMaze, and AntSoccer environments, encompassing both low (2-4D) and high (15-17D, with joint state) dimensional planning problems. The following points summarize empirical findings:

• Success Rate: ECD either matches or outperforms all baseline methods, including CD and CDGS, particularly in long-horizon or high-dimensional environments. It achieves the best success-runtime trade-off.
• Feasibility and Replanning: ECD-produced trajectories require substantially fewer replans (when using model-predictive tracking) to reach goals, indicating better initial plan feasibility and robustness.
• Computational Efficiency: ECD’s inference time is nearly identical to CD. CDGS, which is the only previous method that achieves comparable success, incurs $10\times$ to $20\times$ higher runtime due to iterative resampling.
• Ablation: Using the exact boundary reaction further improves success rates marginally, but the efficient Markov approximation obtains most of the benefit.
• High-Dimensional Environments: The advantage of ECD persists or increases as planning dimension grows and local multimodality or consistency constraints become more severe.

  Figure 6: ECD exhibits rapid commitment to a single feasible global mode during denoising, in contrast to CompDiffuser’s delayed and scattered convergence.

  

  Figure 7: ECD outperforms CD on PointMaze and AntMaze Giant environments across all allowed replan budgets, reflecting higher initial plan quality.

Implications and Future Directions

Theoretically, ECD demonstrates that conservative (energy-based) integration of local chunk denoisers offers correct global mode propagation and eliminates the pathological independence left by heuristic stitching. In practice, it matches the speed of chunk heuristic methods while delivering much higher trajectory feasibility.

Open problems and avenues for future research include:

• Scalable Reaction Approximations: Further study of data-driven or learned surrogate models for the Jacobian boundary reaction in high-dimensional, non-Markovian settings.
• Adaptive Chunking: Dynamic adaptation of chunk size, stride, and overlap based on geometry, uncertainty, or task characteristics.
• Real-world Deployment: Application to manipulation, locomotion, and interactive settings with perception noise and temporal delays.
• Nonlinear Energy Bridge Models: Generalization beyond quadratic bridge energy, possibly leveraging learned nonparametric or contrastive chunk compatibility models.

Conclusion

ECD introduces an energy-based composition strategy for diffusion planning that addresses the key theoretical and practical limitations of heuristic stitching. The conservative energy field and boundary reaction terms restore principled long-horizon consistency, establish a valid global score function, and, through efficient approximations, achieve superior empirical performance over prior art without significant inference-time overhead.

---