A long-standing challenge in AI is to develop agents capable of solving a wide range of physical tasks and generalizing to new, unseen tasks and environments. A popular recent approach involves training a world model from state-action trajectories and subsequently use it with a planning algorithm to solve new tasks. Planning is commonly performed in the input space, but a recent family of methods has introduced planning algorithms that optimize in the learned representation space of the world model, with the promise that abstracting irrelevant details yields more efficient planning. In this work, we characterize models from this family as JEPA-WMs and investigate the technical choices that make algorithms from this class work. We propose a comprehensive study of several key components with the objective of finding the optimal approach within the family. We conducted experiments using both simulated environments and real-world robotic data, and studied how the model architecture, the training objective, and the planning algorithm affect planning success. We combine our findings to propose a model that outperforms two established baselines, DINO-WM and V-JEPA-2-AC, in both navigation and manipulation tasks. Code, data and checkpoints are available at https://github.com/facebookresearch/jepa-wms.

翻译：人工智能领域长期存在的一个挑战是开发能够解决广泛物理任务并泛化至未见任务与环境的智能体。近期流行的方法涉及从状态-动作轨迹训练世界模型，随后结合规划算法解决新任务。规划通常在输入空间中进行，但近期一系列方法引入了在世界模型学习到的表示空间中进行优化的规划算法，其核心假设是通过抽象无关细节可实现更高效的规划。在本研究中，我们将此类模型定义为JEPA-WM，并深入探究该类别算法有效的技术选择。我们提出对多个关键组件进行系统性研究，旨在确定该框架内的最优方法。通过仿真环境与真实机器人数据的实验，我们分析了模型架构、训练目标与规划算法如何影响规划成功率。综合研究结果，我们提出的模型在导航与操作任务中均优于DINO-WM和V-JEPA-2-AC两个基准模型。代码、数据与模型检查点公开于https://github.com/facebookresearch/jepa-wms。