基于策略引导轨迹扩散的世界模型

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