TensorX

TensorX · Paper Blog

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

Xin Zhang · 2026-10-05

World ModelsJEPAVisual PlanningHierarchical Reinforcement LearningRobot Learning

本文提出了首个端到端无重构、无奖励的分层JEPA世界模型H-JEPA,通过在多时间尺度上自发抽象出由粗到细的独立潜空间,成功将长程规划成功率提升数倍并大幅削减测试期算力开销。

💡 核心创新点

传统的视觉规划任务中,单时间尺度潜空间世界模型(如单层 JEPA 或基于像素重构的扩散/自回归模型)面临着根本性的结构瓶颈。H-JEPA(Hierarchical Joint-Embedding Predictive Architecture) 打破了**“所有时序步长共享单一扁平潜空间”**的经典假设,首次提出了端到端联合训练的分层非重构世界模型。

独立潜空间与时间抽象联合演化

以往的分层架构(例如 HWM)虽然在规划时跨越不同时序跨度,但高层与底层预测器依然共用相同的特征空间,导致高层表征被迫保留底层的细节高频噪声。H-JEPA 的核心机制差异在于:每一个层级都拥有独立的编码器、动作编码器与自回归时序预测器。高层模型直接输入低层的潜状态窗口,并以粗粒度步长(Stride)向前预测自身层级所对应的长程潜空间状态。

分层抽象支持多时间尺度规划,在Visual AntMaze上显著提升成功率并节省算力

预测瓶颈诱导的选择性特征丢弃

这种机制从根本上改变了表征学习的动力学。JEPA 依赖未来预测误差的反向传播来优化编码器。当高层预测器跨越更大的环境步长时,不可预测的局部高频细节(如四足机器人的腿部高频抖动、背景随机传送的干扰物)会在高层编码器中被自发过滤,仅保留全局位置和地图拓扑等慢速变化的核心状态。

高层表征自发过滤快速高频干扰,保留低频核心状态

🚀 重大突破:在 Visual AntMaze 复杂导航任务中,H-JEPA 将规划成功率从扁平单层模型的 18% 跃升至 73%,同时消耗的规划器测试期浮点计算量(FLOPs)降低了一个数量级。


🔍 背景与动机

在基于视觉的长程机器人操作与导航中,直接在像素空间或单尺度潜空间中进行轨迹优化存在两大行业痛点:

  1. 误差复合爆炸:长程规划若使用精细的步长逐步展开,单步累积的预测误差会随着展开步数指数级发散,使长程动作搜索完全脱轨。
  2. 目标匹配与低级控制的表征冲突:如果潜空间充满了密集的微观细节(例如机械臂每个关节的高频振动或纹理细节),基于欧氏距离的目标函数将充斥着局部极值与平坦区域,无法为长程目标提供平滑的梯度引导。

一个生活化的比喻是:人类在规划从纽约前往旧金山的旅程时,大脑绝不会在神经元中以毫秒为单位去推演双腿迈步的微小动作。高层规划只关注城市级别的航线(跨度数千公里),中层规划关注交通工具切换,只有底层的脊髓和小脑网络才以毫秒级频率生成具体的肌电控制信号。

数据集特征频率分离比与选择性抽象程度的正相关关系

如图所示,只有当环境中各状态因子的特征频率存在至少 2× 以上的时间尺度分离时,多层预测模型才能自发涌现出这种层次分明的抽象能力。


📊 数据集与实验环境

本文在四大模拟环境与一个大规模真实机器人操控数据集上进行了全面评估:

