Improving Diffusion Planners by Self-Supervised Action Gating with Energies

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 integrates 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.

We present Self-supervised Action Gating with Energies (SAGE), a method for improving diffusion-based planners in offline reinforcement learning. Diffusion planners generate many possible future action sequences and select the one predicted to achieve the highest reward, but they can still choose trajectories that are unrealistic or dynamically inconsistent with the environment. SAGE addresses this issue by introducing a separate feasibility signal that evaluates whether the early part of a generated plan is locally executable. The method learns this signal entirely from offline data using self-supervised predictive representations and an action-conditioned latent dynamics model. At inference time, SAGE assigns an energy score to each candidate plan based on latent prediction consistency and combines this with existing value estimates to filter out implausible trajectories. SAGE is modular, requires no retraining of the planner, and consistently improves robustness and performance across locomotion, navigation, and robotic manipulation benchmarks.