世界模型的典型框架与分类

世界模型是映射世界当前状态的函数和一个动作到对下一个状态的预测.观察和行动的定义因实施而异。

世界模型为智能系统提供了许多重要功能。学习有效的世界模型需要了解物理学、因果关系、空间智能等。当我们在现实世界中部署人工智能系统时，这些尤为重要。

世界模型在不同领域有不同的应用。虽然重要的研究集中在用于训练游戏代理的视频游戏环境中，但它们最重要的应用在于现实世界的具身人工智能。机器人应用需要世界模型来理解其物理环境并与之交互，预测其行为的结果，并适应不断变化的条件。这些模型帮助机器人执行需要规划的复杂任务。

世界模型代表了人工智能研究的一个令人兴奋但仍处于新兴阶段的领域。“世界模型”一词的定义仍然有些松散。在这篇博文中，我们将讨论世界模型的不同实现。

2.生成世界模型

世界模型在物理人工智能中具有直接应用，特别是在机器人和自动驾驶汽车中。这些系统需要交互式数据来展示行动如何导致特定结果，这与仅观察静态数据的传统视觉模型不同。

机器人技术的一个重大挑战是与语言模型相比，训练数据的可用性有限。虽然最有价值的数据来自现实世界与物理代理的交互，但收集这些数据既昂贵又耗时。我们还需要大量和多样化的环境，以便智能体能够很好地学习概括。此类世界模型的目标是生成真实且多样化的环境，以训练强大的智能体。

Genie

谷歌的Genie是一个基于互联网数据训练的“生成交互环境”（Genie-2 博客文章）。它能够根据文本或图像提示生成交互式环境。我们将解释模型的每个组件。

时空Transformer

由于它是在视频上训练的，因此 Genie 使用时空转换器 （ST-Transformer） 架构来实现大多数模型组件。ST-Transformer 模块具有三个组件：时间注意力、空间注意力和单个前馈模块。空间注意力是单个帧的补丁之间的自我注意力。时间关注的是单个补丁时间步长。这使用因果注意力。这种架构随帧数线性而不是二次方缩放，这使得它可以有效地处理视频。

ST-Transformer 架构

潜在动作模型 （LAM）

在世界模型中，我们需要将动作与观测之间的转换相关联。由于作在互联网规模数据上不可用，因此该模型学习“潜在作”。该模型使用编码器和解码器。编码器的输入是所有先前的帧和下一帧.编码器输出潜在动作.解码器使用这些潜在动作和前一帧来预测下一帧.这是使用 VQ-VAE 物镜实现的。使用尺寸为 8 的小代码本以实现可控性。小码本尺寸也是一个瓶颈。这迫使表示仅存储有关更改的信息，而不是对观察本身进行编码。

在 2D 视频上训练时，帧之间的变化可以用紧凑的代码来描述。想想向左、向右移动、跳跃等。在复杂的现实世界中，这可能更困难。由于潜在动作使用一个小密码本，并且解码器可以访问先前的帧，因此模型被迫学习可以表示为动作的帧之间的差异。

视频分词器

另一个 VQ-VAE 经过训练，可以从视频中生成代币。这里的区别在于输入和输出是相同的，只是视频帧。

动力学模型

动态模型采用先前的视频令牌和潜在作，并预测未来的视频令牌。这也使用了 ST-Transformer 架构。由于这是一个因果转换器，因此在每个时间步长该模型使用当前的潜在嵌入预测下一个时间步长的嵌入.在训练时，这是并行的，因此映射到.

源

MaskGIT 方法增强了该模型的性能。由于我们处理的是图像/视频帧，表示潜在嵌入的网格。一次性对这些嵌入进行传统的自回归预测可能不稳定。在此设置中，如果单个标记被错误预测，则所有后续预测都会受到影响。MaskGIT 通过启用预测校正的迭代过程解决此问题。下一帧中的标记可以双向相互关注。

MaskGIT 源代码

