一句话定位
把预训练视频世界模型 cosmos-predict2 原样搬进机器人策略学习:不改架构,只把动作、本体感知、未来状态价值当作新的”帧”塞进同一套视频扩散序列里做单阶段微调,一个模型同时兼任策略、世界模型和价值函数三个角色,在 LIBERO(98.5%)、RoboCasa(67.1%)拿下新纪录,在真实双臂 ALOHA 任务上平均得分 93.6 分超过微调过的 pi0-5、openvla-oft。
背景与定位
这篇论文回应的是视频模型做机器人策略这条路线里一个具体的复杂度问题:此前的做法要么先把视频模型在机器人数据上微调、再单独训一个动作模块或逆动力学模型来产生动作(两阶段,架构上要加新组件),要么从头训一个统一的视频-动作模型(不依赖预训练视频先验,平白丢掉了大规模视频预训练学到的时空物理直觉)。作者是 Moo Jin Kim(openvla-oft、pi0-5 的作者之一,本项目在 NVIDIA 实习期间与 Stanford 合作完成)等人,团队横跨 NVIDIA 和 Stanford,论文录用于 ICLR 2026。
Cosmos Policy 选择的路径是第三种:直接拿预训练好的 cosmos-predict2-2B-Video2World 模型,不加任何新架构组件,只做一次机器人示范数据上的微调,让模型用它本来就有的视频扩散学习机制去同时学动作、未来状态、未来价值。作者的假设是:视频扩散模型本就擅长对复杂、高维、多峰分布建模,也能生成上百帧的时序连贯视频,这套机制天然适合把动作也当成”一种模态”塞进去学。
和同期的世界模型 + 价值函数路线相比(FLARE 用可学习的未来 token 挂在扩散 Transformer 序列上、SAILOR 用独立的世界模型和奖励模型配 MPPI 做迭代搜索、Latent Policy Steering 用光流做跨本体动作表示再单训一个价值函数来引导策略),Cosmos Policy 的差异化在于策略、世界模型、价值函数三者共享同一套权重、同一个训练目标,而不是分开训练几个模块再拼起来。
模型架构
底座是 cosmos-predict2-2B-Video2World:一个基于 Wan2.1 时空 VAE tokenizer 的连续 token 隐空间视频扩散模型,用 EDM 去噪分数匹配训练,去噪网络是一个 Diffusion Transformer(通过交叉注意力接收 T5-XXL 文本嵌入作为条件,通过自适应层归一化接收噪声水平 σ)。Wan2.1 tokenizer 把 (1+T)×H×W×3 的视频压成 (1+T’)×H’×W’×16 的隐帧序列,其中 T’=T/4、H’=H/8、W’=W/8;首帧不做时间压缩,用来接收单张输入图像作为条件。
核心机制:latent frame injection(隐帧注入)。 不新增任何网络模块,而是把动作、本体感知、未来状态价值都编码成和图像同样形状的隐帧,直接插进原本只装图像的隐扩散序列里。举例:一个配两个固定第三人称相机加一个腕部相机的平台,隐序列里排 11 帧——(1) 一个技术性的空白占位帧(因 VAE 的 1+T/4 时间压缩方案而必须存在)、(2) 机器人本体感知、(3) 腕部相机图像、(4)(5) 两路第三人称图像、(6) 动作 chunk、(7) 未来本体感知、(8) 未来腕部图像、(9)(10) 未来两路第三人称图像、(11) 未来状态价值。这个排列顺序恰好对应 (s, a, s′, V(s′)),从左到右天然支持自回归解码。若只有一路第三人称相机,则整条序列缩到 7 帧。注入方式是把每个模态归一化到 [−1,+1] 后复制填满一个 H′×W′×C′ 的隐帧体积(动作 chunk 形状 K×d_act,先展平再复制 (H′×W′×C′)/(K×d_act) 次);取值时反过来,对整个隐帧体积取平均再反归一化,提取动作/本体感知不需要经过 VAE 解码。
具体到 LIBERO 配置(据 HF 模型卡):输入第三人称相机(agentview)224×224 RGB + 腕部相机 224×224 RGB + 9 维本体感知(2 维夹爪关节 + 3 维末端位置 + 4 维末端四元数);输出 16 步动作 chunk(每步 7 维:6 自由度末端控制 + 1 维夹爪)、t+16 时刻的未来本体感知与两路未来图像、以及一个标量价值。
联合训练目标(策略/世界模型/价值函数三合一)。 每个训练 batch 里,50% 样本来自示范数据集,用来训策略 p(a, s′, V(s′)|s);另外 50% 来自 rollout 数据集,平均分成两半,一半训世界模型 p(s′, V(s′)|s, a),一半训价值函数 p(V(s′)|s, a, s′)。哪一部分隐帧被当作条件、哪一部分被当作去噪目标,决定了当前这一步到底在训三个功能里的哪一个。策略和世界模型都带辅助监督:策略不是只学 p(a|s),而是联合学 p(a, s′, V(s′)|s);世界模型不是只学 p(s′|s, a),而是联合学 p(s′, V(s′)|s, a)。消融显示去掉这些辅助目标会让 LIBERO 平均成功率掉 1.5 个百分点(见评测部分)。
推理时既可以并行解码(action/未来状态/价值一次性生成,速度快,适合不做规划的直接策略部署),也可以从左到右自回归解码(质量更高,适合规划场景,且允许策略与世界模型/价值函数用不同 checkpoint)。
噪声调度专门为动作生成做了调整。 底座模型的原始噪声分布是 EDM 式对数正态 ln(σ)~N(P_mean=1.39, P_std=1.2²),这套分布对视频生成够用,但对动作生成不够——扩散采样从 σ_max=80 开始,原始分布在高 σ 区训练信号太弱,导致早期去噪不准、误差逐步累积。Cosmos Policy 训练时改用混合分布:0.7 概率采样自原始对数正态、0.3 概率采样自 [1.0, 85.0] 均匀分布,拉长高噪声区的训练权重(0.7/0.3 这个比例没有调参,只是保留接近原分布)。推理时进一步把最低噪声水平从原来的 σ_min=0.002 提高到 σ_min=4(σ_max 仍是 80),因为发现越接近 σ≈0 的最后几步去噪反而因信噪比太低而不准。
规划机制:policy model + planning model 的双模型部署。 训练完基础 Cosmos Policy(称为 policy model)之后,再拿它在自己 rollout 出来的数据上继续微调出一个 planning model(90% 的 batch 用来训世界模型和价值函数,只有 10% 训策略),planning model 专门负责更准的未来状态和价值预测。测试时用 best-of-N:policy model 采样 N 个候选动作,planning model 对每个候选做集成预测——世界模型查询 3 次、每次未来状态再查询价值函数 5 次,一共 15 个价值预测,用”多数均值”(先判断多数预测成功还是失败,再在多数组内取均值,比朴素平均更抗离群值)聚合,选预测价值最高的动作执行完整个 chunk。value function 训练时可以用输入掩码在 V(s′)(遮住 s,a,需要先由世界模型预测未来状态)和 Q(s,a)(遮住 s′,免世界模型的模型无关版本)之间切换,论文对比发现前者(model-based)优于后者(model-free)。
数据
三个任务套件,示范数据全部来自目标平台本身,没有额外的”无动作”互联网视频联合训练:
- LIBERO 仿真基准(单臂 Franka Emika Panda):Spatial / Object / Goal / Long(即 LIBERO-10)四个套件,每套件 10 个任务 × 50 条示范 = 500 条,四套件合计 2000 条示范(GitHub README 训练说明确认为”2000 demos”)。策略训练只用过滤掉失败重放的示范,世界模型和价值函数训练用未过滤的全集(含失败样本)。
- RoboCasa 仿真基准(单臂 Franka Panda,24 个厨房操作任务):官方提供 50 条人类遥操作示范 + 额外 1000 条 MimicGen 合成示范,但为了体现数据效率,Cosmos Policy 只用 50 条人类遥操作示范(24 任务合计 1200 条,GitHub README 确认),不用 MimicGen 合成数据,对照组基线普遍用 300 条起。
- ALOHA 真实双臂机器人(两只 ViperX 300S + 3 相机:1 俯视 + 2 腕部,控制频率从 50Hz 降到 25Hz 以省算力):4 个任务合训一个策略——“put X on plate”(80 条示范,测语言跟随)、“fold shirt”(15 条,测长程接触密集操作)、“put candies in bowl”(45 条,测多峰抓取序列)、“put candy in ziploc bag”(45 条,测毫米级精度操作),合计 185 条示范。
- 失败样本处理:LIBERO/RoboCasa 仿真里约 10%-20% 的示范重放会失败(人类遥操作误差所致),这部分失败样本被保留进 rollout 数据集用于训练世界模型和价值函数,只在策略训练时被过滤掉;ALOHA 真实环境遥操作更仔细,不存在失败示范,此时 rollout 数据集与示范数据集相等。
- Rollout 数据(用于规划微调):直接策略评测阶段已经在所有方法上跑出 505 条 rollout,额外为”put candy in ziploc bag”任务专门收集 143 条 Cosmos Policy 自己的 rollout(该任务因手臂自遮挡导致相机可观测性低、动作容差窄至毫米级,世界模型特别难学),合计 648 条 rollout 用来微调出 planning model。
- HuggingFace 上公开的训练数据集:
nvidia/LIBERO-Cosmos-Policy(约 64.3 万行)、nvidia/RoboCasa-Cosmos-Policy(约 297.6 万行)、nvidia/ALOHA-Cosmos-Policy(555 行)。
训练方法
基础训练是对 cosmos-predict2-2B-Video2World 的全参数微调(不冻结任何权重),三个任务套件各自独立训练一版 checkpoint:
| 套件 | 训练步数 | GPU | 全局 batch size | 训练时长 | 动作 chunk | 执行步数 |
|---|---|---|---|---|---|---|
| LIBERO | 40K 步 | 64× H100 | 1920 | 48 小时 | 16 步 | 全 chunk(16) |
| RoboCasa | 45K 步 | 32× H100 | 800 | 48 小时 | 32 步 | 半 chunk(16) |
| ALOHA | 50K 步 | 8× H100 | 200 | 48 小时 | 50 步(2 秒 @ 25Hz) | 全 chunk(50) |
40K 步后 LIBERO 的单步 L1 训练损失:动作 0.012、未来本体感知 0.007、未来腕部图像隐帧 0.068、未来第三人称图像隐帧 0.036、价值 0.007(损失量级说明模型学动作/本体感知/价值比学图像快得多)。RoboCasa 和 ALOHA 的对应损失量级相近(RoboCasa:0.016/0.007/0.084/0.048/0.007;ALOHA:0.010/0.008/0.097/0.085/0.007)。
规划阶段的微调是在 policy model 基础上,用 648 条 rollout 数据,按 90%(世界模型+价值函数各占一半)/10%(策略)的比例继续训练,得到 planning model,再和原始 policy model 组成”双模型部署”用于 best-of-N 搜索(见架构部分)。
基线的公平对比设置:在 ALOHA 上为了公平比较,π_0.5、π_0、OpenVLA-OFT+ 都用同样的算力预算(8× H100,48 小时)微调,但因框架吞吐不同,实际完成的梯度步数差异很大——π_0.5 和 π_0 跑了 40 万步(batch 256),OpenVLA-OFT+ 跑了 3.2 万步(batch 96,梯度累积 4)。参数量小得多(约 1.5 亿,对比 Cosmos Policy 等模型的 2B-7B)的 Diffusion Policy 从零训练,单张 H100 跑 48 小时,相当于 190 个 epoch / 7.2 万步(batch 350)。
Infra(训练 / 推理工程)
训练算力:LIBERO 用 64× H100(全参数微调 2B 模型,48 小时),RoboCasa 用 32× H100(48 小时),ALOHA 用 8× H100(48 小时)——均为 80GB H100。精度方面 HF 模型卡注明”仅测试过 BF16 精度推理”,训练精度论文未单独披露。
推理显存(单卡,不含规划):LIBERO 任务 6.8 GB、RoboCasa 任务 8.9 GB、ALOHA 任务 6.0 GB;加上 best-of-N 规划的 ALOHA 任务:串行最低 10.0 GB(单卡),并行推荐 N 张卡各 10.0 GB(GitHub README/HF 模型卡口径一致)。
推理延迟(单张 H100,生成一个完整动作 chunk 的墙钟时间):5 步去噪并行生成动作+未来状态+价值,LIBERO/RoboCasa 用时 0.61 秒;ALOHA 用 10 步去噪,0.95 秒;降到 1 步去噪时 0.16 秒(此时 RoboCasa 平均成功率仍有 66.4%,相比 5 步的 67.1% 只降 0.7 个百分点,速度快了近 4 倍)。
规划(best-of-N)推理:真实 ALOHA 机器人上用 N=8 的 best-of-N 搜索,配 8 张并行 H100(每张负责一个搜索分支)——policy model 先生成 8 个候选动作 chunk(每个 10 步去噪),planning model 对每个候选生成 3 个未来状态集成预测(每个 5 步去噪),再对每个未来状态生成 5 个价值集成预测(每个 5 步去噪),共 15 个价值预测,整个搜索过程总耗时 4.9 秒(8 张 H100 并行)。作者在讨论部分明确承认这个延迟是当前方法应用到动态任务或移动任务上的主要瓶颈。
软件栈:HuggingFace Transformers 作为运行时,官方标注的兼容硬件是 NVIDIA Hopper(H100),HF 模型卡”Test Hardware”一栏同时列出 H100 与 A100。
评测 benchmark
LIBERO(4 套件 × 500 试验 × 3 随机种子,合计 6000 次试验):
| 方法 | Spatial | Object | Goal | Long | 平均 |
|---|---|---|---|---|---|
| Diffusion Policy | 78.3 | 92.5 | 68.3 | 50.5 | 72.4 |
| Dita | 97.4 | 94.8 | 93.2 | 83.6 | 92.3 |
| π0 | 96.8 | 98.8 | 95.8 | 85.2 | 94.2 |
| UniVLA | 96.5 | 96.8 | 95.6 | 92.0 | 95.2 |
| π0.5 | 98.8 | 98.2 | 98.0 | 92.4 | 96.9 |
| OpenVLA-OFT | 97.6 | 98.4 | 97.9 | 94.5 | 97.1 |
| CogVLA | 98.6 | 98.8 | 96.6 | 95.4 | 97.4 |
| Cosmos Policy | 98.1 | 100.0 | 98.2 | 97.6 | 98.5 |
RoboCasa(24 任务 × 50 试验 × 3 随机种子,合计 3600 次试验,评测场景含未见过的物体实例,5 个场景中 2 个含未见过的装修风格):
| 方法 | 每任务训练示范数 | 平均成功率 |
|---|---|---|
| GR00T-N1 | 300 | 49.6 |
| DP-VLA | 3000 | 57.3 |
| GR00T-N1.5 | 300 | 64.1 |
| Video Policy | 300 | 66.0 |
| FLARE / GR00T-N1.5+HAMLET | 300 | 66.4 |
| Cosmos Policy | 50 | 67.1 |
真实 ALOHA 机器人(4 任务合计 101 次试验,含分布内与分布外两类初始条件):
| 方法 | 分布内均分 | 分布外均分 | 全部均分 |
|---|---|---|---|
| Diffusion Policy | 39.2 | 24.1 | 33.6 |
| OpenVLA-OFT+ | 67.1 | 51.7 | 62.0 |
| π0 | 81.3 | 71.7 | 77.9 |
| π0.5 | 87.8 | 92.5 | 88.6 |
| Cosmos Policy | 96.3 | 89.3 | 93.6 |
分任务看,Cosmos Policy 在”put X on plate”(100.0)和”put candy in ziploc bag”(高精度任务,85.4 分,远超 π0.5 的 61.5)上领先明显;“fold shirt”上和 π0.5、OpenVLA-OFT+ 打平(均约 99.5);唯独”put candies in bowl”(多峰抓取)上 π0.5 反而更高(95.2 对 89.6)。分布外场景里 π0.5 略胜 Cosmos Policy(92.5 对 89.3)。
消融(LIBERO):去掉辅助监督(策略/世界模型不再联合预测未来状态和价值)平均成功率从 98.5 掉到 97.0(−1.5);把预训练视频权重换成随机初始化再训同样步数,掉到 94.6(−3.9)。从零训练的版本在 ALOHA “fold shirt” 上只拿 80.8 分(比完整版低 18.7 分),且动作明显更”抖”,作者因担心长期损伤机械臂而中止了对该变体的进一步评测。
逐层消融(RoboCasa):从完整目标(67.1%)开始逐步去掉组件——去掉价值函数训练样本(66.6%)→再去掉世界模型训练样本(64.0%)→再去掉策略训练里的辅助价值监督(62.5%)→最后去掉策略训练里的辅助未来状态和价值监督、只剩纯动作预测(44.4%),证明”让策略联合预测未来状态”是效果提升里最关键的一环。
规划实验:在 ALOHA 两个更难的任务(“put candies in bowl”、“put candy in ziploc bag”)上,加入 best-of-N 规划后平均得分比不做规划的基础策略提高 12.5 分;model-based(V(s′))规划优于 model-free(Q(s,a))规划,作者归因于后者在 rollout 数据有限时更难学准、且输入维度更高更容易过拟合。
创新点与影响
Cosmos Policy 的核心贡献是把”用预训练视频模型做机器人策略”这条路线里普遍存在的”多阶段训练 + 新增架构组件”的复杂度直接砍掉:用同一套隐扩散序列、同一套训练目标,让动作、未来状态、未来价值和图像平起平坐地被当作视频帧来生成,不新增任何网络模块。这让一个 2B 参数的通用视频基础模型经过单阶段、单机器人平台的微调后,在 LIBERO 和 RoboCasa 上双双刷新纪录,并且在数据量远少于对手(RoboCasa 只用 50 条示范对比基线的 300 条起)的情况下做到——说明视频模型的预训练先验本身就是强初始化,不需要额外的大规模动作标注数据来兜底。同一份权重同时充当策略、世界模型、价值函数,也让”用 rollout 数据继续微调出更准的世界模型和价值函数、再用 best-of-N 做测试时规划”这件事变得直接:不需要维护三套独立模块。
作者自己指出的局限也很明确:model-based 规划推理慢(真实机器人上一次搜索约 5 秒),难以用于动态或移动任务;规划要练准需要相当量的 rollout 数据,超出示范分布之外的泛化仍不够,如何用更少 rollout 达到同等效果是开放问题;当前只做了一层 best-of-N 搜索,把世界模型预测时域拉长、搜索树做深会是后续方向。
原始链接
- arXiv: https://arxiv.org/abs/2601.16163
- PDF: https://arxiv.org/pdf/2601.16163
- 项目主页: https://research.nvidia.com/labs/dir/cosmos-policy/
- GitHub: https://github.com/nvlabs/cosmos-policy
- Hugging Face(模型 + 训练数据集合): https://huggingface.co/collections/nvidia/cosmos-policy
- 总结视频: https://youtu.be/V2qdFD9n5BM
一手源存档(sources/)
- nvidia-cosmos-policy—project-page.md — 官方项目主页(评测表格、代码/HF/视频链接)
- nvidia-cosmos-policy—github-readme.md — GitHub README(系统要求、训练算力配置、推理代码示例)
- nvidia-cosmos-policy—hf-card.md — HuggingFace
nvidia/Cosmos-Policy-LIBERO-Predict2-2B模型卡(输入输出维度、训练配置、推理硬件) - arXiv 原文 PDF(2601.16163),不入 git,可通过上方 arXiv 链接重下