西交大 Fast-LeWorldModel:用动作前缀并行预测,加速视觉世界模型规划

Fast-LeWM 用动作前缀并行预测替代逐步 rollout,让视觉世界模型规划更快更稳。

原文标题:世界模型太慢?西交大提出Fast LeWorldModel:用「动作前缀并行预测」让动态估计加速4倍

原文作者:机器之心

冷月清谈:

机器之心介绍了西安交通大学提出的 Fast LeWorldModel(Fast-LeWM)。传统 LeWorldModel 在规划时需要对候选动作序列逐步自回归 rollout,不仅速度慢,也容易让误差沿时间步累积。Fast-LeWM 将预测接口改为“动作前缀并行预测”:模型一次性预测执行 1 到 H 个动作前缀后对应的未来潜变量,并在训练中对每个前缀状态进行密集监督。实验显示,在 Two-Room、Reacher、PushT、OGBench-Cube 等环境中,Fast-LeWM 平均成功率由 LeWM 的 85.8% 提升至 90.5%,加入 self-consistency 评分后达到 92.0%;动态模块耗时从 31.4 秒降至 8.0 秒,完整 CEM 求解时间从 54.4 秒降至 28.3 秒。文章认为,这项工作的重点不只是加速,而是提示世界模型的“预测接口”设计可能会显著影响规划效率与稳定性。

怜星夜思:

1、世界模型做规划时,逐步 rollout 和并行预测未来,哪个更接近真实智能?
2、Fast-LeWM 的 self-consistency 评分,会不会只是让模型更保守,而不是真的更准?
3、动作前缀并行预测如果用到真实机器人上,最大的坑可能是什么?
4、这类世界模型加速方法,对自动驾驶、游戏 AI 或通用机器人哪个方向最先受益?

原文内容


