Fast LeWorldModel论文笔记

约 3527 字大约 12 分钟

2026-07-16

一句话省流

Fast-LeWorldModel 提出了一种基于 Action Prefix Prediction（动作前缀预测）的并行潜空间世界模型，用动作前缀替代传统的一步步自回归状态预测，在保持 LeWorldModel 规划框架不变的前提下，显著降低了长时预测误差和规划时间，同时提升了任务成功率

LeWM 风格的规划仍然成本高昂，它的动力学模型本质上是局部的：它只从当前的潜状态 $z_{t}$ 和动作 $a_{t}$ 预测下一个潜状态 $\hat{z}_{t+1}$ 。这造成了两个实际限制。首先，候选序列的评估变得缓慢，因为自回归展开必须逐步生成整个想象的潜变量轨迹，重复执行动作编码和潜变量预测的过程。其次，在早期或中间想象状态引入的误差可能会传播到后续预测中，使得展开轨迹随着视野（horizon）的增加而变得越来越不可靠。

Fast LeWorldModel (Fast-LeWM)，这是一个围绕动作前缀（action-prefix）预测构建的快速潜变量世界模型，预测在执行编码动作序列的前缀后所达到的潜变量，从而允许模型直接且并行地评估这些前缀结果给定当前的视觉潜变量 $z_{t}$ 和动作前缀 $a_{t:t+k-1}=(a_{t},...,a_{t+k-1})$ ，Fast-LeWM 预测执行该前缀后达到的未来潜变量：

z_{t+k}=G_{Fast-LeWM}(z_{t}, a_{t:t+k-1}), k=1,...,H

Fast-LeWM 使用了动作前缀编码器和并行潜变量预测器

我们指出了 LeWM 中的局部单步转换接口是一个关键瓶颈，它导致了自回归展开过程缓慢，以及潜变量预测误差在长视野下的累积。
我们提出了 Fast-LeWM，这是一种快速潜变量世界模型，它将潜动力学建模从单步转换重新定义为动作前缀预测。这使得密集的前缀级监督成为可能，并将累积的动作效应转化为直接的动力学训练目标。
我们在相同的协议下，对 LeWM 的所有规划任务评估了 Fast-LeWM。Fast-LeWM 将平均成功率从 85.8% 提高到了 90.5%，将动力学模块加速了 3.9 倍（从 31.4 秒缩短到 8.0 秒），将完整的 CEM 求解时间减少了 48.0%（从 54.4 秒缩短到 28.3 秒），并显著降低了开环预测误差及其在长视野下的增长速度。

方法

无奖励的潜变量世界模型 (Reward-Free Latent World Models)

考虑一个由观察-动作轨迹组成的离线、无奖励数据集 $\mathcal{D}=\{\tau^{(n)}\}_{n=1}^{N}$ ，其中 $\tau=\{(o_{t},a_{t})\}_{t=1}^{T}$ (1)。其中 $o_t$ 是像素观测， $a_t \in \mathbb{R}^{d_a}$ 是连续动作目标是学习一个潜变量世界模型，使其能够支持无奖励标签的目标条件规划。

LeWorldModel 中的自回归展开 (Autoregressive Rollout in LeWorldModel)

更详细可以查看 LeWorldModel

常见的潜变量世界模型设计（包括 LeWM 风格的动力学）是学习一个局部的状态转移预测器，该预测器使用预测下一个潜变量的损失进行训练： $\hat{z}_{t+1}=F_{\phi}(z_{t},a_{t})$ , (3)， $\mathcal{L}_{pred}^{1step}=\vert{}\vert{}\hat{z}_{t+1}-z_{t+1}\vert{}\vert{}_{2}^{2}$ (4)。遵循无重建的 JEPA 风格世界模型，在一批潜嵌入 Z 上添加如 SIGReg 的抗崩塌正则化项： $\mathcal{L}_{1step}=\mathcal{L}_{pred}^{1step}+\lambda SIGReg(Z)$ (5)。给定候选序列 $a_{t:t+H-1}$ ，单步模型通过自回归展开来估计终端潜变量： $\hat{z}_{t+k}=F_{\phi}(\hat{z}_{t+k-1},a_{t+k-1})$ ， $k=1,...,H$ (6)。并通过以下方式对候选序列进行打分： $C_{AR}(a_{t:t+H-1})=\vert{}\vert{}\hat{z}_{t+H}-z_{g}\vert{}\vert{}_{2}^{2}$ (7)。

主要缺点是在评估候选序列时需要较长的自回归链。

快速 LeWorldModel (Fast LeWorldModel)

Fast-LeWM 通过减少预测的未来状态之间的顺序依赖性来解决上述限制，直接从观察到的锚点潜变量（anchor latent）和导致该状态的动作前缀预测出每一个未来状态。设计动作前缀编码器以及并行潜变量预测器

