Advantage-Guided Diffusion for Model-Based Reinforcement Learning

📄 arXiv: 2604.09035v1 📥 PDF

作者: Daniele Foffano, Arvid Eriksson, David Broman, Karl H. Johansson, Alexandre Proutiere

分类: cs.AI, cs.LG

发布日期: 2026-04-10

💡 一句话要点

提出Advantage引导的扩散模型(AGD-MBRL)，提升基于扩散模型的模型强化学习性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模型强化学习 扩散模型 优势函数 轨迹生成 机器人控制

📋 核心要点

  1. 基于扩散模型的MBRL虽然能缓解误差累积，但现有引导方法要么忽略价值信息，要么因奖励短视而效果不佳。
  2. AGD-MBRL利用智能体的优势估计引导扩散过程，使采样集中于具有更高长期回报的轨迹，从而改进策略。
  3. 实验表明，AGD-MBRL在MuJoCo控制任务上显著提升了样本效率和最终回报，优于现有方法。

📝 摘要（中文）

基于自回归世界模型的模型强化学习(MBRL)存在误差累积问题，而扩散世界模型通过联合生成轨迹片段来缓解这个问题。然而，现有的扩散引导要么仅依赖策略而忽略价值信息，要么基于奖励而导致短视。本文提出了Advantage引导的扩散模型(AGD-MBRL)，它利用智能体的优势估计来引导反向扩散过程，从而使采样集中在预期产生更高长期回报的轨迹上。我们开发了两种引导方法：(i) Sigmoid Advantage Guidance (SAG)和(ii) Exponential Advantage Guidance (EAG)。我们证明了通过SAG或EAG引导的扩散模型允许我们对轨迹进行重加权采样，权重随着状态-动作优势的增加而增加，这意味着在标准假设下策略得到改进。此外，我们表明与无引导的扩散模型相比，AGD-MBRL生成的轨迹遵循改进的策略（即具有更高的价值）。AGD通过引导状态分量同时保持动作生成策略条件化，与PolyGRAD风格的架构无缝集成，并且不需要更改扩散训练目标。在MuJoCo控制任务（HalfCheetah、Hopper、Walker2D和Reacher）上，AGD-MBRL在样本效率和最终回报方面优于PolyGRAD、在线Diffuser风格的奖励引导以及无模型基线（PPO/TRPO），在某些情况下提升幅度高达2倍。这些结果表明，优势感知引导是扩散模型MBRL中短视问题的简单有效的补救措施。

🔬 方法详解

问题定义：基于扩散模型的模型强化学习旨在学习一个世界模型，该模型能够生成环境的未来状态和奖励。现有的方法，例如直接使用奖励作为扩散过程的引导信号，容易陷入短视问题，因为它们只关注生成窗口内的即时奖励，而忽略了长期回报。此外，一些方法仅依赖策略信息，而忽略了价值信息，导致引导效果不佳。

核心思路：AGD-MBRL的核心思路是利用智能体的优势函数来引导扩散过程。优势函数表示在给定状态下采取某个动作相对于平均策略的优势程度。通过将优势函数作为引导信号，AGD-MBRL可以鼓励扩散模型生成具有更高长期回报的轨迹，从而避免短视问题，并提升策略性能。

技术框架：AGD-MBRL的整体框架包括以下几个主要模块：1) 扩散模型：用于生成轨迹片段。2) 策略网络：用于生成动作。3) 价值网络：用于估计状态的价值。4) 优势函数估计器：利用策略网络和价值网络来估计状态-动作的优势。5) 扩散引导器：利用优势函数来引导扩散过程，生成具有更高优势的轨迹。AGD-MBRL与PolyGRAD风格的架构无缝集成，通过引导状态分量，同时保持动作生成策略条件化。

关键创新：AGD-MBRL的关键创新在于提出了两种基于优势函数的扩散引导方法：Sigmoid Advantage Guidance (SAG)和Exponential Advantage Guidance (EAG)。SAG和EAG通过不同的方式将优势函数融入到扩散过程中，从而引导扩散模型生成更优的轨迹。此外，AGD-MBRL不需要更改扩散模型的训练目标，可以直接应用于现有的扩散模型。

关键设计：AGD-MBRL的关键设计包括：1) 优势函数的估计方法：使用时序差分学习或蒙特卡洛方法来估计优势函数。2) SAG和EAG的具体形式：SAG使用Sigmoid函数将优势函数映射到[0, 1]区间，EAG使用指数函数来放大优势函数的影响。3) 扩散引导的强度：通过调整引导系数来控制优势函数对扩散过程的影响。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，AGD-MBRL在MuJoCo控制任务（HalfCheetah、Hopper、Walker2D和Reacher）上显著优于现有方法，包括PolyGRAD、在线Diffuser风格的奖励引导以及无模型基线（PPO/TRPO）。在某些情况下，AGD-MBRL的样本效率和最终回报提升幅度高达2倍。这些结果验证了优势感知引导在扩散模型MBRL中的有效性。

🎯 应用场景

AGD-MBRL可应用于各种需要长期规划和决策的强化学习任务，例如机器人控制、自动驾驶、游戏AI等。通过利用优势函数引导扩散过程，AGD-MBRL可以提升智能体的学习效率和最终性能，使其能够更好地适应复杂环境并完成各种任务。该方法在样本效率和最终性能上的提升，使其在资源受限的场景下具有重要价值。

📄 摘要（原文）

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.