研究者经600次实验揭示JEPA世界模型的长任务规划能力

无需额外训练即可解决长时程任务

当前AI领域正围绕世界模型应该是什么展开激烈讨论。大部分注意力都集中在生成式世界模型上，这类模型通过生成未来（通常以像素形式）来预测未来。包括Yann LeCun在内的许多科学家认为这是徒劳的。大多数像素对于决定该做什么并不重要，而且很多像素根本无法预测。[1]

提出的替代方案是JEPA，即联合嵌入预测架构。它包含三个部分：编码器、预测器和规划器。编码器将每一帧观察转换为学习到的嵌入——一个描述场景的紧凑向量。预测器接收该嵌入加上一个动作，预测下一帧的嵌入。然后规划器在预测器内部尝试数百个动作序列，选出预测终态最接近目标嵌入的那个。整个循环在嵌入空间中运行，因此模型只需学习场景中对决定动作重要的部分。

今年，包括LeCun在内的一个团队发布了LeWM [2]，这是一个JEPA世界模型，其主要基准是Push-T [3]——一个仿真环境，机器人将一个灰色T形方块推过桌面，直到T形块停留在目标图像所示的绿色目标位置。我花了过去几个月在它上面运行了600次实验。

规划器收到这样一张图像，必须让场景与之匹配。灰色T是方块，蓝色圆盘是机器人的推杆，绿色T是画在桌面上的目标位置。这张图像中的一个细节对成功率的影响，竟然超过了用仿真物理替换预测器。

我首先想知道的是LeWM受其预测能力的限制有多大，因此我首先将预测器替换为模拟器的完美物理，同时保持其他一切不变。然后我给两个版本都设置了长时程任务，目标在75步之外，并配以相同的子目标：来自专家演示的帧，每25步一个。原版LeWM成功率为78%。使用真实物理的版本成功率为77%。在它训练时所涉及的那类路径上，LeWM的预测在规划方面已经和完美物理一样好。

这意味着JEPA世界模型的下一个提升来自于如何使用训练好的架构，而这正是我600次实验的大部分所在。

在没有对模型进行一次重新训练的情况下，我发现了四种新的规划方法。前两种方法合在一起将LeWM在长时程任务上的成功率从22%提升到85%，大约相当于一个被赋予专家演示理想子目标的规划器。另外两种方法将LeWM扩展到更难的目标，其中一种——仅仅改变了目标的评分方式——将成功率提升了60个百分点，是替换为完美物理所带来提升的三倍。

按照论文原样使用，LeWM在长时程任务（目标在75步之外）上的成功率为22%。[2] 当目标在25步之外时，它的成功率达到89.7%，与精确物理版本的90.5%相差不到一个百分点。

所以该模型在25步范围内是准确的，但在75步时就吃力了。通常的解释是小的预测误差在长规划中累积，这就需要更好的模型。我在同样的250个长时程任务上对此进行了对照测试。

250个长时程任务，目标在75步之外 精确物理增加了20个百分点。将任务拆分为25步的片段——每个片段都是LeWM已经几乎完美处理的长度——增加了近60个百分点。LeWM的有效规划视野约为25步，因此很明显，要走得更远，就需要将这种长度的片段串联起来。

LeWM的规划视野约为25步 每25步一个子目标……三次短搜索80% 通过子目标串联25步搜索，使每次搜索都保持在LeWM准确的范围内。 这就是分层规划：高层选择子目标，低层达成它们。这被广泛认为是世界模型的主要开放问题之一，而在Push-T上，它的价值是替换为完美物理的三倍。 额外的搜索不能替代它。将规划器运行四次并保留成本最低的计划，花费了四倍的计算量，却只增加了不到一个百分点，因为四次运行几乎总是一起成功或一起失败。 这个实验中的子目标来自专家演示，而部署的机器人不会有这种演示。第二种方法消除了这种依赖。 LeWM是在任务的专家演示上训练的，其编码器可以嵌入这些演示的每一帧，给出约460,000个嵌入状态，我将它们连接成一个图。每个状态链接到其自身演示中紧随其后的状态，以及嵌入空间中接近的其他演示中的状态。为了求解长时程任务，规划器在这个图中搜索从当前状态到目标的路径，并使用该路径上的状态作为子目标。 这只有在嵌入空间中的距离具有物理意义时才有效，而它确实有。取一个机器人能在25步内到达的子目标和一个它无法到达的子目标，LeWM的嵌入距离在87%的情况下将可达的那个排在更近的位置。因此嵌入空间同时充当了机器人可达范围的映射。 第一个版本有一个缺陷。在两个不同演示中邻近状态之间的跳转与真实步进具有相同的成本，因此搜索将跳转串联成方块永远无法跟随的路径（1,015个尝试的计划中有318个使用了六次或更多跳转）。修复方法是在图中每次演示之间的跳转上增加成本。仅使用LeWM自身对过去演示的嵌入，在同样的250个长时程任务（目标在75步之外）上，成功率从62%提升到85%，这与使用专家子目标的规划器得分相同。这已经好得不能再好了，对吧？

通过演示记忆进行规划 起点目标 跳转无成本 · 8次跳转 · 62% 每次跳转有成本 · 1次跳转 · 85%

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）