Unified World Models:
视频与动作扩散联合预训练用于大规模机器人数据集

01 动机

机器人学习的核心挑战在于：带动作标注的数据十分稀缺，而无标注的视频却海量存在。现有方法要么专注于行为克隆（忽略了大量无动作视频），要么单独训练视频预测模型再桥接到策略，难以在同一框架中统一这两类监督信号。此外，将动力学模型与策略分开训练会导致信息割裂，错失互相强化的机会。

"We propose Unified World Models (UWM), a framework that integrates an action diffusion process and a video diffusion process within a unified transformer architecture, where each modality is independently controlled by its own diffusion timestep."

02 方法

UWM 在同一个扩散 Transformer 中对 action 序列与未来观测的 latent 同时去噪。训练时对每个模态独立采样噪声时间步（t_a 与 t_o'），使模型暴露于全部噪声水平组合；推断时将某模态的时间步固定为边界值，即可将其"边缘化"（设为 T）或"条件化"（设为 0），从而选择所需的推断函数。

独立扩散时间步——核心创新

传统方法（如 PAD）将动作与观测的扩散时间步耦合为同一个值，导致两种模态只能同步加噪/去噪，无法独立控制。UWM 的训练目标为：

"L(θ) = E[w_a · ||ε_a − f_θ(a_{t_a}, o'_{t_o'}, o, t_a, t_o')||² + w_o' · ||ε_o' − f_θ(a_{t_a}, o'_{t_o'}, o, t_a, t_o')||²]"

通过令 t_a 与 t_o' 独立均匀采样，模型学会在任意噪声水平组合下同时去噪，从而在推断时：

• Policy p(a|o)：令 t_o' = T（将观测边缘化），仅对动作去噪。
• Video Prediction p(o'|o)：令 t_a = T（将动作边缘化），仅对未来帧去噪。
• Forward Dynamics p(o'|o, a)：令 t_a = 0（以干净动作为条件），预测未来观测。
• Inverse Dynamics p(a|o, o')：令 t_o' = 0（以干净未来观测为条件），推断动作。

架构细节

预训练与 Co-training

预训练阶段：在 2000 条来自 DROID 数据集的带动作 Franka 机器人轨迹上训练 100K 步，覆盖多样化任务与场景以增强泛化能力。Co-training 阶段：额外混入 2000 条无动作视频轨迹，训练时将 t_a 固定为 T 以掩蔽动作监督，从而充分利用无标注视频。

03 实验

实验分为两大平台：（1）真实 Franka 机器人，5 项操作任务（Stack-Bowls、Block-Cabinet、Paper-Towel、Hang-Towel、Rice-Cooker），设置 in-distribution（ID）与 out-of-distribution（OOD，改变光照/背景/杂乱物）两类评估；（2）LIBERO 仿真基准，90 任务预训练 + 10 任务微调。对比基线：Diffusion Policy（DP）、PAD（联合时间步版本）、GR1（确定性视频-动作 Transformer）。

真实机器人——分任务成功率（ID / OOD）

LIBERO 仿真基准

OOD 分类实验（Stack-Bowls 为例）

在 Stack-Bowls 任务的 OOD 条件下（光照 L1/L2、背景 B1/B2、杂乱物 C1/C2），UWM co-trained 版本在"背景"与"杂乱物"类 OOD 成功率达 70%，而 Diffusion Policy 基线仅为 40%（来自 Table IV）。

前向动力学预测

UWM 的前向动力学模式（令 t_a = 0）能准确预测机器人与物体的未来姿态，验证了模型对物理动力学的内在理解。逆向动力学模式（令 t_o' = 0）在轨迹追踪任务中的成功率高于直接策略模式。

消融实验

图 10 对比了从头训练与在 UWM 预训练模型基础上微调的效果：随着微调数据量增加，预训练版本的成功率曲线始终高于从头训练的 Diffusion Policy，表明 UWM 的预训练具有更强的可扩展性（"UWM scales more effectively with pretraining than DP"）。Register tokens 的消融同样确认了其对多模态特征共享的实质贡献。

04 局限性