# Improving Diffusion Planners by Self-Supervised Action Gating with Energies

arXiv: https://arxiv.org/abs/2603.02650 · HTML: https://mlsanity.com/paper/2603.02650/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for offline RL, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 46 of its 55 references that are in the map.
- First posted on arXiv 3 Mar 2026.

## Placed in

- Offline RL and imitation learning from data / [Diffusion and flow policies for offline RL](https://mlsanity.com/offlinerl/diffusion-and-flow-policies-for-offline-rl/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 30 of 46)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Self-supervised Learning: Generative or Contrastive](https://mlsanity.com/paper/2006.08218/) (2020)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies](https://mlsanity.com/paper/2304.10573/) (2023)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [Efficient Diffusion Policies for Offline Reinforcement Learning](https://mlsanity.com/paper/2305.20081/) (2023)
- [Learning to Poke by Poking: Experiential Learning of Intuitive Physics](https://mlsanity.com/paper/1606.07419/) (2016)
- [Consistency Models as a Rich and Efficient Policy Class for Reinforcement Learning](https://mlsanity.com/paper/2309.16984/) (2023)
- [Diffusion Models for Reinforcement Learning: A Survey](https://mlsanity.com/paper/2311.01223/) (2023)
- [S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning](https://mlsanity.com/paper/2103.06326/) (2021)
- [Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model](https://mlsanity.com/paper/2401.10700/) (2024)
- [Diffusion Policies for Out-of-Distribution Generalization in Offline Reinforcement Learning](https://mlsanity.com/paper/2307.04726/) (2023)
- [DiffuserLite: Towards Real-time Diffusion Planning](https://mlsanity.com/paper/2401.15443/) (2024)
- [Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning](https://mlsanity.com/paper/2407.20109/) (2024)

## Built on by (most cited first, 1 of 1)

- [Refining Compositional Diffusion for Reliable Long-Horizon Planning](https://mlsanity.com/paper/2605.03075/) (2026)

## Abstract (arXiv)

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.