用动作条件预测一致性诊断JEPA世界模型

# 用动作条件预测一致性诊断JEPA世界模型

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

Guo An1,∗, Zijing Wu2,∗, Honghua Dong1,∗ Yuhao Yan3, Zixuan Gui4, Haochong Chen4, Shanzhao Ruan5 Xiang Wang2, Yurong Ling6,†, Qi Tian6,1,†  
1华为  
2中国科学技术大学  
3浙江大学  
4清华大学  
5哈尔滨工业大学  
6广东省人工智能与数字经济实验室（深圳）

###### 摘要

联合嵌入预测架构（JEPAs）学习在紧凑的潜在空间中进行预测的世界模型，而非在像素空间中预测，从而降低对无关外观进行建模的压力。然而，这并不能保证对视觉扰动具有鲁棒性：扰动仍可能改变编码后的表征，并影响后续的、以动作条件为基础的预测。双模拟（bisimulation）精确地刻画了这一需求：仅当两个观测的动作条件后果一致时，它们才应被视为同一状态。在这一标准的指导下，我们引入了动作条件预测一致性（ACPC），这是一种诊断指标，用于衡量一个干净历史记录与其视觉受扰动版本在相同动作序列下向前展开后的发散程度。我们证明，该发散程度为扰动引起的多步预测误差和规划器成本变化提供了上界。在逐对ACPC的基础上，我们定义了两个互补的度量：不变性半径（IR）概括了干净与受扰动版本展开轨迹的散布程度，而分离率（SR）则检查不同状态在展开后是否仍然可区分。在四个视觉控制任务上的实验表明，逐对ACPC能够预测扰动引起的预测误差和成本变化。在LeWM上，IR–SR筛查能够跨任务迁移，并且该联合诊断在模糊和缩放扰动下仍然有效。PLDM在另一种架构下表现出相似的诊断趋势。代码可在https://github.com/Anguo-star/acpc-diagnostics获取。

††脚注：∗同等贡献。†通讯作者。联系：Guo An,[email protected]

## 1 引言

世界模型通过预测候选动作的结果来支持控制[14,15]。联合嵌入预测架构（JEPAs）预测目标表征而非重建像素，从而避免了对再现无关视觉细节的显式需求[23,1,3,42,24]。然而，仅凭这种设计并不能确定哪些视觉差异对控制是重要的。此类表征应当对与任务无关的视觉扰动不敏感，同时保留那些需要不同动作的情境之间的差异。这一要求与双模拟背后的直觉相呼应：双模拟通过动作条件下的后果而非视觉外观来定义状态等价性[9,56]。然而，仅靠编码器距离并不能显示扰动如何通过预测进行传播。在多步展开过程中，预测器可能放大或缩小初始表征差异。编码器级别和单步转移的诊断并不能直接测量这种展开级别的效应。因此，我们提出：当一个干净历史记录及其受扰动视图在相同动作序列下向前展开时，它们的预测轨迹相距多远？

为了测量这种展开级别的效应，我们引入了动作条件预测一致性（ACPC）。ACPC将干净历史记录及其受扰动视图在相同动作序列下向前展开，并测量它们的预测轨迹之间的距离。图1提供了如何使用IR汇总该逐对距离、并由SR补充的概览。然而，仅凭低ACPC并不能排除表征坍缩的可能：一个恒定表征会使每次成对展开都相同。因此，我们使用*不变性半径*（IR）在历史记录间汇总ACPC，其中较低的值表示对视觉扰动的敏感性较低。*分离率*（SR）检查不同状态在展开后是否仍然可区分，较高的值表示更好的分离性。

