Unified World Models: 视频与动作扩散联合预训练用于大规模机器人数据集Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets
在单一 Transformer 中耦合 action diffusion 与 video diffusion，解锁四种推断模式Coupling action diffusion and video diffusion in a single Transformer to unlock four inference modes
Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · University of Washington & Toyota Research Institute
UWM（Unified World Models）将 action diffusion 与 video diffusion 融合在同一个 Transformer 架构中，通过对每个模态独立采样扩散时间步来决定该模态是"生成"还是"条件化"，从而用一个模型同时实现策略、前向动力学、逆向动力学与视频预测四种推断模式。在真实机器人操作与 LIBERO 仿真基准上，UWM 预训练后微调的成功率显著超越 Diffusion Policy，对视觉扰动的鲁棒性提升尤为突出。UWM (Unified World Models) fuses action diffusion and video diffusion within a single Transformer architecture; by independently sampling a diffusion timestep for each modality, it decides whether that modality is "generated" or "conditioned on", so that a single model realizes four inference modes at once — policy, forward dynamics, inverse dynamics and video prediction. On real-robot manipulation and on the LIBERO simulation benchmark, UWM fine-tuned after pretraining attains success rates that clearly surpass Diffusion Policy, with the gain in robustness to visual perturbations being especially pronounced.
机器人学习的核心挑战在于：带动作标注的数据十分稀缺，而无标注的视频却海量存在。现有方法要么专注于行为克隆（忽略了大量无动作视频），要么单独训练视频预测模型再桥接到策略，难以在同一框架中统一这两类监督信号。此外，将动力学模型与策略分开训练会导致信息割裂，错失互相强化的机会。The core challenge of robot learning is that action-annotated data is extremely scarce while unlabeled video exists in vast quantities. Existing approaches either focus on behavior cloning (ignoring the large body of action-free video), or train a video prediction model separately and then bridge it to a policy, making it hard to unify these two kinds of supervision signal within one framework. Moreover, training the dynamics model and the policy separately fragments information and forfeits the chance for the two to reinforce each other.
"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."
图 1：UWM 概览。通过对 action 与 observation 分别采样扩散时间步（t_a、t_o'），模型在训练时同时接触"生成动作"与"生成未来观测"两类目标；推断时只需冻结其中一个模态的时间步为 0 或 T，即可切换至策略、视频预测、前向动力学或逆向动力学四种模式——全程使用同一套参数。Figure 1: UWM overview. By sampling separate diffusion timesteps for action and observation (t_a, t_o'), the model is exposed during training to both objectives at once — "generate the action" and "generate the future observation"; at inference one need only freeze the timestep of one modality at 0 or T to switch among the four modes of policy, video prediction, forward dynamics and inverse dynamics — all with the same set of parameters.
+20%真实机器人 OOD 成功率（相对 Diffusion Policy 最大提升）Real-robot OOD success rate (largest gain relative to Diffusion Policy)
0.79LIBERO 仿真均值成功率（UWM），DP 为 0.71LIBERO simulation mean success rate (UWM); DP is 0.71
4 种4 types推断模式（policy / video / forward / inverse dynamics）inference modes (policy / video / forward / inverse dynamics)
2000无动作视频数据（DROID）用于 co-training 进一步提升性能action-free video data (DROID) used for co-training to further improve performance
02方法Method
UWM 在同一个扩散 Transformer 中对 action 序列与未来观测的 latent 同时去噪。训练时对每个模态独立采样噪声时间步（t_a 与 t_o'），使模型暴露于全部噪声水平组合；推断时将某模态的时间步固定为边界值，即可将其"边缘化"（设为 T）或"条件化"（设为 0），从而选择所需的推断函数。UWM denoises the action sequence and the latent of the future observation simultaneously inside a single diffusion Transformer. During training a noise timestep is sampled independently for each modality (t_a and t_o'), exposing the model to every combination of noise levels; at inference, fixing the timestep of a modality at a boundary value either "marginalizes" it (set to T) or "conditions" on it (set to 0), thereby selecting the desired inference function.
图 2：UWM 训练与推断管线。左侧：预训练阶段在带动作标注的 robot 轨迹上同时学习动作与观测扩散；co-training 阶段将无动作视频的 t_a 固定为 T（即掩蔽动作）来利用额外视频数据。右侧：通过调整 t_a 与 t_o' 的取值，同一模型可分别执行四种推断。Figure 2: UWM training and inference pipeline. Left: in the pretraining stage, action diffusion and observation diffusion are learned jointly on action-annotated robot trajectories; in the co-training stage, t_a of action-free video is fixed at T (i.e. the action is masked) so that extra video data can be exploited. Right: by adjusting the values of t_a and t_o', the same model can carry out four different inferences.
独立扩散时间步——核心创新Independent diffusion timesteps — the core innovation
传统方法（如 PAD）将动作与观测的扩散时间步耦合为同一个值，导致两种模态只能同步加噪/去噪，无法独立控制。UWM 的训练目标为：Conventional approaches (such as PAD) couple the diffusion timesteps of action and observation into one and the same value, so the two modalities can only be noised/denoised in lockstep and cannot be controlled independently. The training objective of UWM is:
通过令 t_a 与 t_o' 独立均匀采样，模型学会在任意噪声水平组合下同时去噪，从而在推断时：By letting t_a and t_o' be sampled independently and uniformly, the model learns to denoise under any combination of noise levels, so that at inference:
Policy p(a|o)：令 t_o' = T（将观测边缘化），仅对动作去噪。Policy p(a|o): set t_o' = T (marginalizing the observation) and denoise the action only.
Video Prediction p(o'|o)：令 t_a = T（将动作边缘化），仅对未来帧去噪。Video Prediction p(o'|o): set t_a = T (marginalizing the action) and denoise the future frames only.
Forward Dynamics p(o'|o, a)：令 t_a = 0（以干净动作为条件），预测未来观测。Forward Dynamics p(o'|o, a): set t_a = 0 (conditioning on the clean action) and predict the future observation.
Inverse Dynamics p(a|o, o')：令 t_o' = 0（以干净未来观测为条件），推断动作。Inverse Dynamics p(a|o, o'): set t_o' = 0 (conditioning on the clean future observation) and infer the action.
图 3：UWM 单个 Transformer Block。观测编码（ResNet-18 + 冻结 SDXL VAE）与动作 token 共同输入，通过 Adaptive Layer Normalization（AdaLN）注入各自的扩散时间步。额外引入的随机初始化 register tokens 不产生直接输出，但显著改善跨模态特征共享。Figure 3: a single UWM Transformer block. The observation encoding (ResNet-18 + frozen SDXL VAE) and the action tokens are fed in together, and their respective diffusion timesteps are injected through Adaptive Layer Normalization (AdaLN). The additionally introduced, randomly initialized register tokens produce no direct output, yet markedly improve cross-modal feature sharing.
架构细节Architecture details
观测编码Observation encoding
当前帧：ResNet-18 编码为紧凑特征向量Current frame: encoded by ResNet-18 into a compact feature vector
未来帧：冻结 SDXL VAE 编码为 28×28×4 latent（224×224×3 输入）Future frame: encoded by a frozen SDXL VAE into a 28×28×4 latent (224×224×3 input)
两个时间步（t_a、t_o'）均通过 AdaLN 注入每一个 Transformer blockBoth timesteps (t_a, t_o') are injected into every Transformer block via AdaLN
Register Tokens
随机初始化，不对应任何真实 token 的预测目标Randomly initialized; they correspond to the prediction target of no real token
充当动作与观测模态之间的"中间人"，使两类特征能相互交换信息They act as an "intermediary" between the action and observation modalities, letting the two kinds of feature exchange information with each other
消融实验证明其对多模态特征共享有实质贡献Ablation studies prove they make a substantive contribution to multimodal feature sharing
预训练与 Co-trainingPretraining and co-training
预训练阶段：在 2000 条来自 DROID 数据集的带动作 Franka 机器人轨迹上训练 100K 步，覆盖多样化任务与场景以增强泛化能力。Co-training 阶段：额外混入 2000 条无动作视频轨迹，训练时将 t_a 固定为 T 以掩蔽动作监督，从而充分利用无标注视频。Pretraining stage: train for 100K steps on 2000 action-labeled Franka robot trajectories from the DROID dataset, covering diverse tasks and scenes so as to strengthen generalization. Co-training stage: additionally mix in 2000 action-free video trajectories and fix t_a at T during training to mask out action supervision, thereby making full use of unlabeled video.
03实验Experiments
实验分为两大平台：（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）。The experiments span two platforms: (1) a real Franka robot with 5 manipulation tasks (Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel, Rice-Cooker), evaluated under both in-distribution (ID) and out-of-distribution (OOD, with altered lighting / background / clutter) settings; (2) the LIBERO simulation benchmark, with 90-task pretraining + 10-task fine-tuning. Baselines for comparison: Diffusion Policy (DP), PAD (the joint-timestep version) and GR1 (a deterministic video-action Transformer).
图 5：真实机器人任务配置。5 项任务的初始状态（第一行）、成功完成（第二行）与 OOD 配置（第三行，含异常光照、背景替换与额外障碍物）。Figure 5: real-robot task configuration. Initial states of the 5 tasks (first row), successful completions (second row) and OOD configurations (third row, with unusual lighting, background replacement and extra obstacles).
图 6：各方法在真实机器人任务上的平均成功率（含 ID 与 OOD）。UWM（预训练 + co-training）在所有设置下均超越竞争基线，co-training 引入无动作视频后 OOD 性能提升尤为显著。Figure 6: mean success rate of each method on the real-robot tasks (ID and OOD included). UWM (pretraining + co-training) surpasses the competing baselines under every setting, and the OOD gain is especially marked once co-training brings in action-free video.
LIBERO 仿真基准LIBERO simulation benchmark
方法Method
平均成功率 ± 95% CIMean success rate ± 95% CI
GR1
0.58 ± 0.14
PAD
0.57 ± 0.19
Diffusion Policy
0.71 ± 0.12
UWM（本文）UWM (ours)
0.79 ± 0.11
OOD 分类实验（Stack-Bowls 为例）OOD breakdown experiment (Stack-Bowls as an example)
在 Stack-Bowls 任务的 OOD 条件下（光照 L1/L2、背景 B1/B2、杂乱物 C1/C2），UWM co-trained 版本在"背景"与"杂乱物"类 OOD 成功率达 70%，而 Diffusion Policy 基线仅为 40%（来自 Table IV）。Under the OOD conditions of the Stack-Bowls task (lighting L1/L2, background B1/B2, clutter C1/C2), the co-trained UWM reaches a success rate of 70% on the "background" and "clutter" OOD categories, whereas the Diffusion Policy baseline reaches only 40% (from Table IV).
图 9：OOD 分类设置可视化。从左至右依次为正常场景、异常光照（L1/L2）、背景替换（B1/B2）、额外杂乱物（C1/C2）。UWM 在背景与杂乱物类 OOD 上的鲁棒性改善最为突出。Figure 9: visualization of the OOD breakdown settings. From left to right: the normal scene, unusual lighting (L1/L2), background replacement (B1/B2) and extra clutter (C1/C2). UWM's robustness improvement is most pronounced on the background and clutter OOD categories.
前向动力学预测Forward dynamics prediction
UWM 的前向动力学模式（令 t_a = 0）能准确预测机器人与物体的未来姿态，验证了模型对物理动力学的内在理解。逆向动力学模式（令 t_o' = 0）在轨迹追踪任务中的成功率高于直接策略模式。The forward dynamics mode of UWM (setting t_a = 0) accurately predicts the future poses of the robot and of objects, confirming the model's intrinsic understanding of physical dynamics. The inverse dynamics mode (setting t_o' = 0) achieves a higher success rate than the direct policy mode on trajectory-tracking tasks.
图 8：前向动力学预测可视化。给定初始观测与动作序列，UWM 准确预测出机器人末端执行器和操作对象的未来状态，展示了其对环境动力学的建模能力。Figure 8: visualization of forward dynamics prediction. Given an initial observation and an action sequence, UWM accurately predicts the future states of the robot end-effector and of the manipulated object, demonstrating its ability to model environment dynamics.
消融实验Ablation studies
图 10 对比了从头训练与在 UWM 预训练模型基础上微调的效果：随着微调数据量增加，预训练版本的成功率曲线始终高于从头训练的 Diffusion Policy，表明 UWM 的预训练具有更强的可扩展性（"UWM scales more effectively with pretraining than DP"）。Register tokens 的消融同样确认了其对多模态特征共享的实质贡献。Figure 10 compares training from scratch with fine-tuning on top of the UWM pretrained model: as the amount of fine-tuning data grows, the success-rate curve of the pretrained version stays consistently above that of Diffusion Policy trained from scratch, indicating that UWM pretraining has stronger scalability ("UWM scales more effectively with pretraining than DP"). The ablation of register tokens likewise confirms their substantive contribution to multimodal feature sharing.
图 10：从头训练 vs. 预训练微调的扩展曲线。随着微调数据量增大，UWM 预训练的收益持续高于 Diffusion Policy，显示出更优的数据效率。Figure 10: scaling curves of training from scratch vs. fine-tuning from pretraining. As the amount of fine-tuning data grows, the benefit of UWM pretraining stays consistently above Diffusion Policy, showing better data efficiency.
04局限性Limitations
说明：以下局限性部分来自论文 Section VII 的明确表述，部分为根据设计推断，已作标注。Note: Some of the limitations below are explicitly stated in Section VII of the paper, others are inferred from the design; both are labeled accordingly.
仿真 OOD 提升幅度小于真实世界（论文明确指出）OOD gains in simulation are smaller than in the real world (explicitly stated in the paper)
论文指出仿真环境中 OOD 设置的性能增益小于真实机器人，"potentially due to simpler dynamics in current simulations"。现有仿真器的物理保真度有限，导致多样化视频数据的潜在价值未被充分挖掘。The paper points out that the performance gain under OOD settings is smaller in simulation than on the real robot, "potentially due to simpler dynamics in current simulations". The limited physical fidelity of existing simulators means the potential value of diverse video data is not fully exploited.
需要预先对齐的动作与视频数据（设计推断）Requires pre-aligned action and video data (inferred from the design)
预训练阶段仍依赖带动作标注的 robot 轨迹，无动作视频仅作为 co-training 补充。若动作与观测数据来自不同机器人或跨具身形态，对齐难度大，限制了对完全异构数据的利用。The pretraining stage still relies on action-annotated robot trajectories, with action-free video serving only as a co-training supplement. If the action and observation data come from different robots or across embodiments, alignment is difficult, which limits the use of fully heterogeneous data.
计算开销与推断速度（设计推断）Computational cost and inference speed (inferred from the design)
Diffusion 推断本身需要多步迭代去噪，扩散 Transformer 的参数量也较大。论文在附录中仅简短提及计算需求，未给出详细的推断延迟数据，实时控制场景下的适用性尚需验证。Diffusion inference itself requires multi-step iterative denoising, and the diffusion Transformer also has a fairly large parameter count. The paper mentions computational requirements only briefly in the appendix and gives no detailed inference-latency figures, so its applicability to real-time control scenarios still needs to be verified.
失效模式分析有限（设计推断）Limited failure-mode analysis (inferred from the design)
OOD 实验仅在光照、背景、杂乱物三类场景下系统评估，对更复杂扰动（如大幅度姿态变化、新型物体类别）的鲁棒性尚未深入分析。The OOD experiments are systematically evaluated only under the three categories of lighting, background and clutter; robustness to more complex perturbations (such as large pose changes or novel object categories) has not yet been analyzed in depth.
