Advantage-Guided Diffusion for Model-Based Reinforcement Learning
On the map: World models → Latent dynamics models and model-based RL
also in 1 more
Abstract (arXiv)
Model-based reinforcement learning (MBRL) with autoregressive world models suffers from compounding errors, whereas diffusion world models mitigate this by generating trajectory segments jointly. However, existing diffusion guides are either policy-only, discarding value information, or reward-based, which becomes myopic when the diffusion horizon is short. We introduce Advantage-Guided Diffusion for MBRL (AGD-MBRL), which steers the reverse diffusion process using the agent's advantage estimates so that sampling concentrates on trajectories expected to yield higher long-term return beyond the generated window. We develop two guides: (i) Sigmoid Advantage Guidance (SAG) and (ii) Exponential Advantage Guidance (EAG). We prove that a diffusion model guided through SAG or EAG allows us to perform reweighted sampling of trajectories with weights increasing in state-action advantage-implying policy improvement under standard assumptions. Additionally, we show that the trajectories generated from AGD-MBRL follow an improved policy (that is, with higher value) compared to an unguided diffusion model. AGD integrates seamlessly with PolyGRAD-style architectures by guiding the state components while leaving action generation policy-conditioned, and requires no change to the diffusion training objective. On MuJoCo control tasks (HalfCheetah, Hopper, Walker2D and Reacher), AGD-MBRL improves sample efficiency and final return over PolyGRAD, an online Diffuser-style reward guide, and model-free baselines (PPO/TRPO), in some cases by a margin of 2x. These results show that advantage-aware guidance is a simple, effective remedy for short-horizon myopia in diffusion-model MBRL.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 41 of its 48 references that are in the map. First posted on arXiv 10 Apr 2026.
Builds on
41 of 48 references.
Closest ancestors · 11
Same direction, newest first
- 2024Policy-Guided Diffusion9 Apr
- 2023World Models via Policy-Guided Trajectory Diffusion13 Dec
- 2023A Generalist Dynamics Model for Control18 May
- 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
- 2023Mastering Diverse Domains through World Models10 Jan
- 2022Transformers are Sample-Efficient World Models1 Sep
- 2020Mastering Atari with Discrete World Models5 Oct
- 2018Recurrent World Models Facilitate Policy Evolution4 Sep
- 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models30 May
- 2018Model-Ensemble Trust-Region Policy Optimization28 Feb
- 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning8 Aug
Other references · 30
- 2017Attention Is All You Needcited by 59,932
- 2020Denoising Diffusion Probabilistic Modelscited by 20,851
- 2017Proximal Policy Optimization Algorithmscited by 14,441
- 2021Diffusion Models Beat GANs on Image Synthesiscited by 8,185
- 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
- 2022Classifier-Free Diffusion Guidancecited by 5,717
- 2022Flow Matching for Generative Modelingcited by 5,426
- 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusioncited by 2,730
- 2015Trust Region Policy Optimizationcited by 2,503
- 2016Asynchronous Methods for Deep Reinforcement Learningcited by 2,218
- 2022RePaint: Inpainting using Denoising Diffusion Probabilistic Modelscited by 1,419
- 2022Planning with Diffusion for Flexible Behavior Synthesiscited by 1,023
18 more
- 2017Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithmcited by 723
- 2022Is Conditional Generative Modeling all you need for Decision-Making?cited by 504
- 2019Model-Based Reinforcement Learning for Ataricited by 495
- 2018Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Reviewcited by 466
- 2023IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policiescited by 307
- 2019Benchmarking Model-Based Reinforcement Learningcited by 158
- 2023AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Plannerscited by 147
- 2023MADiff: Offline Multi-agent Learning with Diffusion Modelscited by 73
- 2024Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Modelcited by 62
- 2020Plannable Approximations to MDP Homomorphisms: Equivariance under Actionscited by 61
- 2025Efficient Online Reinforcement Learning for Diffusion Policycited by 61
- 2019Combating the Compounding-Error Problem with a Multi-step Modelcited by 52
- 2019Learning to Combat Compounding-Error in Model-Based Reinforcement Learningcited by 39
- 2023Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learningcited by 37
- 2024Diffusion Spectral Representation for Reinforcement Learningcited by 15
- 2025Prior-Guided Diffusion Planning for Offline Reinforcement Learningcited by 13
- 2023Value function estimation using conditional diffusion models for controlcited by 10
- 2025Adversarial Diffusion for Robust Reinforcement Learningcited by 4
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
- 2026Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization
- 2025Diffusion Models for Reinforcement Learning: Foundations, Taxonomy, and Development
- 2025Diffusion Guidance Is a Controllable Policy Improvement Operator
- 2023Diffusion Models for Reinforcement Learning: A Survey
- 2026How Does the Lagrangian Guide Safe Reinforcement Learning through Diffusion Models?
- 2026Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF
- 2026Beyond Imitation: Reinforcement Learning Fine-Tuning for Adaptive Diffusion Navigation Policies