Advantage-Guided Diffusion for Model-Based Reinforcement Learning

• URL: http://arxiv.org/abs/2604.09035v1
• Date: Fri, 10 Apr 2026 06:53:25 GMT
• Title: Advantage-Guided Diffusion for Model-Based Reinforcement Learning
• Abstract summary: Advantage-Guided Diffusion for MBRL steers the reverse diffusion process using the agent's advantage estimates.<n>We show that the trajectories generated from AGD-MBRL follow an improved policy compared to an unguided diffusion model.
• Score: 38.18017161791996
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Joint Alignment and Distillation for Video Generation via Sample-Guided Distribution Matching [36.00855152229835]
  We propose a unified, single-stage optimization framework grounded in Distribution Matching.<n>In the standard DM framework, distillation updates the model via a gradient direction that minimizes the gap between the real and fake models.<n>We introduce DM-Align, which derives a complementary gradient direction to guide the model toward human-preferred samples.
  arXiv  Detail & Related papers  (2026-09-03T04:46:46Z)
• Step Back to Move Forward: Reflection-Aware Preference Optimization for Visual Generation [12.58348084349179]
  Reflection-Aware GRPO is a new RL-based preference alignment framework for diffusion generative models.<n>We introduce Diffusion Reflection, which rectifies intermediate sampling trajectories by inverting the diffusion process with a weak estimator.<n>We also introduce Counterfactual Path Synthesis to implicitly distill these rectified trajectories into the policy, enabling the model to internalize the benefits of search-based exploration.
  arXiv  Detail & Related papers  (2026-09-03T04:37:19Z)
• MAR-GRPO: Stabilized GRPO for AR-diffusion Hybrid Image Generation [24.618644100413018]
  Reinforcement learning (RL) has been successfully applied to autoregressive (AR) and diffusion models.<n> extending RL to hybrid AR-diffusion frameworks remains challenging due to interleaved inference and noisy log-probability estimation.<n>In this work, we study masked autoregressive models (MAR) and show that the diffusion head plays a critical role in training dynamics.
  arXiv  Detail & Related papers  (2026-04-08T11:30:35Z)
• Data-regularized Reinforcement Learning for Diffusion Models at Scale [99.01056178660538]
  We introduce Data-regularized Diffusion Reinforcement Learning ( DDRL), a novel framework that uses the forward KL divergence to anchor the policy to an off-policy data distribution.<n>With over a million GPU hours of experiments and ten thousand double-blind evaluations, we demonstrate that DDRL significantly improves rewards while alleviating the reward hacking seen in RLs.
  arXiv  Detail & Related papers  (2025-12-03T23:45:07Z)
• TAG:Tangential Amplifying Guidance for Hallucination-Resistant Diffusion Sampling [53.61290359948953]
  Tangential Amplifying Guidance (TAG) operates solely on trajectory signals without modifying the underlying diffusion model.<n>We formalize this guidance process by leveraging a first-order Taylor expansion.<n> TAG is a plug-and-play, architecture-agnostic module that improves diffusion sampling fidelity with minimal computational addition.
  arXiv  Detail & Related papers  (2025-10-06T06:53:29Z)
• Plug-and-Play Prompt Refinement via Latent Feedback for Diffusion Model Alignment [54.17386822940477]
  We introduce PromptLoop, a plug-and-play reinforcement learning framework that incorporates latent feedback into step-wise prompt refinement.<n>This design achieves a structural analogy to the Diffusion RL approach, while retaining the flexibility and generality of prompt-based alignment.
  arXiv  Detail & Related papers  (2025-10-01T02:18:58Z)
• DiffusionNFT: Online Diffusion Reinforcement with Forward Process [99.94852379720153]
  Diffusion Negative-aware FineTuning (DiffusionNFT) is a new online RL paradigm that optimize diffusion models directly on the forward process via flow matching.<n>DiffusionNFT is up to $25times$ more efficient than FlowGRPO in head-to-head comparisons, while being CFG-free.
  arXiv  Detail & Related papers  (2025-09-19T16:09:33Z)
• VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL [28.95582264086289]
  VAlue-based Reinforced Diffusion (VARD) is a novel approach that first learns a value function predicting expection of rewards from intermediate states.<n>Our method maintains proximity to the pretrained model while enabling effective and stable training via backpropagation.
  arXiv  Detail & Related papers  (2025-05-21T17:44:37Z)
• Prior-Guided Diffusion Planning for Offline Reinforcement Learning [5.819784482811376]
  Prior Guidance (PG) is a novel guided sampling framework that replaces the standard Gaussian prior of a behavior-cloned diffusion model.<n>We present an efficient training strategy that applies behavior regularization in latent space, and empirically demonstrate that PG outperforms state-the-art diffusion policies and planners across diverse long-horizon offline RL benchmarks.
  arXiv  Detail & Related papers  (2025-05-16T05:39:02Z)
• Diffusion Classifier-Driven Reward for Offline Preference-based Reinforcement Learning [45.95668702930697]
  We propose a novel preference-based reward acquisition method: Diffusion Preference-based Reward (DPR)<n>DPR directly treats step-wise preference-based reward acquisition as a binary classification and utilizes the robustness of diffusion classifiers to infer step-wise rewards discriminatively.<n>We also propose Diffusion Preference-based Reward (C-DPR), which conditions on trajectory-wise preference labels to enhance reward inference.
  arXiv  Detail & Related papers  (2025-03-03T03:49:38Z)
• Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review [63.31328039424469]
  This tutorial provides a comprehensive survey of methods for fine-tuning diffusion models to optimize downstream reward functions. We explain the application of various RL algorithms, including PPO, differentiable optimization, reward-weighted MLE, value-weighted sampling, and path consistency learning.
  arXiv  Detail & Related papers  (2024-07-18T17:35:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.