实验环境配置详情

  • FourRoomDistractors:连续 2D 四房间导航任务。观测分辨率为 224×224,动作维度为 2D 速度命令。环境内包含一个以泊松分布(平均间隔 35 步)随机传送的干扰球。训练集包含 3750 条轨迹(共 150 万步转移),单层训练使用 100 万步转移。
  • Visual AntMaze:基于 OGBench 的中等难度四足蚂蚁迷宫任务。观测包括 64×64 RGB 第三视角图像和 27 维本体感受状态(关节角度与速度),动作输出为 8 维扭矩控制。训练集采用 explore 与 stitch 策略各 12500 条轨迹的等比例混合,共计 1005 万步转移。
  • OGBench Cube:机械臂抓取与放置立方体任务。图像为 224×224,动作维度为 5D(末端位移、偏航角与夹爪)。训练集包含 9800 条脚本轨迹(共 196 万步转移)。
  • Push-T:连续 2D 目标推块操作。图像输入 224×224,动作维度为 2D 智能体目标位移。训练集包含 18685 条专家轨迹(共 234 万步转移)。
  • DROID 真实机器人操作数据集:包含 74530 个真实场景演示视频(220084 个视角,256×256 分辨率),动作为 7 维末端执行器相对位姿变化量(xyz、姿态角与夹爪)。

训练超参数与硬件规格

  • 硬件环境:所有模型均在 NVIDIA H200 GPU 集群上训练,采用 bfloat16 混合精度。
  • 优化器与调度:AdamW,权重衰减 1e-3(DROID 上为 1e-4),梯度裁剪阈值 1.0。
  • 学习率:底层编码器峰值学习率为 5e-5,高层为 1e-5(DROID 任务两层统一为 5e-4),采用 1% 步数线性预热结合余弦退火。
  • 批量大小:模拟任务批次大小为 128 条轨迹片段,DROID 批次大小为 256。

🛠️ 核心机制与研究方法

H-JEPA 包含由下至上的多层感知编码与由上至下的分层引导规划两大核心支柱。

H-JEPA的多时间尺度架构与端到端训练流程

模块一:时序分层表征构建与跨层池化机制

设环境原始观测为 oto_t,原始基元动作为 ata_t。第一层(底层)JEPA 包含基于 ViT-Tiny 的图像编码器 E(1)E^{(1)},将输入图片编码为全局特征向量 zt(1)=E(1)(ot)z_t^{(1)} = E^{(1)}(o_t)。对于上层 ℓ>1\ell > 1,模型通过跨层聚合建立粗时间尺度:

zt(ℓ)=E(ℓ)(zt⋅sℓ:t⋅sℓ+wℓ(ℓ−1)),at(ℓ)=A(ℓ)(at⋅sℓ+wℓ−1:(t+1)⋅sℓ+wℓ−1(ℓ−1))z^{(\ell)}_t = E^{(\ell)}\left(z^{(\ell-1)}_{t\cdot s_\ell : t\cdot s_\ell+w_\ell}\right), \quad a^{(\ell)}_t = A^{(\ell)}\left(a^{(\ell-1)}_{t\cdot s_\ell+w_\ell-1 : (t+1)\cdot s_\ell+w_\ell-1}\right)

在公式 (1) 中,sℓs_\ell 表示上层的降采样步长(Stride),wℓw_\ell 为时序窗口跨度。状态编码器 E(ℓ)E^{(\ell)} 将下层的窗口表征池化为抽象状态;动作编码器 A(ℓ)A^{(\ell)} 则聚合跨越 sℓs_\ell 步的底层动作,形成高阶宏观动作(Macro-Action)。在本文的标准实验中,所有高层状态编码器均取 wℓ=1w_\ell = 1,即点对点抽象映射。

模块二:端到端非坍塌预测目标与 SIGReg 正则化

每一层 ℓ\ell 均配备一个自回归 Transformer 预测器 F(ℓ)F^{(\ell)}。在给定前序上下文潜状态及对应的动作嵌入序列时,通过单步教师强制(Teacher-forced)损失预测未来的潜状态:

Lpred(ℓ)=1cℓ∑τ=1cℓ∥F(ℓ)(z<τ(ℓ),a<τ(ℓ))−zτ+1(ℓ)∥22\mathcal{L}_{\text{pred}}^{(\ell)} = \frac{1}{c_\ell} \sum_{\tau=1}^{c_\ell} \left\| F^{(\ell)}(z^{(\ell)}_{<\tau}, a^{(\ell)}_{<\tau}) - z^{(\ell)}_{\tau+1} \right\|_2^2

