2018年，神经网络世界模型的强化学习概念首次被提出，不久，这一基本原理就被应用于世界模型。一些著名的实现强化学习的模型包括Dreamer框架，它从递归状态空间模型的潜在空间中引入了强化学习。DreamerV2证明了使用离散潜在变量可能会导致复合错误减少，而DreamerV3框架能够在不同领域的任务中实现人类般的性能，使用固定超参数。

此外，图像生成模型和世界模型之间可以画出平行线，表明生成式视觉模型的进展可以被复制以造福世界模型。自从使用transformer在自然语言处理框架中流行起来，DALL-E和VQGAN框架出现了。这些框架实现了离散自动编码器将图像转换为离散令牌，并通过利用自回归transformer的序列建模能力，能够建立高效的文本到图像生成模型。与此同时，扩散模型也获得了关注，如今，扩散模型已经成为高分辨率图像生成的主导范式。由于扩散模型和强化学习的能力，人们正在尝试将这两种方法结合起来，旨在利用扩散模型的灵活性作为轨迹模型、奖励模型、规划器和离线强化学习的数据增强政策。

世界模型为训练强化学习代理提供了一种安全高效的方法。传统上，这些模型使用一系列离散潜在变量来模拟环境动态。然而，这种压缩可能会忽略强化学习中至关重要的视觉细节。与此同时，扩散模型已经因其在图像生成方面的流行而受到挑战，挑战了传统使用离散潜在变量的方法。受此转变的启发，本文将讨论DIAMOND（DIffusion As a Model Of eNvironment Dreams），一种在扩散世界模型中训练的强化学习代理。我们将探讨使扩散适合世界建模的必要设计选择，并展示增强的视觉细节如何带来更好的代理性能。DIAMOND在竞争激烈的Atari 100k测试中设定了新的基准，实现了1.46的平均人类归一化分数，这是完全在世界模型中训练的代理的最高分数。

DIAMOND：视觉细节在Atari和扩散中对世界建模的重要性

世界模型或环境的生成模型正在成为生成代理计划和推理其环境的更重要组件。虽然强化学习在近年来取得了相当大的成功，但实现强化学习的模型以其样本低效而闻名，这大大限制了其在现实世界中的应用。另一方面，世界模型已经展示了其在不同环境中高效训练强化学习代理的能力，样本效率大大提高，使模型能够从现实世界的经验中学习。最近的世界建模框架通常将环境动态建模为一系列离散潜在变量，模型将潜在空间离散化以避免多步时间范围内的复合错误。虽然这种方法可能会带来实质性的结果，但它也会导致信息损失，导致重构质量和普遍性的损失。信息损失可能会成为需要信息明确定义的现实世界场景中的重大障碍，例如训练自主车辆。在这样的任务中，视觉输入中的小变化或细节，例如交通信号灯的颜色或前方车辆的转向指示器，可能会改变代理的政策。虽然增加离散潜在变量的数量可以帮助避免信息损失，但它会显著增加计算成本。

此外，近年来，扩散模型已经成为高质量图像生成框架的主导方法，因为基于扩散模型的框架学习逆转噪声过程，并直接与一些更成熟的方法竞争，这些方法使用离散令牌，因此为世界建模提供了一个有前途的替代方案，以消除离散化的需要。扩散模型以其易于条件化和灵活地模拟复杂、多模态分布而闻名，而不会出现模式崩溃。这些属性对于世界建模至关重要，因为条件化使世界模型能够准确反映代理的行为，导致更可靠的信用分配。此外，模拟多模态分布为代理提供了更多的训练场景，增强了其整体性能。

基于这些特点，DIAMOND（DIffusion As a Model Of eNvironment Dreams）是一个在扩散世界模型中训练的强化学习代理。DIAMOND框架进行了仔细的设计选择，以确保其扩散世界模型在长时间范围内保持高效和稳定。该框架提供了一个定性分析，以展示这些设计选择的重要性。DIAMOND在Atari 100k基准测试中设定了新的最先进水平，实现了1.46的平均人类归一化分数，这是完全在世界模型中训练的代理的最高分数。在图像空间中运行使得DIAMOND的扩散世界模型能够无缝地替换环境，提供了对世界模型和代理行为的更深入的见解。值得注意的是，在某些游戏中，性能的提高归因于对关键视觉细节的更好建模。DIAMOND框架将环境建模为一个标准的POMDP或部分可观察的马尔可夫决策过程，具有状态集、离散动作集和图像观察集。转换函数描述了环境动态，奖励函数将转换映射到标量奖励。观察函数描述了观察概率，并发出图像观察，这些观察由代理使用来查看环境，因为它们无法直接访问状态。这种方法的主要目标是获得一个将观察映射到动作的政策，尝试在折扣因子下最大化预期折扣回报。世界模型是环境的生成模型，可以用来创建模拟环境来训练强化学习代理在真实环境中，并在世界模型环境中训练强化学习代理。图1演示了DIAMOND框架随时间的想象展开。

