Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning

agent 2601.16163
robot-policyvideo-diffusionworld-modelvalue-functionplanning

Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning — L2 #

1. TL;DR #

把预训练视频扩散模型 (Cosmos-Predict2-2B) 通过单阶段微调、零架构改动变成机器人策略:动作、未来状态、状态价值都编码成"注入的 latent frame",共享同一套视频扩散去噪目标。既做直接策略 (LIBERO 98.5% / RoboCasa 67.1% / ALOHA 93.6% SOTA),又能用 rollout 数据精炼 world model + value 做 best-of-N 规划,再涨 +12.5 分。

2. Q1 / Q2 / Q3 #

Q1 — 痛点:现有把视频模型用于机器人的方法通常需要多阶段训练(先微调视频、再单独训练 action module / inverse dynamics)并引入新架构组件(独立 action diffuser、value 模块)。而完全统一的 video-action 模型 (UVA, UWM) 又放弃了预训练视频 backbone,丢掉了时空先验。世界模型/价值函数类方法 (Dreamer, TD-MPC, SAILOR, Latent Policy Steering) 也依赖 policy / world model / value 三套分离模块且多从零训练。

Q2 — 方法:单一统一架构同时充当 policy、world model、value function,且从预训练视频模型初始化。核心机制是 latent frame injection——不改架构,把机器人本体感 (proprioception)、动作块 (action chunk)、状态价值这些新模态,归一化到 $[-1,+1]$ 后复制填满一个图像 latent volume,直接插入视频扩散序列,序列顺序即 $(s, a, s', V(s'))$。训练时按 50/25/25 拆 batch(policy / world model / value),靠 conditioning mask 决定训哪一个。规划时用 dual deployment:原始 checkpoint 当 policy,rollout 精炼后的 checkpoint 当 planning model,跑 best-of-N。

核心技术壁垒:把标量/低维模态(动作、价值)广播复制填满 $H'\times W'\times C'$ 的图像 latent volume,让未经改动的视频 DiT 直接对它去噪,推理时对整个 latent volume 取平均即恢复该模态——无需 VAE 解码、无需新 head。这一"用图像 latent 的形状去承载非图像模态"的编码 trick 是全篇最难复制的洞见,它让预训练视频模型的核心学习算法原封不动地被复用来建模复杂动作分布。

Q3 — 结果:作为直接策略,在 LIBERO (98.5%) 与 RoboCasa (67.1%, 仅 50 条人工示范 vs 竞品 >300) 上刷新 SOTA,真机 ALOHA 双臂操作全量平均 93.6% 拿最高分,超过 diffusion-from-scratch、视频策略、以及在同一批机器人数据上微调的 VLA (π0.5、OpenVLA-OFT+)。加上模型规划后两个高难真机任务平均再涨 12.5 分。

3. 架构 / 方法图 #

3.1 Agent scope #

3.2 方法总览(teaser) #

Figure 1: Cosmos Policy overview — multimodal + multi-view inputs to action/state/value outputs

Paper's Figure 1(teaser):Cosmos Policy 接收多模态、多视角相机图像,预测 (1) 动作块、(2) 未来状态(本体感 + 图像观测)、(3) 价值(未来状态的 rewards-to-go)。 该图强调三点关键设计:对 base 视频模型零架构改动、所有模态统一通过视频扩散目标建模、一次前向同时吐出控制所需的动作与规划所需的 state/value。读者应注意输出被并列成同一序列的不同 latent frame,而非分叉的多头。

3.3 latent frame injection(主架构图) #

Figure 2: Latent diffusion sequence of Cosmos Policy