**图 1：** ACPC、IR和SR的概念概览。(a) 一条记录的干净历史及其视觉受扰动视图构成成对输入。(b) 一个冻结的世界模型对两个视图进行编码，并在相同的记录动作序列下向前展开H步。ACPC测量两个预测轨迹之间的距离。在脆弱情形下，扰动引起的散布增大，并与其他状态（灰色）的表征重叠；在鲁棒情形下，成对视图保持接近，同时其他状态保持分离。示意图仅为示意：只有面板内的邻近性有意义，位置不进入任何报告的指标。(c) 不变性半径（IR）汇总跨历史记录的扰动诱导展开散布，而分离率（SR）则衡量选定的不同状态展开在超过该散布后仍保持分离的频率。只有当IR较低且SR较高时，检查点才通过诊断筛查；通过是针对所评估的视觉偏移和状态标签的，并不代表整体鲁棒性的认证。

我们证明ACPC为扰动引起的多步预测误差变化提供了上界。当从两个视图评估相同的候选动作序列时，ACPC还为预测成本的变化提供了上界（命题1和2）。这些上界对每个被评估的成对输入都成立，且不需要分布性或平滑性假设。在LeWM[26]上的实验表明，ACPC能够提供关于预测误差变化以及视觉扰动导致交叉熵方法（CEM）[21]选择不同计划时产生的额外预测成本的信息。IR和SR共同帮助识别在多个任务和扰动类型下视觉扰动表现良好的检查点。在PLDM[36]上，更好的视觉扰动下规划性能通常也伴随着较低的IR和较高的SR，这表明在第二种世界模型架构下呈现出相同的定性诊断模式。

本文的贡献如下：

1. 我们引入了*动作条件预测一致性*（ACPC），它比较干净历史记录及其受扰动视图在相同动作序列下的预测轨迹。IR在跨历史记录间汇总敏感性，而SR检查不同状态是否仍可区分（第3节）。
2. 我们证明了视觉扰动能改变多步预测误差的*逐样本上界*。当从两个视图评估相同的候选动作序列时，我们还限定了预测成本的变化（命题1和2）。
3. 我们在四个控制任务和三种视觉扰动下评估了ACPC、IR和SR，并检验了它们在第二种世界模型架构上的行为。ACPC为扰动引起的多步预测误差和规划成本变化提供预测性信息，而IR和SR有助于识别能保持控制性能的已训练模型（第4节）。

## 2 相关工作

#### 世界模型。

世界模型预测环境在动作作用下的演变，并利用这些预测进行控制。DreamerV3[14]通过想象轨迹进行学习，而TD-MPC2[15]使用学习到的潜在动力学进行规划。联合嵌入预测架构（JEPAs）通过在表征空间中预测目标而非重建像素来学习表征[23,1,3,2]。LeWM[26]是我们的主要模型系列，它将联合嵌入目标与动作条件下的潜在预测相结合。PLDM[36]提供了另一种联合嵌入动力学架构，我们也在其上评估ACPC。早期工作分析了联合嵌入预测目标如何有利于缓慢特征[37]。相关的训练目标在PBL[12]中从历史和动作预测未来的潜在观测，或在SPR[33]中在增强下匹配预测表征和目标表征。后续工作研究了自预测学习中的坍缩[39]以及动作条件的影响[19]。理论分析探讨了潜在预测何时可以忽略无关特征或偏好可预测、高影响力的特征[42,24]，以及何时可以恢复线性变换下的潜在状态[20]。seq-JEPA[10]进一步研究了已知视图变换下的不变性和等变性。这些工作解释了预测性表征和动力学是如何学习的；我们评估的是已训练预测器在面对视觉输入扰动时的响应。

#### 视觉扰动鲁棒性。

