无需高斯分布：用于JEPA世界模型的对比逆动态方法

# 无需高斯分布：基于对比式逆动力学的JEPA世界模型
来源：https://arxiv.org/html/2608.17542
###### 摘要

联合嵌入预测架构（JEPAs）通过预测未来嵌入来学习世界模型，但其目标函数允许编码器退化为常数解，因此每个实际系统都必须添加防塌陷机制\(13 (https://arxiv.org/html/2608.17542#bib.bib1);2 (https://arxiv.org/html/2608.17542#bib.bib3);5 (https://arxiv.org/html/2608.17542#bib.bib2);4 (https://arxiv.org/html/2608.17542#bib.bib4)\)。LeWorldModel（LeWM）使用SIGReg正则化器防止塌陷，该正则化器强制隐变量分布匹配各向同性高斯分布：通过规定表征必须呈现的形态来稳定表征，而与其建模的环境无关。我们认为防塌陷压力可以来自转移动作本身。对比式掩码转移动作建模（AC-MTM）保留了LeWM的前向隐变量预测目标，并添加了仅在训练阶段使用的逆动力学头，该头通过Action-NCE进行训练：每个隐变量转移必须在批次中的其他动作中识别出产生该动作的行动，这是一个塌陷编码器无法完成的判别任务。逆分支在训练后被丢弃，使得测试时的编码、前向预测、规划和计算与LeWM完全相同。在四个标准像素控制任务上，在匹配的规划协议下，AC-MTM可以从头开始稳定训练，并在平均性能上匹配SIGReg。在更困难的多物体OGBench视觉场景任务上，结果与规定的几何形状成为瓶颈一致：AC-MTM达到80.0±2.0%的成功率，而SIGReg为58.0±2.0%，在每个训练随机种子上提升了20-24个百分点。单个50集随机策略运行给出了52%的基线估计。因此，对比式逆动力学提供了一种无需分布先验的防塌陷信号，无需目标网络、停止梯度、预训练编码器或重建目标，我们分析了其成立的动作空间和可观测性假设。我们将代码开源在 https://github.com/jackboyla/action-contrastive-jepa。

## 1 引言

一个有用的世界模型应能从经验中学习，预测动作的后果，并在没有奖励标签或人工设计状态的情况下支持规划。联合嵌入预测架构（JEPAs）提供了一条吸引人的路径：编码器将感觉观测映射到紧凑的隐变量，预测器根据动作建模未来的隐变量\(13 (https://arxiv.org/html/2608.17542#bib.bib1);24 (https://arxiv.org/html/2608.17542#bib.bib5);19 (https://arxiv.org/html/2608.17542#bib.bib6)\)。由于预测发生在表示空间中，模型无需重建像素。我们的实验使用图像，但目标本身只需要成对的观测和动作。

参考图1：AC-MTM的训练和推理路径。训练时，观测到的动作是前向预测器的输入，也是对比式逆任务的正标签。测试时，CEM采样候选动作，仅使用不变的编码器和前向预测器；逆动力学/Action-NCE分支被丢弃。
核心困难在于*塌陷*：如果每个观测都映射到相同的隐变量，那么下一隐变量的预测会 trivially 完美，但该表征对规划毫无用处。现有系统通过使用预训练编码器、停止梯度和教师-学生启发式方法，或显式的分布正则化器来约束表征以防止塌陷\(2 (https://arxiv.org/html/2608.17542#bib.bib3);5 (https://arxiv.org/html/2608.17542#bib.bib2);4 (https://arxiv.org/html/2608.17542#bib.bib4)\)。LeWM\(14 (https://arxiv.org/html/2608.17542#bib.bib8)\)将这条路线精炼到最简形式：一个端到端的像素JEPA，使用下一嵌入预测加上SIGReg进行训练，SIGReg是来自LeJEPA\(3 (https://arxiv.org/html/2608.17542#bib.bib7)\)的正则化器，它将嵌入投影到随机方向上，并推动每一个一维边缘分布向高斯分布靠拢。这将多术语的PLDM配方简化为单一的有效系数。

SIGReg原理合理，但它通过规定一个*全局的、各向同性高斯*隐变量几何形状来稳定世界模型，编码器必须满足这一规定，而不管所建模的环境如何。\(14 (https://arxiv.org/html/2608.17542#bib.bib8)\)自己也指出这种规定可能是LeWM在低内在维度环境中表现不佳的原因。这引出了我们的问题：*与其施加一个动力学从未要求的分布，能否从转移动作本身导出防塌陷信号？*

感觉运动世界模型（SMWM）在相同的LeWM设置中用一个独立的逆动作MSE正则化器回答了这个问题\(11 (https://arxiv.org/html/2608.17542#bib.bib9)\)。我们的重点是对比式表述：一个机会水平的塌陷边界，其相对于逆MSE的可靠性，一个跨任务固定的损失权重，以及一个更难的Scene压力测试。

我们研究对比式掩码转移动作建模（AC-MTM）。与LeWM一样，模型从当前隐变量和动作预测未来的隐变量。仅在训练阶段，一个逆头额外将每个相邻的隐变量对映射到一个动作查询，该查询必须在批次中的其他动作中识别出观测到的动作，这是一个InfoNCE风格的目标函数\(7 (https://arxiv.org/html/2608.17542#bib.bib13);20 (https://arxiv.org/html/2608.17542#bib.bib14)\)，我们称之为Action-NCE。一个常数编码器会将每个转移映射到相同的查询，因此动作识别无法超越机会水平（第3节）：塌陷将前向目标的全局最优解转化为一个在逆任务上无法超越机会水平的表征。由此产生的防塌陷压力是动力学固有的；编码器被迫保留恰好区分不同动作效果的信息。整个逆分支在训练后被丢弃，使得部署的模型与LeWM相同。我们将AC-MTM作为受控的LeWM修改版本进行评估，与SIGReg、一个非对比式逆回归消融版本（MTM-MSE）以及一个对比式前向预测替代方案（AC-CPC）进行比较。得出三个结论：

\(1\) 基于动力学的防塌陷机制有效，无需分布先验，无需测试时改变。在LeWM的单阶段配方下从头开始训练，AC-MTM在所有任务和随机种子上从未塌陷，并在标准四任务套件（TwoRoom、Reacher、PushT、OGBench-Cube）上匹配SIGReg，在两个任务上获胜，一个任务上打平，并输掉了PushT。规划器从不调用逆头：评估使用与LeWM相同的编码器、自回归预测器、隐变量距离和CEM规划器。

\(2\) 在更难的Scene任务上，结果与规定的几何形状成为瓶颈一致。在OGBench视觉场景任务上（单个机械臂控制抽屉、窗户、按钮和一个可移动的立方体），在匹配的轨迹-目标MPC协议下，SIGReg在三个随机种子上的平均成功率降至58.0%，而AC-MTM保持在80.0%，在三个训练种子上分别提升了24、20和22个百分点。单个50集随机策略运行给出了52%的基线估计。逆回归消融版本也显著超过了SIGReg（75.3%），因此优势来自于动力学导出的信号本身，而非某种损失形式。两个模型都避免了塌陷；该结果可能结合了更有用的转移动作几何形状和更准确的短时域隐变量动力学。附录B审计了该结果。

\(3\) 逆信号的对比形式使其可靠。朴素的逆回归（MTM-MSE）在TwoRoom、Cube和Scene上表现强劲，但在Reacher的三个随机种子中的两个上出现塌陷；Action-NCE消除了这种双峰性，但在一个长时域压力测试上损失了3.8个百分点。两个逆变体在PushT上都落后于SIGReg，在该任务中，识别动作会低估弱受控的物体状态，我们将在第4.5节分析这一局限性。

我们的贡献是：(i) AC-MTM，一种面向端到端JEPA世界模型的无分布、动力学固有的防塌陷机制，附带简单的机会水平下界以证明塌陷不可取；(ii) 在五个像素控制任务上的受控三随机种子比较，一直匹配到规划器，将训练信号隔离为唯一的变量；(iii) 探针、压力测试和反事实惊奇分析，区分“不塌陷”与“支持规划”，包括对逆动力学防塌陷何时应与不应成立的明确说明。

## 2 背景与定位

#### 无奖励的隐变量规划。
我们考虑离线轨迹，包含像素观测和动作 τ=\(o0,a0,o1,...\)\\tau=\(\\mathbf\{o\}\_\{0\},\\mathbf\{a\}\_\{0\},\\mathbf\{o\}\_\{1\},\dots\)，没有奖励或最优性假设。训练期间，at\\mathbf\{a\}\_\{t\}是将ot\\mathbf\{o\}\_\{t\}带到ot\+1\\mathbf\{o\}\_\{t\+1\}的记录连续命令。测试时，没有可用的动作标签：一个基于交叉熵方法（CEM）的模型预测控制求解器采样、评分并重新拟合候选动作序列的分布，以最小化到编码目标的隐变量距离\(19 (https://arxiv.org/html/2608.17542#bib.bib6);14 (https://arxiv.org/html/2608.17542#bib.bib8)\)。

#### 从PLDM到LeWM。
PLDM使用几个相互作用的项构成的VICReg衍生目标来防止塌陷\(19 (https://arxiv.org/html/2608.17542#bib.bib6);5 (https://arxiv.org/html/2608.17542#bib.bib2)\)。LeWM将其替换为预测加SIGReg，使所有一维隐变量边缘分布匹配高斯分布（受Cramér–Wold定理启发）\(3 (https://arxiv.org/html/2608.17542#bib.bib7);14 (https://arxiv.org/html/2608.17542#bib.bib8)\)。这种简化是与我们工作进行比较的正确基准：LeWM的贡献不仅在于性能，更在于证明稳定的端到端像素JEPA世界模型可以比早期的配方简单得多。我们保持这一路线，同时将剩余的全局分布匹配项替换为局部的、对比式的转移动作信号。

#### 一类无分布信号。
MTM-MSE是我们的SMWM风格的逆回归基线\(11 (https://arxiv.org/html/2608.17542#bib.bib9)\)：它从一个相邻隐变量对回归出观测到的动作。SMWM建立了这种基本的逆MSE机制，并针对每个环境调整其损失权重；我们的受控基线使用一个在Reacher上选择的系数，并跨任务固定。为了区分“不塌陷”和“规划良好”，我们还评估了AC-CPC，一个动作条件对比式*前向*目标，它在批次负样本中识别真实的未来隐变量，而不是识别动作。所有变体都是插入到相同LeWM编码器、预测器和规划器的训练信号；表1总结了每个信号对表征的要求。

表1：本文比较的防塌陷机制。AC-MTM在训练时使用负样本，但不规定全局隐变量边缘分布，也不改变测试时的规划器。

## 3 对比式掩码转移动作建模

令观测ot\\mathbf\{o\}\_\{t\}在时刻t的隐变量嵌入为zt=encθ\(ot\)\\mathbf\{z\}\_\{t\}=\\mathrm\{enc\}\_\{\\theta\}\(\\mathbf\{o\}\_\{t\}\)。为简洁起见，我们写一个单步预测器；实现继承了LeWM的因果隐变量历史预测器。前向（规划）任务是

z^t\+1=fwdφ\(zt,at\),Lfwd=‖z^t\+1−zt\+1‖22。\\hat\{\\mathbf\{z\}\}\_\{t\+1\}=\\mathrm\{fwd\}\_\{\\phi\}\(\\mathbf\{z\}\_\{t\},\\mathbf\{a\}\_\{t\}\),\\qquad\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}=\\big\\\|\\hat\{\\mathbf\{z\}\}\_\{t\+1\}\-\\mathbf\{z\}\_\{t\+1\}\\big\\\|\_\{2\}^\{2\}。(1)
AC-MTM添加了一个仅在训练阶段使用的逆动力学头。对于展平批次/窗口中的N个转移中的每一个，它生成一个动作查询

a^i=invψ\(zi,zi\+1\)。\\hat\{\\mathbf\{a\}\}\_\{i\}=\\mathrm\{inv\}\_\{\\psi\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{z\}\_\{i\+1\}\)。(2)
我们使用负平方距离对每个观测到的动作aj\\mathbf\{a\}\_\{j\}作为转移i的候选进行评分，得到一个基于批次内动作的InfoNCE风格分类损失\(7 (https://arxiv.org/html/2608.17542#bib.bib13);20 (https://arxiv.org/html/2608.17542#bib.bib14)\)：

sij=−‖a^i−aj‖22τda,LNCE=−1N∑i=1Nlogexp⁡\(sii\)∑j=1Nexp⁡\(sij\),s\_\{ij\}=\-\\frac\{\\\|\\hat\{\\mathbf\{a\}\}\_\{i\}\-\\mathbf\{a\}\_\{j\}\\\|\_\{2\}^\{2\}\}\{\\tau d\_\{a\}\},\\qquad\\mathcal\{L\}\_\{\\mathrm\{NCE\}\}=\-\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\log\\frac\{\\exp\(s\_\{ii\}\)\}\{\\sum\_\{j=1\}^\{N\}\\exp\(s\_\{ij\}\)\},(3)
其中dad\_\{a\}是动作维度。总目标是

LAC−MTM=Lfwd\+λLNCE,λ=0\.30,τ=0\.10。\\mathcal\{L\}\_\{AC\-MTM\}=\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}\+\\lambda\\mathcal\{L\}\_\{\\mathrm\{NCE\}\},\\qquad\\lambda=0\.30,\\quad\\tau=0\.10。(4)
没有SIGReg项。固定系数是通过有界Reacher稳定性扫描选择的，然后在所有任务中保持不变使用。

“掩码”指的是转移动作元组内的因子预测，而非图像块掩码。前向任务隐藏zt\+1\\mathbf\{z\}\_\{t\+1\}并从\(zt,at\)\(\\mathbf\{z\}\_\{t\},\\mathbf\{a\}\_\{t\}\)预测它；逆任务隐藏at\\mathbf\{a\}\_\{t\}并从\(zt,zt\+1\)\(\\mathbf\{z\}\_\{t\},\\mathbf\{z\}\_\{t\+1\}\)预测它。我们在每个批次上同时优化两个任务，而不是每次采样一个掩码。

公式3中的候选集是批次中已存在的N=B⁡\(T−1\)N\{=\}B\(T\{\-\}1\)个观测到的动作块；任何数据集动作或记忆库也可以替代，但出于几个原因我们使用批次内动作。首先，这些块已驻留在GPU上，因此负样本仅增加一个N×NN\\times N距离矩阵。其次，候选aj\\mathbf\{a\}\_\{j\}是*原始*动作，而非编码器输出，因此它们不会给encθ\\mathrm\{enc\}\_\{\\theta\}带来梯度；扩大候选池只会改变判别任务的难度，而不改变梯度路径。第三，目标是防塌陷，批次内池已经通过公式5的\\(\\log N\\)下限来实施这一目标。全局采样主要会添加容易的负样本（远离预测、贡献可忽略梯度的块），同时增加假负样本的比率，因为控制动作会重复（接近零或饱和的块），而重复的“负样本”会惩罚正确的预测。因此，更大的候选池增加了计算和标签噪声，以锐化该方法不需要的一种动作检索特性。

#### 为何Action-NCE反对塌陷。
仅靠前向预测无法区分有用的表示和常数表示：如果encθ\(o\)=c\\mathrm\{enc\}\_\{\\theta\}\(\\mathbf\{o\}\)=c对每个观测都成立，且fwdφ\(c,a\)=c\\mathrm\{fwd\}\_\{\\phi\}\(c,\\mathbf\{a\}\)=c，则Lfwd=0\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}=0。Action-NCE将这种退化解转变为失败的分类问题。塌陷下，每个转移都给逆头相同的输入对\(c,c\)\(c,c\)，因此动作分类器的每一行都相同。然后模型被迫为批次中的所有N个正样本分配一个固定的概率向量pp。由于每个候选动作恰好是正确标签一次，平均损失满足

−1N∑i=1Nlogpi≥logN,\\-\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\log p\_\{i\}\\geq\\log N,(5)
等式仅在机会分类器时成立。因此塌陷