H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

Embodied AI and Robotics World Models Task Decomposition ML RL SSLRO
利用潜在世界模型进行长时程规划,需要在不同时间尺度和抽象层级之间进行推理。现有的、任务无关的JEPA(联合嵌入预测架构)世界模型,要么仅在单一时间尺度上进行预测与规划,要么虽支持多时间跨度,却仍共享同一个潜在空间。我们提出H-JEPA——一种端到端训练分层化动作条件JEPA模型的完整方案:该分层结构中,每一层级均在其各自学习所得的潜在空间中,预测更远期的状态演化。规划过程采用自顶向下的方式进行:顶层负责优化朝向最终目标的整体进展;而每一层级所生成的预测结果,则作为其下一层级规划器的子目标。当数据中蕴含的因素以显著分离的时间尺度演化时,高层级会主动舍弃快速变化且难以预测的细节,仅保留与任务相关、演化较慢的状态信息。在四个模拟的导航与操作环境中,分层式规划性能均优于扁平化(单层)JEPA模型;尤其在Visual AntMaze任务中,采用三层分层结构后,任务成功率从18%大幅提升至73%,同时所消耗的规划器计算资源反而更低。消融实验表明,上述性能提升主要归因于两方面:一是时间维度上的分解能力,二是高层级所生成的目标表征更具语义抽象性与任务导向性。此外,在逆动力学监督(inverse-dynamics supervision)下,该方法可拓展至DROID数据集中的多种真实机器人视频,同样展现出分层结构在降低规划器计算开销的同时,显著提升离线规划保真度的优势。
Long-horizon planning with latent world models requires reasoning across timescales and levels of abstraction. Existing task-agnostic JEPA world models predict and plan at a single timescale or with multiple horizons in one shared latent space. We introduce H-JEPA, an end-to-end recipe for training a hierarchy of action-conditioned JEPAs in which each level predicts farther ahead in its own learned latent space. Planning proceeds top-down: the top level optimizes progress toward the goal, and each level's predictions become subgoals for the planner below it. When factors in the data evolve at separated timescales, higher levels discard fast, unpredictable detail and retain slower task-relevant state. Across four simulated navigation and manipulation environments, hierarchical planning improves over a flat JEPA; on Visual AntMaze, a three-level hierarchy raises success from 18% to 73% using less planner compute. Ablations attribute these gains to both temporal decomposition and higher-level goal representations. With inverse-dynamics supervision, the approach extends to diverse real-robot videos from DROID, where hierarchy improves offline planning fidelity at lower planner compute.
许愿