训练时增强改进了DrQ[52]和DrQ-v2[51]中的视觉控制，而SODA[16]在辅助表征学习目标中使用增强。ViGMO[29]和VIBR[7]在潜在空间或价值空间中学习不变性。ReOI[5]则改为在测试时修改观测，以减少视觉模型预测控制中的分心物敏感性。其他方法改变世界模型的训练方式。TPC[27]偏好时间上可预测的信息，DreamerPro[6]用原型预测替代像素重建。任务信息抽象[8]和去噪MDP[45]使用任务或奖励信息将有用状态与分心物分离。Iso-Dream[28]将可控与不可控动力学分开，其他方法将时间掩蔽与双模拟结合[38]，推断视觉上相似分心物的动作[46]，或使用任务感知重建[18]。这些方法反映了一个更广泛的原则：鲁棒性应去除无关的视觉变化，而不丢弃控制所需的区分。双模拟通过奖励和动作条件下的转移形式化行为相似性[9,56,35]，包括一个仅基于转移的无奖励形式[40]。价值等价模型和价值感知的模型学习目标则专注于价值预测或规划所需的信息[11,44]。最接近的工作MWM[50]在训练期间强制执行动作条件下的展开一致性。ACPC则测量冻结模型在相同动作序列下干净与受扰动的潜在展开是如何分化的，而不需要重新训练它。

#### 世界模型诊断工具。

模型精度与控制性能并不总是同步变化：一步似然可能对下游性能缺乏良好的指导[22]。这种不匹配推动了将模型学习和评估与下游决策联系起来的方法[48]。因此，世界模型评估应反映模型的预期用途，包括长展开、规划和策略评估[55]。一组方法测试或强制学习的动力学是否尊重动作。ATM[4]比较编码和预测转移中的动作信息，而Delta-JEPA[58]从潜在差异中解码动作。ACID[34]在规划成本中加入逆动力学循环一致性。Future-Compatible[31]检查生成的未来是否与其声明的动作一致，而World Models as Group Actions[47]测试恒等、逆和复合结构。相关的诊断暴露长展开中的运动学失败[32]，或使用冻结的逆动力学探针测试潜在表征在视觉损坏下是否保留动作信息[53]。

第二组方法评估下游失败。WMAttack[13]搜索针对闭环世界模型智能体的视觉攻击，而ARB4WM[57]则针对其策略、价值和潜在动力学。其他工作比较预测的多步潜在转移与环境产生的转移[41]，或比较预测的计划成本与真实计划成本[54]。最后，CARRL[25]在有界观测不确定性下证书了基于Q的动作选择，CROP[49]通过平滑化证书了每状态动作和有限时域奖励上界。这些证书覆盖策略决策或回报，而我们的上界仅针对所评估的成对输入和固定候选池保留优胜者或精英集。更广泛地说，现有诊断评估的是动作内容、展开有效性、模型误差或攻击敏感性。ACPC则问的是：一个保持任务性质的视觉扰动如何在共享动作展开中传播，而无需真实未来。

## 3 作为诊断工具的动作条件预测一致性

视觉扰动可能改变世界模型预测的轨迹。ACPC通过在一个动作序列下将干净历史记录及其受扰动版本向前展开并比较两条轨迹来测量这种变化。我们首先定义这种逐对测量。然后我们证明它为预测误差和候选成本的扰动诱导变化提供了上界，并在此过程中引入IR和SR。

### 3.1 成对测量

令 τ = (o_0, a_0, o_1, a_1, ...) 为一条记录的历史轨迹，令 d 为一个视觉扰动变换。受扰动历史记为 d(τ) = (d(o_0), a_0, d(o_1), a_1, ...)。给定一个冻结世界模型，该模型包含编码器 enc 和动力学模型 dyn，我们可以从历史中编码得到潜在状态 s_t = enc(o_0:t)，随后使用相同的动作序列 a_0, a_1, ... 在该潜在状态下展开 H 步。对干净历史和受扰动历史分别执行此操作，得到两条潜在轨迹：

s_t, s_{t+1}, ..., s_{t+H}  
s'_t, s'_{t+1}, ..., s'_{t+H}

其中 s'_{t+k} 是从受扰动历史编码和预测的潜在状态。ACPC 被定义为这两条轨迹之间的平均距离：

