Go back

Unified World Models

Edit page

论文精读：Unified World Models

论文： Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets
作者： Chuning Zhu 等，University of Washington 与 Toyota Research Institute
发表： RSS 2025；项目主页还标注其获得 ICML 2025 Building Physically Plausible World Models Workshop Best Paper。(Weird Lab)

一、这篇论文到底做了什么？

一句话概括：

UWM 在一个 Diffusion Transformer 中同时对“机器人动作”和“未来图像”做扩散，并为两种模态设置独立的 diffusion timestep，从而让同一个模型既能当策略，也能当前向动力学、逆动力学和视频预测模型。

传统机器人学习通常把下面几个模型分开训练：

Policy:p(a∣o)\text{Policy:}\quad p(a\mid o) Forward Dynamics:p(o′∣o,a)\text{Forward Dynamics:}\quad p(o'\mid o,a) Inverse Dynamics:p(a∣o,o′)\text{Inverse Dynamics:}\quad p(a\mid o,o') Video Prediction:p(o′∣o)\text{Video Prediction:}\quad p(o'\mid o)

其中：

UWM 的核心观点是：这四个分布并非四个独立问题，它们都可以从同一个联合条件分布中得到：

p(a,o′∣o).p(a,o'\mid o).

因此作者不再分别训练四个模型，而是训练一个统一的多模态扩散模型。(arXiv)

二、为什么只做 Behavior Cloning 不够？

普通 Diffusion Policy 学的是：

a∼p(a∣o).a\sim p(a\mid o).

它只需要回答：

“看到当前图像，专家会做什么？”

但一条机器人轨迹中还包含非常重要的信息：

“做了这个动作以后，环境发生了什么变化？”

也就是：

o′∼p(o′∣o,a).o'\sim p(o'\mid o,a).

如果模型只预测动作，那么轨迹中的时间动力学信息没有被显式利用。作者认为，这也是普通模仿学习在大规模、多任务数据上预训练效率不够高的原因之一。

与此同时，互联网上有大量视频，或者机器人运行过程中有大量没有动作标注的视频。普通策略无法直接利用这些数据，因为它们缺少 aa；但视频预测模型仍然可以从 (o,o′)(o,o') 中学习环境变化。UWM 正是希望同时利用：

Drobot(oi,ai,oi′)i=1N\mathcal D_{\text{robot}} {(o_i,a_i,o_i')}_{i=1}^N

和

Dvideo(oi,oi′)i=1M.\mathcal D_{\text{video}} {(o_i,o_i')}_{i=1}^M.

(arXiv)

三、基础：扩散模型在这里做什么？

对于任意变量 xx，标准前向扩散过程为：

xtαˉtx0+1−αˉtϵ,ϵ∼N(0,I).x_t \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon, \qquad \epsilon\sim\mathcal N(0,I).

当 t=0t=0 时：

xt≈x0,x_t\approx x_0,

也就是变量基本没有被破坏。

当 t=Tt=T 时：

xT≈ϵ∼N(0,I),x_T\approx \epsilon\sim\mathcal N(0,I),

原始变量的信息几乎完全消失。

模型学习预测加入的噪声：

ϵθ(xt,t)≈ϵ.\epsilon_\theta(x_t,t)\approx\epsilon.

对应的训练目标为：

LdiffE[∣ϵθ(xt,t)−ϵ∣22].\mathcal L_{\text{diff}} \mathbb E \left[ \left| \epsilon_\theta(x_t,t)-\epsilon \right|_2^2 \right].

推理时从纯高斯噪声 xTx_T 开始，不断去噪，最后得到数据样本 x0x_0。

四、UWM 最核心的设计：两个独立 timestep

UWM 同时扩散动作 aa 和未来观测 o′o'：

ataαˉtaa+1−αˉtaϵa,a_{t_a} \sqrt{\bar\alpha_{t_a}}a + \sqrt{1-\bar\alpha_{t_a}}\epsilon_a, oto′′=αˉto′ o′+1−αˉto′ ϵo′.o'_{t_{o'}} = \sqrt{\bar{\alpha}_{t_{o'}}}\,o' + \sqrt{1-\bar{\alpha}_{t_{o'}}}\,\epsilon_{o'}.

关键是：

ta≠to′t_a\neq t_{o'}

而且二者在训练时独立采样：

ta∼U(0,T),to′∼U(0,T).t_a\sim U(0,T), \qquad t_{o'}\sim U(0,T).

模型输入为：

(o,,ata,,oto′′,,ta,,to′),\left( o,, a_{t_a},, o'_{t_{o'}},, t_a,, t_{o'} \right),

同时预测动作噪声和图像噪声：

(ϵ^aθ,ϵ^o′θ)sθ(o,ata,oto′′,ta,to′).\left( \hat\epsilon_a^\theta, \hat\epsilon_{o'}^\theta \right) s_\theta \left( o,a_{t_a},o'_{t_{o'}},t_a,t_{o'} \right).

最终损失是：

L(θ)E[wa∣ϵ^aθ−ϵa∣22+wo′∣ϵ^o′θ−ϵo′∣22].\mathcal L(\theta) \mathbb E \left[ w_a \left| \hat\epsilon_a^\theta-\epsilon_a \right|_2^2 + w_{o'} \left| \hat\epsilon_{o'}^\theta-\epsilon_{o'} \right|_2^2 \right].

论文中默认：

wa=wo′=1.w_a=w_{o'}=1.

由于 tat_a 和 to′t_{o'} 独立采样，模型会看到完整的二维噪声组合：

(ta,to′)∈0,…,T2.(t_a,t_{o'}) \in {0,\ldots,T}^2.

而不是只看到共享 timestep 下的一条对角线：

(t,t).(t,t).

(arXiv)

五、为什么 timestep 可以充当“模态 Mask”？

这是全文最需要理解的地方。

1. t=0t=0：把该变量作为条件

当：

ta=0t_a=0

时：

ata=a,a_{t_a}=a,

动作是干净的，因此模型可以把动作作为已知条件。

同理，当：

to′=0t_{o'}=0

时，未来观测 o′o' 是干净的，可以作为已知条件。

2. t=Tt=T：近似把该变量边缘化掉

当：

to′=Tt_{o'}=T

时：

oT′≈N(0,I).o'_T\approx\mathcal N(0,I).

此时输入的 oT′o'_T 几乎不再包含具体未来状态的信息，因此在预测动作时：

E[ϵa∣o,ata,oT′]≈E[ϵa∣o,ata].\mathbb E[ \epsilon_a \mid o,a_{t_a},o'_T ] \approx \mathbb E[ \epsilon_a \mid o,a_{t_a} ].

这相当于把 o′o' 积分掉：

p(a∣o)∫p(a,o′∣o),do′.p(a\mid o) \int p(a,o'\mid o),do'.

因此：

将某一模态的 timestep 固定为 TT，相当于把该模态完全加噪并“遮住”，进而近似实现边缘化。

严格来说这只是近似成立，因为实际扩散过程中的 αˉT\bar\alpha_T 不一定严格等于零。但只要最终噪声足够接近标准高斯，这个近似通常成立。

(arXiv)

六、同一个模型如何实现四种推理？

推理模式timestep 设置已知条件去噪对象对应分布
Policyto′=Tt_{o'}=T当前观测 oo动作 aap(a∣o)p(a\mid o)
Video Predictionta=Tt_a=T当前观测 oo未来观测 o′o'p(o′∣o)p(o'\mid o)
Forward Dynamicsta=0t_a=0当前观测 oo、动作 aa未来观测 o′o'p(o′∣o,a)p(o'\mid o,a)
Inverse Dynamicsto′=0t_{o'}=0当前观测 oo、目标未来观测 o′o'动作 aap(a∣o,o′)p(a\mid o,o')

(arXiv)

1. Policy

把未来观测完全加噪：

to′=T,oT′∼N(0,I),t_{o'}=T, \qquad o'_T\sim\mathcal N(0,I),

然后只对动作进行反向扩散：

aT→aT−1→⋯→a0.a_T\rightarrow a_{T-1}\rightarrow\cdots\rightarrow a_0.

得到：

a0∼p(a∣o).a_0\sim p(a\mid o).

此时它就是普通机器人策略。

2. Video Prediction

把动作完全加噪：

ta=T,aT∼N(0,I),t_a=T, \qquad a_T\sim\mathcal N(0,I),

然后对未来图像去噪：

oT′→oT−1′→⋯→o0′.o'_T\rightarrow o'_{T-1}\rightarrow\cdots\rightarrow o'_0.

得到：

o0′∼p(o′∣o).o'_0\sim p(o'\mid o).

这允许 UWM 使用没有动作标注的视频。

3. Forward Dynamics

给定干净动作：

ta=0,ata=a,t_a=0, \qquad a_{t_a}=a,

然后对未来图像去噪：

o′∼p(o′∣o,a).o'\sim p(o'\mid o,a).

它回答：

“如果执行这一段动作，环境接下来会变成什么样？”

4. Inverse Dynamics

给定干净的目标未来观测：

to′=0,oto′′=o′,t_{o'}=0, \qquad o'_{t_{o'}}=o',

然后对动作去噪：

a∼p(a∣o,o′).a\sim p(a\mid o,o').

它回答：

“想从当前状态到达这个未来状态，应该执行什么动作？”

论文正文在描述逆动力学时有一句“对 next observations 做 reverse diffusion”，但从公式和整个推理过程看，实际去噪的是动作，这应当是文字笔误。

七、为什么独立 timestep 比 PAD 的共享 timestep 更重要？

PAD 同样联合预测动作和未来图像，但它对两种模态使用共享 timestep：

ta=to′=t.t_a=t_{o'}=t.

这样训练时只能看到：

(0,0),(1,1),…,(T,T).(0,0),(1,1),\ldots,(T,T).

却看不到两个非常关键的组合：

(ta,to′)=(0,T),(t_a,t_{o'})=(0,T),

以及：

(ta,to′)=(T,0).(t_a,t_{o'})=(T,0).

然而：

UWM 使用独立 timestep 后，训练会覆盖整个二维平面：

(0,0)(0,1)⋯(0,T) (1,0)(1,1)⋯(1,T) ⋮⋮⋮ (T,0)(T,1)⋯(T,T).\begin{matrix} (0,0)&(0,1)&\cdots&(0,T)\ (1,0)&(1,1)&\cdots&(1,T)\ \vdots&\vdots&&\vdots\ (T,0)&(T,1)&\cdots&(T,T) \end{matrix}.

所以它不只是联合生成 p(a,o′∣o)p(a,o'\mid o)，还专门学习了各种条件分布和边缘分布。作者将其描述为学习动作与未来观测之间的“causal relationship”。(arXiv)

不过，“causal”这个用词需要谨慎。模型学习到的是：

p(o′∣o,a)p(o'\mid o,a)

和

p(a∣o,o′),p(a\mid o,o'),

它确实比单纯的 p(a∣o)p(a\mid o) 更了解动作和状态变化之间的结构关系，但仅凭专家轨迹上的观察数据，并不能严格保证识别出了因果干预分布：

p(o′∣o,do⁡(a)).p(o'\mid o,\operatorname{do}(a)).

因此更准确的说法是：

UWM 学到了 action-conditioned predictive dynamics，而不是完成了严格的因果识别。

八、模型架构

UWM 是一个多模态 Diffusion Transformer。

1. 当前观测 oo

当前图像不是扩散对象，而是条件信息。

每个相机画面先经过 ResNet-18：

o⟶ResNet-18ho.o \stackrel{\text{ResNet-18}}{\longrightarrow} h_o.

然后将：

ho,e(ta),e(to′)h_o,\quad e(t_a),\quad e(t_{o'})

拼接起来，通过 AdaLN 调制每个 Transformer block。

2. 未来观测 o′o'

未来图像先经过冻结的 SDXL VAE：

o′⟶SDXL VAEzo′.o' \stackrel{\text{SDXL VAE}}{\longrightarrow} z_{o'}.

原始图像尺寸为：

224×224×3,224\times224\times3,

latent 尺寸为：

28×28×4.28\times28\times4.

随后使用时空 patch：

4×4×24\times4\times2

将未来图像 latent 转成 token。

3. 动作

动作序列长度为：

Ha=16.H_a=16.

每个动作通过浅层 MLP 映射成 action token。

模型推理出 16 个动作，但部署时只执行前 8 个，然后重新观察并规划：

a1,…,a8⏟执行,a9,…,a16⏟丢弃,\underbrace{a_1,\ldots,a_8}_{\text{执行}} ,\quad \underbrace{a_9,\ldots,a_{16}}_{\text{丢弃}},

这仍然属于 receding-horizon control。

4. Register tokens

动作 token 和图像 patch token 差异很大。作者加入 8 个不对应任何输出的 register tokens：

[r1,…,r8].[r_1,\ldots,r_8].

Transformer 输入为：

[action tokens;image tokens;register tokens].[ \text{action tokens}; \text{image tokens}; \text{register tokens} ].

最终 register token 被丢弃，但在中间层可以充当动作和图像之间交换信息的共享工作区。

模型主体配置为：

(arXiv)

消融实验也支持 register 的作用：

模型Book-CaddySoup-Cheese
8 registers0.880.90
4 registers0.830.86
无 registers0.810.85
Cross-attention conditioning0.780.86

(arXiv)

九、没有动作标注的视频怎么训练？

对于机器人数据：

(o,a,o′),(o,a,o'),

正常独立采样：

ta,to′∼U(0,T).t_a,t_{o'}\sim U(0,T).

对于无动作视频：

(o,o′),(o,o'),

作者手动设置：

ta=T,t_a=T,

并使用随机高斯噪声填充缺失动作：

aT∼N(0,I).a_T\sim\mathcal N(0,I).

因此图像分支学习：

p(o′∣o),p(o'\mid o),

共享 Transformer 则可以从视频中学习物体、环境以及时间变化模式。(arXiv)

这里有一个值得注意的实现细节：论文称 action loss 仍然会在包含视频样本的整个 batch 上计算。由于视频样本中的动作是随机噪声，而且 ta=Tt_a=T，这一部分 action loss 更像是在重建输入噪声，并没有提供真实的动作语义监督。真正有用的信号主要来自未来图像去噪和共享表示。

十、实验结果

1. 真实机器人

作者使用 DROID 风格的 Franka Panda 机器人平台，预训练 100K steps，然后在五个任务上微调：

UWM 在五个任务的 ID 和 OOD 设置中都优于 DP、PAD 和 GR1。例如：

任务UWM 预训练UWM 加视频DP
Stack-Bowls ID0.860.920.48
Stack-Bowls OOD0.760.840.36
Block-Cabinet OOD0.600.720.26
Hang-Towel OOD0.640.760.28
Rice-Cooker ID0.600.650.35

视频共同训练几乎总能进一步提高 UWM，而 GR1 加入视频后在部分任务上反而下降。(arXiv)

2. LIBERO 仿真

在 LIBERO-90 上预训练，然后在 5 个 LIBERO-10 任务上微调并进行分布外测试：

方法平均成功率
UWM0.79 ± 0.11
DP0.71 ± 0.12
GR10.58 ± 0.14
PAD0.57 ± 0.19

(arXiv)

提升不是特别巨大，但比较稳定。

3. 分类 OOD 测试

作者针对光照、背景、大小物体干扰等六种 OOD 条件进行测试。

总成功次数：

任务UWM + 视频UWM 仅机器人预训练DP
Stack-Bowls21/3015/3012/30
Block-Cabinet15/308/306/30

这组结果比普通 ID 测试更有说服力，因为它显示视频预测辅助目标主要改善的是视觉鲁棒性，而不仅仅是拟合训练动作。(arXiv)

4. Forward Dynamics 与 Inverse Dynamics

作者展示了 UWM 可以根据动作预测未来机器人和物体位置。

在 inverse dynamics 轨迹跟踪中，给定与专家轨迹相同的时间限制：

方法Book-CaddySoup-Cheese
普通 Policy0.370.33
Inverse Dynamics0.730.45

这说明当明确提供目标未来图像 o′o' 时：

p(a∣o,o′)p(a\mid o,o')

比没有目标条件的：

p(a∣o)p(a\mid o)

更容易紧密跟踪参考轨迹。(arXiv)

十一、最关键的消融：提升来自“未来动力学”还是图像重建？

作者比较了三种预训练目标：

目标Stack-BowlsBlock-Cabinet
重建未来观测 o′o'0.860.76
重建当前观测 oo0.700.66
DP，不重建图像0.480.60

如果只是重建当前图像，模型也能学习更好的视觉特征；但预测未来图像效果更好。

这说明增益不完全来自额外的视觉 reconstruction loss，而部分来自学习：

(o,a)→o′(o,a)\rightarrow o'

所包含的时间动力学。(arXiv)

此外，从零训练时 UWM 与 DP 表现相近，但经过大规模预训练后 UWM 提升更明显。这意味着它最核心的价值不是成为更强的单任务策略，而是：

在多任务、大数据预训练阶段，更充分地利用轨迹中的动态监督。

(arXiv)

十二、这篇论文最强的地方

1. 统一方式非常简洁

它没有设计四套 decoder，也没有专门增加复杂的 task token，而只是引入：

ta,to′t_a,\quad t_{o'}

两个独立 timestep，就统一了四个分布。

2. 同一份机器人数据得到额外监督

普通 DP 从一条轨迹中只学习：

o→a.o\rightarrow a.

UWM 同时学习：

o→a,o\rightarrow a, (o,a)→o′,(o,a)\rightarrow o', (o,o′)→a,(o,o')\rightarrow a, o→o′.o\rightarrow o'.

即使没有额外数据，监督密度也更高。

3. 无动作视频的接入方式自然

视频样本不用预测 latent action，也不需要手工估计 optical flow 或关键点，只需要把缺失动作设置为完全噪声。

4. 方法特别适合作为预训练框架

实验表明，它从零训练时未必显著强于 DP，但随着预训练数据增多，优势逐渐出现。这与机器人 foundation model 的数据扩展方向相符。(arXiv)

十三、局限和需要谨慎的地方

1. 它还不是一个真正用于规划的 world model

虽然 UWM 能预测：

p(o′∣o,a),p(o'\mid o,a),

但部署时主要仍然直接调用 policy mode：

p(a∣o).p(a\mid o).

模型不会生成多个候选动作、模拟多个未来、比较奖励，再选择最优动作。

因此它更接近：

world-model-regularized diffusion policy

或者：

使用动力学辅助预训练的统一策略模型

而不是完整的 model-based planning system。

2. “因果理解”的证据不充分

独立 timestep 确实让模型学习更多条件组合，但论文没有进行明确的干预实验、反事实实验或因果变量解耦。

所以它证明的是：

但还没有严格证明模型获得了可识别的 causal representation。

3. PAD 对比存在一定混杂因素

论文的 PAD baseline 不仅使用共享 timestep，还采用了不同的当前图像 conditioning 方式；作者将干净图像与 noisy future observation 沿通道拼接，而 UWM 使用 ResNet + AdaLN。(arXiv)

因此 UWM 对 PAD 的巨大提升可能同时来自：

论文没有给出一个完全控制其他变量、只比较：

ta=to′t_a=t_{o'}

与

ta⊥to′t_a\perp t_{o'}

的核心消融。这是实验中比较明显的缺口。

4. Internet video 的提升仍然有限

互联网视频共同训练结果为：

数据Stack-BowlsBlock-Cabinet
Robot data + robot videos0.920.84
Robot data + internet videos0.880.80
Robot data only0.860.76

互联网视频有正向作用，但明显弱于同域机器人视频。这说明方法解决了“没有动作标注”的问题，却尚未真正解决：

(arXiv)

5. 未来预测仍然较短且可能有伪影

作者明确承认，forward dynamics 生成结果有时包含视觉伪影，可能影响后续规划；模型也尚未充分利用大规模人类视频，未来需要更密集的视频预测。(arXiv)

十四、放在 World Action Model 脉络中的定位

从你的 WAM 阅读方向看，UWM 很重要，因为它体现了一个关键转变：

Policy only⟶Joint action-world modeling.\text{Policy only} \quad\longrightarrow\quad \text{Joint action-world modeling}.

但它仍然属于较早期、较轻量的 World Action Model：

维度UWM
动作预测有
未来视觉预测有
动作条件世界模型有
无动作视频预训练有
显式语言指令基本没有
多步 imagination planning没有
奖励或 value model没有
通过预测未来选择动作没有
长期记忆或 latent memory没有
大规模 human-video embodiment transfer很有限

因此 UWM 的核心贡献不是：

“让机器人先想象未来，再决定动作。”

而是：

“用一个统一生成模型学习动作与未来，并把 world modeling 当作大规模策略预训练的额外监督。”

十五、最终评价

这篇论文最大的创新可以浓缩成一句话：

Independent modality diffusion timestepscontrollable conditioning and marginalization\boxed{ \text{Independent modality diffusion timesteps} \text{controllable conditioning and marginalization} }

它非常巧妙地把：

p(a∣o),p(o′∣o),p(o′∣o,a),p(a∣o,o′)p(a\mid o),\quad p(o'\mid o),\quad p(o'\mid o,a),\quad p(a\mid o,o')

统一在一个 Diffusion Transformer 中。

从方法设计看，它是一个很漂亮、很适合扩展的框架；从实验看，它较充分证明了“未来动力学预测有利于机器人策略预训练”和“无动作视频有助于 OOD 鲁棒性”。但它尚未证明 world model 可以真正用于在线规划，其因果理解表述也略强于实验支持。

我对它的定位是：

UWM 是从 Diffusion Policy 向 World Action Model 过渡的一篇关键工作。其核心价值在于统一预训练，而非显式 imagination-based planning。

Edit page
Share this post on:
Share this post via WhatsAppShare this post on FacebookShare this post on XShare this post via TelegramShare this post on PinterestShare this post via email
Previous Post
Unified Video Action Model（UVA）
Next Post
Video Prediction Policy - A Generalist Robot Policy with Predictive Visual Representations