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.

翻译：人工智能领域长期存在的挑战是开发能够解决广泛物理任务并泛化至未见任务和环境的智能体。近期一种流行方法是从状态-动作轨迹中训练世界模型，随后结合规划算法解决新任务。规划通常在输入空间中进行，但最近一类方法引入了在世界模型的 learned 表示空间中优化的规划算法，其优势在于通过抽象无关细节实现更高效规划。本研究将这类模型定义为 JEPA-WMs，并探究使其生效的技术选择。我们提出对多个关键组件的综合研究，以寻找该类方法中的最优方案。实验采用仿真环境和真实机器人数据，分析了模型架构、训练目标和规划算法对规划成功的影响。结合研究发现，我们提出的模型在导航和操作任务中均优于 DINO-WM 和 V-JEPA-2-AC 两项基准。代码、数据和模型检查点见 https://github.com/facebookresearch/jepa-wms。