Improving Diffusion Planners by Self-Supervised Action Gating with Energies

TL;DR

Self-supervised Action Gating with Energies (SAGE), an inference-time re-ranking method that penalises dynamically inconsistent plans using a latent consistency signal, improves the performance and robustness of diffusion planners.

Abstract

Diffusion planners are a strong approach for offline reinforcement learning, but they can fail when value-guided selection favours trajectories that score well yet are locally inconsistent with the environment dynamics, resulting in brittle execution. We propose Self-supervised Action Gating with Energies (SAGE), an inference-time re-ranking method that penalises dynamically inconsistent plans using a latent consistency signal. SAGE trains a Joint-Embedding Predictive Architecture (JEPA) encoder on offline state sequences and an action-conditioned latent predictor for short horizon transitions. At test time, SAGE assigns each sampled candidate an energy given by its latent prediction error and combines this feasibility score with value estimates to select actions. SAGE can integrate into existing diffusion planning pipelines that can sample trajectories and select actions via value scoring; it requires no environment rollouts and no policy re-training. Across locomotion, navigation, and manipulation benchmarks, SAGE improves the performance and robustness of diffusion planners.

Figures (11)

• Figure 1: The SAGE framework. (a) Learn a predictive JEPA state representation from masked state windows with an EMA teacher. (b) Train an action-conditioned latent predictor whose prediction error defines a transition energy. (c) At test time, score diffusion-generated candidate plans with this energy and gate value-based selection toward locally feasible actions
• Figure 2: SAGE energy localises feasibility violations. Per-step latent-consistency energy on an offline episode with clean and after corrupted action segment; corruption induces a sharp, local spike (shaded interval). Full D4RL results in Appendix \ref{['app:addexp_diagnostics']} Figure \ref{['fig:traj_all']}.
• Figure 3: 100 trajectories example sampled from Maze2d, MCSS can selects wall-crossing or out-of-bounds trajectories (left), while SAGE’s prefix energy filtering and soft penalty suppresses these failures without collapsing trajectory diversity (right).
• Figure 4: Average performance as a function of the prefix window length $K$ across four evaluation domains.
• Figure 5: Average performance as a function of the keep rate $\mathcal{P}$ across four evaluation domains.