在训练过程中，根据伯努利分布随机屏蔽标记，采样概率在 0.5 到 1 之间，模型预测这些屏蔽标记。这种训练方法可以在推理时进行迭代解码。下一帧初始化为零，动态模型根据上下文进行预测。然后，动力学模型将此输出作为输入进行处理。由于预测是离散标记值的软最大值，因此我们可以确定哪些标记的置信度较低。后续迭代中的模型仅更新这些低置信度预测。这在 25 个步骤中重复，其中重新生成输出令牌，直到达到高置信度。这种方法类似于扩散，但使用离散而不是连续的标记。掩蔽方法使动力学模型的预测更加稳健，并减少了视频数据中的冗余，从而提高了学习效率。

训练和推理

首先分别训练潜在动作模型和视频分词器。他们使用 VQ-VAE 物镜进行训练。然后使用 LAM 和视频分词器来训练动态模型。有一个停止渐变，以便冻结潜在动作模型。

在推理时，LAM 被丢弃，因为我们可以直接定义作。模型会提示单个初始帧，该帧已标记化。在每一步中，用户为潜在动作选择一个整数。初始帧标记和潜在作用于生成下一个潜在标记。此生成的令牌将附加到输入 .然后，用户可以选择其他作。在每一步中，分词器解码器都可以将潜在嵌入映射到可以显示给用户的图像。

视频分词器和 LAM 非常相似。一个显着的区别是 LAM 使用更小的码本大小。重点是拥有一组可控的作（8 个代码）。视频分词器使用更大的码本大小 （1024） 来更有效地重建视频。LAM 能够使用较小的码本大小，因为它一次只尝试编码一帧。而视频分词器对整个视频进行标记化，然后将其全部重建在一起。

---

Genie 通过在不需要显式作的情况下对大规模数据进行训练，实现了重大突破。相反，它学习一组有限的潜在动作，这些动作能够与生成的环境进行交互。

通过这些潜伏的动作，用户可以控制新生成的游戏。虽然这种方法很强大，因为它可以在没有基本实况作的情况下对大量数据集进行训练，但学习到的作也有局限性。由于作是发现的而不是定义的，因此无法强制执行特定的映射。例如，对于 4 个潜在动作，我们不能保证动作 1 表示“向左移动”，动作 2 表示“向右移动”。模型可能会学习对角线运动。这种潜在动作的使用在一定程度上削弱了 Genie 作为真实世界模型的分类，因为传统世界模型旨在了解观察结果如何根据明确定义的现实世界动作而变化。

Cosmos

由英伟达发布的Cosmos是“世界基础模型”。它是在包含 2000 万小时视频的大型数据集上进行训练的。与其他作品不同，Cosmos 专注于物理 AI 和现实世界数据，而不是模拟视频游戏环境。它专为机器人和自动驾驶汽车等物理人工智能应用而构建。我们将介绍该模型的不同组件。

分词器

有效的世界建模需要有效地对视频进行标记化。它们提供离散和连续分词器。它们都以相同的方式进行训练，但离散分词器使用有限标量量化 （FSQ） 来离散化标记。

分词器模型使用时间因果编码器-解码器架构来映射视频到“代币视频”在空间和时间上都更加紧凑。为了训练这一点，他们使用简单的 L1 重建损失。在训练的第二阶段，他们使用光流损耗来提高视频帧之间的平滑度，并使用革兰氏矩阵损耗来提高帧的清晰度。对抗性损失用于进一步增强这一点。

世界模型

与分词器一样，世界模型也针对连续和离散令牌实现。但是，模型架构差异更大。扩散用于连续标记，而自回归模型则使用离散标记进行学习。

扩散模型

他们使用基于 DiT（扩散变压器）的架构。连续标记受到高斯噪声的扰动，DiT 模型对这些标记进行去噪。该模型通过交叉注意力以文本为条件。文本提示提供有关用于去噪的视频的其他信息。在推理时，文本可用于生成新视频。该模型还可以使用输入图像帧进行调节以生成视频延续。

自回归

自回归世界模型使用离散视频令牌在下一个令牌预测目标上进行训练。此培训分多个阶段进行：

第 1 阶段：在给定前 17 帧的上下文的情况下预测未来帧的标记。

阶段 1.1：上下文长度增加到 34 帧。请注意，每一帧都会产生大量标记。

第 2 阶段：使用文本条件进行训练。

扩散模型用作解码器，以实现更高质量的视频生成。此解码器是单独训练的。

源

---

