强化学习在近年来取得了许多成功，但其样本效率低下，限制了其在现实世界中的应用。世界模型作为一种环境生成模型，为解决这一问题提供了希望。它可以作为模拟环境，以更高的样本效率训练强化学习智能体。

目前，大多数世界模型通过离散潜变量序列来模拟环境动态。然而，这种压缩成紧凑离散表示的方法可能会忽略对强化学习至关重要的视觉细节。

与此同时，扩散模型已经成为图像生成领域的主导方法，挑战了传统的离散潜变量建模方法。受此启发，研究人员提出了一种名为DIAMOND（环境梦境扩散模型）的新方法，它是一种在扩散世界模型中训练的强化学习智能体。DIAMOND在设计上做出了关键选择，以确保扩散模型在长时间范围内的高效性和稳定性。

DIAMOND在著名的Atari100k基准测试中取得了1.46的平均人类标准化得分，这是完全在世界模型中训练的智能体的最佳成绩。此外，在图像空间中操作的优势在于，扩散世界模型可以直接替代环境，从而更好地理解世界模型和智能体的行为。研究人员发现，某些游戏性能的提升源于对关键视觉细节的更好建模。

DIAMOND的成功得益于EDM（Elucidating the Design Space of Diffusion-based Generative Models）框架的选择。与传统的DDPM(Denoising Diffusion Probabilistic Models)相比，EDM在较少的去噪步骤下表现出更高的稳定性，避免了模型在长时间范围内出现严重的累积误差。

此外，DIAMOND还展示了其扩散世界模型可以作为交互式神经游戏引擎的能力。通过在87小时的静态《反恐精英:全球攻势》游戏数据上进行训练，DIAMOND成功生成了一个可交互的Dust II地图神经游戏引擎。

未来，DIAMOND可以通过整合更先进的记忆机制，例如自回归Transformer，来进一步提升其性能。此外，将奖励/终止预测整合到扩散模型中也是一个值得探索的方向。

论文地址：https://arxiv.org/pdf/2405.12399