在公式 (2) 中,cℓc_\ell 为各层预测器的上下文长度。为了防止非对比自监督学习中潜向量坍塌为常数的平凡解,H-JEPA 在每层引入了 SIGReg(Sketched Isotropic Gaussian Regularizer)规范化项:

L(ℓ)=Lpred(ℓ)+λℓSIGReg(Z(ℓ))\mathcal{L}^{(\ell)} = \mathcal{L}_{\text{pred}}^{(\ell)} + \lambda_\ell \text{SIGReg}\left(Z^{(\ell)}\right)

在公式 (3) 中,Z(ℓ)Z^{(\ell)} 代表展平批次与时序维度的潜嵌入矩阵,λℓ\lambda_\ell 为各层的正则化惩罚系数。该正则化项利用经验特征函数与标准高斯分布的距离,在无需负样本或显式重构的前提下强制潜空间各向同性,保证了特征梯度的稳定性。

自顶向下的分层规划流程,下层受上层预测子目标的引导

模块三:自顶向下的跨抽象层梯度规划算法

在推理阶段,给定初始帧 o0o_0 与目标帧 ogo_g,各层编码器自底向上计算各层初始状态 z0(ℓ)z_0^{(\ell)} 与目标状态 g(ℓ)g^{(\ell)}。规划过程自顶向下递推:最顶层 LL 在宏观时间跨度上利用梯度下降直接优化宏观动作序列,最小化其终点与抽象目标的欧氏距离:

a0:HL−1(L),∗=arg⁡min⁡a0:HL−1(L)∥z^HL(L)−g(L)∥22a^{(L),*}_{0:H_L-1} = \arg\min_{a^{(L)}_{0:H_L-1}} \left\| \hat{z}^{(L)}_{H_L} - g^{(L)} \right\|_2^2

在公式 (4) 中,z^HL(L)\hat{z}_{H_L}^{(L)} 是顶层预测器在其规划时域 HLH_L 内展开的最终预测状态。优化完成后,顶层轨迹被用作下层的子目标(Subgoals)。下层 ℓ\ell 将其展开预测状态提升(Lift)至上一层的潜空间:

z~i(L+1)=E(ℓ+1)(z^isℓ+1:isℓ+1+wℓ+1(ℓ)),i=1,…,Kℓ+1\tilde{z}^{(L+1)}_i = E^{(\ell+1)}\left(\hat{z}^{(\ell)}_{is_{\ell+1} : is_{\ell+1}+w_{\ell+1}}\right), \quad i = 1, \dots, K_{\ell+1}

随后,下层 ℓ\ell 通过最小化提升状态与上层规划子目标之间的加权距离来求解其动作:

a0:Hℓ−1(ℓ),∗=arg⁡min⁡a0:Hℓ−1(ℓ)[∥z~Kℓ+1(ℓ+1)−z^Kℓ+1(ℓ+1),∗∥22+β∑i=1Kℓ+1−1∥z~i(ℓ+1)−z^i(ℓ+1),∗∥22]a^{(\ell),*}_{0:H_\ell-1} = \arg\min_{a^{(\ell)}_{0:H_\ell-1}} \left[ \left\| \tilde{z}^{(\ell+1)}_{K_{\ell+1}} - \hat{z}^{(\ell+1),*}_{K_{\ell+1}} \right\|_2^2 + \beta \sum_{i=1}^{K_{\ell+1}-1} \left\| \tilde{z}^{(\ell+1)}_i - \hat{z}^{(\ell+1),*}_i \right\|_2^2 \right]

在公式 (6) 中,Kℓ+1K_{\ell+1} 为子目标数量,β≥0\beta \ge 0 调节中间子目标与终点目标的相对权重。这一优化逐层向下递归,直到底层输出环境可执行的基元动作。

模块四:抑制真实场景慢特征坍塌的逆动力学建模(IDM)

