Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
一句话总结
Delta-JEPA 在基于 JEPA 的潜在世界模型中引入了潜在差分动作解码器(LDAD),通过相邻观测的潜在表示之差重建实际执行的动作,要求潜在变化保留动作信息,从而抑制表示坍塌,并鼓励从同一状态出发的不同动作产生可区分的潜在位移。与此同时,潜在前向动态预测器根据当前表示和动作预测下一时刻的表示,为规划提供轨迹展开模型。两个目标共同进行端到端训练,无需重建观测像素,使模型学习到更适合规划、对动作更加敏感的潜在动态。
摘要
现有问题:学习用于规划的视觉世界模型,需要紧凑且对动作保持敏感的潜在动态表示。然而,无需重建的联合嵌入目标可能导致表示坍塌,使其失去对动作的敏感性。
Delta-JEPA,一种端到端、无需重建观测的世界模型,通过引入潜在差分动作解码器(Latent Difference Action Decoder,LDAD),增强潜在空间中的前向预测。LDAD 根据连续观测之间的潜在位移,重建实际执行的动作。
这种位移层面的监督直接约束了状态转移的几何结构:相邻观测的嵌入若发生坍塌,就会丢失动作信息;同时,该监督鼓励不同动作产生可区分的潜在变化,从而支持基于轨迹展开的规划。
Delta-JEPA 仅使用潜在表示预测和动作重建,无需像素重建或分布匹配正则项。
引言
联合嵌入预测架构(JEPA)直接预测紧凑的未来表示,而非未来的像素,因此为潜在空间中的世界建模提供了一个尤其有吸引力的基础。 这种高效性也带来了一个重大挑战:如果仅使用潜在表示预测目标进行端到端训练,基于 JEPA 的世界模型很容易坍塌为平凡的常量表示。在这种情况下,模型虽然取得了看似很低的预测损失,却破坏了规划所需的表示结构。
提出 Delta-JEPA:一种以潜在差分动作解码器(Latent Difference Action Decoder,LDAD)为核心的端到端潜在世界模型。LDAD 不再通过拼接潜在状态来重建动作,而是根据潜在差分 Δzt=zt+1−zt 预测实际执行的动作。
这种位移层面的逆向目标能够促进对动作敏感的潜在动态,而这对于规划至关重要: 如果从同一个潜在状态出发,不同动作产生的下一时刻嵌入无法区分,那么世界模型就无法表示由动作控制的下一状态转移,潜在空间中的轨迹展开也就无法为规划提供有效信息。 如果从同一状态出发,不同动作能够产生可区分的下一状态嵌入,那么潜在表示就具有更好的可控性。
通过要求从 Δzt 中恢复动作,LDAD 鼓励不同动作产生可区分的潜在位移和下一状态嵌入,同时抑制动作预测依赖特定状态的线索,而非状态转移本身。
Delta-JEPA 采用一个简单的双目标方案来训练这一机制:潜在空间中的前向预测负责建模给定动作下的未来表示,LDAD 则使动作引起的潜在位移能够用于预测对应动作。 规划需要通过潜在空间中的轨迹展开来评估候选动作序列,模型必须能够区分不同动作会如何推动环境演化。
本文的主要贡献如下:
- 对动作敏感的潜在动态: 我们提出 LDAD,这是一种基于位移的逆向目标,通过约束潜在状态转移在动作上具有可区分性,缓解表示坍塌。
- 双目标训练框架: 我们构建了 Delta-JEPA,一种仅使用潜在空间前向预测和基于 LDAD 的动作重建进行训练的端到端潜在世界模型。
- 实验验证: 我们在多种连续控制任务上评估了 Delta-JEPA,结果表明,它在提升规划性能的同时,也学到了对动作更加敏感的潜在动态。
方法
问题定义
遵循无监督潜在世界模型的标准范式,关注无奖励、离线设置下的世界模型学习问题。给定一个离线数据集:
D={(o1,a1,…,oT)},
其中包含由高维原始图像观测 ot∈RC×H×W 和连续动作 at∈Rda 交替组成的轨迹。关键在于,D 不包含任何任务特定的奖励信号,且由任意的、未知的行为策略采集而成。
目标是在不重建像素、不使用任务奖励的情况下,学习一个紧凑的潜在表示空间 Z⊆Rd,以及一个对动作敏感的潜在动态预测器。
Delta-JEPA 概述
图1:Delta-JEPA 框架概览。原始观测 ot 和 ot+1 通过共享编码器映射为潜在表示 zt 和 zt+1。在前向路径中,动力学预测器根据 zt 和动作 at 预测后续表示 z^t+1,并由预测损失 Lpred 引导。同时,潜在差分动作解码器接收潜在位移 Δzt 以重建动作 a^t,并由动作损失 Laction 监督。这种基于位移的动作监督促使由动作引起的潜在差异具有可区分性,整个框架通过 L=Lpred+λLaction 进行端到端优化。
如图 1 所示,Delta-JEPA 包含两个相互耦合的目标。 潜在前向动态预测器学习根据当前表示和动作预测下一时刻的表示,为规划提供所需的轨迹展开模型。 潜在差分动作解码器(LDAD)则对相邻潜在状态之间的位移施加逆向动力学约束,要求能够从该位移中恢复引发这次状态转移的动作。
这两个目标共同训练出一个端到端、无需重建观测的世界模型,抑制表示坍塌为对动作不敏感的形式,并促进对动作敏感的下一状态预测。
潜在前向动态预测器
编码器 fθ 将每个观测 ot 映射为潜在表示 zt=fθ(ot)。以 zt 和动作 at 为条件,动态预测器 Pϕ 估计下一潜在状态:
z^t+1=Pϕ(zt,at).(1)
其中,z^t+1 表示预测的下一潜在状态。
使用潜在空间中的均方预测损失来训练编码器和预测器:
Lpred=∥z^t+1−zt+1∥22,(2)
其中,zt+1=fθ(ot+1) 是由同一个编码器生成的目标表示。
虽然式(2)能够实现无需重建观测的动态学习,但单独使用这一目标会产生退化解:编码器和预测器可以通过坍塌为近乎恒定的表示来降低损失。 即使预测损失很小,这样的解也几乎无法保留规划所需的信息。LDAD 通过对相邻潜在状态之间的差分施加以动作为依据的约束,来解决这一失效问题。
潜在差分动作解码器
LDAD 对相邻潜在状态之间的差分施加逆向动力学约束。给定两个编码后的观测 zt 和 zt+1,我们将潜在位移定义为:
Δzt=zt+1−zt.(3)
随后,解码器根据该位移预测实际执行的动作:
a^t=DΘ(Δzt),(4)
其中,DΘ 表示动作解码器,a^t 表示预测动作。解码器使用均方动作重建损失进行端到端训练:
Laction=∥a^t−at∥22.(5)
动作监督的位移机制
图2:LDAD 诱导的动作敏感潜在几何示意图。在没有位移级动作监督的情况下(左上),从同一潜在状态 zt 出发的不同动作可能产生相似的下一嵌入。LDAD 计算每个位移 Δzt(i)=zt+1(i)−zt,解码动作 a^t(i),并用 Laction=∥a^t−at∥22 对其进行监督(下)。这促使以动作为条件的转移在潜在空间中占据可区分的方向和终点(右上)。
图 2 所示,左上图展示了一种对动作不敏感的潜在空间几何结构:从同一个 zt 出发,不同动作可能产生彼此接近的下一时刻嵌入,使得这些潜在状态转移难以根据动作加以区分。
LDAD 通过下方所示的解码流程来解决这一问题。对于每次状态转移,它计算位移 Δzt(i)=zt+1(i)−zt,预测对应动作 a^t(i),并以实际执行的动作 at(i) 为目标优化重建损失。 由于解码器只能观测到位移,要成功恢复动作,局部状态转移的几何结构就必须编码实际执行的动作,从而促使不同动作引起的位移变得可区分。
右上图展示了这种监督的预期效果:不同动作产生彼此分离的状态转移方向和下一时刻嵌入。 这种几何结构对于规划尤为重要。当不同候选动作产生相似的潜在终点时,轨迹展开几乎无法为比较这些动作的后果提供依据,因此可能导致规划器选择效果难以区分或错误的动作。 相反,彼此分离的动作条件状态转移,使候选轨迹展开更能体现动作的控制作用,并为规划提供更多信息。
图5:两条 Two-Room 潜在轨迹的 PCA 可视化,二者初始状态相近但终点不同,按训练轮次分别展示 Delta-JEPA(上)和 LeWM(下)。蓝色和橙色表示两条轨迹,颜色深浅表示从早期状态(浅色)到后期状态(深色)的时间进展。
图6:Two-Room 上动作条件预测器响应的 PCA 可视化。我们采样 512 个起始历史,并在将最后一个动作替换为每个候选动作时保持历史表示 zt 固定。对于每个候选动作,我们可视化其相对于零动作预测的预测位移。每个半透明点对应一个起始历史在一个候选动作下的响应,每个编号标记表示该候选动作在全部 512 个历史上的平均响应。
基于位移的动作解码的作用
这种基于位移的逆向目标从三个方面影响所学习的表示:
- 抗表示坍塌。 动作重建目标抑制编码器将连续观测映射为近乎相同的潜在向量。如果相邻观测的表示发生坍塌,那么 Δzt 就无法提供有效信息,DΘ 也就无法恢复实际执行的动作。
- 减少对单一状态线索的依赖。 标准的逆向动力学解码器通过拼接潜在状态来预测动作,即 a^t=DΘ([zt,zt+1])。这种形式可能允许模型走捷径:由于前向预测器在预测 z^t+1 时接收了 at,学习到的目标表示 zt+1 可能包含与动作相关的线索,使逆向解码器无需充分建模状态转移本身,就能恢复 at。LDAD 通过仅以相对位移 Δzt 作为解码器的输入来降低这一风险,使动作重建必须依靠相邻潜在状态之间的变化,而非特定状态的线索。
- 为规划提供对动作敏感的潜在动态。 对于规划而言,潜在表示必须支持以动作为条件的潜在轨迹展开。LDAD 鼓励从同一个潜在状态出发的不同动作产生可区分的潜在位移和下一状态嵌入。因此,可以通过候选动作各自产生的不同潜在轨迹来比较这些动作,为动作选择提供更有信息量的预测。
多步动作解码
使用 Transformer 骨干网络实现 DΘ,并将 LDAD 扩展为多步动作解码,以捕捉更长时间范围内的时序结构。给定时间跨度 N≥1,解码器使用长时间跨度的潜在位移,重建从 t 到 t+N 这一时间区间内的动作序列:
{a^τ}τ=tt+N−1=DΘ(zt+N−zt).(6)
多步 LDAD 动作解码器采用带有 N 个可学习动作查询的 Transformer。位移 zt+N−zt 通过自适应层归一化(Adaptive Layer Normalization,AdaLN)注入每个查询,随后由 Transformer 各层输出 N 个重建的连续动作。这一多步扩展在潜在空间中更长的时间区间上施加了以动作为依据的动态约束。
联合优化与端到端训练
将框架的整体训练目标定义为由前向预测损失和动作重建损失组成的联合损失:
L=Lpred+λLaction,(7)
其中,λ>0 是用于平衡两项损失的超参数。
Delta-JEPA 仅使用两个目标:潜在表示预测学习以动作为条件的动态,动作重建则使局部潜在状态转移对动作保持敏感。该方法无需冻结编码器、停止梯度传播的分支或分布匹配正则项。
实验
主要围绕三个问题展开:Delta-JEPA 能否改善规划、改进是否来自差分动作解码、学到的潜在变化是否具有动作和物理意义
实验设置
任务: 四个视觉连续控制环境,覆盖导航(Two-Room)、机械臂到达(Reacher)、接触式推动(Push-T)和三维物体操作(OGB-Cube)。 训练与比较: 使用无奖励离线数据,对比 LeWM、Sub-JEPA 和 PLDM;各方法均从头训练 50 个 epoch,主结果使用 3 个随机种子。 模型与评价:编码器和预测器沿用 LeWM 的架构及评价协议,主要考察规划成功率;Delta-JEPA 额外加入动作解码器,默认重建 5 步动作,动作损失权重 λ=10。
实验设计与结果
| 实验 | 设计与原因 | 结果及含义 |
|---|---|---|
| 规划性能比较 | 在四个任务中与基线比较,检验表示学习的改进能否转化为实际控制收益。 | 四个任务的平均成功率均最高。相对各任务最强基线,Two-Room、Push-T、OGB-Cube 分别提高约 6.27、4.54、15.14 个百分点;Reacher 仅小幅领先,不能据此强调明显优势。 |
| 动作损失权重消融 | 在 Push-T 中改变 λ,检验动作重建是否必要,以及监督强度的影响。 | 去掉动作重建时模型近乎坍塌、规划几乎失败;权重过小也无效,适中时表现较稳定,过大则下降。该实验中 λ=50 最好,说明需要平衡两个目标。 |
| 差分输入 vs. 状态拼接 | 仅将解码器输入从 [zt,zt+1] 换成 Δzt,隔离“使用差分”本身的作用。 | 差分输入在四个任务中均更好,Push-T 提升最大(12.60 个百分点)。这是支持论文核心设计最直接的消融证据。 |
| 解码目标消融 | 在 Reacher 中,将动作替换为末端位置变化、关节位置变化或两者拼接,检验收益是否只是来自预测任意状态变化。 | 原始动作最好,关节位置变化接近;末端位置变化较差,拼接也没有改善。说明监督信号与可控动态的关联,比增加监督信息更重要。 |
| 潜在空间是否坍塌 | 对不同训练阶段的 Push-T 表示做 PCA,观察表示是否逐渐形成结构。 | 表示从集中分布逐渐扩展,并形成更清晰的结构,为缓解坍塌提供定性支持。 |
| 动态是否对动作敏感 | 一方面观察初始状态相近、终点不同的轨迹;另一方面固定历史表示、只改变动作,直接检查预测器响应。 | Delta-JEPA 的轨迹随时间逐渐分离;不同动作产生方向更清晰的预测变化,较大动作通常对应较大变化。LeWM 的动作响应则更集中、重叠。后一个实验更直接支持“预测器使用了动作信息”。 |
| 物理状态与状态变化探测 | 冻结编码器,用线性/MLP 探针分别从 zt 预测真实状态、从 Δzt 预测真实状态变化,检查表示是否保留物理信息。 | Two-Room 上两类探测均明显领先;附录中,OGB-Cube 的多项状态变化探测也表现突出。但其他任务存在反例,不能概括为所有物理属性都更好。 |
| 注意力可视化 | 查看编码器是否关注任务相关实体,辅助解释表示学到了什么。 | 注意力集中于智能体、T 形物体、方块及夹爪等区域;不同层关注的实体有所不同。说明模型在没有像素重建和物体监督时,也能形成任务相关的视觉关注。 |
总结
提出了 Delta-JEPA,一种无需重建观测的潜在世界模型。该模型通过潜在差分动作解码,直接对潜在位移中的动作信息施加监督。通过从 Δzt=zt+1−zt 中重建动作,Delta-JEPA 能够抑制表示坍塌,并鼓励不同动作产生可区分的潜在状态转移,以支持规划;同时,它保留了一个仅由潜在表示预测和动作重建组成的简单训练目标。 对潜在差分施加监督,是学习紧凑、抗表示坍塌且适用于规划的世界模型的一项有效原则。
局限性
- 状态变化不一定能唯一确定动作,多步解码尤其如此
LDAD 要求从 Δzt 恢复动作,但现实中可能存在不同动作产生相同观测变化的情况,例如受到障碍阻挡、执行器饱和,或者动作效果暂时不可见。此时,即使表示准确保留了环境状态,也未必能恢复实际动作。
论文的多步版本进一步要求从 zt+N−zt 恢复整段动作序列。不同动作序列可能有相同的起点和终点,却经过不同路径,因此仅凭端点差分未必能确定中间动作。原文没有专门检验这种歧义对训练的影响,也没有比较不同解码跨度 N 的效果。
- 差分输入减少了某些捷径,但没有保证学到真实动态
使用 zt+1−zt 确实改变了解码器可利用的信息,但差分中仍可能包含与动作相关的视觉线索,而非真正的物理变化。例如,模型可能主要编码容易预测动作的智能体局部变化,却忽略物体运动。
附录中,Delta-JEPA 在 Push-T 上对智能体运动的探测很好,但对物块位置和角度变化的探测不如部分基线。这与“表示更偏向动作相关信息”的解释相容,还不足以证明它全面学到了更准确的环境动态。
- 抗坍塌是经验效果,缺少对退化解的充分约束
动作重建能抑制所有状态都映射为同一向量,但它不保证潜在空间具有良好的尺度和完整的信息结构。
例如,将编码器输出缩小,再相应放大动作解码器的输入权重,在函数表达能力允许时,仍可能保留动作预测能力,同时降低潜在预测误差。实际训练可能受到架构和优化机制的限制,但论文没有分析这类尺度退化。因此,更稳妥的表述是:实验表明它能缓解坍塌,而非从理论上排除了坍塌及其他退化表示。
- 核心机制的证据还有缺口
固定历史、改变动作的实验说明预测器确实会响应动作,但响应明显,不等于响应正确。模型也可能对不同动作产生清晰、稳定,却与真实后果不一致的变化。
如果进一步比较“预测变化”和“实际执行动作后的真实变化”,并考察多步误差累积,就能更直接地验证世界模型的动态准确性。同样,PCA 图中表示分散只能提供定性支持,无法单独证明表示质量或有效维度更好。
- 验证范围和超参数分析仍有限
实验覆盖四个仿真任务,支持了导航和操作中的有效性,但尚未验证真实机器人、明显的视觉干扰、分布外动作或更复杂的部分可观测场景。对动作损失权重的分析也只在 Push-T 上进行,多步解码跨度则缺少消融。
另外,Reacher 上的优势仅为约 0.33 个百分点,且结果存在波动,因此“四个任务均取得最高平均值”不能理解为“四个任务都有明确优势”。