图2：Fast-LeWM训练流程。 当前观测 (o_t) 被编码为隐状态 (z_t)。动作前缀编码器（Action-Prefix Encoder）接收一个状态-动作令牌序列，该序列由将 (z_t) 映射得到的状态令牌作为第0个令牌，后接动作令牌 ((a_{t}, ..., a_{t+H-1})) 构成。在因果掩码机制下，动作位置的输出形成前缀令牌 ((p_{t, 1}, ..., p_{t, H}))，其中 (p_{t, k}) 对应前缀 ((a_{t}, ..., a_{t+k-1}))，而第0个输出令牌被丢弃。并行预测器（Parallel Predictor）利用前缀令牌和锚定隐状态，在单次前向传播中直接生成所有未来隐状态 ((\hat{z}{t+1}, ..., \hat{z}))，并通过密集多视野损失（dense multi-horizon loss）以编码后的未来观测为监督进行训练。与孤立的一步转移学习不同，密集动作前缀预测直接监督多个视野，迫使模型学习由动作序列引发的多视野状态演化，同时减少向后续视野的误差传播。

对于候选动作序列 $a_{t:t+H-1}$ ，每个视野 k 对应于前缀 $a_{t:t+k-1}=(a_{t},a_{t+1},...,a_{t+k-1})$ ， $k=1,...,H$ (8)。与此前缀关联的目标是在执行了这 k 个动作后达到的潜状态，即 $z_{t+k}$ 将每个前缀映射到一个学习到的前缀 token： $p_{t,k}=E_{\psi}^{(k)}(a_{t},...,a_{t+k-1})$ (9)，其中 $p_{t,k}$ 总结了前 k 个动作的累积效应。然后，潜变量预测器使用锚点潜变量和每个前缀 token 来预测对应的未来潜变量： $\hat{z}_{t+k}=G_{\phi}(z_{t},p_{t,k})$ ， $k=1,...,H$ (10)

在训练期间，对所有前缀 token 的密集监督迫使模型学习动作效应如何在不同视野上累积在展开期间，所有动作前缀可以被集中处理，并能并行地生成它们对应的未来状态

动作前缀编码器 (Action-Prefix Encoder)

动作前缀编码器的作用是输出特定视野的前缀 token $p_{t,1:H}$ ，其中处于视野 k 的 token 仅对应于长度为 k 的动作前缀，同时能避免未来动作的信息泄漏到较短的前缀中。给定动作序列 $a_{t:t+H-1}=(a_{t},a_{t+1},...,a_{t+H-1})$ (11)，编码器会输出前缀 token 的密集序列 $p_{t,1:H}=E_{\psi}(a_{t:t+H-1})=(p_{t,1},p_{t,2},...,p_{t,H})$ (12)。

为了以最小的开销提供这一上下文，我们将当前的潜变量 $z_{t}$ 通过一个轻量级的 MLP（多层感知机）映射为一个状态 token，并将其作为第 0 个 token 附加在动作 token 序列的最前面。加上这个状态 token 后，该编码器可以被视为一种条件形式： $p_{t,k}=E_{\psi}^{(k)}(a_{t},...,a_{t+k-1}\vert{}z_{t})$ (14)

并行潜变量预测器 (Parallel Latent Predictor)

给定当前的潜变量和各个前缀 token，预测器会并行地估计所有前缀视野的未来潜变量： $\hat{z}_{t+1:t+H}=G_{\phi}(z_{t},p_{t,1:H})$ (15)。等价于： $\hat{z}_{t+k}=G_{\phi}(z_{t},p_{t,k})$ ， $k=1,...,H$ (16)

所有查询的视野只需共享一次动作前缀编码的前向传递和一次并行潜变量预测的前向传递，无需去重复迭代单步的下一状态预测，这有助于减少误差积累。

密集前缀预测目标 (Dense Prefix Prediction Objective)

对于一个训练片段 $(o_{t},a_{t},o_{t+1},a_{t+1},...,a_{t+H-1},o_{t+H})$ (17)，我们将当前和未来的观测编码为 $z_{t+i}=f_{\theta}(o_{t+i})$ ， $i=0,...,H$ (18)。动作前缀编码器和预测器将产生前缀级的预测结果 $\hat{z}_{t+1:t+H}$ 。每个动作前缀都会接收其自己的潜变量目标： $\mathcal{L}_{prefix}=\frac{1}{H}\sum_{k=1}^{H}\vert{}\vert{}\hat{z}_{t+k}-z_{t+k}\vert{}\vert{}_{2}^{2}$ (19)