在 DROID 等背景高度多变的多场景数据集中,纯自监督预测模型极易遭遇慢特征坍塌(Slow-Feature Collapse):编码器倾向于只编码跨帧不变的静态背景(因为背景极度可预测),而丢弃受动作控制的微小机械臂末端。为解决这一难题,模型在各层引入了逆动力学建模分支:

Lidm(ℓ)=1T−1∑t=1T−11da∥I(ℓ)(zt(ℓ),zt+1(ℓ))−sg[at(ℓ)]∥22\mathcal{L}_{\text{idm}}^{(\ell)} = \frac{1}{T-1} \sum_{t=1}^{T-1} \frac{1}{d_a} \left\| I^{(\ell)}\left(z^{(\ell)}_t, z^{(\ell)}_{t+1}\right) - \text{sg}\left[a^{(\ell)}_t\right] \right\|_2^2

在公式 (7) 中,I(ℓ)I^{(\ell)} 为逆动力学网络头,dad_a 为动作维度,sg[⋅]\text{sg}[\cdot] 表示梯度截断。强制两步潜状态反推中间动作,为编码器保留受控主体提供了理论上的下界保护。

💡 关键技巧:在多层架构中,高层逆动力学的拟合目标是池化后的宏观动作 at(ℓ)a_t^{(\ell)},切断该目标的梯度可防止高层动作表示退化。


📈 结果与深度分析

模拟环境的规划性能与算力对比

在四大模拟任务中,H-JEPA 与单层扁平基线 LeWM 以及单潜空间分层基线 HWM 进行了系统对比:

跨算力预算与分层深度的梯度规划性能对比

  • 帕累托效率飞跃:在 Visual AntMaze、FourRoom 和 Cube 任务中,增加层级数量持续将“成功率-算力”前沿推向左上方。在 Visual AntMaze 上,两层 H-JEPA 成功率达到约 40%,三层 H-JEPA 达到 73.3%,而单层 LeWM 的成功率仅为 18.0%。
  • 超越单潜空间分层:相比在单一潜空间内分层的 HWM,H-JEPA 在 AntMaze 上的成功率几乎是 HWM 的两倍,证明了分层独立潜表征带来的抽象增益远超单纯的时序切分。

消融实验:目标代价函数几何结构改变

为了解耦“时序解耦”与“抽象代价空间”对规划的独立贡献,作者进行了仅改变评估代价空间的消融实验:保持底层单层动力学与规划器不变,仅仅将终点距离评估提升至高层潜空间。

实验表明,在 AntMaze 中,单层模型使用投影至第二层和第三层的代价函数后,成功率从基准的 16.7% 分别跃升至 20.7% 和 22.0%。高层潜空间有效消除了迷宫局部死胡同的伪极值,形成了沿通道单调上升的平滑梯度。

DROID 真实环境拓展与失败分析

DROID数据集上的逆动力学必要性消融与规划保真度

在 DROID 真实机器人操作中,评估使用路径 Fréchet 保真度(相对于零动作基准):

  • 缺乏 IDM 损失的单层模型由于慢特征坍塌,保真度归零(低于静止基线);引入 IDM 后,单层保真度恢复至 +34%。
  • 结合两层 H-JEPA 结构后,保真度进一步提升至 +40%,并在规划算力上比单层模型节约数倍。

⚠️ 局限性:在动作序列极短的环境(如 Push-T,平均轨迹仅 125 步)中,随着时序跨度加倍,高层模型可提取的有效训练片段急剧萎缩(四层模型仅能利用 14% 的数据转移),导致深层模型性能退化。此外,当前模型未融合机器人腕部相机视点,对细微抓取接触的建模仍显不足。

未来工作将聚焦于将 H-JEPA 扩展到无动作标签的互联网视频预训练、结合自然语言指定高层子目标,以及在变长自适应时序步长(Adaptive Striding)下构建动态层次世界模型。针对实体物理机器人的闭环控制部署亦是极具前景的落地路线。针对复杂长程任务,H-JEPA 成功证明了自监督分层世界模型在统一时间抽象与低级控制上的巨大潜力。

北京市昌平区探索星信息技术及软件开发工作室

京ICP备2026059466号