Genie: Generative Interactive Environments
公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）
目录
0. 摘要
2. 方法
2.1. 模型组件
2.2. 推理：可控行动视频生成
3. 实验
0. 摘要
我们介绍 Genie，这是第一个通过无监督方式从未标记的互联网视频中训练出来的生成式交互环境。通过文本、合成图像、照片，甚至素描，该模型可以被提示生成各种可控行动的虚拟世界。拥有 11B 参数，Genie 可以被视为基础世界模型（foundation world model）。它包括一个时空视频分词器（spatiotemporal video tokenizer）、一个自回归动力学模型（autoregressive dynamics model），以及一个简单而可扩展的潜在行动模型（latent action model）。尽管在训练中没有任何地面真实行动标签或其他通常在世界模型文献中找到的领域特定要求，Genie 使用户能够在生成的环境中逐帧执行操作。此外，所得到的学习潜在行动空间有助于训练代理程序模仿来自未见过视频的行为，为未来训练通用代理程序打开了道路。
项目页面：https://sites.google.com/view/genie-2024/home
2. 方法
Genie 是一个从仅视频数据中训练出来的生成式交互环境。在这一部分，我们从开始讲解模型的主要组成部分之前，先介绍一些基本概念。
Genie 架构中的几个组件基于 Vision Transformer（ViT）（Dosovitskiy等，2021；Vaswani等，2017）。值得注意的是，Transformer 的二次存储成本对于可能包含多达 O(10^4) 个 token 的视频提出了挑战。因此，我们在所有模型组件上采用了一种内存高效的 ST-transformer 架构（受Xu等人（2020）启发，见图 4），在模型容量和计算约束之间取得平衡。
与传统的 Transformer 不同，其中每个标记都与其他所有标记有关，ST-transformer 包含 L 个时空块，其中交错使用空间和时间注意层，然后是标准的前馈层（FFW）作为注意块。空间层中的自注意力关注每个时间步内的 1 × H × W token，而时间层中的自注意力关注跨越 T 个时间步的 T × 1 × 1 token。与序列 Transformer 类似，时间层假设具有因果结构和因果掩码。至关重要的是，在我们的架构中，计算复杂性的主导因素（即空间注意力层）与帧数成线性关系而不是二次关系，使其在具有持续动态交互的视频生成中更加高效。此外，请注意在 ST 块中，我们在空间和时间组件之后只包含一个 FFW，省略了后空间（post-spatial） FFW，以便提高模型的其他组件，我们观察到这显著改善了结果。
2.1. 模型组件
如图3所示，我们的模型包含三个关键组件：1）一个潜在行动模型，推断每一对帧之间的潜在行动𝒂，2）一个视频分词器，将原始视频帧转换为离散 token 𝒛，以及 3）一个动力学模型，给定潜在行动和过去的帧标记，预测视频的下一帧。该模型按照标准自回归视频生成流程进行两个阶段的训练：首先训练视频分词器，然后用于动力学模型。然后，我们同时训练潜在行动模型（直接从像素中获得）和动力学模型（在视频标记上）。
潜在行动模型（Latent Action Model，LAM）。为实现可控制的视频生成，我们使每个未来帧的预测依赖于前一帧的行动。然而，在互联网视频中很少有此类行动标签，并且获得行动注释可能成本高昂。相反，我们以完全无监督的方式学习潜在行动（见图 5）。
首先，一个编码器将所有先前的帧 𝒙_(1:𝑡) = (𝑥1, · · · 𝑥𝑡) 以及下一帧 𝑥_(𝑡+1) 作为输入，并输出相应的一组连续的潜在行动 ˜𝒂_(1:𝑡) = (˜𝑎1, · · · ˜𝑎𝑡)。然后，一个解码器将所有先前的帧和潜在行动作为输入，并预测下一帧 ˆ𝑥_(𝑡+1)。
为了训练模型，我们利用基于 VQ-VAE 的目标（van den Oord等，2017），这使我们能够将预测的动作数量限制在一个小的离散代码集中。我们将 VQ 码书的词汇大小 |𝐴|（即可能的潜在行动的最大数量）限制为一个小的值，以允许人类可玩性，并进一步强制可控性（在我们的实验中使用 |𝐴| = 8）。由于解码器只能访问历史和潜在行动，˜𝑎𝑡 应编码过去和未来之间对解码器成功重构未来帧最有意义的变化。请注意，这个解码器仅在 LAM 训练信号中存在。实际上，除了 VQ 码书之外，整个 LAM 在推断时被丢弃，并用用户的行动替换。
我们利用我们的 ST-transformer 架构来构建潜在行动模型。在时间层中的因果掩码允许我们将整个视频 𝒙_1:𝑇 作为输入，并生成每一帧之间的所有潜在行动 ˜𝒂_1:𝑇−1。
视频分词器（Video Tokenizer）。借鉴之前的工作（Gupta等，2023；Villegas等，2023；Yan等，2023），我们将视频压缩为离散标记，以降低维度并实现更高质量的视频生成（见图 6）。我们再次使用 VQ-VAE，它以视频的T帧 𝒙_1:𝑇 = (𝑥1, 𝑥2, · · · , 𝑥𝑇 ) ∈ ℝ^(𝑇×𝐻×𝑊×𝐶) 作为输入，在整个视频序列上生成每个帧的离散表示 𝒛_1:𝑇 = (𝑧1, 𝑧2, · · · , 𝑧𝑇 ) ∈ 𝕀^(𝑇×𝐷)，其中 𝐷 是离散潜在空间的大小。使用标准的 VQ-VAE 目标训练分词器，覆盖整个视频序列。
与先前专注于仅在分词阶段进行空间压缩的工作不同（Gupta等，2023；Hong等，2022；Wu等，2022），我们在编码器和解码器中都使用 ST-transformer 来合并编码中的时间动态，从而提高了视频生成的质量。由于 ST-transformer 的因果性质，每个离散编码 𝑧𝑡 包含来自视频 𝒙_1:𝑡 中所有先前看到的帧的信息。Phenaki（Villegas等，2023）也使用了一个具有时间感知的分词器，C-ViViT，但这种架构的计算密集性较高，因为其成本随着帧数的增加呈二次增长。相比之下，我们基于 ST-transformer 的分词器（STViViT）在成本上更加高效，其成本的主导因素与帧数成线性关系。
动力学模型（Dynamics Model）。动力学模型是一个仅解码的 MaskGIT（Chang等，2022）transformer（见图7）。在每个时间步 𝑡 ∈ [1, 𝑇]，它接收标记化视频 𝒛_1:𝑡−1 和 stopgrad 潜在行动 ˜𝒂_1:𝑡−1，并预测下一帧标记 ˆ𝑧𝑡。我们再次使用 ST-transformer，其因果结构使我们能够使用来自所有 (𝑇 − 1) 帧的标记 𝒛_1:𝑇−1 和潜在行动 ˜𝒂1:𝑇−1 作为输入，并生成所有下一帧的预测 ˆ𝒛2:𝑇。该模型通过预测标记 ˆ𝒛_2:𝑇 和地面实况标记 𝒛_2:𝑇 之间的交叉熵损失进行训练。在训练时，我们根据在 0.5 和 1 之间均匀采样的伯努利分布掩蔽率随机掩蔽输入标记 𝒛_2:𝑇−1。请注意，训练世界模型的常见做法，包括基于 transformer 的模型，是将时间 𝑡 的行动连接到相应的帧上（Micheli等，2023；Robine等，2023）。然而，我们发现将潜在行动视为对潜在行动和动力学模型的加法嵌入有助于提高生成的可控性。
2.2. 推理：可控行动视频生成
我们现在描述如何在推理时使用Genie进行可控制的视频生成（见图8）。玩家首先用一幅图像 𝑥1提示模型，作为初始帧。该图像通过视频编码器进行分词，得到 𝑧1。然后，玩家通过在 [0，|𝐴|) 范围内选择任何整数值来指定离散的潜在行动 𝑎1。动力学模型获取帧标记 𝑧1 和相应的潜在行动˜𝑎1，通过使用离散输入 𝑎1 在 VQ 码书中索引得到，然后预测下一帧标记 𝑧2。这个过程在自回归的方式下重复进行，因为行动继续传递给模型，同时标记被解码成视频帧 ˆ𝒙2:𝑇，使用分词器的解码器。请注意，我们可以通过将模型传递给起始帧和从视频中推断出的行动，重新生成数据集中的实况视频，或者通过改变行动生成完全新的视频（或轨迹）。
3. 实验
模型扩展。在固定的视频分词器和行动模型架构的情况下，我们训练了一系列参数从 40M 到 2.7B 的动力学模型。图 9 显示我们的架构在模型参数上呈现出良好的扩展性，每次规模增加都对应着最终训练损失的一致减小。这是我们的方法从规模调整中受益的强烈迹象，而我们主要的 Genie 模型也充分利用了这一点。
Batch 扩展。我们还研究了 Batch 大小的调整效果，考虑了一个具有 2.3B 模型的 Batch 大小为128、256 和 448 的情况，相当于 1.9M、3.8M 和 6.6M 个 token。如图 9 所示，增加 Batch 大小导致了类似的效果。
我们在一个程序生成的 2D 平台游戏环境 CoinRun（Cobbe等，2020）的难和易两个设置中进行评估，并与具有专家行动作为上界的 oracle 行为克隆（BC）模型以及具有随机行动作为下界进行比较（见图 15）。基于 LAM 的策略在只有 200 个专家样本用于适应的情况下就实现了与 oracle 相同的得分，尽管几乎可以确定它以前从未见过 CoinRun。这证明了学到的潜在行动在迁移方面是一致且有意义的，因为从潜在到实际的映射不包含关于当前观察的任何信息。