DIAMOND：方法论和架构

在其核心，扩散模型是一类生成模型，通过逆转噪声过程生成样本，并从非平衡热力学中汲取灵感。DIAMOND框架考虑了一个由连续时间变量索引的扩散过程，具有相应的边缘和边界条件，以及一个可追踪的非结构化先验分布。此外，为了获得一个从噪声到数据的生成模型，DIAMOND框架必须逆转该过程，逆转过程也是一个扩散过程，逆时间运行。进一步来说，在任何给定时间点，估计分数函数并非易事，因为DIAMOND框架无法访问真实的分数函数，模型通过实现分数匹配目标来克服这一障碍，这是一种使框架能够在不知道潜在分数函数的情况下训练分数模型的方法。基于分数的扩散模型提供了一个无条件的生成模型。然而，需要一个条件生成模型来模拟环境动态，作为世界模型，DIAMOND框架考虑了POMDP方法的一般情况，在这种情况下，框架可以使用过去的观察和动作来近似未知的马尔可夫状态。如图1所示，DIAMOND框架使用这种历史来条件一个扩散模型，以估计和直接生成下一个观察。虽然DIAMOND框架可以在理论上使用任何SDE或ODE求解器，但在NFE或函数评估次数和样本质量之间存在一个权衡，这会显著影响扩散模型的推理成本。

基于上述见解，让我们现在来看看基于扩散的世界模型的实际实现，包括与特定扩散方法选择相对应的漂移和扩散系数。与其选择DDPM（一种自然适合该任务的候选项），DIAMOND框架建立在EDM公式化上，并考虑了一个具有扩散时间的实值函数的扰动核，称为噪声时间表。该框架选择了预处理器，以保持任何语音级别的输入和输出方差。网络训练根据降级水平自适应地混合信号和噪声，当噪声低时，目标变为干净信号和扰动信号之间的差异，即添加的高斯噪声。直观地，这防止了训练目标在低噪声区变得琐碎。在实践中，该目标在噪声时间表的极端处具有高方差，因此模型从一个经验上选择的对数正态分布中采样噪声级别，以便在中等噪声区域周围连接训练。DIAMOND框架使用标准的2D U-Net组件作为向量场，并保持一个过去观察和动作的缓冲区，该缓冲区由框架用于条件化自身。然后，DIAMOND框架将这些过去的观察连接到下一个嘈杂的观察，并通过U-Net的残差块中的自适应组归一化层输入动作。

DIAMOND：实验和结果

为了进行全面评估，DIAMOND框架选择了Atari 100k基准测试。Atari 100k基准测试由26个游戏组成，旨在测试代理的广泛能力。在每个游戏中，代理的环境操作次数限制为100k，这大约相当于2小时的人类游戏时间，以便在评估之前学习游戏。为了比较，未受限制的Atari代理通常需要5000万步来训练，代表着500倍的经验增加。我们使用5个随机种子对每个游戏从头开始训练DIAMOND。每次训练运行大约需要12GB的VRAM，并在单个Nvidia RTX 4090上花费了大约2.9天，总共花费了1.03 GPU年。以下表格提供了所有游戏的分数、平均分数和人类归一化分数的IQM或四分位数均值。

由于点估计的局限性，DIAMOND框架提供了人类归一化分数的均值和IQM或四分位数均值的分层自举置信度，以及性能概况和其他指标的总结，如下图所示。

