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

机器之心报道

编辑：小舟、陈萍

在图像生成领域占据主导地位的扩散模型，开始挑战强化学习智能体。

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

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

受这种范式转变的推动，来自日内瓦大学、爱丁堡大学、微软研究院的研究者联合提出一种在扩散世界模型中训练的强化学习智能体 —— 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

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

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

方法介绍

实验

为了全面评估 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 想象的轨迹通常具有更高的视觉质量，并且更符合真实环境。

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