快速潜在世界模型：基于动作前缀预测的高效视觉规划方法 Fast LeWorldModel
通过动作前缀并行预测替代自回归rollout，在提升视觉规划成功率的同时大幅降低计算开销
前置知识
Joint-Embedding Predictive Architectures (JEPAs)
JEPAs是一类无需重建像素的预测架构，其核心思想是预测未来或被遮蔽的嵌入表示而非原始像素。与传统的自编码器不同，JEPAs通过对比学习或匹配预测嵌入与真实嵌入来训练，避免了像素重建计算开销，同时防止表示崩溃。典型的JEPAs包括I-JEPA、V-JEPA等，它们通过最大化预测嵌入与目标嵌入之间的相似性来学习有意义的视觉表示。
本文的LeWorldModel和Fast-LeWM都基于JEPAs框架，理解JEPAs对于理解为什么可以跳过像素重建直接在潜在空间进行规划至关重要。
Cross-Entropy Method (CEM)
CEM是一种基于采样的优化算法，常用于模型预测控制中的动作序列搜索。它通过迭代地从初始分布中采样候选动作序列，评估每个候选的目标函数值，然后选择表现最好的样本重新拟合分布参数，重复这个过程直到收敛。CEM在连续控制问题中表现良好，因为它能够处理非凸优化问题并且对噪声具有鲁棒性。
本文使用CEM作为规划器来搜索最优动作序列，理解CEM的工作原理有助于理解为什么规划过程中的动力学模型评估次数如此频繁，从而凸显Fast-LeWM加速的重要性。
SIGReg正则化
SIGReg是一种防止潜在表示崩溃的正则化方法，其全称为Signed Information Regularization。它通过在潜在空间中引入高斯分布约束来防止编码器将所有输入映射到相同的表示。具体来说，SIGReg鼓励潜在向量的分布在统计上接近标准高斯分布，通过最大化潜在表示与高斯先验之间的互信息来实现。这种方法在JEPAs风格的架构中特别重要，因为它们缺乏像素重建的监督信号。
Fast-LeWM保留了LeWM使用的SIGReg正则化，理解这个技术有助于理解如何在没有像素重建的情况下维持潜在空间的多样性和信息量。
开环预测与闭环规划
开环预测是指给定初始状态和动作序列，模型预测整个轨迹而不考虑实际执行过程中的反馈。闭环规划则是指在规划过程中考虑实际观测到的状态，并根据反馈调整动作。开环预测误差会随着预测步数的增加而累积，这是由于模型的不完美导致的误差传播。在模型预测控制中，通常使用开环预测来评估候选动作序列，然后选择最优序列执行，形成闭环控制。
本文的核心问题之一就是开环预测误差的累积，Fast-LeWM通过动作前缀预测来缓解这个问题，理解开环与闭环的区别对于理解论文的动机和贡献至关重要。
研究动机
现有的LeWorldModel风格的世界模型在视觉规划中存在两个关键瓶颈。首先，LeWM使用局部一步转移动力学模型，要评估一个候选动作序列就必须通过自回归rollout逐步生成整个潜在轨迹，这意味着要重复调用动作编码器和潜在预测模块。在Two-Room任务上，这种动力学模块的rollout时间达到31.4秒，占整个CEM求解时间54.4秒的大部分。其次，由于每一步预测都依赖于前一步的预测结果，早期或中间状态的预测误差会传播到后续预测中，随着规划horizon的增长，累积误差使得rollout变得越来越不可靠。当使用CEM优化时，需要评估大量候选序列，这两个问题被进一步放大。
本文的目标是本文的目标是设计一个快速的潜在世界模型，能够在保持或提升规划成功率的同时，大幅降低规划过程中的计算开销和累积误差。具体来说，希望减少动力学模块的评估时间，从LeWM的31.4秒降低到更合理的水平，同时降低开环预测误差并减缓其在长horizon下的增长速度。
与已有工作不同的是，以往的研究主要集中在改进视觉表示学习目标或改进规划算法本身，但很少有人关注动力学模型的评估接口本身。本文的独特切入角度是重新设计动力学模型的查询接口，从传统的单步局部预测改为动作前缀预测。这种接口级别的改变不是简单地扩大动作块长度，而是从根本上改变了模型如何学习和使用动作序列信息，使得每个未来状态都可以直接从观测到的初始状态和对应的动作前缀预测，而不需要通过中间的想象状态。
核心方法
Fast-LeWM的核心思想是用动作前缀预测替代传统的自回归rollout。给定当前潜在状态和一个候选动作序列，模型将这个动作序列编码成多个前缀token，每个token总结了从开始到某个horizon的动作累积效果。然后，一个并行预测器使用当前潜在和所有前缀token一次性预测所有对应的未来潜在状态。这种方法在训练时对每个前缀都有监督，迫使模型学习动作如何在不同horizon上累积影响。在规划时，所有horizon可以并行评估，既减少了计算次数，又避免了误差的顺序累积。
核心创新点是将动力学模型的预测单元从单步转移改为动作前缀。LeWM风格的动力学模型本质上是局部的，它只预测下一个潜在状态，要获得多步预测就必须复合这些局部转移。Fast-LeWM直接预测执行动作前缀后到达的潜在状态。这种改变的本质区别在于，不同horizon的预测不再依赖于前面预测的潜在状态，而是都锚定在观测到的初始潜在状态上，因此预测之间是独立的，可以并行计算。密集的前缀级监督使模型直接学习多horizon的状态演化，而不是只拟合单步状态转移。
方法步骤详情
Fast-LeWM的训练包含三个主要步骤。第一步，动作前缀编码：将当前潜在通过一个轻量级MLP映射为状态token，作为第0个token，然后与动作token序列拼接。一个因果Transformer对这个状态-动作token序列进行处理，在因果掩码下，第k个位置的输出只能看到从开始到该位置的动作，因此形成前缀token。第二步，并行潜在预测：预测器使用当前潜在和所有前缀token，在一个前向传播中生成所有未来潜在。第三步，密集前缀预测损失：对每个horizon，计算预测潜在与真实潜在的L2距离，形成损失，加上SIGReg正则化得到总损失。在规划时，使用相同的CEM协议，但评估候选序列时直接使用最后一个前缀token预测终端潜在，而不需要逐步rollout。可选的自一致性项可以在规划时添加，通过比较直接预测和通过中间前缀预测的终端潜在来提供额外的模型一致性信号。
技术新颖性
Fast-LeWM的技术新颖性主要体现在三个方面。第一，接口级别的重新设计：从单步局部预测改为动作前缀预测，这不是架构的简单修改，而是动力学模型使用方式的基本改变。第二，密集多horizon监督：传统方法通常只监督终端预测或下一个状态，而Fast-LeWM对每个前缀都有监督，迫使模型学习连续的状态演化。第三，并行预测架构：动作前缀编码和并行预测器的组合使得所有horizon可以在一个前向传播中评估，这与LeWM的顺序rollout形成鲜明对比。此外，状态token的引入为动作前缀编码器提供了当前场景配置的上下文，使得相同动作在不同初始状态下可以产生不同的效果预测。
实验结果
Fast-LeWM在四个目标条件规划任务上都取得了优于或等于LeWM的性能。在Two-Room任务上，成功率从87%提升到98%，在Reacher上从86%提升到88%，在PushT上保持96%，在OGBench-Cube上从74%提升到80%。平均成功率从85.8%提升到90.5%，加上自一致性项后进一步提升到92.0%。更重要的是，Fast-LeWM在Two-Room任务上将动力学模块的rollout时间从31.4秒降低到8.0秒，加速了3.9倍，整个CEM求解时间从54.4秒降低到28.3秒，减少了48%。开环预测实验显示，Fast-LeWM在所有四个任务上的初始潜在误差都明显低于LeWM，并且误差增长速度也更慢。物理状态探针实验在PushT任务上进行，Fast-LeWM在MLP探针下对智能体位置、块位置和块角度的预测都达到了最优性能，表明潜在空间保留了更丰富的物理状态信息。消融实验显示，简单地扩大LeWM的动作块长度会导致性能严重下降，而终端only的Fast-LeWM虽然优于扩大动作块的方法但仍弱于完整模型，证明密集前缀监督的必要性。去掉状态token后性能略有下降，表明当前状态上下文对动作效果预测的重要性。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Two-Room视觉规划 | 成功率 | 98% | LeWM 87% | +11个百分点 |
| Two-Room规划效率 | CEM求解时间 | 28.3秒 | LeWM 54.4秒 | -48% |
| Two-Room规划效率 | 动力学模块时间 | 8.0秒 | LeWM 31.4秒 | -74.5% |
| 四个任务平均 | 规划成功率 | 90.5% | LeWM 85.8% | +4.7个百分点 |
| PushT物理状态探针 | 智能体位置MLP探针相关性 | 1.000 | LeWM 0.998 | +0.002 |
局限与改进
作者承认Fast-LeWM仍然依赖于潜在空间的准确性，如果视觉编码器无法捕获任务相关的信息，规划性能会受到影响。此外，虽然Fast-LeWM减少了rollout误差累积，但它仍然是一个开环预测模型，在实际部署中可能需要结合在线适应或模型不确定性估计来进一步提高鲁棒性。从架构角度看，动作前缀编码器引入了额外的Transformer参数，虽然总参数量与LeWM相当，但推理时的内存占用可能略高。另一个潜在的局限性是，Fast-LeWM假设动作序列是预先采样的，这对于离散动作空间可能不如连续动作空间自然。
独立分析的弱点
Fast-LeWM的一个潜在弱点是动作前缀编码器的计算复杂度随着horizon线性增长，对于非常长的规划horizon可能会成为瓶颈。改进方向可以是使用层次化或递归的前缀编码方式，或者引入注意力机制来动态选择最重要的前缀。另一个弱点是Fast-LeWM仍然使用固定的horizon，对于需要更长规划的任务可能不够。可以通过自适应horizon选择或多尺度预测来解决这个问题。此外，当前的自一致性项需要额外的模型调用，可以通过蒸馏或一致性正则化来避免。潜在空间的可解释性也是一个可以改进的方向，虽然探针实验显示物理信息被保留，但显式地建模物理约束可能会进一步提高规划性能。
未来方向
作者提出的未来工作方向包括将Fast-LeWM扩展到其他视觉规划任务，如导航和操作，以及结合在线学习来适应环境变化。基于论文成果，可以延伸的研究方向包括：将动作前缀预测与其他规划算法结合，如树搜索或梯度优化，探索更高效的规划策略；研究多模态动作前缀编码，以处理语言指令或图像目标；探索层次化前缀预测，在多个时间尺度上建模动作效果；将Fast-LeWM与模型不确定性估计结合，实现风险感知的规划；研究Fast-LeWM在部分可观测环境中的表现，结合历史信息进行预测。
复现评估
论文提供了详细的实验设置和超参数，代码已在GitHub开源，模型参数量为17.9M，与LeWM的18.0M相当。实验使用NVIDIA 4090 GPU，训练批次大小在大多数环境为128，在Cube环境为32以获得更稳定和快速的收敛。所有模型训练10个epoch，匹配LeWM的协议。由于Fast-LeWM不依赖视觉历史，所有环境的历史大小设置为1。数据预处理、图像分辨率、优化器设置和评估设置都遵循LeWM，除非另有说明。整体来说，复现难度中等，需要类似的GPU资源和LeWM的数据集，但开源代码和详细配置应该能够保证复现。
论文图表