Cosmos 包括另外两种世界模型分类。Text2World 将文本提示映射到视频。Video2World 模型将文本提示和视频作为输入。这些模型生成视频的延续。

在 Cosmos 中，动作被称为扰动或.在 WFM 训练期间，这些扰动只是描述视频的文本提示。文本在继续视频时有效地编码动作，从而允许进行可推广的世界模型训练。在微调期间，可以将作编码为更具体的格式。例如，在对相机控制进行微调时，模型会使用表示相机位姿的普吕克坐标进行调节。该模型学习生成与这些坐标匹配的视频。还可以合并其他模型，以根据特定于任务的作来调节 WFM。

在世界模型预训练期间，不使用显式作。该模型本质上是作为视频生成模型。然而，这些模型后来经过微调，以合并针对不同用例的显式作，将它们转换为真实的世界模型。

源

Cosmos 最近发布。看看“世界基础模型”是否会在未来几个月内得到更广泛的开发和使用，将会很有趣。目前还没有模型在这种规模的物理 AI 数据上进行训练。

代理的世界模型

之前的工作侧重于生成可用于训练智能自主代理的交互式环境。现在，我们将探索使用世界模型来开发代理本身的方法。这对于利用所有生成的环境至关重要。我们希望创建能够在不同环境和任务中表现良好的自主代理，尤其是在数据有限的情况下。世界模型为这一挑战提供了有价值的解决方案。

世界模型是基于模型的强化学习的重要组成部分。在 MBRL 中，动力学模型根据当前状态和动作预测下一个状态。有不同的方法可以将其用于规划。通常，动力学模型是从手头任务的轨迹数据集、现实世界或模拟中学习的。正在预测的状态是特定于任务的表示。

动力学模型可以使用世界模型来实现。这是当模型在大量且多样化的离线数据集上进行训练时，而不是特定于任务的轨迹。世界模型支持学习单个动力学模型，该模型可以推广到多个任务。这对于机器人技术很重要，因为您希望机器人能够执行在训练期间未优化的事情。对于这种概括，重要的是对世界模型进行各种任务的训练，这样它就不会关注特定于任务的信息。

Ha 等人，2018 年

这篇相对较早的论文展示了世界模型在强化学习背景下的优势。他们在不同的 OpenAI Gym 环境中训练世界模型以学习有效的策略。

V：视觉模型是学习潜在表示的 VAE来自环境的观察（图像）。

M：记忆RNN学习跨时间的表示，并在潜在空间中进行预测。这是建模.下一个潜在表示是根据当前表示、动作和隐藏状态预测的。这是使用混合密度模型的概率分布，因此我们可以从中采样。

C：控制器从隐藏状态和当前潜在表示中预测动作。这可以是一个轻量级线性模型。

训练步骤

1. 在随机策略下生成推出
2. 火车在框架上
3. 火车关于随机策略的潜在表示和作用
4. 优化控制器
5. 重启以了解新策略的推出

我们还可以使用世界模型在潜在空间中生成新的推出。这是通过对并使用它来训练控制器。这种方法的一个问题是，该政策可以利用世界模式中的差距。的随机性帮助训练更健壮的控制器。

这项早期工作为世界模型的后续研究奠定了重要基础，展示了它们创建样本效率更高、广义学习系统的潜力。

Dreamer V3

Hafner 等人 2023 年的论文《Mastering Diverse Domains through World Models》介绍了 DreamerV3 模型。这一工作旨在训练世界模型以用于基于模型的强化学习。这种方法利用两个阶段：世界模型学习和行为者批评学习。

世界模型

世界模型作为循环状态空间模型 （RSSM） 实现，该模型是使用带有帧和动作的轨迹来学习的。

序列模型：该模型根据先前的隐藏状态、潜在表示和动作预测下一个隐藏状态。

编码器：此模型生成时间步长的潜在表示从观测值（帧）中。这也利用了隐藏状态.

动力学预测器：这预测隐藏状态的潜在表示。

奖励/继续预测器,：这些模型预测奖励并继续发出信号。RL 环境设置为在每个时间步长都有一个标量值。

译码器：这从潜在表示预测观察系。解码器不是规划所必需的，但使用它进行训练可确保从观测中捕获信息。该模型还可以可视化计划的轨迹。编码器和解码器共同形成一个离散的自动编码器。

