Hamiltonian JEPA：具有继承式控制状态的动作条件世界模型

https://t.co/izoGpqysAw

来源： https://arxiv.org/html/2609.33497

作者： Ameen Ali、Lior Wolf，The Blavatnik School of Computer Science，Tel Aviv University 通讯信息： {tamimzoabi, ameenali023}@mail.tau.ac.il, [email protected]

摘要

仅凭像素进行规划，需要的不仅仅是一个稳定且可预测的潜空间。规划器所打分的状态，还必须按照动作如何驱动系统来组织。联合嵌入预测架构（Joint-embedding predictive architectures，JEPAs）通过预测未来表征来避免像素重建，但现有的动作条件 JEPA 要求同一个嵌入同时承担感知与控制两种职能。我们提出 H-JPA，将两者分离。一个较宽的感知编码会被 Bures-Wasserstein 先验正则化到尺度良好的各向同性几何上；而该编码的一个固定正交切片即为控制状态，它直接继承编码的协方差，自身不需要任何目标函数。该状态在相位条件的耗散端口哈密顿（port-Hamiltonian）动力学下演化，其输入端口具有正交列。端口逆一致性（Port-inverse consistency，PIC）通过该端口的转置把已执行的动作回读出来。我们证明这一读出正是沿端口方向投影的展开误差，因此 PIC 是对预测误差的一种无参数重加权，而不是额外的动作解码器。将读出与端口解绑会破坏这一恒等关系，并损失掉一半的增益。在至多 10¹⁰ 训练轮后，H-JPA 在四个基于像素的控制基准上达到或超过了无需重建的基线方法，包括采用动作解码的 Delta-JPA；其最大增益出现在 OGB-Cube 上（91.9 对 79.3，单位为百分比）。我们在 PushT 和 OGB-Cube 上的消融实验分离了结构化预测器、PIC、预测时程、状态秩以及防坍缩先验各自的贡献。

1 引言

