UWM 阅读笔记
Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets — Chuning Zhu et al., RSS 2025
摘要
UWM 找到了一个极其优雅的方式来统一策略学习和世界建模：给视频和动作各自一个独立的扩散时间步。就这一个改动，同一个模型在推理时只需要调两个数字，就能变成四种东西：纯策略、纯世界模型、逆动力学模型、正向动力学模型。一个 checkpoint 四种用法，无需改任何架构。
更妙的是，这个设计让模型自然地消化没有动作标注的纯视频数据：只需把动作侧的时间步设为最大噪声（"我对动作一无所知"），模型就自动退化成只学视频预测。在真实机器人上，UWM 比 Diffusion Policy 高出最多 20 个百分点；协同训练 2000 条无标注视频后，OOD 泛化能力再次大幅提升。而最反直觉的发现来自消融：独立时间步本身就贡献了 12% 的性能提升，超过任何其他单项设计选择。
一、Motivation
机器人学习面临一个数据困境：高质量示范轨迹（操作员用遥操作设备采集、每条都带精确动作标注的那种）收集成本极高，100 条可能就要大半天，1000 条需要数周甚至更久。然而互联网上有海量的视频，人做饭、组装家具、操作工具、整理房间，里面蕴含着丰富的物理动力学知识：物体被推会怎么滑、液体倒出来会怎么流、柔性物体抓握时会怎么变形。问题是这些视频只有像素，没有机器人关节角度，也没有末端执行器位姿。
之前的工作大致走了三条路，各有各的问题。第一条路是只用标注数据：Diffusion Policy 完全不碰视频，标注数据用得高效，但海量无标注视频里的物理知识就白白浪费了。第二条路是从视频里反推动作标注（inverse dynamics labeling），代价大、误差也大，标注质量直接制约了策略质量。第三条路是联合训练视频和动作，PAD 走的就是这条，但它给所有模态共用同一个扩散时间步，训练时图像噪声和动作噪声永远绑定在一起，推理时想只输出动作、不预测图像都做不到。始终没有一个方案能同时利用两种数据，又让策略和世界模型真正共享知识。
UWM 的切入点很巧妙：扩散模型的核心就是"在不同噪声水平之间穿梭"。如果给每种模态一个独立的噪声水平控制，那"模态缺失"就等价于"把它的噪声设到最大"。不需要特殊处理，不需要 mask，不需要改架构，噪声本身就是最自然的"我不知道"。
二、现存问题
- 标注数据稀缺：遥操作采集每条轨迹需要人工操控数分钟，100 条就是几小时的人力，扩展到数千条需要团队级投入。
- 视频数据有物理知识但无动作：YouTube 上无穷无尽的操作视频只有像素序列，没有任何机器人动作标注。如何利用这些数据是 open problem。
- PAD 的耦合时间步：PAD 给所有模态共用一个扩散时间步 $t$，训练时图像和动作的噪声水平永远绑定，推理时就无法"只去噪动作、忽略图像"，没法灵活切换推理模式。
- 策略和世界模型割裂：传统方案里策略 $p(a|o)$ 和世界模型 $p(o'|o,a)$ 是两个独立的模型，不共享特征提取，也无法互相增强。
- PAD 的条件注入方式：PAD 用 channel-wise 拼接做条件注入（把观测像素直接和噪声拼在一起），在复杂真实场景下效果很差。UWM 的实验表明 PAD 在多数真实世界任务上成功率接近 0%。
四、方法详解
Figure 2：UWM 训练与推理。左：联合训练，独立采样两个时间步。中：策略推理（marginalize 观测）。右：逆动力学推理（condition on 观测）。
4.1 核心设计：独立时间步
UWM 的噪声预测网络接收五个输入：
$$s_\theta(o, a_{t_a}, o'_{t_{o'}}, t_a, t_{o'})$$关键在于 $t_a$ 和 $t_{o'}$ 是独立采样的。训练时模型见过所有可能的 $(t_a, t_{o'})$ 组合，从 (0, 0) 到 (T, T)。推理时只需要固定其中一个到边界值，就能得到不同的模型：
4.2 训练目标
$$\ell(\theta) = \mathbb{E}_{(o,a,o') \sim \mathcal{D},\; t_a, t_{o'} \sim \mathcal{U}(0,T),\; \epsilon_a, \epsilon_{o'} \sim \mathcal{N}(0,I)} \left[ w_a \|\epsilon_a^\theta - \epsilon_a\|_2^2 + w_{o'} \|\epsilon_{o'}^\theta - \epsilon_{o'}\|_2^2 \right]$$4.3 消化无标注视频
对于没有动作标注的视频数据，直接设 $t_a = T$，把动作换成纯噪声 $\epsilon_a \sim \mathcal{N}(0, I)$。此时训练 loss 自然退化成只有观测侧的项，模型只学"世界怎么变化"，不需要知道"是什么动作导致了变化"。整个过程不需要特殊的 mask 或数据处理流程。
4.4 架构细节
Figure 3：UWM DiT 架构。底部输入：当前观测 $o$、带噪动作 $a_{t_a}$、带噪下一帧 $o'_{t_{o'}}$、register tokens。两个时间步 $\epsilon_{t_a}$, $\epsilon_{t_{o'}}$ 通过 AdaLN 注入。
- 图像编码：全分辨率 (224, 224, 3) 通过冻结的 SDXL VAE 压成 (28, 28, 4) latent，再用 (4, 4, 2) 的 patch 切成 transformer token。多相机视角（$n_c$ 个）的帧堆叠处理。
- 动作编码：action chunk（长度 $h_a$）的每个时间步通过浅层 MLP 编码成 token。
- Register token：额外加了一组随机初始化的 learnable token（推理时丢弃），充当跨模态信息交换的"缓冲区"。消融显示去掉后性能下降 6%。
- 时间步注入：$t_a$ 和 $t_{o'}$ 各自通过正弦编码，然后拼接后通过 AdaLN（Adaptive Layer Normalization）注入 transformer 每一层。AdaLN 比 concatenation 好 15%。
五、实验结果
5.1 真实机器人（DROID 数据集）
2000 条 DROID 轨迹预训练，另外 2000 条无标注视频用于协同训练。5 个真实世界任务，在 ID（同分布）和 OOD（新环境/新物体）下评估。
关键发现：UWM 预训练版就比 DP 平均高出约 20 个百分点；加上 2000 条无标注视频协同训练后（UWM-Co），在所有 OOD 场景下再提升 4-12 个百分点。而 PAD 在多数真实任务上接近 0%，channel 拼接的条件注入在复杂真实场景下完全崩溃。
5.2 LIBERO 仿真
LIBERO-90 上预训练，LIBERO-10 上 OOD 测试。UWM 平均 0.79（±0.11），DP 为 0.71，PAD 为 0.57，GR-1 为 0.58。仿真下提升幅度比真实世界小，可能是因为仿真的视觉条件更简单。
5.3 消融实验：每个组件值多少分
最关键的发现：独立时间步贡献了 12% 的性能提升，超过观测预测（8%）和 register token（6%），仅次于 AdaLN（15%）。这说明独立时间步不只是工程上的便利，它从根本上让模型学到了更好的联合表征。
5.4 逆动力学模式
在 LIBERO 轨迹跟踪任务上，逆动力学模式（给定真实的 $o'$，反推动作）比策略模式更准：Book-Caddy 0.73 vs 0.37，Soup-Cheese 0.45 vs 0.33。这证明模型确实在不同推理模式下表现出不同能力，四种模式不是噱头。
六、总结
UWM 用一个极简的设计（独立时间步 + AdaLN）统一了策略、世界模型、逆动力学和正向动力学，证明了同一个模型的四种推理模式之间的知识是共享增强的。协同训练无标注视频的能力让 UWM 在数据利用效率上远超只用标注数据的方法。
留下的问题：仿真场景下提升幅度有限，$w_a$ 和 $w_{o'}$ 权重需要针对任务调参，计算开销也没有详细分析。
七、Insight
在扩散模型里，"缺失模态"和"设时间步为 T"是同一件事。
这不是工程 trick，而是关于扩散模型表达能力的一个深层洞察。当 $t = T$ 时信号完全被噪声覆盖，等效于"这个模态不存在"。一旦接受这个等价关系，模型设计的自由度就打开了：你可以把任意数量的模态塞进同一个扩散过程，只需给每个模态一个独立的时间步控制，缺什么就把什么设成 $T$。这个设计模式后来被 DreamZero（Flash 调度）、X-WAM（ANS）、Motus（UniDiffuser 调度）等后续工作广泛采用和变体。
八、关键引用
[Section III-B] "Simply setting combinations of tₐ and t_{o'} allows for flexible inference of policies, dynamics models, inverse models, and video prediction from the same model."
只需要设置 $t_a$ 和 $t_{o'}$ 的组合，就能从同一个模型灵活推理出策略、动力学模型、逆模型和视频预测。
[Section III-D] "For video-only samples, action timestep is fixed at tₐ = T with actions imputed as random noise ϵₐ ~ N(0,I), enabling the same loss computation."
对纯视频数据，把动作时间步固定为 $T$、动作用纯噪声填充，就能用同一个 loss 训练。不需要任何特殊处理。
[Section IV-D5] "Fixed joint timestep [coupled diffusion] reduces performance by ~12%, validating decoupling necessity."
用固定的耦合时间步（PAD 的方式）性能下降约 12%，证明了时间步解耦是必要的而非可选的。
[Section IV-D1] "When trained from scratch without pretraining, UWM and DP perform comparably. However, UWM scales more effectively with pretraining."
不做预训练时 UWM 和 DP 差不多。但 UWM 在利用大规模预训练数据时拉开了差距，统一架构的优势在规模上才显现。
九、Q&A
时间步是否独立。PAD 给所有模态共用一个扩散时间步 $t$，训练时图像和动作的噪声水平永远相同。UWM 给每个模态独立的 $t_a$ 和 $t_{o'}$，训练时随机采样所有组合。看似微小的改动却带来了 12% 的性能差异，因为独立时间步让模型学到了更灵活的条件分布表示。
$t_{o'} = T$ 意味着下一帧观测完全被高斯噪声覆盖，对模型来说等价于"我对未来画面一无所知"。此时模型只能依赖当前观测 $o$ 来去噪动作 $a$。训练时模型见过了 $t_{o'} = T$ 的情况（因为时间步是均匀采样的，$T$ 是合法值），所以推理时这个设定是自洽的。本质上利用了扩散模型"噪声最大 = 无信息"的数学性质。
Concatenation（PAD 的方式）把观测像素直接拼到噪声 latent 的 channel 维度上，这改变了输入的分布，破坏了预训练权重学到的特征结构。AdaLN 通过调节每一层的 scale 和 shift 参数来注入条件信息，不改变输入形状，预训练的 attention 模式完整保留。这在真实世界的复杂视觉条件下差异巨大。
Register token 是一组额外的 learnable token，在 self-attention 中参与计算但推理时丢弃。作者认为它充当了跨模态信息交换的"缓冲区"：动作 token 和图像 latent 维度与语义差异很大，直接做 attention 可能不够高效，register token 提供了一个中间地带，类似 ViT 中 register token 消除 attention sink 的作用。去掉后性能下降 6%。
视频数据包含了各种环境变化（光照、背景、物体外观、相机角度），模型通过预测这些变化学到了对视觉干扰鲁棒的表征。Table IV 的分类 OOD 实验最清楚：Stack-Bowls 在不同光照 (L1/L2)、背景 (B1/B2)、杂物 (C1/C2) 条件下，UWM-Co 达到 21/30，UWM-Pre 只有 15/30，DP 只有 12/30。2000 条额外视频的边际收益极高。
UWM 论文指出 PAD 使用原始像素级的 channel-wise 拼接做条件注入，在真实世界的高方差视觉条件下完全崩溃。仿真环境视觉条件规范，PAD 还能工作（LIBERO 上 0.57）；但到了真实世界，光照变化、背景杂乱，channel 拼接的特征就被干扰噪声淹没了。
这说明 UWM 的优势来自"统一架构在大规模数据上的 scaling 效率"，而非架构本身的即时优势。数据少时策略学习本身就是瓶颈，世界模型目标帮不上忙；数据多时世界模型目标提供了额外的监督信号，让共享特征提取器学到更通用的表征。和"预训练越大 fine-tune 越强"是同一个道理。
论文没有报告具体延迟。不过策略模式下（$t_{o'} = T$）只需要去噪动作，不需要预测图像，理论上比 PAD 快（PAD 必须同时 denoise 图像和动作）。参数量约 300M，在现代 GPU 上 DDPM 采样应该可以达到 5-10Hz 量级，但论文没提这个细节，可能没做专门优化。