---

通过这些模块，我们可以使用离线轨迹训练世界模型。训练不同模块的损失有三种不同的损失。

预测损失训练解码器、奖励预测器和继续预测器（对应于方程中的三个项）。和训练编码器和动力学模型，类似于 VQ-VAE 码本和承诺损失（VAE 博客）。动力学损失训练动力学预测器从与编码器输出匹配的隐藏状态中预测潜在表示，而表示损失训练编码器生成与动力学预测器输出匹配的潜在表示。这些损失被裁剪，因此当其他两个损失得到合理优化时，学习将重点放在预测损失上。这也可以防止编码器在表示中包含不相关的信息，从而使控制更加困难。序列模型不直接包含在任何这些损失中，而是在我们通过轨迹生成隐藏状态时接收梯度。

这些损失针对轨迹的每个时间步长进行加权和计算，以一次训练世界模型的所有模块。

演员评论家学习

一旦我们使用离线轨迹训练世界模型，我们就可以在模拟的在线轨迹上使用行为者-批评者学习。这是一种流行的 RL 方法。参与者 （） 预测给定当前状态要采取的作。批评者 （） 预测一个州未来奖励的预期总和。

演员可以根据世界模型中想象的轨迹进行训练，而不是与真实环境交互。世界模型和 actor 模拟轨迹以及奖励和继续信号。这个阶段就是梦想。批评者接受训练以学习奖励的总和，而行为者则学习最大化回报的政策。这使得通过模拟交互学习有效的策略成为可能。

演员和评论家同时接受来自世界模型的想象轨迹和存储在回放缓冲区中的真实轨迹的训练。虽然世界模型提高了样本效率，但仍需要一些真实的交互来微调特定任务。想象轨迹与真实轨迹的比率以及损失的权重是可调的超参数。

Dreamer 使用通用世界模型，能够更有效、更高效地学习各种任务和环境的策略。Ha 等人 2018 和 Dreamer 的一个缺点是他们只使用带有奖励的轨迹进行训练。众所周知，此类数据是有限的。我们希望能够使用没有动作和奖励的视频进行训练，甚至使用未标记的静态图像。

DINO-WM模型

此方法使用预训练的 DINOv2 SSL 模型（此处介绍）来引导世界模型。DINO 模型经过非常大规模的训练，以生成用于视觉理解的图像表示。此方法将此功能合并为世界模型的组件，而不是尝试通过世界模型本身来学习它。

作者还选择避免重建图像，而是预测潜在表示。这遵循 JEPA 框架。

该环境被建模为部分观察到的马尔可夫决策过程 POMDP，定义为（观察、行动、概率）。与其他方法一样，此世界模型有三个组成部分：

观察模型学习潜在表示。过渡模型根据先前的观察和行动预测未来的表示。解码器模型重建图像观察，这是可选的（可视化/生成）。

观察模型

观察模型经过训练以学习视觉表示。该模型旨在推广到许多任务和环境。为此，他们使用预训练的 DINOv2 模型。该模型是在互联网规模的图像数据上训练的，因此可以很好地泛化。在时间步长它对图像进行编码到补丁嵌入中形状的哪里是补丁的数量，而是嵌入维度。

过渡模型

过渡模型被训练为仅解码器的转换器模型。给定过去的潜在因素，该模型预测未来的潜在因素。通过注意力掩码，由只关注同一索引的嵌入。这就是潜在状态关注.表示上下文长度，即模型可以回顾的步数。观测模型中的斑块之间存在注意力，但过渡模型独立查看每个斑块。

为了合并动作，将 K 维向量连接到输入中嵌入的每个补丁。这是从原始动作表示映射而来的，原始动作表示是使用 MLP 定义动作的连续向量。作的表示形式会根据环境或任务而变化。例如，它可以是一个 2D 向量，例如 （-0.1， 2.1），描述如何推动对象。这些作包含在训练数据集中。该模型使用这些轨迹进行教师强制训练。

其他元数据（例如本体感觉信息）也可以映射到嵌入并连接起来。该模型使用简单的 L2 损失进行训练：

译码器