本文第一作者为西安交通大学硕士生高云天,通讯作者为西安交通大学教授许翔宇,其研究方向涵盖世界模型、三维视觉与具身智能(个人主页:https://xuxy09.github.io/)


在视觉规划与具身智能中,“世界模型” 被认为是智能体走向通用决策能力的核心组件:在真正执行动作之前,先在潜在空间中 “想象未来”,再选择最优行为。


但在视觉规划里,这个 “想象” 过程往往很慢


以 LeWorldModel(LeWM)为例,它在规划时有一个重要瓶颈:每评估一条候选动作序列,模型都要一步步自回归 rollout。也就是说,LeWM 先预测下一步 latent,再把预测出的 latent 输入 dynamics model,继续预测下一步:


图片


这种方式有两个问题:一是规划慢,CEM 需要反复评估大量候选动作序列;二是误差会沿 imagined trajectory 累积,早期预测偏一点,后面可能越滚越偏。


针对这一瓶颈,西安交通大学研究团队提出了 Fast LeWorldModel(Fast-LeWM),试图从根本上改变世界模型的预测方式:从 step-by-step rollout 变成 trajectory-level parallel prediction。



  • 论文标题:Fast LeWorldModel

  • 作者:Yuntian Gao, Xiangyu Xu

  • 单位:西安交通大学

  • 论文链接:https://arxiv.org/pdf/2606.26217

  • 项目主页:https://fast-lewm.github.io/

  • 代码:https://github.com/Yuntian-Gao/Fast-LeWorldModel


它的核心思想非常直接:不再用一步转移模型反复 rollout,而是把一段动作序列的不同前缀作为预测单元,直接并行预测执行这些动作前缀后到达的未来潜变量,并且通过密集的监督迫使模型学会状态随着不同动作序列的演化过程,而非状态的单步转移。


换句话说,模型不再问:“执行下一个动作后会怎样?”,而是直接问:“执行 1 个、 2 个、…… H 个动作后,分别会到达什么状态?”


实验显示,在与 LeWM 相同的规划协议下,Fast-LeWM 将平均成功率从 85.8% 提升到 90.5%;加入自一致性约束后进一步提升到 92.0%。同时模型的 rollout 中的动态模块耗时从 31.4s 降至 8.0s,完整 CEM 求解时间从 54.4s 降至 28.3s



Fast-LeWM 的 pipeline



Fast-LeWM 的方法由三部分组成。


第一步,视觉编码器把当前观测和未来观测映射到 latent space:


图片


其中当前 latent z_t 是模型输入,未来 latent 图片 是训练监督目标。


第二步是 Action-Prefix Encoder。它把候选动作序列通过 causal Transformer 编码成一组 prefix tokens,每个 token 对应一个不同长度的动作前缀:



其中,第 k 个 prefix token 只包含前  k 个动作的信息:


图片


考虑到同样的动作序列在不同初始位置、物体状态和接触关系下对于动作会产生不同的后果,实际实现中,Fast-LeWM 还会把当前 latent z_t  映射成 state token,放在动作 token 序列最前面,为动作的编码提供上下文信息。


图片


第三步,Parallel Latent Predictor 使用当前 latent 和全部 prefix token,一次性输出所有未来 latent


图片


训练时,Fast-LeWM 对每一个前缀的预测都施加监督,而不只是监督最终状态:



最终目标保留 SIGReg 防坍塌正则:


图片


这也是 Fast-LeWM 区别于 LeWM 的关键:模型不只学习状态的局部变化,还要学习动作逐步累积时状态如何连续变化。


Planning 时:

基于动作前缀的快速 rollout


在测试阶段,Fast-LeWM 仍然沿用 LeWM 的 CEM 规划协议,对于 CEM 中第 m 条候选动作序列,LeWM 需要通过一步动态模型沿 imagined latents 逐步推进,最终得到终点预测 。相比之下,Fast-LeWM 将动作前缀作为 rollout 单元:对于长度为 H 的候选动作序列,模型通过对应的 prefix token 图片,直接从当前 latent 建立到未来 latent 图片 的预测路径,因此候选序列的代价函数:



论文这种动作前缀的设计还额外带来了一个可选的 self-consistency scoring:模型一方面可以直接从长度为 H 的动作前缀预测终点;另一方面也可以先预测一个中间 latent,再从中间 latent 继续预测剩余时域。两种终点预测之间的差异被作为一致性惩罚项:



其中,β 控制 self-consistency 项的权重。当 β=0,Fast-LeWM 退化为只使用 goal distance 的 CEM 打分;当 β>0 时,CEM 可以选择那些在不同 prefix 分解下预测结果一致的候选动作序列,进一步提升规划稳定性。


成功率提升,

规划时间近乎减半


实验沿用 LeWM 的 goal-conditioned latent planning 协议,在 Two-Room、Reacher、PushT、OGBench-Cube 四个环境上评测: 



结果显示,Fast-LeWM 在四个任务上的平均成功率从 LeWM 的 85.8% 提升到 90.5%;加入动作前缀预测带来的额外 Self-Consistency 规划项后进一步达到 92.0%


效率提升更明显。在相同 CEM budget、单张 NVIDIA 4090 上,Fast-LeWM 的 dynamics time 从 31.4s 降到 8.0s,加速约 4 倍,其中包含动作编码和 predictor 预测时间。完整 CEM solve time 从 54.4s 降到 28.3s,减少 48.0%


另外开环情况下,Fast-LeWM 想象未来时的 latent 误差和误差随 Horizon 的增长率也更小:



消融实验:

不是简单把动作块变长就行



作者进一步通过消融实验验证了 Fast-LeWM 各个组件的作用。首先一个看似直接的加速方式是:把 LeWM 的动作块变长,让一次 transition 覆盖更长时间。作者构造了 Long-Action LeWM,将原本 action encoding 从 5 个 primitive actions 改为 25 个 primitive actions, 结果效果并不好,Terminal-only Fast-LeWM 只监督最终 latent,不监督中间 prefix latent,表现优于 Long-Action LeWM,但仍低于完整模型。这说明 action prefix 本身已经是更有效的长时域表示,但 dense prefix supervision 对学习连续状态演化仍然关键。


作者还发现,去掉 state token 后,模型在多个任务上性能下降。这进一步说明了动作编码需要提供有效的上下文信息。


总结


Fast-LeWM 针对世界模型在规划阶段的关键瓶颈,提出了 action-prefix prediction 机制,将传统的一步自回归 rollout 改为并行多时域潜变量预测。


在相同 LeWM 评测协议下,Fast-LeWM 将平均规划成功率从 85.8% 提升到 90.5%,加入 self-consistency 后达到 92.0%;同时将动态模块耗时从 31.4s 降到 8.0s,将完整 CEM 求解时间从 54.4s 降到 28.3s。


更深层意义:

世界模型的瓶颈不在 “模型”,而在 “接口”


这项工作的核心启示并不只是加速,更本质的是它表明,对于面向规划的世界模型而言,动态模型的接口设计本身可能与表征学习目标同样重要。相比一步步预测 “下一个 latent”,直接预测动作前缀导致的多时域未来状态,或许是让视觉世界模型走向高效规划的一条更直接的路径。这可能意味着,世界模型正在从 “逐步想象未来”,走向 “并行生成未来”。



© THE END

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com


关于“哪个方向最先受益”,我投游戏 AI。一是仿真环境可控,二是可以大规模采数据,三是动作空间和奖励定义比较清楚。世界模型在游戏里先跑通,再往机器人和自动驾驶迁移,现实阻力会小很多。

1 个赞

真实机器人上最大的问题可能不是模型慢,而是世界太脏。桌面不平、夹爪打滑、光照变了、物体被碰歪一点点,latent 里都可能变成大麻烦。Fast-LeWM 加速了“大脑想象”,但“眼睛和手”的误差还得有人擦屁股。

1 个赞

如果从模型校准角度看,self-consistency 更像是一种内部不确定性估计。它不能保证预测等于真实世界,但能筛掉模型自己都前后矛盾的候选轨迹。至于是否过度保守,需要看任务中探索性动作的比例,以及 β 权重怎么调。

1 个赞

我倒觉得 self-consistency 有点像考试验算,不一定让你更聪明,但能帮你少写离谱答案。直接预测终点是一种说法,拆成中间状态再预测终点是另一种说法,两边对不上,那这条动作序列大概率不太靠谱。

我选通用机器人,尤其是桌面操作、仓储分拣这种。因为这类任务很需要“想一下再动手”,同时规划速度又直接影响操作效率。自动驾驶当然也需要预测未来,但安全认证和长尾场景太硬核,落地周期可能更长。