世界模型也扩散！训练出的智能体竟然不错

世界模型提供了一种以安全且样本高效的方式训练强化学习智能体的方法。近期，世界模型主要对离散潜在变量序列进行操作来模拟环境动态。

然而，这种压缩为紧凑离散表征的方法可能会忽略对强化学习很重要的视觉细节。另一方面，扩散模型已成为图像生成的主要方法，对离散潜在模型提出了挑战。

这种范式转变的推动，来自日内瓦大学、爱丁堡大学、微软研究院的研究者联合提出一种在扩散世界模型中训练的强化学习智能体——DIAMOND（DIffusion As a Model Of eNvironment Dreams）。

• 论文地址：https://arxiv.org/abs/2405.12399
• 项目地址：https://github.com/eloialonso/diamond
• 论文标题：Diffusion for World Modeling: Visual Details Matter in Atari

在Atari 100k基准测试中，DIAMOND+获得了1.46的平均人类归一化得分（HNS）。这意味着在世界模型中训练的智能体的SOTA水平完全可以在世界模型中训练的智能体的SOTA水平进行完全训练。该研究提供了稳定性分析来说明DIAMOND的设计选择对于确保扩散世界模型的长期高效稳定是必要的。

除了在图像空间中操作的好处是使扩散世界模型能够成为环境的直接代表，从而提供对世界模型和智能体行为更深入的了解。特别地，该研究发现某些游戏中性能的提高源于对关键视觉细节的更好建模。

方法介绍

接下来，本文介绍了 DIAMOND，这是一种在扩散世界模型中训练的强化学习智能体。具体来说，研究者基于 2.2 节引入的漂移和扩散系数 f 和 g，这两个系数对应于一种特定的扩散范式选择。此外，该研究还选择了基于 Karras 等人提出的 EDM 公式。

首先定义一个扰动核，，其中， 是一个与扩散时间相关的实值函数，称为噪声时间表。这对应于将漂移和扩散系数设为 和。

接着使用 Karras 等人（2022）引入的网络预处理，同时参数化公式（5）中的，作为噪声观测值和神经网络 预测值的加权和： 

得到公式（6）

其中为了简洁定义，包含所有条件变量。 

预处理器的选择。选择预处理器和，以保持网络输入和输出在任何噪声水平 下的单位方差。 是噪声水平的经验转换，  由  和数据分布的标准差  给出，公式为

 结合公式 5 和 6，得到训练目标：

该研究使用标准的 U-Net 2D 来构建向量场，并保留一个包含过去 L 个观测和动作的缓冲区，以此来对模型进行条件化。接下来他们将这些过去的观测按通道方式与下一个带噪观测拼接，并通过自适应组归一化层将动作输入到 U-Net 的残差块中。正如在第 2.3 节和附录 A 中讨论的，有许多可能的采样方法可以从训练好的扩散模型中生成下一个观测。虽然该研究发布的代码库支持多种采样方案，但该研究发现欧拉方法在不需要额外的 NFE（函数评估次数）以及避免了高阶采样器或随机采样的不必要复杂性的情况下是有效的。 

实验

为了全面评估 DIAMOND，该研究使用了公认的 Atari 100k 基准测试，该基准测试包括 26 个游戏，用于测试智能体的广泛能力。对于每个游戏，智能体只允许在环境中进行 100k 次操作，这大约相当于人类 2 小时的游戏时间，以在评估前学习玩游戏。作为参考，没有限制的 Atari 智能体通常训练 5000 万步，这相当于经验的 500 倍增加。研究者从头开始在每个游戏上用 5 个随机种子训练 DIAMOND。每次运行大约使用 12GB 的 VRAM，在单个 Nvidia RTX 4090 上大约需要 2.9 天（总计 1.03 个 GPU 年）。 

表 1 比较了在世界模型中训练智能体的不同得分：

图 2 中提供了平均值和 IQM（ Interquartile Mean ）置信区间：

结果表明，DIAMOND 在基准测试中表现强劲，超过人类玩家在 11 个游戏中的表现，并达到了 1.46 的 HNS 得分，这是完全在世界模型中训练的智能体的新纪录。该研究还发现，DIAMOND 在需要捕捉细节的环境中表现特别出色，例如 Asterix、Breakout 和 Road Runner。

为了研究扩散变量的稳定性，该研究分析了自回归生成的想象轨迹（imagined trajectory），如下图 3 所示：

该研究发现有些情况需要迭代求解器将采样过程驱动到特定模式，如图 4 所示的拳击游戏：

如图 5 所示，与 IRIS 想象的轨迹相比，DIAMOND 想象的轨迹通常具有更高的视觉质量，并且更符合真实环境。

感兴趣的读者可以阅读论文原文，了解更多研究内容。