一、研究背景与问题核心问题联合嵌入预测架构JEPA是学习世界模型的有力框架但现有方法极其脆弱容易发生表示坍缩即模型将所有输入映射为相同表示。为防止坍缩现有方法依赖复杂的多目标损失、指数移动平均、预训练编码器或辅助监督导致训练复杂且不稳定。二、主要贡献作者提出了LeWorldModel (LeWM)具有以下创新点特性LeWM训练方式端到端从原始像素直接学习损失项数仅2项预测损失 正则化损失可调超参数仅1个正则化权重λ防坍缩机制SIGReg正则化理论上可证明依赖技巧无停止梯度、无EMA、无预训练编码器模型规模1500万参数单GPU数小时完成训练三、方法论模型架构编码器ViT-Tiny将图像观测映射为低维潜表示预测器6层Transformer以动作和当前潜状态为输入预测下一潜状态训练目标预测损失最小化预测嵌入与真实下一嵌入的MSESIGReg正则化通过随机投影Epps-Pulley检验统计量强制潜嵌入服从各向同性高斯分布从而防止坍缩规划方法使用交叉熵方法CEM在潜空间优化动作序列采用模型预测控制MPC策略仅执行第一个动作后重新规划四、实验评估4.1 控制任务性能图6环境Push-T2D操作、OGBench-Cube3D操作、Reacher2D到达、Two-Room2D导航结果在Push-T和Reacher上优于PLDM和DINO-WM在OGBench-Cube上与DINO-WM接近后者略微领先在简单的Two-Room上表现欠佳SIGReg在高维空间强制高斯先验与低内在维度环境不匹配4.2 效率优势图3规划速度比DINO-WM快48倍完整规划在1秒内完成编码观测使用的token数比DINO-WM少约200倍4.3 训练稳定性图18-19LeWM的损失曲线平滑单调收敛PLDM的七项损失曲线噪声大且非单调训练不稳定五、物理理解评估5.1 潜空间探测表1训练线性和MLP探测器从潜嵌入预测物理量位置、角度等LeWM优于PLDM与DINOv2预训练表示竞争力相当5.2 违反预期测试图10对物理扰动物体瞬间传送产生显著高惊奇峰值对视觉扰动颜色变化反应弱得多证明模型捕捉了物理连续性而非仅仅是视觉特征5.3 涌现特性时间潜路径拉直图17LeWM的潜轨迹在训练中自发变得更直且优于带有显式时间平滑正则化的PLDM解码可视化图8即使未使用重建损失潜表示仍保留足够视觉信息六、消融实验关键发现消融项结论SIGReg投影数性能不敏感无需调参嵌入维度低于184时性能下降超过后饱和编码器架构ViT vs ResNet两种均有效ViT略优预测器dropoutp0.1最佳过低或过高均下降加入重建损失反而损害控制性能七、局限性与未来方向短视野规划当前仅限于短时域需要分层世界建模处理长时域数据依赖需要覆盖充分的离线数据集低复杂度环境下SIGReg可能不适用动作标签依赖需要显式动作标注未来可通过逆动态模型学习动作表示大规模预训练潜力在更大视频数据上预训练可能进一步提升性能LeWM是第一个能以简单双项损失稳定端到端训练的JEPA世界模型在控制性能、训练速度、规划效率和物理理解方面均达到或超越现有方法且将超参数从6个减至1个大幅降低了研究和应用门槛。这里是自己的论文阅读记录感兴趣的话可以参考一下如果需要阅读原文的话可以看这里如下所示官方项目地址在这里如下所示项目地址在这里如下所示数据集发布地址在这里如下所示摘要联合嵌入预测架构JEPA为在紧凑的潜空间中学习世界模型提供了一个有吸引力的框架。然而现有方法仍然很脆弱依赖于复杂的多目标损失、指数移动平均、预训练编码器或辅助监督来避免表示坍缩。在本工作中我们介绍了LeWorldModelLeWM这是第一个能够仅使用两个损失项就从原始像素进行稳定端到端训练的JEPA模型一个下一嵌入预测损失和一个强制潜嵌入服从高斯分布的正则化项。与现有唯一的端到端替代方案相比这可将可调的损失超参数从六个减少到一个。LeWM拥有1500万个参数可在单张GPU上在几小时内完成训练其规划速度比基于基础模型的世界模型快高达 48×同时在各种2D和3D控制任务中保持竞争力。除了控制任务我们还展示了LeWM的潜空间通过物理量探测编码了有意义的物理结构。惊奇评估证实该模型能可靠地检测物理上不合理的事件。图2潜世界模型方法的特点。方法按训练范式分组。端到端方法PLDM同时学习编码器和预测器直接处理像素不依赖预训练表示或诸如停止梯度或指数移动平均之类的启发式技巧但需要许多超参数且缺乏形式化的抗坍缩保证。基于基础模型的方法DINO-WM通过冻结预训练的基础视觉编码器来避免坍缩但牺牲了端到端学习。任务特定方法Dreamer, TD-MPC在训练期间需要奖励信号或特权状态访问。LeWM解决了每类方法的局限性它是端到端的、任务无关的、基于像素的、无需重建和奖励的并且仅需一个超参数同时具有可证明的抗坍缩保证。1 引言人工智能的一个核心目标是开发能够在不同任务和环境中通过单一、统一的学习范式获取技能的智能体——这种范式直接作用于来自其周围环境的感官输入——无需手工设计的状态表示或领域特定的校准。视觉非常适合这一目标相机廉价且可扩展从像素学习使得从原始感官输入到动作的完全端到端训练成为可能[1]。世界模型WMs是一类强大的方法[2]它们学习预测环境中动作的后果。一旦成功世界模型允许智能体仅通过其世界模型即在想象空间中进行规划和自我改进。这在离线设置中尤其有价值此时智能体必须从固定数据集中学习无法与环境交互——利用模型生成合成经验并评估反事实动作序列[3, 4]。最近一种流行的学习世界模型的方法是联合嵌入预测架构JEPA[5]。JEPA不试图对环境的所有方面进行建模而是专注于捕捉预测未来状态所需的最相关特征。具体来说JEPA学习将观测编码为紧凑的低维潜空间并通过预测未来观测的潜表示来建模时间动态。然而尽管概念上简单现有的JEPA方法非常容易发生坍缩。在这种失败模式中模型将所有输入映射到几乎相同的表示以琐碎地满足时间预测目标导致表示不可用。因此防止坍缩是训练JEPA模型的核心挑战之一。许多有影响力的工作提出了解决此问题的方法。然而这些方法通常依赖于启发式正则化、多目标损失函数、外部信息源或架构简化例如使用预训练编码器。在实践中这些策略常常引入额外的不稳定性或显著增加训练复杂性。为了克服这些限制我们提出了LeWorldModelLeWM这是第一种无需启发式方法、原则性强且简单参见图2的从原始像素学习稳定JEPA模型的方法。此外LeWM可以在单张GPU上训练降低了研究的门槛。我们在2D和3D环境中的一系列操作、导航和 locomotion 任务上评估了LeWM。此外我们通过在潜空间中进行有针对性的探测和惊奇量化评估来探究其直观的物理理解。总体而言我们的主要发现和贡献是我们提出了一种端到端的JEPA方法用于在单张GPU上从原始像素学习潜世界模型。该方法依赖于一个简单且稳定的双目标函数该函数在不同架构和超参数选择下保持稳健同时支持高效的对数超参数搜索。图3固定计算量下的规划时间和性能。左图50次运行的平均规划时间比较。使用比DINO-WM少 ∼200× 的token对观测进行编码使LeWM能够达到与PLDM相当的规划速度同时比DINO-WM快高达 ∼50×。中-右图相同计算预算固定FLOPs下的规划性能。在Push-T中图和OGBench-Cube右图上LeWM显著优于DINO-WM。有关规划设置的详细信息请参见附录D。LeWM在多种2D和3D任务中以紧凑的1500万参数模型取得了强劲的控制性能超越了现有的端到端基于JEPA的方法同时以显著更低的成本与基于基础模型的世界模型保持竞争力规划速度提升高达 48×。我们通过潜空间中的物理量探测和用于检测非物理轨迹的违反预期测试评估了潜空间中的物理理解。2 相关工作世界模型旨在从数据中学习环境动态的预测模型使智能体能够在想象中推理未来状态。世界模型的一个主要类别包括生成式方法这些方法在像素空间中显式地建模环境动态。这些动作条件下的生成式模型通过基于过去状态和动作生成未来观测充当学习到的模拟器。生成式世界模型已成功应用于模拟现有游戏类环境。例如IRIS [3]、DIAMOND [6]、Δ-IRIS [7]、OASIS [8] 和 DreamerV4 [4] 模拟了Minecraft、Counter-Strike和Crafter等环境提高了强化学习中策略的样本效率。其他方法生成全新的交互式模拟器例如Genie [9] 和HunyuanWorld [10]而学习到的模拟器也已应用于机器人策略评估[11]。重要的是许多生成式世界模型假设可以访问包含奖励信号的数据集从而能够联合建模动态和与价值相关的信息以用于下游强化学习。相比之下我们关注的是无奖励设置这对应于JEPA系列工作中所考虑的设置该系列旨在从观测数据中学习通用的、任务无关的世界模型而不依赖奖励监督。JEPA是一个用于学习世界模型的框架该模型在紧凑的低维潜空间中预测系统的动态演变。自LeCun [5]引入以来JEPA方法有了显著发展主要区别在于它们的目标任务以及用于学习非坍缩表示的策略。一个重要的研究方向是将JEPA应用于自监督表示学习通过预测掩码输入块的潜嵌入。例如针对图像的I-JEPA [12]、针对视频的V-JEPA [13, 14] 以及针对医疗数据的Echo-JEPA和Brain-JEPA [15, 16]。这些方法通常采用目标编码器的指数移动平均EMA以及停止梯度SG更新来稳定训练并防止表示坍缩。然而对EMA和SG的理论理解仍然有限因为它们通常不对应于良好定义的目标函数的极小化[17]。第二个研究方向是将JEPA方案用于动作条件下的潜世界建模。一些方法依赖预训练编码器来获取表示[14, 18–20]。这避免了坍缩但将表示的表达能力限制于所使用的预训练编码器。相比之下PLDM [21, 22] 使用VICReg [23] 和额外的正则化项端到端地学习表示但代价是已知的训练不稳定性和可扩展性限制[24]。一些工作通过结合辅助信号或架构组件如本体感觉输入或动作解码器进一步提高了稳定性[18, 19]。在本工作中我们提出了一种稳定方法用于使用一个简单的双项损失从原始像素端到端训练JEPA一个关于未来嵌入的预测目标和一个强制嵌入服从高斯分布的正则化目标[25]。使用潜动态进行规划。世界模型[26]开创了从高维观测的紧凑潜表示中直接学习策略的先河。一些工作利用学习到的潜动态模型通过强化学习训练策略[27-29, 4]。在这些方法中生成式世界模型充当一个模拟器在其中可以“想象”出轨迹从而使得策略优化主要在潜空间的想象中进行。训练完成后策略直接执行世界模型在测试时不再需要。最近的工作则转向在测试时使用模型预测控制MPC直接在潜空间中进行规划[30-33, 18, 22]。与基于想象力的策略学习不同这些方法在线使用世界模型来预测候选动作序列的结果并在执行过程中迭代优化它们。因此模型在运行时仍处于控制回路中实现了适应性决策但增加了计算需求。3 方法LeWorldModel在本节中我们介绍LeWorldModelLeWM。我们首先描述了从离线数据学习潜世界模型的简化训练流程包括数据集、模型架构和训练目标。然后我们解释了如何通过学习到的模型利用模型预测控制MPC进行潜空间规划来实现决策。3.1 学习潜世界模型离线数据集。我们考虑一个完全离线且无奖励的设置。LeWorldModel仅从未标注的观测-动作轨迹中训练无法访问奖励信号或任务规范。此设置与JEPA系列工作[18, 14]一致旨在从观测数据中学习通用的、任务无关的世界模型。我们的目标不是针对特定任务优化行为而是学习能够捕捉环境动态并在之后可用于控制或适应各种任务的表示。通过预测损失编码器被激励去学习预测器可预测的表示。然而仅此损失会导致表示坍缩产生一个将所有输入映射到常数表示的琐碎解。为防止这种行为我们引入了一个抗坍缩正则化项以促进嵌入空间中的特征多样性。具体来说我们采用Sketched-Isotropic-Gaussian Regularizer (SIGReg) [25]因为它简单、可扩展且稳定。SIGReg鼓励潜嵌入与各向同性高斯目标分布匹配。算法1. LeWorldModel训练过程的伪代码。像素观测被编码为潜嵌入预测器通过预测以动作为条件的下一步嵌入来估计动态。模型使用下一步嵌入预测损失和逐步的SIGReg正则化项进行端到端优化以防止表示坍缩。该方法仅引入两个训练超参数SIGReg中使用的随机投影数 M 和正则化权重 λ。除非另有说明我们使用 M1024 个投影和 λ0.1。在实践中我们观察到投影数量对下游性能影响可忽略不计见第4节和附录G使得 λ 成为唯一需要调整的有效超参数。这大大简化了超参数选择因为可以使用简单的二分搜索以对数复杂度高效优化 λ。我们不使用停止梯度、指数移动平均或额外的稳定化技巧。梯度通过损失的所有组件传播所有参数以端到端的方式进行联合优化从而形成一个精简且易于实现的训练流程。训练逻辑总结在算法1中。3.2 潜空间规划我们采用模型预测控制MPC策略仅执行规划出的前 K 个动作然后根据更新后的观测重新规划。关于规划策略的更多细节见附录D。图5用于评估的环境。左图Push-T一个2D操作任务智能体必须将方块推向目标配置常用作机器人基准测试。中图(1)OGBench-Cube一个视觉上更丰富的3D操作环境机器人手臂与立方体交互以到达目标位置。中图(2)Two-Room一个简单的2D导航环境智能体在房间之间移动以到达目标位置。右图Reacher一个2关节手臂需要在2D平面中到达目标配置的任务。所有环境都具有连续动作空间。环境和数据集的更多细节见附录E。4 潜空间规划性能4.1 规划评估设置环境。我们在多种任务上评估LeWM包括导航、运动规划和操作涵盖二维和三维环境所有这些环境如图5所示。关于数据集生成和环境的更多细节见附录E。基线。我们将LeWM的性能与多个基线进行比较DINO-WM和PLDM两种最先进的基于JEPA的方法一个目标条件的行为克隆策略GCBC以及两种目标条件的离线强化学习算法GCIVL和GCIQL。在这些基线中PLDM与我们的设置最接近因为它也端到端地直接从像素观测学习世界模型。然而它依赖于从VICReg准则派生的七项训练目标这引入了训练不稳定性并增加了超参数调整的复杂性。相比之下DINO-WM使用DINOv2 [41]作为特征编码器来建模动态以减轻表示坍缩但其原始公式还结合了其他模态如本体感觉输入为了公平比较除非另有说明我们从DINO-WM中排除了本体感觉信息。基线的其他实现细节附录C和评估设置附录F.1在附录中提供。对于每种方法我们在所有环境中保持超参数固定。4.2 迈向高效的世界模型规划我们在图6中报告了规划性能。在更具挑战性的规划任务上LeWM相较于PLDM有所改进在PushT上实现了 18% 更高的成功率同时与DINO-WM保持竞争力。值得注意的是在PushT上即使DINO-WM可以访问额外的本体感觉信息仅使用像素的LeWM也超越了它这证明了LeWM捕捉潜在任务相关数量的能力。此外在比较规划加速图3时LeWM实现了 48× 更快的规划时间完整规划在一秒内完成同时在各项任务中保持了具有竞争力的性能。对于固定的规划设置此规划时间在不同环境中保持一致缩小了与实时控制的差距。我们在图6中报告了规划性能。在更具挑战性的规划任务上LeWM优于PLDM在PushT上实现了 18%18% 更高的成功率同时与DINO-WM保持竞争力。值得注意的是在PushT上即使DINO-WM可以访问额外的本体感觉信息仅使用像素的LeWM也超越了它这证明了LeWM捕捉潜在任务相关数量的能力。有趣的是在最简单的环境Two-Room上LeWM表现较差。一个可能的解释是该数据集的低多样性和低内在维度使得编码器难以在高维潜空间中匹配由SIGReg强加的各向同性高斯先验这可能导致潜表示结构较差。这突显了SIGReg正则化在非常低复杂度的环境中的潜在局限性。此外在比较规划加速图3时LeWM实现了 48×48× 更快的规划时间完整规划在一秒内完成同时在各项任务中保持了具有竞争力的性能。对于固定的规划设置此规划时间在不同环境中保持一致缩小了与实时控制的差距。图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正则化鼓励高维潜空间中的高斯分布而环境的内在维度要低得多。4.3 迈向稳定的世界模型训练消融实验。我们对LeWM的几个设计选择进行了消融研究。首先我们分析了SIGReg对其内部参数即随机投影数和积分节点数的敏感性。性能在很大程度上不受这些数量影响表明它们不需要仔细调整。因此正则化权重 λλ 仍然是唯一有效的超参数。由于只需要调整一个超参数可以使用简单的二分策略 (O(log⁡n))高效地进行网格搜索而PLDM需要多项式时间 (O(n6))搜索。我们还研究了嵌入维度的影响。虽然表示维度必须足够大方法才能表现良好但性能超过某个阈值后会迅速饱和表明该方法对编码器容量的精确选择具有鲁棒性。此外我们通过将默认的ViT编码器替换为ResNet-18骨干网络表8来检查编码器架构的影响。LeWM在使用两种架构时都取得了具有竞争力的性能表明它在很大程度上与视觉编码器的选择无关。所有消融实验的细节见附录G。训练曲线。我们在图18中报告了LeWM在PushT上的训练损失曲线在图19中报告了PLDM的曲线。LeWM的双目标函数表现出平滑且单调的收敛预测损失稳步下降而SIGReg正则化项在训练早期急剧下降然后趋于平稳表明潜分布迅速接近各向同性高斯目标。相比之下PLDM的七项目标函数在其多个损失组件中表现出噪声和非单调行为。这些观察突显了LeWM的一个关键优势通过将训练目标缩减为仅两个表现良好的项训练变得显著更稳定消除了平衡来自多个正则化器的竞争梯度的需要。5 量化LeWM中的物理理解在本节中我们通过从潜嵌入中学习提取物理量或测量世界模型检测物理变化的能力来评估LeWM潜空间所捕捉动态的质量。5.1 潜空间的物理结构探测物理量。作为物理理解的第一个度量我们评估哪些物理量可以从LeWM的潜表示中恢复。我们训练线性和非线性探测器从给定的嵌入中预测感兴趣的物理量。Push-T环境的结果见表1。我们的方法始终优于PLDM同时与大型预训练模型如DINOv2产生的表示保持竞争力。我们在附录F.2中提供了其他环境上的探测结果。表1Push-T上的物理潜空间探测结果。LeWM始终优于PLDM同时与DINO-WM保持竞争力。DINO-WM在某些属性上的强探测性能可能源于其基础模型预训练DINOv2编码器在数量级更多的数据约 124M124M 图像上训练覆盖了更加多样的分布这可能使其默认在其嵌入中捕捉到一些物理属性。解码潜空间。为了进一步评估潜表示中捕获的信息我们在图8中报告了由解码器生成的图像该解码器经过训练可从单个潜嵌入192维重建像素观测。尽管训练期间从未使用重建解码器能够从学习到的表示中恢复视觉场景证实了低维且紧凑的潜空间保留了关于底层物理状态的足够信息。解码器架构的细节见附录D。可视化潜空间。我们进一步使用t-SNE可视化潜空间的结构。图9提供了PushT环境中潜空间的定性可视化。可视化表明学习到的表示捕捉了环境的空间结构在潜空间中保留了邻域关系和相对位置。时间潜路径拉直。受神经科学中的时间拉直假说[42]启发我们测量了整个训练过程中连续潜速度向量之间的余弦相似度公式9。我们发现在PushT上LeWM的潜轨迹在训练过程中作为一种纯粹涌现现象变得越来越直没有任何显式正则化鼓励这种行为参见图17。值得注意的是尽管PLDM采用了专门的时间平滑正则化项LeWM仍实现了比PLDM更高的时间直线度。我们在附录H中详述了我们的发现。5.2 违反预期框架量化物理理解的另一种方法是检测对学习到的世界模型的违反能力。受发展心理学中使用的违反预期VoE范式以及最近在机器学习中采用的方法[43-45]的启发该框架评估模型是否对违反学习到的物理规律的事件分配更高的惊奇度。遵循先前工作我们通过测量模型预测的未来观测与实际观测之间的差异来量化惊奇度。我们在三个环境中评估此框架图7PushT和OGBench-Cube上的预测器展开。我们可视化了由LeWM在给定上下文和动作序列情况下生成的解码潜计划。每次展开使用三个图像观测作为上下文将其编码为潜表示。以动作序列为条件预测器以开环方式自回归生成未来的潜状态。所有预测的潜状态使用训练期间未使用的解码器解码为图像。产生的想象展开与真实观测非常吻合表明潜表示有效地捕捉了整体场景结构和基本的环境动态。然而一些更精细的细节未被LeWM完全捕捉例如OGBench-Cube中末端执行器的角度。更多的展开结果见图11。TwoRoom、PushT和OGBench Cube。对于每个环境我们引入两种类型的扰动。第一种是视觉扰动其中一个物体的颜色在轨迹过程中突然改变。第二种是物理扰动其中一个或多个物体被传送到随机位置违反了预期的场景物理连续性。图10显示LeWM始终对包含物理违反的帧分配比其未扰动对应帧更高的惊奇度。关于VoE的更多细节见附录F.3。6 结论本工作介绍了LeWorldModelLeWM一种用于学习环境潜世界模型的稳定端到端方法。LeWM是一种联合嵌入预测架构它使用编码器将图像观测映射到潜空间并使用预测器通过预测以动作为条件的未来嵌入来建模嵌入空间中的时间动态。在各种连续控制环境中仅使用原始像素输入LeWM在数据效率、规划时间、训练时间和稳定性方面优于先前的方法同时保持了具有竞争力的最终任务性能。训练的稳定性和简单性源于明确鼓励潜嵌入遵循各向同性高斯分布以避免坍缩。总体而言LeWM为现有的潜世界模型方法提供了一种可扩展的替代方案提供了原则性的训练动态以及可解释和涌现的表示属性。局限性与未来工作。尽管取得了这些有希望的成果但几个局限性指明了重要的研究方向。首先当前的潜世界模型规划仍局限于短视野。分层世界建模代表了解决长视野图8训练过程中的解码器可视化。随着训练的进行潜表示越来越多地捕获重建视觉场景所需的信息尽管训练期间未使用重建损失。在训练早期解码图像对应于慢特征这一现象此前已有报道[21]。图9使用LeWM获得的PushT环境潜空间可视化。左侧通过在x-y平面上移动智能体和方块获得状态网格。右侧使用t-SNE可视化这些状态的嵌入。推理和规划的一个有前途方向。其次我们的方法仍然依赖于具有足够交互覆盖范围的离线数据集这可能收集成本高昂或困难。特别是有限的数据多样性可能会影响SIGReg正则化在内在维度非常低的简单环境中的有效性因为在这样的环境中在高维潜空间中匹配各向同性高斯先验变得具有挑战性。在大规模和多样化的自然视频数据集上进行预训练可以提供强大的表示先验并减少对领域特定数据的依赖。最后当前的端到端潜世界模型依赖于动作标签来预测未来状态这同样可能获取成本高昂。一个有前途的方向是通过逆动态建模学习未来的动作表示这可能会减少对显式动作注释的需求。图10在三个环境中的违反预期评估。每个图显示了沿三条轨迹的模型惊奇度一条未扰动的参考轨迹一条视觉扰动的轨迹物体颜色突然改变以及一条物理扰动的轨迹一个或多个物体被传送到随机位置。传送违反了物理连续性并产生了明显的惊奇峰值而未扰动轨迹保持低基线。在所有三个环境中传送扰动的惊奇度显著更高配对t检验p0.01而对于立方体颜色扰动增加较弱且不显著表明模型对物理扰动比对视觉扰动更敏感。从左到右环境分别为TwoRoom、PushT和OGBench Cube。