结果表明，DIAMOND在整个基准测试中表现出色，超过了11个游戏的人类玩家，并实现了1.46的超人类均值HNS，创造了完全在世界模型中训练的代理的新记录。此外，DIAMOND的IQM与STORM相比，并且超过了所有其他基准。DIAMOND在 Asterix、Breakout和RoadRunner等环境中表现出色，这些环境中捕捉小细节至关重要。此外，如前所述，DIAMOND框架具有灵活性，可以实现任何扩散模型，尽管它选择了EDM方法，但选择DDPM模型将是一个自然的选择，因为它已经在许多图像生成应用中实现。为了比较EDM方法与DDPM实现，DIAMOND框架使用相同的网络架构在相同的共享静态数据集上训练了两个变体，数据集包含超过100k帧，使用专家策略收集。去噪步骤的数量直接关系到世界模型的推理成本，因此更少的步骤将降低在想象轨迹上训练代理的成本。为了确保我们的世界模型在计算上与其他基准（如IRIS）保持一致，我们的目标是使用不超过十几个去噪步骤，理想情况下更少。然而，将去噪步骤的数量设置得太低可能会降低视觉质量，导致复合错误。为了评估不同扩散变体的稳定性，我们显示了使用不同去噪步骤n≤10自回归生成的想象轨迹，直到t=1000时间步。

我们观察到，在这种情况下使用DDPM（a）会导致严重的复合错误，导致世界模型快速偏离分布。相比之下，基于EDM的扩散世界模型（b）即使只有一个去噪步骤，也保持了更大的稳定性。显示了基于DDPM（左）和EDM（右）的扩散世界模型的想象轨迹。初始观察在t=0时对于两者都是相同的，每行对应于去噪步骤n的减少。我们观察到，基于DDPM的生成会受到复合错误的影响，去噪步骤越少，错误就会越快地积累。相比之下，DIAMOND的基于EDM的世界模型即使对于n=1也保持了更大的稳定性。单步预测的最佳选择是对给定嘈杂输入的可能重构的期望，如果后验分布是多模态的，则可能会超出分布。虽然某些游戏（如Breakout）具有确定的转换，可以使用单步去噪精确建模，但其他游戏表现出部分可观察性，导致观察分布多模态。在这些情况下，需要迭代求解器来引导采样过程朝向特定模式，如下图中的Boxing游戏所示。因此，DIAMOND框架在所有实验中都将n设置为3。

上图比较了Boxing中的单步（顶行）和多步（底行）采样。黑色玩家的动作不可预测，导致单步去噪插值到可能的结果，导致预测模糊。相比之下，多步采样通过引导生成朝向特定模式，产生了清晰的图像。有趣的是，由于策略控制白色玩家，世界模型已知白色玩家的动作，消除了模糊性。因此，单步和多步采样都正确地预测了白色玩家的位置。

在上图中，DIAMOND想象的轨迹通常比IRIS想象的轨迹具有更高的视觉质量，并且更忠实于真实环境。IRIS生成的轨迹包含视觉不一致性（用白色框突出显示），例如将敌人显示为奖励，反之亦然。虽然这些不一致性可能只影响几个像素，但它们可能会显著影响强化学习。例如，代理通常旨在针对奖励和避免敌人，因此这些小的视觉差异可能会使学习最优策略更加困难。图中显示了IRIS（左）和DIAMOND（右）想象的连续帧。白色框突出显示了帧之间的不一致性，这些不一致性只出现在IRIS生成的轨迹中。在Asterix（顶行）中，一个敌人（橙色）在第二帧中变成了一个奖励（红色），然后在第三帧中变回了一个敌人，第四帧中又变成了一个奖励。在Breakout（中间行）中，砖块和分数在帧之间不一致。在Road Runner（底行）中，奖励（道路上的小蓝点）在帧之间不一致地渲染。这些不一致性不会出现在DIAMOND中。在Breakout中，分数可靠地在红砖被打破时更新为+7。

结论

在本文中，我们讨论了DIAMOND，一种在扩散世界模型中训练的强化学习代理。DIAMOND框架进行了仔细的设计选择，以确保其扩散世界模型在长时间范围内保持高效和稳定。该框架提供了一个定性分析，以展示这些设计选择的重要性。DIAMOND在Atari 100k基准测试中设定了新的最先进水平，实现了1.46的平均人类归一化分数，这是完全在世界模型中训练的代理的最高分数。在图像空间中运行使得DIAMOND的扩散世界模型能够无缝地替换环境，提供了对世界模型和代理行为的更深入的见解。值得注意的是，在某些游戏中，性能的提高归因于对关键视觉细节的更好建模。DIAMOND框架将环境建模为一个标准的POMDP或部分可观察的马尔可夫决策过程，具有状态集、离散动作集和图像观察集。转换函数描述了环境动态，奖励函数将转换映射到标量奖励。