ACPC(τ, d) = (1/H) * Σ_{k=1}^{H} dist(s_{t+k}, s'_{t+k})

其中 dist 是潜在空间中的距离度量，例如欧氏距离。该值越低，说明视觉扰动对预测轨迹的影响越小。由于该定义不依赖于真实未来的观测，ACPC 可以在没有环境交互的情况下被计算，适用于日志数据。

### 3.2 预测误差与计划成本的上界

我们建立两个命题，将ACPC与下游影响联系起来。

**命题1（预测误差变化界）。** 令 e(τ) 为模型在干净历史上从时间 t+1 到 t+H 的多步预测误差，e(d(τ)) 为受扰动历史对应的误差。假设预测误差是潜在轨迹的 L-Lipschitz 函数，则

|e(d(τ)) - e(τ)| ≤ L * ACPC(τ, d)。

**命题2（计划成本变化界）。** 令 C(τ, A) 为模型在干净历史下对候选动作序列集 A 中的每条序列预测出的计划成本，C(d(τ), A) 为受扰动历史下的对应成本。如果成本函数是潜在轨迹的 L'-Lipschitz 函数，则

max_{a∈A} |C(d(τ), a) - C(τ, a)| ≤ L' * ACPC(τ, d)。

这两个命题表明，ACPC 从样本层面限制了视觉扰动对预测误差和计划成本的影响。该上界对所评估的每一对历史-扰动组合均成立，且不需要分布性假设或平滑性假设。

### 3.3 不变性半径与分离率

ACPC 只测量单个干净-受扰动对的分化。然而，仅仅低 ACPC 也可能来自表征坍缩：一个将所有历史映射到同一潜在状态的模型会使任何成对展开完全相同。因此，我们需要一个互补的度量来确保不同状态在展开后仍然可区分。

**不变性半径（IR）。** 给定一组干净历史记录 {τ_i} 和扰动 d，IR 定义为：

IR(d) = max_i ACPC(τ_i, d)

或者更稳健地，定义为所有成对ACPC值的分位数（例如，第90百分位）。IR 总结了扰动诱导的展开散布：IR 越低，模型对视觉扰动越不敏感。

**分离率（SR）。** 给定从不同初始状态抽取的一对历史记录 (τ_i, τ_j)，其中 i ≠ j，我们使用相同的动作序列展开两者，并计算两条展开轨迹之间的距离。如果该距离大于 IR 的一个倍数（例如，max(IR, ε)），则称它们被“分离”。SR 是这些被选定的不同状态展开中仍然分离的比例：

SR(d) = (1/N_pairs) * Σ_{i<j} 1[dist(τ_i, τ_j rollout) > κ * IR(d)]

其中 κ 是一个超参数，ε 是一个避免除零的常数。SR 越高，说明即使存在视觉扰动，不同状态在展开后也能保持可区分性。

一个检查点只有在 IR 低且 SR 高时才能通过诊断筛查。该筛查是针对所评估的视觉扰动和状态标签的，并不等同于在所有条件下的鲁棒性证书。

## 4 实验

我们在四个视觉控制任务上评估 ACPC、IR 和 SR，并使用两种世界模型架构：LeWM 和 PLDM。我们回答以下问题：

1. 逐对 ACPC 是否能够预测扰动引起的预测误差和规划成本变化？
2. IR 和 SR 是否能够识别在视觉扰动下保持控制性能的已训练模型？
3. 该诊断在不同任务、扰动类型和架构之间是否具有可迁移性？

### 4.1 设置

**任务。** 我们使用四个来自 DeepMind Control Suite 的视觉控制任务：Cartpole Swing-up、Cheetah Run、Walker Walk 和 Finger Spin。每个任务使用 64×64 RGB 观测。

**扰动。** 我们考虑三种视觉扰动：随机抖动（random jitter）、模糊（blur）和缩放（resize）。扰动强度被设置为在不清除任务语义的前提下显著改变观测。

**模型。** 我们使用 LeWM 作为主要模型系列。对于每个任务，我们训练多个检查点，这些检查点在干净环境下的规划性能各异。我们的实验中还包括一组使用不同训练设置（例如，不同的数据增强强度）得到的检查点，以产生 IR/SR 的多样性。此外，我们还使用 PLDM 来验证诊断趋势在不同架构下的稳定性。

**评估协议。** 对每个检查点，我们从回放缓冲区中采样 T=200 条长度为 H=5 的干净历史。对每个历史 τ_i，我们采样一个扰动强度（从验证集中选择）并构造受扰动历史 d(τ_i)。我们使用相同的动作序列从两条历史展开模型，计算 ACPC。IR 取所有 ACPC 值的第 90 百分位。对于 SR，我们从干净历史中采样 500 对不同的初始状态，使用相同的动作序列展开，并检查展开轨迹之间的距离是否超过 κ*IR。我们同时测量扰动引起的预测误差变化（在环境真值潜在状态上计算）和规划成本变化（使用 CEM 选择的计划成本）。

### 4.2 逐对 ACPC 预测下游影响

图2展示了逐对 ACPC 值与扰动引起的多步预测误差变化之间的关系。在所有任务上，ACPC 与预测误差变化之间呈正相关。对于规划成本变化，我们也观察到了类似的正相关，尤其是在 CEM 因扰动而选择了不同计划时。这表明 ACPC 作为扰动传播的度量，确实携带了关于预测和规划如何受到影响的信息。

### 4.3 IR–SR 筛查识别鲁棒检查点

图3显示了在每个任务上，不同检查点的 IR、SR 与其在扰动下的规划性能之间的关系。规划性能被定义为在受扰动环境中使用 CEM 执行时的平均回报（归一化到 [0,1]）。我们观察到，具有低 IR 和高 SR 的检查点往往具有更高的扰动下规划性能。有些检查点在干净环境下表现良好但 IR 较高且 SR 较低，这些检查点在扰动下性能显著下降。IR–SR 筛查的准确性在不同的扰动类型（模糊、缩放）中保持一致。

### 4.4 跨架构的通用性

在 PLDM 模型上重复相同的评估程序后，我们观察到类似的定性结果：更好的扰动下规划性能通常伴随较低的 IR 和较高的 SR。尽管 PLDM 的潜在空间维度、训练目标与 LeWM 不同，ACPC 诊断仍然能够区分鲁棒检查点与脆弱检查点。

### 4.5 消融与敏感性

我们测试了 IR 的分位数选择、SR 中的 κ 值以及展开长度 H 对诊断结果的影响。结果显示，IR 在第 80-95 百分位之间取值时，筛查结果基本稳定。κ 在 1.0-2.0 范围内对排序影响很小。H 越大，对长程传播的敏感性越高，但 IR 和 SR 的相对排序在 H=3 到 H=10 之间保持一致。

## 5 结论

我们提出了动作条件预测一致性（ACPC），用于诊断世界模型对视觉扰动的鲁棒性。ACPC 测量干净历史记录与其视觉受扰动视图在相同动作序列下展开后的轨迹差异。我们证明它可以为扰动引起的多步预测误差和计划成本变化提供样本级上界。此外，通过不变性半径（IR）和分离率（SR），我们的诊断同时关注低敏感性和可区分性，避免了单一距离度量的坍缩陷阱。在包括 LeWM 和 PLDM 在内的多种架构和任务上的实验表明，ACPC 及其派生的 IR–SR 筛查能够识别在视觉扰动下保持控制性能的模型检查点。这些发现说明，即使不依赖真实未来观测，我们也可以对世界模型的视觉鲁棒性进行有意义的评估。未来的工作可以将 ACPC 作为训练目标，以主动提升视觉扰动下的世界模型鲁棒性，并探索在更复杂环境中的扩展。