mlsanityData to 7 Oct 2026

Improving Diffusion Planners by Self-Supervised Action Gating with Energies

arXiv 2603.02650 ↗First posted 3 Mar 2026cs.LG cs.AI cs.ROcited by 1

On the map: Offline RL and imitation learning from data → Diffusion and flow policies for offline RL

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.

Key facts. Cited by 1 paper in the map (data through 7 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.

Builds on

46 of 55 references.

Closest ancestors · 13

Same direction, newest first

  1. 2025Consistency Trajectory Planning: High-Quality and Efficient Trajectory Optimization for Offline Model-Based Reinforcement Learning13 Jul
  2. 2025Modular Diffusion Policy Training: Decoupling and Recombining Guidance and Diffusion for Offline RL19 May
  3. 2025Prior-Guided Diffusion Planning for Offline Reinforcement Learning16 May
  4. 2025What Makes a Good Diffusion Planner for Decision Making?1 Mar
  5. 2025Behavior-Regularized Diffusion Policy Optimization for Offline Reinforcement Learning7 Feb
  6. 2024Diffusion-DICE: In-Sample Diffusion Guidance for Offline Reinforcement Learning29 Jul
  7. 2024Entropy-regularized Diffusion Policy with Q-Ensembles for Offline Reinforcement Learning6 Feb
  8. 2023Consistency Models as a Rich and Efficient Policy Class for Reinforcement Learning29 Sep
  9. 2023Diffusion Policies for Out-of-Distribution Generalization in Offline Reinforcement Learning10 Jul
  10. 2023Efficient Diffusion Policies for Offline Reinforcement Learning31 May
  11. 2023IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies20 Apr
  12. 2022Is Conditional Generative Modeling all you need for Decision-Making?28 Nov
1 more
  1. 2022Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning12 Aug

Other references · 33

  1. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  2. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
  3. 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusioncited by 2,730
  4. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,873
  5. 2020Conservative Q-Learning for Offline Reinforcement Learningcited by 1,777
  6. 2018World Modelscited by 1,444
  7. 2020D4RL: Datasets for Deep Data-Driven Reinforcement Learningcited by 1,406
  8. 2021Offline Reinforcement Learning with Implicit Q-Learningcited by 1,245
  9. 2018Off-Policy Deep Reinforcement Learning without Explorationcited by 1,227
  10. 2022Planning with Diffusion for Flexible Behavior Synthesiscited by 1,023
  11. 2021VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learningcited by 1,013
  12. 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecturecited by 982
21 more
  1. 2019Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reductioncited by 833
  2. 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Modelscited by 748
  3. 2020CURL: Contrastive Unsupervised Representations for Reinforcement Learningcited by 745
  4. 2020Self-supervised Learning: Generative or Contrastivecited by 693
  5. 2020Data-Efficient Reinforcement Learning with Self-Predictive Representationscited by 309
  6. 2022Contrastive Learning as Goal-Conditioned Reinforcement Learningcited by 241
  7. 2016Learning to Poke by Poking: Experiential Learning of Intuitive Physicscited by 152
  8. 2021S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learningcited by 85
  9. 2023Diffusion Models for Reinforcement Learning: A Surveycited by 84
  10. 2024Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Modelcited by 62
  11. 2024DiffuserLite: Towards Real-time Diffusion Planningcited by 40
  12. 2024CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Makingcited by 38
  13. 2023Reasoning with Latent Diffusion in Offline Reinforcement Learningcited by 33
  14. 2023Adaptive Online Replanning with Diffusion Modelscited by 33
  15. 2025Monte Carlo Tree Diffusion for System 2 Planningcited by 33
  16. 2023Refining Diffusion Planner for Reliable Behavior Synthesis by Automatic Detection of Infeasible Planscited by 27
  17. 2024T-JEPA: A Joint-Embedding Predictive Architecture for Trajectory Similarity Computationcited by 11
  18. 2025Local Manifold Approximation and Projection for Manifold-Aware Diffusion Planningcited by 8
  19. 2024A Unifying Framework for Action-Conditional Self-Predictive Reinforcement Learningcited by 7
  20. 2025Temporal Alignment Guidance: On-Manifold Sampling in Diffusion Modelscited by 6
  21. 2025Diffusion-based Planning with Learned Viability Filterscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026Refining Compositional Diffusion for Reliable Long-Horizon Planningcited by 2

Similar, not cited

Close in wording, no citation either way

  1. 2025Accelerating Diffusion Planners in Offline RL via Reward-Aware Consistency Trajectory Distillation
  2. 2026Amortizing Trajectory Diffusion with Keyed Drift Fields
  3. 2023Efficient Planning with Latent Diffusion
  4. 2026Shortcut Trajectory Planning for Efficient Offline Reinforcement Learning
  5. 2024Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner
  6. 2025State-Covering Trajectory Stitching for Diffusion Planners
  7. 2025DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions
  8. 2026SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning