AliveRiver

外观

LeWorldModel论文笔记

约 5794 字大约 19 分钟

World ModelLeWorldModel

2026-07-12

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

摘要以及引言部分

LeWorldModel (LeWM)是第一个能够从原始像素稳定进行端到端训练的 JEPA JEPA 学习将观察结果编码到一个紧凑、低维的潜在空间中，并通过预测未来观察的潜在表示来对时间动态进行建模 LeWorldModel (LeWM)，这是第一个能够不使用启发式方法，从原始像素端到端学习稳定 JEPA 的方法，且该方法原则性强且简单

方法：LeWorldModel

图1图1：LeWorldModel训练流程。 给定帧观测 (o_{1: T}) 和动作 (a_{1: T})，编码器将帧映射为低维潜在表示 (z_{1: T})。预测器通过从当前潜在状态 (z_{t}) 和动作 (a_{t}) 自回归地预测下一潜在状态 (z_{t+1}) 来建模环境动态。编码器和预测器通过均方误差（MSE）预测损失进行联合优化。LeWM 不依赖任何训练技巧，如停止梯度、指数移动平均或预训练表示。为防止平凡坍塌，SIGReg 正则项强制潜在嵌入服从高斯分布，以促进特征多样性。为便于处理，潜在嵌入被投影到多个随机方向上，并对每个一维投影应用正态性检验。汇总这些统计数据可促使整体嵌入分布逼近各向同性高斯分布。

首先描述用于从离线数据中学习潜在世界模型的精简训练过程，包括数据集、模型架构和训练目标。然后，我们解释如何利用学习到的模型，通过使用模型预测控制 (MPC) 的潜在规划来进行决策

LeWorldModel 通过编码器将图像映射到潜在空间，利用预测器学习潜在空间中的环境动力学，并结合 SIGReg 正则化避免表示塌缩；在推理阶段，固定世界模型参数，通过 CEM 与 MPC 在潜在空间中优化动作序列，实现目标导向的规划与控制。

学习潜在世界模型 (Learning the Latent World Model)

离线数据集

考虑一个完全离线且无奖励的环境。LeWorldModel 仅从没有标注的观察和动作轨迹中训练，无法访问奖励信号或任务规范 目标不是针对特定任务优化行为，而是学习能够捕捉环境动态的表示，并且随后可以将其控制或适应到多种任务中

训练数据由长度为 TT 的轨迹组成，包含原始像素观察 o1:To_{1:T} 和相关的动作 a1:Ta_{1:T}。轨迹是离线从没有最优性要求的行为策略中收集的；它们可以是伪专家的，也可以是探索性的，只要它们能够充分覆盖环境的动态即可

模型架构 (Model Architecture)

LeWM 建立在两个组件之上：编码器 (encoder) 和预测器 (predictor)，编码器将给定的帧观察 oto_{t} 映射到紧凑、低维的潜在表示 ztz_{t}。预测器通过在给定潜在嵌入 ztz_{t} 和动作 ata_{t} 的情况下预测下一帧观察的嵌入 zt+1z_{t+1}，在潜在空间中对环境动态进行建模 编码器： zt=encθ(ot)z_{t}=enc_{\theta}(o_{t}) 预测器： z^t+1=predϕ(zt,at)\hat{z}_{t+1}=pred_{\phi}(z_{t},a_{t})

编码器实现为视觉 Transformer (ViT) 预测器是一个包含 6 层、16 个注意力头和 10% dropout 的 Transformer（约 1000 万参数）。预测器将 NN 个帧表示的历史作为输入，并使用时间因果掩码自回归地预测下一个帧表示，以避免查看到未来的嵌入

训练目标 (Training Objective)

目标是学习对预测未来有用的潜在表示，即建模环境的动态

LeWorldModel 的训练目标是两项的总和：预测损失和正则化损失。 完整的 LeWM 训练目标定义为：LLeWM≜Lpred+λSIGReg(Z)\mathcal{L}_{LeWM}\triangleq\mathcal{L}_{pred}+\lambda SIGReg(Z) (3)

预测损失 Lpred\mathcal{L}_{pred}（教师强制）计算连续时间步长的预测嵌入之间的误差： Lpred≜∣∣z^t+1−zt+1∣∣22\mathcal{L}_{pred}\triangleq\vert{}\vert{}\hat{z}_{t+1}-z_{t+1}\vert{}\vert{}_{2}^{2},z^t+1=predϕ(zt,at)\hat{z}_{t+1}=pred_{\phi}(z_{t},a_{t}) (1)

单独使用预测损失，会导致表示崩溃，为了防止这种行为，采用了草图各向同性高斯正则化器(Sketched-Isotropic-Gaussian Regularizer, SIGReg) SIGReg(Z)≜1M∑m=1MT(h(m))SIGReg(Z)\triangleq\frac{1}{M}\sum_{m=1}^{M}T(h^{(m)}) (2)

该方法只引入了两个训练超参数：SIGReg 中使用的随机投影数量 MM 以及正则化权重 λ\lambda。实践中观察到投影数量对下游任务的性能影响可以忽略不计，这使得 λ\lambda 成为唯一需要调整的有效超参数。

潜在规划 (Latent Planning)

图4图4：使用LeWorldModel的潜在规划。 给定初始观测 (o_1) 和目标 (o_g)，图2中学习到的世界模型在LeWM潜在空间中进行规划。初始状态嵌入 (z_1) 和目标嵌入 (z_g) 由编码器获得。随后，预测器向前展开未来潜在状态直至视界 (H)。最终预测状态与目标嵌入之间的潜在代价引导求解器优化动作序列。这一预测-优化循环重复进行，直到收敛到一个良好的规划候选。与之相对，在测试时则直接使用模型预测控制（MPC）[30–33, 18, 22]在潜在空间中进行规划。

模型根据以下规则预测潜在转换：
z^t+1=predϕ(z^t,at)\hat{z}_{t+1}=pred_{\phi}(\hat{z}_{t},a_{t})z^1=encθ(o1)\hat{z}_{1}=enc_{\theta}(o_{1})

规划是通过优化动作序列，以最小化终端潜在目标匹配度指标来执行的： C(z^H)=∣∣z^H−zg∣∣22\mathcal{C}(\hat{z}_{H})=\vert{}\vert{}\hat{z}_{H}-z_{g}\vert{}\vert{}_{2}^{2}zg=encθ(og)z_{g}=enc_{\theta}(o_{g}) (4) 其中 z^H\hat{z}_{H} 是展开结束时的预测潜在状态，zgz_{g} 是目标观察 ogo_{g} 的潜在嵌入。世界模型参数在规划期间保持固定。此过程对应于一个有限视野最优控制问题： a1:H∗=arg mina1:HC(z^H)a_{1:H}^{*}=arg~min_{a_{1:H}}\mathcal{C}(\hat{z}_{H}) (5) 使用交叉熵方法 (Cross-Entropy Method, CEM) 解决这个问题,迭代地选择最佳计划并用最佳计划的统计数据更新采样分布的参数

潜在规划性能 (Latent Planning Performance)

规划评估设置 (Planning evaluation setup)

环境 (Environments)

图5图5：用于评估的环境。 左侧：Push-T，一个二维操控任务，智能体需将木块推向目标配置，常作为机器人学基准任务。中图（1）：OGBench-Cube，一个视觉更丰富的三维操控环境，机械臂与立方体交互以到达目标位置。中图（2）：Two-Room，一个简单的二维导航环境，智能体在房间之间移动以到达目标位置。右侧：Reacher，一个两关节臂需要在二维平面上到达目标配置的任务。所有环境均具有连续动作空间。

基线 (Baselines)

将 LeWM 的性能与几个基线方法进行了比较：DINO-WM 和 PLDM

迈向高效的世界模型规划 (Towards Efficient Planning with WMs)

图6图6：各环境下的规划性能。 结果分别展示于Two-Room（左侧）、Reacher（中图1）、PushT（中图2）和OGBench-Cube（右侧）。在Push-T和Reacher上，LeWM始终优于PLDM和DINO-WM。在OGBench-Cube上，DINO-WM略优于LeWM，这可能是因为该环境视觉复杂度更高且具有3D特性，使得编码器训练更具挑战性。在较为简单的Two-Room环境中，PLDM和DINO-WM的表现优于LeWM，这或许可以解释为：SIGReg正则化鼓励高维潜在空间中的高斯分布，而该环境的本征维度要低得多。

LeWorldModel 在复杂规划任务上取得了优于 PLDM、接近甚至超过 DINO-WM 的性能，同时实现了最高 48 倍的规划加速；但在简单环境中，SIGReg 的高斯分布约束可能限制表示能力，导致规划性能下降。

迈向稳定的世界模型训练 (Towards Stable Training of World Models)

消融实验表明，LeWorldModel 对 SIGReg 参数、潜在表示维度和编码器架构均具有较好的鲁棒性，仅需调节正则化权重 λ 即可；同时，由于训练目标仅包含预测损失和 SIGReg 两项，其训练过程相比 PLDM 更稳定、收敛更平滑且调参成本更低。

量化 LeWM 中的物理理解 (Quantifying Physical Understanding in LeWM)

潜在空间的物理结构 (Physical Structure of the Latent Space)

探测物理量 (Probing physical quantities)

表1：Push-T上的物理潜在属性探测结果。 LeWM始终优于PLDM，同时与DINO-WM保持竞争力。DINO-WM在部分属性上表现强劲，可能源于其基础模型预训练：DINOv2编码器训练所用的数据量高出两个数量级（约1.24亿张图像），且覆盖的分布范围更为广泛，这可能使其嵌入默认就能捕获某些物理属性。 表1

作者想验证一个问题：LeWorldModel 学到的 latent 是否真正包含了环境中的物理信息？

为此，他们采用了 Probe（探测器） 方法： 冻结世界模型的 Encoder，不再训练； 从 Encoder 输出的 latent 中，训练一个很简单的预测器（Linear 或 MLP）； 看能否恢复出真实的物理状态，例如： 智能体位置（Agent Location） 方块位置（Block Location） 方块角度（Block Angle） 如果一个简单的探测器就能准确预测这些属性，说明 latent 已经编码了这些物理信息。

解码潜在空间 (Decoding Latent Space)

图10图10：训练过程中的解码器可视化。 随着训练推进，潜在表示逐渐捕获重构视觉场景所需的信息，尽管训练过程中并未使用重构损失。训练早期，解码出的图像对应的是慢特征（slow features）

尽管在训练期间从未使用过重建目标，但解码器仍能够从学习到的表示中恢复视觉场景，这证实了低维且紧凑的潜在空间保留了有关底层物理状态的充足信息。

可视化潜在空间 (Visualizing Latent Space)

tu13图13：LeWM在PushT环境中获得的潜在空间可视化。 左侧展示了通过在x-y平面上移动智能体和木块所获得的状态网格；右侧则使用t-SNE对这些状态的嵌入进行了可视化。

该可视化表明，学习到的表示捕获了环境的空间结构，并在潜在空间中保留了邻域关系和相对位置

时间潜在路径拉直 (Temporal Latent Path Straightening)

  1. 什么是“时间潜在路径拉直”？简单来说，现实世界中物体的运动通常是连续、平滑且有惯性的。在人工智能的“脑海”（潜在空间）里，如果模型真正理解了这种物理规律，那么代表物体状态的点在时间上的移动轨迹也应该是一条平滑、笔直的线，而不是随机乱跳的折线 。
  2. 他们是怎么测量的？研究人员通过计算连续时间步上“速度向量”（即状态变化的快慢和方向）的余弦相似度来量化这一点 。如果模型预测的前后两步的方向高度一致，就说明这个轨迹走得很“直” 。
  3. 令人惊讶的发现（涌现现象）研究发现，LeWM 在训练过程中，这种轨迹“变直”的现象是自发产生的（纯粹的涌现现象） 。这意味着，研究人员并没有在代码里写任何强制它去走出直线轨迹的规则（显式正则化），它是自己“悟”出这个道理的 。
  4. 对比优越性这正是最有趣的地方。之前的基线模型 PLDM 煞费苦心地专门加了一个数学规则（时间平滑度正则化项）来逼迫模型走出平滑的轨迹 。然而结果是：LeWM 什么都没刻意做，反而比刻意去做的 PLDM 走出了更直、更平滑的时间轨迹 。
  5. 总结来说：这段话是在高度评价 LeWM 的架构设计。它证明了该模型虽然训练机制简单，但能够自然而然地捕捉到物理世界时间连续性的本质，体现了其在底层逻辑上对物理动态的深刻理解。

期望违背框架 (Violation-of-expectation Framework)

图8图8：三种环境下的违反预期（Violation-of-expectation）评估。 每张图展示了LeWM在三条轨迹上的惊讶度（surprise）：一条无扰动的参考轨迹、一条视觉扰动轨迹（物体颜色突变）和一条物理扰动轨迹（物体瞬移至随机位置）。瞬移违反了物理连续性，并产生显著的惊讶度尖峰，而无扰动轨迹则保持在较低基线水平。在所有环境中，瞬移引起的惊讶度增加均显著（配对t检验，(p<0.01)），而颜色变化引起的增加则较弱且不显著，表明模型对物理扰动的敏感度高于视觉扰动。环境从左至右依次为：TwoRoom、PushT、OGBench Cube。

  1. 核心思想：什么是“期望违背” (VoE)？想象你在看一场魔术：如果魔术师让一个小球在空中平稳飞行，你觉得很正常；但如果小球突然凭空消失，又瞬间在另一个地方出现，你会感到非常“惊讶”，因为这违背了你脑海中根深蒂固的物理连续性常识。心理学家经常用这种“打破预期”的方法来测试婴儿是否建立起了物理常识。现在，研究人员把这套方法搬到了 AI 身上，用来测试世界模型能否察觉出“不符合物理规律”的事件 。
  2. 量化指标：如何计算 AI 的“惊奇度” (Surprise)？AI 不会有真正的情绪，它的“惊讶”是用数学误差来表示的。惊奇度 = 模型“脑补”预测的未来画面 与 实际发生的画面 之间的误差差异 。如果事情按常规物理规律发展，AI 预测得准，惊奇度（误差）就很低；如果发生了反常的事情导致 AI 预测彻底失效，惊奇度（误差）就会瞬间飙升。
  3. 实验手段：两种“捣乱”测试为了考察模型，研究人员在三个不同的虚拟测试环境（TwoRoom、PushT 和 OGBench Cube）中，故意给正在运动的物体制造了两种突发意外 ：
    1. 视觉扰动（换皮测试）： 运动中的物体突然改变了颜色 。这仅仅是表面的视觉变化，并没有破坏底层的物理运动轨迹。
    2. 物理扰动（瞬移测试）： 运动中的物体突然“瞬移”到了另一个随机位置 。这严重违背了物理学中的空间连续性（物体必须经过空间路径才能移动，不能瞬间转移）

结论

LeWorldModel (LeWM)，这是一种用于学习潜在世界模型的稳定的端到端方法 。是一种联合嵌入预测架构，其中编码器将图像观察映射到潜在空间中，而预测器则通过在给定动作的条件下预测未来的嵌入，来对时间动态进行建模 。

一些局限性指出了未来的研究方向 。首先，规划仍然局限于较短的时间视野，这促使我们在未来研究分层世界建模，以进行长视野的推理 。其次，我们的方法还依赖于具有足够覆盖率的离线数据集 ；特别是，在简单、低维的环境中，数据的低多样性会削弱 SIGReg 的效果，因为在这种环境中去匹配高维的高斯先验更加困难

LeWM 流程图

LeWorldModel (LeWM) 在推理和交互时的完整执行流程可以总结为以下五个步骤：

  1. 状态与目标的编码 (Encoding State and Goal)
  • 系统接收当前的初始观察画面 o1o_1 和任务想要达到的目标观察画面 ogo_g 。
  • 编码器 (Encoder) 将当前画面 o1o_1 映射压缩为低维的初始潜在状态向量 z1z_1 。
  • 同时，编码器也会将目标画面 ogo_g 映射压缩为目标潜在特征向量 zgz_g 。
  • 在接下来的整个规划过程中，世界模型自身的网络参数保持完全固定 。
  1. 脑海中的自回归推演 (Autoregressive Rollout in Imagination)
  • 系统随机初始化一个候选动作序列 a1:Ha_{1:H} 。
  • 预测器 (Predictor) 从初始潜在状态 z^1=z1\hat{z}_1 = z_1 开始，结合当前的动作，利用时间因果掩码自回归地向后预测未来的潜在状态 。
  • 状态转移的递推规则为 z^t+1=predϕ(z^t,at)\hat{z}_{t+1} = pred_\phi(\hat{z}_t, a_t) 。
  • 预测器会一直向后推演，直到达到预设的规划视野（Planning Horizon）步数 HH 。
  1. CEM 动作优化循环（直至收敛）(CEM Action Optimization Loop)
  • 系统计算推演到最后一小步的预测潜在状态 z^H\hat{z}_H 与目标潜在向量 zgz_g 之间的终端目标匹配度误差成本：C(z^H)=∣∣z^H−zg∣∣22\mathcal{C}(\hat{z}_H) = \vert{}\vert{}\hat{z}_H - z_g\vert{}\vert{}_2^2 。
  • 接下来进入动作优化循环：交叉熵方法 (CEM) 是一种零阶采样优化算法 ，它在当前动作分布中批量采样出许多组候选动作序列，并通过世界模型评估每组动作的成本 。
  • 算法从中筛选出成本最低的前几组“精英 (elites)”动作序列 。
  • 系统利用这些精英序列的统计数据，去更新和修正下一轮动作采样的分布参数（均值和方差） 。
  • 这个“采样-评估-筛选精英-更新分布”的优化循环会不断迭代进行 ，直到动作分布收敛，从而筛选出一个最优的动作计划序列 a1:H∗a_{1:H}^* 。
  1. MPC 滚动执行与自适应重新规划 (MPC Execution and Replanning)
  • 由于自回归推演随着视野 HH 越拉越长，其预测误差会不断累积，产生模型偏差 。为了防止偏离轨道，模型不会盲目将优化出来的 HH 步动作全部执行完 。
  • 此时系统采用模型预测控制 (MPC) 策略：系统仅执行当前最优动作计划序列中的前几个动作，而不是全部执行 （注：在论文的具体控制实验中设置视野 H=5H=5，并滚动执行完 5 步动作 ）。
  • 执行完这部分计划动作后，智能体会作用于真实的物理环境，使环境状态发生改变并返回一个新的真实观察画面 。
  1. 闭环迭代直至完成目标 (Closed-loop Iteration to Goal)
  • 智能体以这个最新的真实观察画面作为全新的初始观察 o1o_1 。
  • 整个流程将重新回到第 1 步（重新生成当前潜在状态、重新在脑海中做 HH 步推演、重新通过 CEM 优化动作序列） 。
  • 这个“规划-执行一小段-重新规划”的闭环预测优化循环将不断往复，直到最终完全达成设定的目标任务 。

LeWM 对比 Dreamer，提升到底在哪？

LeWM 对比 Dreamer，提升到底在哪？

这两者的根本差异在于生成式重建（Generative）与联合嵌入预测（JEPA）的路线之争，这一点在I-JEAP中讲过了。LeWM 的提升并不在于它能在某个特定游戏里玩得比 Dreamer 更好，而在于它更加轻量、通用且无需人工标注的奖励。

  • 彻底摆脱像素重建：Dreamer 是一类生成式世界模型，需要在像素空间中显式建模环境动态（相当于要在脑海里生成未来的每一帧画面）。这种方法不仅计算成本极高，而且模型容易将算力浪费在与物理动态无关的视觉背景细节上。LeWM 则属于 JEPA 架构，它是免重建（Reconstruction-free）的，只在紧凑的低维空间预测未来最相关的特征 。

  • 无需奖励信号，泛化能力强：Dreamer 通常需要依赖环境的奖励信号（Reward）来联合建模动态和价值信息，以此来优化策略 。而 LeWM 的设定完全是“无奖励（Reward-free）”的纯离线观察学习 。这意味着它学到的是一套通用的物理运行规律，不需要像 Dreamer 那样针对特定的单一任务进行训练 。

既然 CEM 又慢又有缺陷，为什么不用 Dreamer 的 Actor？

论文作者在附录中也坦诚地承认了这一点：CEM 作为一种零阶采样方法，深受“维度灾难”的困扰，在大型动作空间中表现挣扎，并且由于非凸性，很容易陷入局部最优 。

放弃高效的 Actor（策略网络）而采用笨拙的 CEM，主要是由应用场景和目标设定的根本差异决定的：

  • 零样本目标规划 vs. 预训练策略：Dreamer 的 Actor-Critic 架构是在世界模型的“想象”中，通过成千上万次的试错来预先训练出一个策略网络。这需要一个明确的奖励函数来指导，且一旦训练完成直接执行策略网络，测试时就不再需要世界模型了。

  • 即时的模型预测控制（MPC）：LeWM 面临的是“目标引导”的离线场景，即临时给你一张目标画面，你需要立刻给出动作走向这个目标 。在测试阶段，通过 MPC 和 CEM，系统可以直接在潜在空间中搜索一组动作，使推演出的最终状态 z^H\hat{z}_H 距离目标状态 zgz_g 最近 。这是一种“随需随测”的规划，不需要提前花费大量时间去预训练一个绑定在特定任务上的 Actor。

  • 极致的模型速度弥补了算法缺陷：虽然 CEM 算法本身比较低效，但因为 LeWM 舍弃了像素重建，模型极其小巧（仅 15M 参数），它的推演速度极快 。这使得即使搭配了笨重的 CEM 采样，其整体规划速度依然比基于基础模型的方法快高达 48 倍，能够在不到一秒的时间内完成规划 。用极其快速的模型正向传播，硬生生拉平了 CEM 采样慢的劣势。