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 获取。