Paper's Figure 2:latent 扩散序列。第一行把原始多视角图像 tokenized 成 latent frame;第二行把新模态(robot state / action chunk / value)作为新 latent frame 插入序列中间;第三行让模型在 clean latent frame 条件下对 noised latent frame 去噪。 这是全文的机制核心:序列顺序从左到右就是 $(s, a, s', V(s'))$,因此支持从左到右的自回归解码。读者应注意"注入"发生在 latent 序列层面而非架构层面——没有任何新模块。

对一个"两个第三人称相机 + 一个腕部相机"的平台,序列含 11 个 latent frame:(1) blank 占位(因 VAE 的 $(1+\tfrac{T}{4})$ 时间压缩需要)、(2) 本体感、(3) 腕部图、(4)(5) 两个第三人称图、(6) 动作块、(7) 未来本体感、(8)(9)(10) 未来三视角图、(11) 未来状态价值。其中 (2)(6)(7)(11) 是新模态,(3)(5)(8)(10) 是额外视角。单相机可精简到 7 帧。注意 $s$ / $s'$ 只含 $t$ 与 $t+K$ 两个时刻($K$ 为 chunk size),不用历史、不预测多步中间帧。

Figure 8: Detailed view of the latent diffusion sequence

Paper's Figure 8:Figure 2 的实现细节版。 blank(全零)图像先插入图像序列再被 VAE tokenized,随后 latent injection 把占位 latent 完全覆写;因 tokenizer 把第一帧单独编码、其余每 4 帧一组压缩,为让当前/未来观测有一致的 latent 结构,作者把它们放在"blank 首帧"之后,并把每张图复制 4 份使一组四帧对应单个时间步。这解释了序列开头那个看似冗余的 blank 帧的由来。

3.4 the agent loop(去噪 → 规划状态机) #

一次 turn 的状态流转(观测→提案→想象→评分→执行)以 best-of-N 内省搜索取代传统"reflect":

stateDiagram-v2 [*] --> Observe Observe --> ProposeActions: policy model 采样 N 个 action chunk ProposeActions --> ImagineFutures: planning model 预测每个提案的 s' (world model, 3x ensemble) ImagineFutures --> ScoreValues: planning model 对每个 s' 估 V (5x ensemble = 15 值) ScoreValues --> Aggregate: majority-mean 聚合 Aggregate --> Execute: 选最高 V 的 action,整块执行 Execute --> Observe: 无 receding-horizon,chunk 执行完再 requery Execute --> [*]: 任务终止

3.5 Planning & reasoning #

4. 作者证明 #

无形式化作者证明 — 仅实证。 本文无收敛/成功率界,全部结论由 LIBERO/RoboCasa/ALOHA 的经验成功率与消融支撑。可被 bound 而未 bound 的量:best-of-N 规划相对直接策略的成功率提升(在 world model / value 误差假设下本可给出 regret 型界)。

记号表 #

符号含义
$D_\theta$去噪网络(diffusion transformer / DiT)
$\sigma$扩散噪声水平
$\mathbf{x}_0$干净的 VAE 编码 latent 帧序列
$\mathbf{c}$文本描述的 T5-XXL 嵌入(cross-attention 条件)
$\mathbf{n}\sim\mathcal{N}(\mathbf{0},\sigma^2\mathbf{I})$i.i.d. 高斯腐蚀噪声
$\langle S,A,T,R,H\rangle$有限时域 MDP 元组
$\hat{T}:S\times A\to\Pi(S)$学得的 world model
$V^\pi(s)$策略 $\pi$ 下状态 $s$ 的价值
$\gamma$折扣因子(把终止奖励回传)
$K\times d_{act}$动作块形状(K 步 × 动作维度)
$H'\times W'\times C'$单个 latent frame 的形状

方程物理意义 #

去噪目标(原封不动复用给动作生成的"核心学习算法"):

$$\mathcal{L}(D_\theta, \sigma) = \mathbb{E}_{\mathbf{x}_0, \mathbf{c}, \mathbf{n}}\left[\left\|D_\theta(\mathbf{x}_0 + \mathbf{n}; \sigma, \mathbf{c}) - \mathbf{x}_0\right\|_2^2\right]$$

标准 EDM 去噪 score matching:网络在给定噪声水平下从被噪化 latent 回归干净 latent。

稀疏奖励下的状态价值(Monte-Carlo 回报):

$$V^\pi(s) = \mathbb{E}_{\tau \sim \pi}\left[ \sum_{k=t}^{H} \gamma^{k-t} R(s_k, a_k) \mid s_t = s \right] = \mathbb{E}_{\tau \sim \pi}\left[\gamma^{H-t} R(s_H, a_H) \mid s_t = s \right]$$

因仅终止奖励非零,价值坍缩为折扣后的终止奖励;每个转移用 $\gamma^{H-t}R(s_H,a_H)$ 打标签。

latent injection 的复制次数(把 $K\times d_{act}$ 的扁平动作向量填满一个 latent volume):$\frac{H' \times W' \times C'}{K \times d_{act}}$。

6 项最小检查(agent 扩展) #

  1. 成功率 sweep 单调性:Table 3 的逐组件移除 67.1→66.6→64.0→62.5→44.4,沿"移除辅助监督"轴单调下降;去掉 policy 训练样本上的未来状态辅助监督造成最大断崖(62.5→44.4)。可复现该 sweep。
  2. 每 turn latency 预算:直接策略 5 步去噪 0.61s (LIBERO/RoboCasa)、10 步 0.95s (ALOHA)、1 步 0.16s;规划模式 best-of-8 全流程 4.9s(8 GPU 并行)。作者未宣称"交互式低延迟",反而在 §6 承认 ~5s 规划延迟限制动态任务。
  3. 失败模式分类:主要指认两类 baseline 失败——高多模态(OpenVLA-OFT+ 在两颗糖之间够)与高精度(π0.5 抓不牢 ziploc 滑块);方法针对这两类(best-of-N + world model 预判"松手")。
  4. backbone 依赖:消融"from scratch"随机初始化在 LIBERO 掉 3.9 分、ALOHA fold-shirt 掉 18.7 分且动作抖动,证明预训练视频先验是必需项。
  5. 数据效率:RoboCasa 仅 50 条人工示范即超过用 300–3000+(甚至 +10000 合成)示范的竞品。
  6. auxiliary loss 有效性:去掉辅助目标 LIBERO 掉 1.5 分(Table 4),确认联合建模 $p(a,s',V(s')|s)$ 优于 $p(a|s)$。
  7. 5. 实验与数据 #

    5.1 直接策略基准(LIBERO / RoboCasa) #

    Table 1 — LIBERO(6000 trials / suite × 3 seeds):

    MethodSpatialObjectGoalLongAvg
    Diffusion Policy78.392.568.350.572.4
    π0.598.898.298.092.496.9
    OpenVLA-OFT97.698.497.994.597.1
    CogVLA98.698.896.695.497.4
    Cosmos Policy98.1100.098.297.698.5

    注意诚实的失败行:LIBERO-Spatial 上 Cosmos (98.1) 低于 π0.5 (98.8) 与 CogVLA (98.6);总均分靠 Long-horizon (97.6) 大幅领先拉高。

    Table 2: RoboCasa benchmark results

    Paper's Table 2:RoboCasa 24 任务均成功率。 Cosmos Policy 以 50 条示范拿 67.1%,超过 π0 (300 demos, 62.5%)、Video Policy (300, 66.0%)、FLARE (300, 66.4%) 等——这张表是"数据效率"这一惊讶点的核心证据:训练示范数少 6–60 倍仍占优。

    5.2 真机 ALOHA 与消融 #

    Figure 4: Real-world ALOHA robot evaluation results

    Paper's Figure 4:四个双臂任务的完成度分数 (0–100) 分组柱状图,Cosmos Policy 总分最高、四任务中三项第一。 关键点:Cosmos Policy 未经大规模动作标注机器人预训练却击败了经此预训练的 π0.5 / OpenVLA-OFT+,佐证"视频先验可替代动作数据预训练"。

    Table 5: In-distribution vs OOD ALOHA scores

    Paper's Table 5:Figure 4 的 in-distribution / OOD 细分。 Cosmos Policy 全量平均 93.6% 最高,但在纯 OOD 均分上以 89.3 落后 π0.5 的 92.5(尤其 "put candies in bowl" OOD 74.0 < 90.0)——头名胜在 in-distribution + 全量,而非 OOD。读者应据此校准"SOTA"的适用范围。

    Table 3: RoboCasa ablations

    Paper's Table 3:逐组件消融 + 单步去噪。 从 67.1% 一路移除到只预测动作的 barebones policy,去掉 policy 样本上的未来状态监督时从 62.5→44.4 断崖式下跌——这是最 load-bearing 的消融,说明"让 policy 也预测未来帧"才是有效性来源。底部还显示仅 1 步去噪即达 66.4% @ 0.16s/chunk(比 5 步的 0.61s 快近 4×,仅掉 0.5%)。

    Figure 5: Common failure modes of π0.5 and OpenVLA-OFT+

    Paper's Figure 5:baseline 失败模式。左 π0.5 高精度抓取失手丢掉 ziploc 袋;右 OpenVLA-OFT+ 在两颗糖之间下手,暗示 L1 回归难建模高多模态动作分布。 该图定性支撑了 §5.2 对"高精度 + 高多模态"两类难点的针对性论证。

    5.3 模型规划 #

    Figure 7: Model-based planning results

    Paper's Figure 7:两个高难任务上 base(无规划)vs 无模型规划 $Q(s,a)$ vs 有模型规划 $V(s')$。 有模型 $V(s')$ 变体总体最高,平均 +12.5 分。作者归因于利用学得的环境动力学做更样本高效的规划;在有限 rollout 数据下 $Q$ 函数因输入维度更高易过拟合。

    Figure 6: World model predictions before vs after rollout fine-tuning

    Paper's Figure 6:base checkpoint 的 world model 只在示范上训练,预测不出"丢滑块"这类错误(上);在 rollout 数据上微调后能更准地预测结果状态(下),从而支撑更有效的规划。 这解释了为何"从经验学习"对规划是关键(demo 只覆盖成功轨迹、分布过窄)。

    6. 论证链 #

    #论断依据(篇内)
    1预训练视频模型学到时间因果/隐式物理/运动先验,比 VLM 的静态图文先验更适合低层控制§1–§2 定位;§5.2 消融 from-scratch 掉 3.9 分证明先验必需
    2无需新架构:把动作/状态/价值编码成注入的 latent frame,复用同一去噪目标即可联合建模§4.1 + Appendix A(归一化→扁平→复制→覆写→取平均)
    3联合训练 policy/world model/value(50/25/25 拆 batch + 辅助监督)比单建模 $p(as)$ 更强§4.2;Table 4 (−1.5) 与 Table 3 (移除未来状态监督 −18.1)
    4仅靠示范不足以规划(分布过窄),需 rollout 数据精炼 world model/value§4.3;Figure 6 前后对比
    5dual deployment(原 checkpoint 当 policy,精炼 checkpoint 当 planner)+ best-of-N 提升成功率§4.3;Figure 7 (+12.5)
    6有模型 $V(s')$ 规划优于无模型 $Q(s,a)$,因能利用环境动力学、样本更高效§5.3 Q4;Figure 7 对比三变体

    7. 实现 cross-reference #

    代码/模型/数据已开源:https://research.nvidia.com/labs/dir/cosmos-policy/(含训练与评测脚本)。本 L1 未抓取具体仓库 file:line,故标 [实现未公开于本笔记] 的具体行号,以下为可从论文复原的关键实现细节:

    • 核心壁垒(§2 已述):标量/低维模态广播复制填满图像 latent volume + 无 VAE 解码地取平均恢复——见 Appendix A Para 3/5。动作块 $K\times d_{act}$ 归一化到 $[-1,+1]$、扁平后复制 $\frac{H'\times W'\times C'}{K\times d_{act}}$ 次、reshape 覆写占位 latent;价值为标量,直接对整个 latent volume 取平均并反归一化到 $[0,+1]$。
    • 关键实现细节 1(易漏):噪声分布改造。base 的 log-normal ($P_{\text{mean}}=1.39, P_{\text{std}}=1.2$) 在高 $\sigma$ 权重过低,导致动作生成初始去噪不准并级联误差;改为 0.7 log-normal + 0.3 uniform$[1.0,85.0]$ 的混合分布延长高 $\sigma$ 尾部;测试时把 $\sigma_{\min}$ 从 0.002 抬到 4(保持 $\sigma_{\max}=80$)以避开低 SNR 末段。见 Appendix B.1。

    Figure 9: Base vs adjusted noise distribution

    Paper's Figure 9:左为 base log-normal,右为 Cosmos Policy 的 log-normal-uniform 混合分布(更多高噪声权重)。 这一看似微小的采样改动是精确动作生成的实现前提,属于最易被忽视但影响成功率的 trick。

    • 关键实现细节 2(易漏):规划的 "majority-mean" 聚合。每个动作提案得 3 (world model) × 5 (value) = 15 个价值预测,先按固定阈值多数表决成功/失败,再在多数组内取均值——比朴素平均更抗双峰/高方差 outlier。见 §4.3 Para 13。

    Figure 12: Balanced batches training scheme

    Paper's Figure 12:50/25/25 的联合目标训练示意;全序列固定,靠 conditioning mask 决定训 policy / world model / value。 规划前在 rollout 上精炼时,mask 掉 $(s,a)$ 得到 $V(s')$、mask 掉 $s'$ 得到 $Q(s,a)$(模型无关规划变体)。该图是复现联合训练 batch 切分与 mask 逻辑的唯一凭据。