Causal World Modeling for Robot Control
让世界模型 不仅能预测视觉未来帧,还能推理对应的动作序列,并用于机器人闭环控制。
环境模型既输出状态,也输出动作
传统方法往往分开训练视觉世界模型与策略模型,而作者提出一种将视觉预测与动作生成同时纳入一个统一自回归框架的方案
1. 自回归视频-动作世界建模(Autoregressive Video--Action World Modeling)
统一序列设计(Video + Action Token)
环境模型同时推导s和a
-
模型将历史视觉 latent 与历史动作 token 交错组成一个统一序列
-
在每个 autoregressive 步骤中预测下一对未来视觉 latent 和对应动作 token
双流 Mixture‑of‑Transformers 结构
Video Stream(VS):用一个完备的VIT模型,将其中的Encoder部分作为环境模型的视觉处理部分
输入为过去时刻VS的输出z,过去时刻的a
z t + 1 : t + K ∼ p θ ( ⋅ ∣ z ≤ t , a < t ) z_{t+1:t+K} \sim p_\theta(\cdot \mid z_{\le t}, a_{<t}) zt+1:t+K∼pθ(⋅∣z≤t,a<t)
Action Stream (AS):另外训练的新模型
a t : t + K − 1 ∼ g ψ ( ⋅ ∣ z ^ t + 1 : t + K , z ≤ t , a < t ) a_{t:t+K-1} \sim g_\psi(\cdot \mid \hat z_{t+1:t+K}, z_{\le t}, a_{<t}) at:t+K−1∼gψ(⋅∣z^t+1:t+K,z≤t,a<t)
2. 世界动态:基于流匹配的扩散(World Dynamics:Flow Matching based Diffusion)
KV-cache
标准因果attention对于一个序列[B,T,N,C]中的第t个时间步:
计算第t步的Q:[B,1,N,C]
重新计算第t步的K和V:[B,[1,t],N,C]
而KVcache的目的在于不用计算t步之前的K和V,即:[B,[1,t-1],N,C]部分不用再计算
计算第t步的Q:[B,1,N,C]
重新计算第t步的K和V:[B,1,N,C]
然后把K和V拼接到Kcache和Vcache中去,维度一样来到[B,[1,t],N,C],用Kcache和Vcache和Q做因果attention
噪声增强
z ~ ≤ t = { ( 1 − s aug ) ϵ + s aug z ≤ t , p , z ≤ t , 1 − p , \tilde{z}{\le t} = \begin{cases} (1 - s{\text{aug}})\,\epsilon + s_{\text{aug}}\,z_{\le t}, p, \\ z_{\le t}, 1 - p, \end{cases} z~≤t={(1−saug)ϵ+saugz≤t,p,z≤t,1−p,
p = 0.5 , s aug ∼ Uniform ( 0.5 , 1.0 ) , ϵ ∼ N ( 0 , I ) p = 0.5,\quad s_{\text{aug}} \sim \text{Uniform}(0.5,\, 1.0),\quad \epsilon \sim \mathcal{N}(0, I) p=0.5,saug∼Uniform(0.5,1.0),ϵ∼N(0,I)
3. FDM-grounded异步
Naive异步:
基于旧的obs + 旧的世界模型输出 + 动作历史预测未来
提出的FDM-grounded异步计算:
用旧的obs + 动作历史 ,以视觉模型推测当前obs,再用旧的obs,动作历史,当前obs的推测值,推测新的动作
4. 模型的反向公式:
L fdm = E t , s , z ^ t + 1 , ϵ ∥ v ψ ( z ~ t + 1 , s , z t , a t , z ~ < t , a ^ < t ∣ c ) − z ˙ t + 1 ( s ) ∥ 2 2 L_{\text{fdm}} = \mathbb{E}{t,s,\hat z{t+1},\epsilon} \Big\| v_\psi(\tilde z_{t+1}, s, z_t, a_t, \tilde z_{<t}, \hat a_{<t} \mid c)- \dot z_{t+1}(s) \Big\|_2^2 Lfdm=Et,s,z^t+1,ϵ vψ(z~t+1,s,zt,at,z~<t,a^<t∣c)−z˙t+1(s) 22
L fdm L_{\text{fdm}} Lfdm: 前向动力学模型(Forward Dynamics Model)的损失
v ψ ( . . . ) v_\psi(...) vψ(...): FDM 中的 velocity field 网络输出
z ~ t + 1 \tilde z_{t+1} z~t+1: 由 FDM 预测的当前下一 latent
s: flow time(扩散时间坐标)
z t z_t zt: 最新真实 latent
a t a_t at: 当前动作
z ~ < t , a ^ < t \tilde z_{<t}, \hat a_{<t} z~<t,a^<t: 视觉模型输出的加噪图像,动作模型输出的动作
z ˙ t + 1 ( s ) \dot z_{t+1}(s) z˙t+1(s): 真实值