对观测模型的输出进行处理，以使用简单的重建损失重建输入图像。解码器损耗不会反向传播到编码器/观测模型。这种方法在经验上被证明更有效，并且完全在潜在空间中保持世界模型的规划能力。

测试时间优化

虽然模型是使用数据集中的轨迹进行训练的，但在测试时，我们希望进行规划，以便获得最佳的作序列。

鉴于目前的观察结果和目标观察，这是期望结果的图像。模型预测控制用于查找实现目标的一系列作。优化了以下成本：哪里,,.

示例 开始观察（左）和目标观察（ 右）来源

此方法使用图像定义目标状态。这仅适用于模拟环境。此方法不直接适用于更复杂的实际应用。

交叉熵法 （CEM）

为了实现 MPC，发现这种方法最有效。该策略表示为形状高斯分布的张量哪里是轨迹中的步数，并且是作嵌入的大小。方法示例该分布的轨迹。世界模型处理每个轨迹并计算最终的潜在状态，该状态用于计算每个轨迹的成本。顶部选择轨迹（不要与动作嵌入大小混淆）。然后使用它们的平均值和标准差来重新定义高斯分布。此迭代过程一直持续到找到可最大程度地降低成本的一系列作。

CEM 与重要性抽样类似，因为它对一组轨迹进行采样，并使用成功的轨迹来更新分布。

梯度下降（Gradient Descent）

这种方法被发现效果较差，但作为未来工作的方向仍然值得注意。在此方法中，作被初始化为学习嵌入。这些作是随机初始化的。成本用于通过梯度下降来更新这些嵌入。

其他方法

视频生成模型

像 Veo 和 Sora 这样的视频生成模型在作为世界模型的分类中占据了一个模糊的空间。虽然他们可以根据文本提示生成逼真的视频并创建现有视频的自然延续，但它们不会模拟世界对行为的反应。尽管人们可以将动作概念化为视频片段，并将其延续概念化为世界模型模拟，但这种方法还不够。我们需要能够模拟世界的模型，而无需输入动作，并按顺序理解多个动作的后果。

JEPA世界模型

I-JEPA 和 V-JEPA（在我的 JEPA 博客文章中进行了讨论），以及类似的自监督学习方法可以解释为世界模型。这些方法通过屏蔽图像或视频的某些部分来工作。掩码本身充当动作。该模型可以访问掩码的位置并预测掩码区域的潜在表示。

源

这项工作介绍了图像世界模型（IWM）。IWM 通过使用各种图像失真来扩展 I-JEPA，超越简单的遮罩，例如亮度、对比度和饱和度的变化。这些失真通常用于对比 SSL 方法。该动作将成为应用于源图像以生成目标图像的扭曲的表示。IWM 使用源潜在和动作来预测目标图像的潜在表示。

由于这部作品只关注图像，因此缺乏时间成分。IWM 建模的世界没有时间维度，限制了其在现实世界中的应用。此外，动作空间对于实际应用并不是特别有用。尽管如此，这种方法有助于开发比完全依赖掩蔽的 I-JEPA 的图像表示更丰富的图像表示。扭曲通常用于使图像表示对它们不变。然而，在这个用例中，我们不是直接学习不变性，而是学习扭曲图像之间的关系。我将这种方法归类为自监督表示学习的巨大改进，而不是世界模型的进步。

值得注意的是，许多其他世界模型架构，例如 DINO-WM，可以被视为联合嵌入谓词架构 （JEPA）。

结论

世界模型是一个具有巨大潜力的重要研究领域。这里描述的作品通常使用简单的动作，例如 2D 视频游戏中的运动。未来，我们需要在不同粒度下对更复杂的动作进行建模。

例如，让我们考虑一个做家务的家用机器人。一个高级行动是选择从哪个家务开始。这需要考虑哪项任务现在最重要，以及它在不久的将来将如何影响家庭。执行杂务需要较低的粒度级别。假设机器人决定洗碗，它必须选择要洗碗的顺序以优化存储。在洗碗时，它必须确定机械运动的顺序才能有效和高效地做到这一点。我们可能还需要为更长的时间范围（例如几天或几个月）进行计划的世界模型。我们的内部世界模型非常强大，这些推理和规划对我们来说非常容易。我们低估了在机器中复制这一点所需的条件。