不仅监督终端结果，还监督部分动作前缀所诱发的中间状态，迫使模型学习随着向序列中附加更多动作，潜状态是如何演变的。

保留了免重建潜变量世界模型使用的 SIGReg 正则化项来防止潜空间坍缩： $\mathcal{L}_{AP}=\mathcal{L}_{prefix}+\lambda SIGReg(Z)$

使用动作前缀进行规划与自一致性 (Planning with Action Prefixes and Self-Consistency)

设 m 为 CEM 采样的候选动作序列的索引，基础的候选代价为 $C_{goal}^{(m)}=\vert{}\vert{}\hat{z}_{t+H}^{(m)}-z_{g}\vert{}\vert{}_{2}^{2}$ (21) 因此，Fast-LeWM 保留了与 LeWM 相同的规划目标，但更改了展开接口

前缀级接口在规划期间还提供了一个可选的自一致性信号：想象一下你规划去某个目的地的路线，这里有两种预测方式：

直飞路线（一步到位）：利用完整的动作前缀，直接从起点预测终点的位置，得到结果 $\hat{z}_{t+H}^{(m)}$ 。
转机路线（中转一次）：先利用部分动作前缀预测走到中间某一步的位置，然后再从这个中间位置出发，预测走到终点的位置，得到另一个结果 $\tilde{z}_{t+H}^{(m)}$ 。

这里的逻辑是：如果模型真的理解了物理世界的规律，那么“直飞”和“转机”预测出的终点应该是同一个地方。如果两个预测结果偏差很大，说明模型对这个动作方案的预测很不稳定、不靠谱。 $C^{(m)}=C_{goal}^{(m)}+\beta\vert{}\vert{}\hat{z}_{t+H}^{(m)}-\tilde{z}_{t+H}^{(m)}\vert{}\vert{}_{2}^{2}$ (22) 就是把这种“偏差值”作为一种惩罚加到了总分里。参数 $\beta$ 就是你对这种“偏差”的容忍度。 $\beta=0$ 代表你不开启这个检查机制，只关心能不能到目标； $\beta$ 越大，代表你越严格，强制要求 AI 选出的方案必须在不同的预测路径下都保持极高的稳定性，不能前后矛盾。

实验

作者主要在四个维度对模型进行了测试：

规划性能与效率测试：测试新方法在多个模拟环境（如控制机械臂、推木块等）中完成任务的成功率和所花的时间。
开环潜变量预测评估：让模型“闭着眼睛脑补”未来多步的状态，测试随着预测步数的增加，模型的想象力会不会严重跑偏。
物理状态探测：检查模型在内部处理信息时，是否真正学到了现实世界的物理规律（比如准确掌握物体的位置和角度）。
消融实验：把新方法拆开，依次去掉各个新组件，以此证明它的强悍表现确实归功于作者提出的“动作前缀预测”等核心设计，而不是碰巧。

一句话总结：又快、又稳、又准。

规划速度大幅提升（快）：新方法的核心动力学模块耗时缩短了将近 4 倍，使得整体任务规划（CEM）的求解时间缩短了将近一半。
任务成功率更高（稳）：在算得大幅变快的同时，它完成任务的平均成功率不仅没有掉，反而比老方法更高。如果开启“自一致性”检查机制，成功率还能进一步提升。
长线预测误差更小（准）：传统的模型预测未来时，步数越多，误差就像滚雪球一样越来越大。新方法因为采用了直接预测“动作前缀”的设计，极大地缓解了这种误差累积，长期预测准确度远超老方法。
物理直觉更好：测试表明，新方法不仅学得快，而且保留了比对照组更丰富的底层物理细节（如更准的位置和角度感知）。

总结

这篇文章的核心逻辑：在LeWM基础上，用因果掩码提取动作前缀（解决时序信息泄露），再用并行网络一步到位（解决速度慢和误差累积）。

并行预测网络（Parallel Latent Predictor）:这个网络并不是复杂的 Transformer，而是一个 6 层的动作调制残差多层感知机（Action-modulated residual MLP）。使用了 AdaLN-zero 调制技术（这种技术在强大的扩散模型如 DiT 中非常经典）。你可以这样理解：把当前的初始状态 $z_t$ 当作画板（主干特征），把前缀 token $p_{t,k}$ 当作画笔（控制信号），通过调整归一化层的缩放和偏移，将动作的影响平滑地注入到初始状态中。

由于每一个未来状态的预测，都只依赖起点 $z_t$ 和对应的前缀 token，它们互相之间完全不需要对方的数据。所以在代码和矩阵运算时，这 $H$ 个任务会被直接打包成一个大矩阵（Batch），扔给那个 6 层 MLP。一次前向传播（One forward pass），GPU 就能把这 $H$ 个时间步的未来状态同时算出来。