近年来，强化学习领域取得了显著成就，但其样本效率不足成为现实应用的一大障碍。世界模型作为一种环境生成模型，为克服这一难题带来了曙光。它能够模拟环境，从而提高强化学习智能体的样本效率。

当前，多数世界模型采用离散潜变量序列来模拟环境变化。但这种将信息压缩为紧凑离散表示的方式，可能会丢失对强化学习至关重要的视觉细节。

在此背景下，扩散模型在图像生成领域取得了主导地位，挑战了传统的离散潜变量建模方法。基于此，研究人员提出了一种名为DIAMOND（环境梦境扩散模型）的新方法，这是一种在扩散世界模型中训练的强化学习智能体。DIAMOND在设计中做出了关键决策，确保了扩散模型在长期运行中的高效与稳定。

image.png

在著名的Atari100k基准测试中，DIAMOND取得了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