Cosmos World Foundation Model Platform for Physical AI
NVIDAI
目录
[0. 摘要](#0. 摘要)
[1. 介绍](#1. 介绍)
[2. 世界基础模型平台](#2. 世界基础模型平台)
[5. 世界基础模型预训练](#5. 世界基础模型预训练)
[6. 世界基础模型后训练](#6. 世界基础模型后训练)
[6.1 用于相机控制的世界基础模型后训练](#6.1 用于相机控制的世界基础模型后训练)
[6.1.1 数据集](#6.1.1 数据集)
[6.1.2 微调](#6.1.2 微调)
[6.1.3. 评估](#6.1.3. 评估)
[6.2. 用于机器人操作的世界基础模型后训练](#6.2. 用于机器人操作的世界基础模型后训练)
[6.2.1 数据集](#6.2.1 数据集)
[6.2.2 微调](#6.2.2 微调)
[6.2.3 评估](#6.2.3 评估)
[9. 结论与讨论](#9. 结论与讨论)
0. 摘要
物理人工智能(Physical AI)需要首先在数字环境中进行训练。它需要自身的数字孪生(策略模型)和世界的数字孪生(世界模型)。在本文中,我们提出了 Cosmos World Foundation Model Platform,以帮助开发者为其物理人工智能系统构建定制的世界模型。我们将世界基础模型定位为一种通用的世界模型,可以通过微调转换为特定下游应用的定制化世界模型。我们的平台包括视频策划管道(video curation pipeline)、预训练的世界基础模型、预训练世界基础模型的后续训练示例以及视频分词器(tokenizer)。为了帮助物理人工智能构建者解决社会中最关键的问题,我们将平台开源,并以宽松的许可协议开放模型权重。
1. 介绍
物理人工智能是一种配备了传感器和执行器(actuator)的人工智能系统:传感器使其能够观察世界,执行器使其能够与世界互动并对其进行修改。物理人工智能的潜力在于将人类从危险、繁重或乏味的体力劳动中解放出来。尽管过去十年中,由于数据和计算能力的扩展,人工智能的多个领域取得了显著进展,但物理人工智能的发展却步履维艰。这主要是因为扩展物理人工智能的训练数据更加困难:所需数据必须包含观察与行动交替的序列,而这些行动会扰动物理世界,可能对系统和世界造成严重破坏。这种情况在人工智能的初期阶段尤为突出,因为此时探索性行为是必要的。世界基础模型(World Foundation Model,WFM)作为物理世界的数字孪生体(digital twin),可供物理人工智能安全交互,是解决数据扩展问题的一种长期追求的方案。
在本文中,我们介绍了用于构建物理人工智能的 Cosmos World Foundation Model (WFM) Platform 。**我们主要关注视觉世界基础模型,其中观察以视频的形式呈现,而扰动可以以多种形式存在。**如图 2 所示,我们提出了一种 "预训练-后训练" 范式,将 WFM 分为预训练 WFM 和后训练 WFM。为了构建预训练 WFM,我们利用大规模视频训练数据集,使模型接触到多样化的视觉体验,从而成为通用模型。为了构建后训练 WFM,我们使用从特定物理人工智能环境中收集的数据集微调预训练WFM,从而生成针对特定物理人工智能配置的专业化 WFM。如图 1 所示,我们展示了预训练和后训练 WFM 的示例结果。
数据决定了人工智能模型的上限。为了构建高上限的预训练 WFM,我们开发了一个视频数据策划管道(video data curation pipeline)。
- 我们使用该管道定位具有丰富动态和高视觉质量的视频内容,以促进学习视觉内容中编码的物理原理。
- 通过该管道,我们从 2000 万小时的视频集中提取了约 1 亿段时长为 2 至 60 秒的视频片段。对于每个片段,我们使用视觉语言模型(VLM)为每 256 帧生成一个视频标题(caption)。
- 视频处理计算成本高昂,因此我们利用现代 GPU 中可用的 H.264 视频编码器和解码器的硬件实现进行解码和转码。
- 我们的视频数据策划管道还利用了许多预训练的图像/视频理解模型。由于这些模型的吞吐量不同,为了最大化可训练视频数据的总体生成吞吐量,我们构建了一个基于 Ray 的编排管道(Moritz等,2017)。详情见第 3 节。
我们探讨了两种构建预训练 WFM 的可扩展方法,这些方法在第 5 节中进行了讨论。
- 这些方法分别是基于 Transformer 的扩散模型和基于 Transformer 的自回归模型。扩散模型通过逐步从高斯噪声视频中去除噪声来生成视频。自回归模型则按照预设顺序,基于过去生成的视频片段逐步生成新的视频。
- 这两种方法都将复杂的视频生成问题分解为更简单的子问题,从而使其更易于处理。
- 我们利用了当前最先进的 Transformer 架构以实现可扩展性。
- 在第5.1节中,我们提出了一种具有强大世界生成能力的基于 Transformer 的扩散模型设计。
- 在第5.2节中,我们提出了一种基于 Transformer 的自回归模型设计用于世界生成。
基于 Transformer 的扩散模型和基于 Transformer 的自回归模型都使用 token 作为视频的表示,其中前者使用向量形式的连续 token,而后者使用整数形式的离散 token。
- 需要注意的是,视频的 token 化过程------即将视频转换为 token 集的过程------并不简单。视频中包含丰富的视觉世界信息。
- 然而,为了促进 WFM 的学习,我们需要将视频压缩为紧凑的 token 序列,同时尽可能保留视频的原始内容,因为 WFM 训练的计算复杂性随 token 数量增长。
- 某种程度上,构建视频 tokenizer 类似于构建视频编解码器。我们开发了一种基于注意力的编码器-解码器架构,用于学习连续和离散 token 的视频 tokenization 方法,具体内容详见第 4 节。
我们在第 6 节中通过微调预训练的 WFM 来生成适用于各种物理人工智能任务的后训练 WFM。
- 在第 6.1 节中,我们微调了预训练的扩散 WFM,使其具备以相机姿态为条件的能力。这种后训练创造了一个可导航的虚拟世界,用户可以通过移动虚拟视点探索生成的世界。
- 在第 6.2 节中,我们将 WFM 微调用于各种机器人任务,这些任务包含视频-动作序列。我们展示了通过利用预训练的 WFM,可以更好地预测机器人执行动作后世界的未来状态。
- 在第 6.3 节中,我们演示了如何微调预训练的 WFM 以用于各种与自动驾驶相关的任务。
我们开发的 WFM 旨在服务于物理人工智能构建者。为了更好地保护开发者在使用世界基础模型时的安全性,我们设计了一个强大的防护系统,包括一个 "预防护" 模块用于阻止有害输入,以及一个 "后防护" 模块用于阻止有害输出。具体细节详见第 7 节。
我们的目标是构建一个世界基础模型平台,帮助物理人工智能开发者推进其系统的发展。为实现这一目标,我们通过 NVIDIA Cosmos 和 NVIDIA Cosmos Tokenizer 分别以 NVIDIA 开放模型许可证提供预训练的世界基础模型和 tokenizer。预训练脚本和后训练脚本将与视频数据策划管道一同在NVIDIA Nemo 框架中提供,帮助开发者定制微调数据集。尽管本文在世界基础模型设计方面取得了若干进展,但这一问题仍远未解决。需要进一步的研究来推动最先进技术的发展。
2. 世界基础模型平台
设 x_{0:t} 为从时间 0 到时间 t 的一系列真实世界视觉观察序列。设 c_t 为对世界的扰动。如图 3 所示,世界基础模型(WFM)是一个模型 W,用于基于过去的观察 x_{0:t} 和当前的扰动 c_t,预测时间 t+1 的未来观察 x_{t+1}。在我们的情况下,x_{0:t} 是一段 RGB 视频,而 c_t 是可以采取多种形式的扰动。它可以是物理人工智能采取的动作、随机扰动、扰动的文本描述等。
5. 世界基础模型预训练
预训练的世界基础模型是通用模型,能够捕捉现实世界物理和自然行为的普遍知识。
- 我们利用两种不同的可扩展深度学习范式,扩散模型和自回归模型,来构建两类世界基础模型。
- 扩散模型和自回归模型都将一个困难的生成问题分解为一系列更易处理的子问题,并且一直在加速生成模型的发展。
- 在扩散模型的情况下,困难的生成问题被分解为一系列去噪问题;在自回归模型的情况下,困难的生成问题被分解为一系列下一 token 预测问题。
- 我们讨论了如何通过各种并行化技术,在我们构建预训练世界基础模型的过程中,利用现代 GPU 扩展这些深度学习范式。
- 我们使用 10000 个 NVIDIA H100 GPU 的集群,在三个月的时间内训练了本文中报告的所有世界基础模型。
6. 世界基础模型后训练
在本节中,我们展示了如何对我们的 Cosmos 世界基础模型进行微调,以支持多种物理 AI 应用。我们包括了以下几个示例:通过相机控制对世界基础模型进行后训练,以实现 3D 可导航视觉世界的生成;在两种不同的机器人配置上进行后训练,以支持两种不同的机器人操作任务;以及通过多视角支持对世界基础模型进行后训练,以训练自动驾驶智能体。
6.1 用于相机控制的世界基础模型后训练
通过以相机姿势为条件,我们将相机控制集成到 Cosmos-1.0-Diffusion-7B-Video2World 中,使其成为一个有效的 3D 世界模拟器。我们将结果的后训练世界基础模型命名为 Cosmos-1.0-Diffusion-7B-Video2World-Sample-CameraCond。我们专注于从单一参考输入图像生成 3D 世界,通过相机控制从指定的相机轨迹生成时间一致且 3D 一致的视频模拟,其中视角的变化与场景的基础 3D 结构对齐。
6.1.1 数据集
- 我们使用 DL3DV-10K(Ling et al., 2024),一个大规模的静态场景视频数据集,来完成这项任务。
- 作为预处理步骤,我们将所有视频分割成包含 256 帧的片段。为了为片段中的所有帧密集地获取相机姿势注释,我们使用 GLOMAP(Pan et al., 2025)对分割后的片段进行结构光束恢复。
- 我们将第一帧的相机姿势设置为单位变换(identity transform),并计算所有后续帧的相对相机姿势。
- 我们还使用一个专有的 VLM 为视频添加标题(caption),以获取描述视频为静态场景的文本提示。
6.1.2 微调
我们通过将采样的潜在嵌入与 Plücker 嵌入(Sitzmann et al., 2021)拼接,来添加相机控制条件,Plücker 嵌入的空间维度与潜在嵌入(latent embeddings)相同。具体而言,给定相机姿势,我们通过以下公式计算 Plücker 坐标:
其中 c 是相机中心位置,d 是每个潜在像素的单位射线方向(其中潜在嵌入被视为降采样后的图像)。所有相机姿势都相对于初始帧是相对的。Cosmos-1.0-Diffusion-7B-Video2World 模型使用的 Cosmos-1.0-Tokenizer-CV8x8x8 具有 8 倍的时间压缩率,因此对于每 8 帧,我们使用第 4 帧的 Plücker 嵌入与相应的潜在表示拼接。
我们将训练视频的输入帧大小调整为 704×1252,并通过反射填充它们到 704×1280。在训练期间,我们采样 57 帧。训练目标和其他超参数与基础扩散世界基础模型训练(第 5.1.3 节)相同。
6.1.3. 评估
我们假设给定了一个世界的单一参考图像,并从该输入图像生成未来的滚动视频(rollout)。我们与 CamCo(Xu et al., 2024),在这种设置下用于相机可控视频生成的最新模型进行比较。为了公平比较,我们使用了同样在 DL3DV-10K(Ling et al., 2024)训练集上微调过的 CamCo 模型。由于我们的后训练世界基础模型生成 57 帧,而 CamCo 只能生成 14 帧,我们比较了相同的 57 帧轨迹,其中我们将 CamCo 的帧数按时间下采样为 4 倍。CamCo 生成的视频分辨率限制为 256 × 256。我们还最大限度地对输入图像和测试帧进行中心裁剪以进行评估。
6.2. 用于机器人操作的世界基础模型后训练
世界基础模型具有作为机器人操作的强大规划者和模拟器的潜力。在这里,我们展示了如何针对两个任务微调我们的预训练世界基础模型:
- 基于指令的视频预测(instruction-based video prediction):输入是机器人当前的视频帧以及文本指令,输出是预测的机器人根据指令执行的操作视频。
- 基于动作的下一帧预测(action-based next-frame prediction):输入是机器人当前的视频帧以及当前和下一帧之间的动作向量,输出是显示机器人执行指定动作结果的下一帧。给定一系列动作,模型可以自回归地运行,预测机器人执行给定动作的视频。
6.2.1 数据集
我们为上述两个任务创建了两个数据集。
对于基于指令的视频预测,我们创建了一个内部数据集,名为 Cosmos-1X 数据集。
- 该数据集包含约 200 小时由 EVE(1x.Tech 的类人机器人)拍摄的自我中心(egocentric)视频,涵盖了多种任务,包括导航、折叠衣物、清洁桌面、捡起物体等。
- 从原始视频中,我们选取了约 12,000 个时长从 1 秒到 9 秒不等的片段。每个片段都带有一个一句话指令,之后通过专有的 VLM 进行上采样。视频的拍摄帧率为 30 FPS,分辨率为 512 × 512。
对于基于动作的下一帧生成,我们使用了一个名为 Bridge(Ebert et al., 2022)的公开数据集,并使用与先前工作(Zhu et al., 2024)相同的配置进行比较。
- Bridge 数据集包括约 20,000 个第三人称视角的视频片段,展示了机器人臂在厨房环境中执行不同任务的视频,分辨率为 320 × 256,帧率为 5 FPS。
- 对于每个视频帧,相应的动作被定义为在抓取器坐标空间中的一个 7 维向量(Δ𝑥, Δ𝑦, Δ𝑧, Δ𝜃𝑟, Δ𝜃_𝑝, Δ𝜃_𝑦, ΔGripper),如 OpenVLA(Kim et al., 2024)中所述。
6.2.2 微调
我们微调了我们的 Cosmos-1.0-Diffusion-7B-Video2World(第5.1节)和 Cosmos-1.0-Autoregressive-5B-Video2World(第5.2节)模型,以支持基于指令的视频预测和基于动作的下一帧预测任务。
对于基于指令的视频预测,我们基于基础世界基础模型构建了两个模型。
- 第一个叫做 Cosmos-1.0-Diffusion-7B-Video2World-Sample-Instruction,第二个叫做 Cosmos-1.0-Autoregressive-5B-Video2World-Sample-Instruction。
- 我们计算指令的 T5 嵌入,并通过交叉注意力将其添加到基础模型的微调中。
对于基于动作的下一帧预测,我们同样基于基础世界基础模型构建了两个模型。
- 第一个叫做 Cosmos-1.0-Diffusion-7B-Video2World-Sample-ActionCond,第二个叫做 Cosmos-1.0-Autoregressive-5B-Video2World-Sample-ActionCond。
- 由于动作是一种在预训练过程中未遇到的新模态,我们在模型内部引入了额外的模块用于条件化。
- 对于 Cosmos-1.0-Autoregressive-5B-Video2World-Sample-ActionCond,我们添加了一个动作嵌入器 MLP,用于将动作向量投影到张量中,然后通过交叉注意力将其集成到模型中。
- 对于 Cosmos-1.0-Diffusion-7B-Video2World-Sample-ActionCond,我们也添加了一个动作嵌入器 MLP 来预测动作到张量,但我们通过将其添加到 DiT 模块的时间戳嵌入中来集成到模型中。
6.2.3 评估
9. 结论与讨论
Cosmos 世界基础模型标志着朝着构建物理世界的通用模拟器迈出了重要一步。本研究概述了我们的综合方法,包括数据整理流程、连续和离散 tokenizer 的设计、扩散和自回归世界基础模型的架构,以及针对多样化下游物理 AI 任务的微调过程。值得注意的是,我们展示了预训练世界模型在关键应用中的适应性,包括 3D 世界导航、机器人操作和自动驾驶系统,这些应用需要 3D 一致性和动作可控性。
局限性。尽管取得了进展,世界基础模型的开发仍处于初期阶段。目前的模型,包括我们的模型,仍不足以作为可靠的物理世界模拟器。我们观察到我们的模型仍然存在一些问题,包括缺乏物体持久性、接触丰富的动力学不准确,以及指令执行的不一致性。此外,生成的视频的真实性并不总是反映出对基本物理原理(如重力、光的相互作用和流体动力学)的遵循。
评估提出了另一个重大挑战。定义强有力的标准让人类评估物理逼真度很困难,因为这种评估往往受到个人偏见、背景和其他主观因素的影响。此外,这些评估可能与下游物理AI任务中使用的度量标准不完全一致。为了解决这些挑战,有前景的方向包括开发由多模态大语言模型(LLMs)驱动的自动化评估器,并利用现有的物理模拟器来实现可重复和交互式评估,从而减少对人工评估的依赖。
自回归与扩散世界基础模型。我们在 3D 一致性(第 5.3.1 节)和机器人视频生成(第 6.2 节)的评估结果表明,基于扩散的世界基础模型目前提供更好的生成质量。通过微调,基于扩散的世界基础模型能够结合多种控制信号,包括相机姿势、末端执行器位置或自动驾驶车辆轨迹,并生成新的格式输出,如多视角视频。然而,基于自回归的世界基础模型具有重要的未开发潜力。它们可以(1)利用大型语言模型(LLMs)的预训练权重来继承广泛的世界知识;(2)通过使用为因果注意力设计的高级推理优化技术实现更快的生成。如果这些能力得到充分实现,自回归世界基础模型可能特别适合需要交互控制或实时处理的应用,例如机器人中的规划和仿真。
重要的是,扩散模型和自回归模型之间的界限并不是严格的。最近的进展表明,带有双向注意力的扩散 Transformer 可以蒸馏成带有因果注意力的学生 Transformer,从而在推理过程中支持 Key-Value 缓存(Yin et al., 2024)。同样,自回归模型可以通过在生成过程中结合局部双向注意力,利用扩散头生成图像(Zhou et al., 2024)。探索这些混合方法及其权衡仍然是一个活跃且有前景的研究领域。我们计划进一步研究这些,并在未来的工作中提供全面的分析。
**论文地址:**https://arxiv.org/abs/2501.03575
**项目页面:**https://github.com/NVIDIA/Cosmos