能够预测自身动作后果的智能体才能进行规划：它在想象中对候选动作序列打分，并执行最优的那个（Ha and Schmidhuber, 2018 (https://arxiv.org/html/2609.33497#bib.bib17)；Hafner et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib18)）。在一个离线交互数据语料上训练后，基于所学世界模型的模型预测控制（model-predictive control）能够抵达训练过程中从未演示过的目标图像，无需奖励函数，也无需在测试时进行梯度更新（Zhou et al., 2025 (https://arxiv.org/html/2609.33497#bib.bib47)；Sobal et al., 2025 (https://arxiv.org/html/2609.33497#bib.bib36)）。规划的质量完全取决于模型对动作的响应。一个模型如果只捕捉到外观、却模糊了不同动作序列之间的差异，就会对所有方案给出相同的排序，因此动作敏感性是所学动力学的核心属性。

从像素学习此类模型时，一个关键问题是该预测什么。基于重建的模型回归未来帧，因而会把容量花在与控制无关的视觉细节上（Watter et al., 2015 (https://arxiv.org/html/2609.33497#bib.bib41)；Hafner et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib18)；Hafner et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib19)）。联合嵌入预测架构（JEPAs）转而预测未来表征（LeCun, 2022 (https://arxiv.org/html/2609.33497#bib.bib23)；Assran et al., 2023 (https://arxiv.org/html/2609.33497#bib.bib3)；Bardes et al., 2024 (https://arxiv.org/html/2609.33497#bib.bib6)），这移除了重建目标，但使训练暴露在表征坍缩（representation collapse）之下。LeWM（Maes et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib27)）针对动作条件世界模型解决了坍缩问题：它把 SIGReg（Balestriero and LeCun, 2025 (https://arxiv.org/html/2609.33497#bib.bib4)）——一种草图化的高斯正则化器——挂在规划器所打分的那个嵌入上。

本文的出发点是这样一个观察：稳定并不等同于控制。匹配边缘嵌入分布可以证明没有任何方向坍缩掉了，但它无法说明动作会推动哪些方向、哪些坐标应当进入规划度量，也无法说明控制需要多少容量。由此可推出三点结论。第一，LeWM 在完整嵌入空间中进行规划，因此用于候选方案排序的度量把与控制相关的变异和那些仅仅满足先验、却不提供任何动作选择信息的维度混在了一起。第二，正则化器与预测损失作用在同一组坐标上，因此每当其他任何压力发生改变时，二者的平衡都必须重新调节。第三，潜空间 MSE 目标存在尺度退化解，因为把目标缩小就能让预测变容易。Sub-JEPA（Zhao et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib45)）发现全空间先验对低维动力学过于严格，将其放松到随机子空间并取得了明显增益，但它的子空间只是供正则化器使用的视角，而非被提取出的状态。

H-JPA（Hamiltonian JEPA）将感知编码与控制状态分离开来，并让状态继承编码的几何结构。一个较宽的编码 (h\in\mathbb{R}^{D}) 通过逐样本的无仿射归一化被放置在球面上，并通过 Bures-Wasserstein 矩惩罚被正则化到各向同性（Gelbrich, 1990 (https://arxiv.org/html/2609.33497#bib.bib13)）。一个固定的正交投影 (U) 定义了状态 (s=U^{\top}h)，其协方差是编码协方差的固定函数，因此不需要额外的状态空间目标。在这个状态上，我们施加一个耗散端口哈密顿步进，其输入端口具有正交列。端口逆一致性（PIC）通过该端口的转置，从实际发生的位移中把已执行的动作回读出来——该转置正是其精确的伪逆。这一读出没有参数、没有自由尺度，且其残差等于沿端口方向投影的展开误差，因此 PIC 将预测误差重新加权，使其偏向动作实际作用的方向。在相同的规划协议下，H-JPA 在 Two-Room、Reacher、PushT 和 OGB-Cube 上达到或超过了重建、全空间先验以及动作解码三类基线方法，其中最大增益出现在 OGB-Cube 上（91.9 对最强基线的 79.3，单位为百分比）。

我们的贡献如下。

(i) 一个双空间世界模型：其中一个分布目标塑造感知，而一个固定的正交映射把协方差传递给规划器面对的状态，并给出了继承协方差的精确刻画（§3.1 (https://arxiv.org/html/2609.33497#S3.SS1) 至 §3.3 (https://arxiv.org/html/2609.33497#S3.SS3)）。

(ii) 对方差趋近于零时防坍缩惩罚的分析，表明 Bures-Wasserstein 惩罚在每个标准差上保持恒定的恢复力，而 SIGReg 与欧氏协方差惩罚都会失去这种力（§3.3 (https://arxiv.org/html/2609.33497#S3.SS3)）。

(iii) 一个相位条件的耗散端口哈密顿预测器，配有正交输入端口；以及 PIC，我们证明它是一种精确的误差重加权，并通过对照实验表明读出与端口之间的绑定关系解释了它一半的增益（§3.4 (https://arxiv.org/html/2609.33497#S3.SS4)、§3.5 (https://arxiv.org/html/2609.33497#S3.SS5)、§4.1 (https://arxiv.org/html/2609.33497#S4.SS1)）。

(iv) 在两个任务上、匹配表征条件下得到的规划、探针与消融结果，其中诊断实验表明规划器在被探针化物理变量之外的方向上对候选方案进行排序（§4 (https://arxiv.org/html/2609.33497#S4)）。

2 相关工作

JEPA 世界模型。 LeCun（2022）(https://arxiv.org/html/2609.33497#bib.bib23) 提出将表征空间中的预测作为世界模型的基础，并由 I-JEPA 和 V-JEPA（Assran et al., 2023 (https://arxiv.org/html/2609.33497#bib.bib3)；Bardes et al., 2024 (https://arxiv.org/html/2609.33497#bib.bib6)）加以验证。LeJEPA（Balestriero and LeCun, 2025 (https://arxiv.org/html/2609.33497#bib.bib4)）论证了各向同性高斯嵌入的合理性，并用 SIGReg 强制其实现。LeWM（Maes et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib27)）将这一研究路线带入了从像素出发的动作条件世界模型；Sub-JEPA（Zhao et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib45)）则在 (K) 个随机子空间中施加高斯约束（(K) 逐基准选取），从而改进了规划。Delta-JEPA（Zhang et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib44)）通过一个作用于潜变量差分的线性解码器引入动作敏感性；感觉运动世界模型（sensorimotor world models，Ivashkov et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib20)）则把逆动力学作为感知目标。

H-JPA 在三个方面与之不同。其子空间是一个单一的固定切片，而这个切片就是控制状态；先验保持在 LeWM 所设置的全空间位置上，状态则通过继承获得该几何；逆向读出正是前向模型自身的输入端口。

用于控制的潜空间动力学模型。 E2C、DVBF、World Models、PlaNet 和 Dreamer 通过像素重建学习潜空间动力学，并常常以预测（或奖励预测）作为目标（Watter et al., 2015 (https://arxiv.org/html/2609.33497#bib.bib41)；Karl et al., 2017 (https://arxiv.org/html/2609.33497#bib.bib21)；Ha and Schmidhuber, 2018 (https://arxiv.org/html/2609.33497#bib.bib17)；Hafner et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib18)；Hafner et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib19)）。DINO-WM 使用冻结的预训练编码器进行规划（Zhou et al., 2025 (https://arxiv.org/html/2609.33497#bib.bib47)），PLDM 则以 VICReg 风格的目标端到端训练一个像素 JEPA（Sobal et al., 2025 (https://arxiv.org/html/2609.33497#bib.bib36)）。H-JPA 既不以像素也不以奖励为目标，并且在没有教师、没有 stop-gradient（Grill et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib15)）的情况下端到端地训练其编码器。它的防坍缩机制属于 Barlow Twins 与 VICReg 一类的矩匹配惩罚（Zbontar et al., 2021 (https://arxiv.org/html/2609.33497#bib.bib42)；Bardes et al., 2022 (https://arxiv.org/html/2609.33497#bib.bib5)），其中欧氏协方差度量被高斯上的 Bures-Wasserstein 度量所取代（Bhatia et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib7)）。

结构化动力学。 哈密顿神经网络（Hamiltonian neural networks）学习一个标量能量，其梯度生成保守流（Greydanus et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib14)），并衍生出受控与广义变体（Zhong et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib46)；Course et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib11)；Chen et al., 2021 (https://arxiv.org/html/2609.33497#bib.bib8)）以及基于图像的生成形式（Toth et al., 2020 (https://arxiv.org/html/2609.33497#bib.bib39)）。该文献针对的是指定的物理状态或重建。H-JPA 则把端口哈密顿形式（van der Schaft and Jeltsema, 2014 (https://arxiv.org/html/2609.33497#bib.bib40)）作为针对规划而训练的、无需重建的表征的归纳偏置，并利用端口正交性消除一个原本会与表征目标发生相互作用的自由尺度。

逆动力学与状态发现。 从状态转移中恢复动作是获取与控制相关特征的经典途径（Agrawal et al., 2016 (https://arxiv.org/html/2609.33497#bib.bib1)；Pathak et al., 2017 (https://arxiv.org/html/2609.33497#bib.bib33)），DeepMDP 与双模拟（bisimulation）目标同理（Gelada et al., 2019 (https://arxiv.org/html/2609.33497#bib.bib12)；Zhang et al., 2021 (https://arxiv.org/html/2609.33497#bib.bib43)）。在外生分块 MDP 中，带有逐步逆预测的前向模型在给定假设下可以识别出一个控制内生状态（Lamb et al., 2023 (https://arxiv.org/html/2609.33497#bib.bib22)），并且已知逐步逆预测在何种情况下才充分（Levine et al., 2024 (https://arxiv.org/html/2609.33497#bib.bib24)）；最近的 JEPA 世界模型中也出现了位移层面的动作解码（Zhang et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib44)；Ivashkov et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib20)）。PIC 将读出绑定到前向输入映射的转置上。把读出绑定到输入映射的转置是语言模型嵌入中的标准做法（Press and Wolf, 2017 (https://arxiv.org/html/2609.33497#bib.bib34)），不过在一般情形下该转置并非该映射的逆（Gu et al., 2026 (https://arxiv.org/html/2609.33497#bib.bib16)）。在 H-JPA 中，端口被约束为正交，这就使其转置恰好成为该端口的精确伪逆。

瓶颈视角与信息瓶颈（information bottleneck）（Tishby et al., 2000 (https://arxiv.org/html/2609.33497#bib.bib38)；Alemi et al., 2017 (https://arxiv.org/html/2609.33497#bib.bib2)）、预测状态表征（predictive state representations）（Littman and Sutton, 2001 (https://arxiv.org/html/2609.33497#bib.bib25)）以及平衡截断（balanced truncation）（Moore, 1981 (https://arxiv.org/html/2609.33497#bib.bib30)）相关。

3 方法

表 1： 我们在研发过程中放宽过的各项设计选择、它们在该架构中分别导致的失败，以及最终所采纳的约束。这些观察构成本文 §3 (https://arxiv.org/html/2609.33497#S3) 中各项约束的动机来源，并与 §4.1 (https://arxiv.org/html/2609.33497#S4.SS1) 中的受控消融实验相互补充。

问题设定。 我们给定的是离线轨迹，包含 RGB 帧 (x_t) 与连续动作，没有奖励，也没有模拟器状态。目标是一个潜空间模型，规划器可以在其中展开候选动作序列，并依据其到某个已编码目标图像的距离对方案进行排序。一个编码器产生感知编码 (h_t)（§3.1 (https://arxiv.org/html/2609.33497#S3.SS1)）。该编码的一个固定线性切片即控制状态 (s_t)，也就是规划器唯一需要打分的量（§3.2 (https://arxiv.org/html/2609.33497#S3.SS2)）。一个分布先验作用于该编码，防止其坍缩（§3.3 (https://arxiv.org/html/2609.33497#S3.SS3)）。一个端口哈密顿预测器在动作驱动下推进该状态（§3.4 (https://arxiv.org/html/2609.33497#S3.SS4)）。训练中额外加入端口逆一致性（PIC）：要求预测器自身的动作输入映射能够从观测到的位移中恢复出已执行的动作（§3.5 (https://arxiv.org/html/2609.33497#S3.SS5)）。

一个模型步可以写成

(h_{t}=E(x_{t}),\qquad s_{t}=U^{\top}h_{t},\qquad v_{t}=s_{t}-s_{t-1},\qquad c_{t}=[s_{t},v_{t}],\qquad \hat{s}_{t+1}=\mathcal{F}_{\theta}(s_{t},\bar{a}_{t}\mid c_{t}).)

(1)

从左往右读式 1 (https://arxiv.org/html/2609.33497#S3.E1)，(E) 是图像编码器，(h_t\in\mathbb{R}^{D}) 是单帧 (x_t) 的编码。矩阵 (U\in\mathbb{R}^{D\times r}) 具有正交列，且从不参与训练；(s_t\in\mathbb{R}^{r})（其中 (r\leq D)）即为控制状态。由于 (s_t) 由单帧计算得到，不携带运动信息，因此我们构造有限差分速度 (v_t\in\mathbb{R}^{r})，并将其与状态拼接成相位 (c_t\in\mathbb{R}^{2r})（方括号表示拼接）。动作 (\bar{a}_t\in\mathbb{R}^{d_a}) 是施加在第 (t) 帧与第 (t+1) 帧之间的归一化控制量。带参数 (\theta) 的预测器 (\mathcal{F}_{\theta}) 在 (\bar{a}_t) 驱动下将 (s_t) 推进一步。记号 (\mathcal{F}_{\theta}(\cdot\mid c_t)) 表示预测器内部的向量场是相位的函数，而被推进的只有状态。速度由相邻状态重新计算得到，初始时令 (v_0=0)，即在起点处取零……