World Models via Policy-Guided Trajectory Diffusion

📄 arXiv: 2312.08533v4 📥 PDF

作者: Marc Rigter, Jun Yamada, Ingmar Posner

分类: cs.LG, cs.AI

发布日期: 2023-12-13 (更新: 2024-03-27)

备注: Published in TMLR, March 2024

💡 一句话要点

提出PolyGRAD：一种基于策略引导轨迹扩散的非自回归世界模型

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

关键词: 世界模型 轨迹预测 扩散模型 强化学习 策略引导

📋 核心要点

  1. 现有世界模型依赖自回归方式，预测误差随轨迹增长而累积，限制了长期预测的准确性。
  2. PolyGRAD利用扩散模型，通过策略梯度引导，一次性生成完整轨迹，避免了自回归模型的误差累积问题。
  3. 实验表明，PolyGRAD在短轨迹预测上优于或媲美现有方法，且计算成本更低，并成功应用于MuJoCo环境。

📝 摘要（中文）

世界模型是开发智能体的重要工具。通过预测一系列动作的结果，世界模型能够利用合成数据，即“在想象中”，通过on-policy强化学习（RL）来优化策略。现有的世界模型是自回归的，它们交替预测下一个状态和从策略中采样下一个动作。预测误差不可避免地随着轨迹长度的增长而累积。本文提出了一种新颖的非自回归世界建模方法，该方法通过扩散模型一次性生成整个on-policy轨迹。我们的方法，策略引导轨迹扩散（PolyGRAD），利用去噪模型以及策略的动作分布的梯度，将初始随机状态和动作的轨迹扩散到on-policy合成轨迹中。我们分析了PolyGRAD、基于分数的生成模型和分类器引导的扩散模型之间的联系。结果表明，除了自回归扩散之外，PolyGRAD在短轨迹的轨迹预测误差方面优于最先进的基线。对于短轨迹，PolyGRAD获得了与自回归扩散相似的误差，但计算需求较低。对于长轨迹，PolyGRAD获得了与基线相当的性能。我们的实验表明，PolyGRAD能够通过在MuJoCo连续控制领域中的想象中进行on-policy RL来训练高性能策略。因此，PolyGRAD为精确的on-policy世界建模引入了一种新的范例，无需自回归采样。

🔬 方法详解

问题定义：论文旨在解决世界模型中长期轨迹预测不准确的问题。现有自回归世界模型在预测未来状态时，需要逐步采样动作并预测状态，导致预测误差随着时间步的增加而累积，限制了模型在长期规划和控制中的应用。

核心思路：论文的核心思路是利用扩散模型生成完整的轨迹，而不是像传统方法那样逐步预测。通过将轨迹视为一个整体，并使用策略梯度引导扩散过程，可以生成与策略一致的、高质量的合成轨迹，从而避免了自回归模型的误差累积问题。

技术框架：PolyGRAD的整体框架包括以下几个主要模块：1) 轨迹扩散模型：使用扩散模型将随机噪声轨迹逐步去噪，生成符合环境动态的轨迹。2) 策略引导：利用策略的动作分布梯度，引导扩散过程生成与策略一致的轨迹。3) 强化学习：使用生成的合成轨迹进行on-policy强化学习，优化策略。

关键创新：PolyGRAD的关键创新在于其非自回归的轨迹生成方式。与传统的自回归模型不同，PolyGRAD一次性生成整个轨迹，避免了误差累积。此外，PolyGRAD还利用策略梯度引导扩散过程，确保生成的轨迹与策略一致。

关键设计：PolyGRAD的关键设计包括：1) 使用去噪扩散概率模型（DDPM）作为轨迹生成器。2) 使用策略的动作分布梯度作为引导信号，指导扩散过程。3) 使用特定的损失函数来训练扩散模型和策略，例如，轨迹预测误差和强化学习奖励。

📊 实验亮点

实验结果表明，PolyGRAD在短轨迹预测方面优于或媲美最先进的基线方法，例如自回归扩散模型，同时计算成本更低。在长轨迹预测方面，PolyGRAD也取得了与基线方法相当的性能。此外，PolyGRAD成功应用于MuJoCo连续控制领域，证明了其在复杂环境中的有效性。

🎯 应用场景

PolyGRAD具有广泛的应用前景，例如机器人控制、游戏AI和自动驾驶等领域。通过在虚拟环境中训练智能体，可以降低训练成本和风险，并提高智能体的泛化能力。此外，PolyGRAD还可以用于生成高质量的合成数据，用于训练其他机器学习模型。

📄 摘要（原文）

World models are a powerful tool for developing intelligent agents. By predicting the outcome of a sequence of actions, world models enable policies to be optimised via on-policy reinforcement learning (RL) using synthetic data, i.e. in "in imagination". Existing world models are autoregressive in that they interleave predicting the next state with sampling the next action from the policy. Prediction error inevitably compounds as the trajectory length grows. In this work, we propose a novel world modelling approach that is not autoregressive and generates entire on-policy trajectories in a single pass through a diffusion model. Our approach, Policy-Guided Trajectory Diffusion (PolyGRAD), leverages a denoising model in addition to the gradient of the action distribution of the policy to diffuse a trajectory of initially random states and actions into an on-policy synthetic trajectory. We analyse the connections between PolyGRAD, score-based generative models, and classifier-guided diffusion models. Our results demonstrate that PolyGRAD outperforms state-of-the-art baselines in terms of trajectory prediction error for short trajectories, with the exception of autoregressive diffusion. For short trajectories, PolyGRAD obtains similar errors to autoregressive diffusion, but with lower computational requirements. For long trajectories, PolyGRAD obtains comparable performance to baselines. Our experiments demonstrate that PolyGRAD enables performant policies to be trained via on-policy RL in imagination for MuJoCo continuous control domains. Thus, PolyGRAD introduces a new paradigm for accurate on-policy world modelling without autoregressive sampling.