Operator-on-F complements value-equivalence: a planning-time diagnostic for latent world models

技术摘要：作为世界模型补充诊断工具的 Operator-on-F

问题陈述
在基于模型的强化学习（MBRL）中，对已学习世界模型的评估传统上依赖于价值等价性（value-equivalence）。这一原则认为，出于规划目的，如果一个模型能准确预测奖励和价值，即使其潜在动力学（latent dynamics）并不完美，该模型也是足够的。然而，作者认为这种方法在对规划至关重要的失效模式面前可能是“沉默”的。具体而言，一个模型可能表现出较低的奖励预测误差和较低的贝尔曼残差（Bellman residual），但其潜在状态展开（pushforwards）却存在显著误差。这些潜在误差会在规划过程中累积，导致策略性能下降，却无法被标准的奖励或价值头指标所检测到。本文指出，需要一种直接衡量模型在选定的可观测子集上，其潜在动力学相对于环境真实动力学的保真度的诊断工具，且该工具应独立于奖励头。

方法论：Operator-on-F
本文引入了 operator-on-F，这是一种通过比较模型的 $k$ 步潜在展开与环境在选定可观测子集 $F$ 上的编码后的下一状态来进行诊断的方法。

• 定义： 令 $z_t = \text{enc}(s_t)$ 为模型的潜在编码。模型的 $k$ 步展开为 $\hat{z}_{t+k} = O_k(z_t, a_{t:t+k-1})$，其中 $O_k$ 是模型自身的学习到的非线性潜在动力学映射。环境的真实情况为 $z'_{t+k} = \text{enc}(s_{t+k})$。
• 探测器（Probe）： 一个线性探测器 $\phi_F: Z \to F$ 将潜在状态映射到可观测子集 $F$。为了确保公平的跨模型比较，$F$ 被选择为一个不同编码器共享的表面。
• 误差指标： 对于每个函数 $\phi \in F$ 和每个锚点（状态-动作-下一状态元组），误差计算为归一化差异：
  $$e_\phi(t) = \frac{|\phi(\hat{z}_{t+k}) - \phi(z'_{t+k})|}{\sigma_\phi}$$
  其中 $\sigma_\phi$ 是跨锚点池的真实推前（pushforward）的标准差。最终指标是聚合了函数和锚点的这些误差的均方根（RMS）。
• F 的组成部分：
  • 价值切片（Value Slice）： $F = \{r, V\}$，使用模型自身的奖励和价值头。
  • 全量 F（Full-F）： 价值切片加上在编码后的下一状态几何结构上的留出 PCA 基（前 16 个方向，经过白化处理）。这捕捉了超越单纯奖励和价值之外的动力学结构。
• 实现： 该指标在模型自身的预测器在规划时界（例如 $k=5$ 步）处进行原生计算。对于没有奖励/价值头的模型（如纯自监督学习模型），比较被限制在观测层级的 $F$ 上。探测器采用岭回归（ridge regression），其正则化通过网格搜索选择，以避免容量与信息之间的混淆。

核心贡献

1. 定义与估计器： 形式化了 operator-on-F 以及一个在其部署时界上直接作用于模型潜在预测器的探测器估计器。
2. 规模扫描分析（Size-Sweep Analysis）： 对发布的 TD-MPC2 (mt80) 检查点在 cheetah-run 任务上的五个参数规模（1M 到 317M）进行了评估。
3. 跨架构比较： 将 TD-MPC2 与从头开始训练的纯自监督学习（SSL）潜在世界模型（LeWM）进行了合理性检查对比。
4. 诊断差异： 证明了 operator-on-F 提供了一种与传统价值等价性代理（贝尔曼残差、奖励误差）不一致的模型质量排序，并且与规划回报（planning return）强相关。

实验结果

• TD-MPC2 规模扫描：

  • 奖励误差： 在所有模型规模下保持在狭窄区间 $[0.028, 0.091]$ 内（约 $3\times$ 变化）。
  • Operator 误差： 跨度很大，从 $0.28$到$2.62$。
  • 317M 失效案例： 最大的模型（317M 参数）表现出最高的 operator 误差（$2.62$，比次高值高出一个数量级）以及规划回报的灾难性崩溃（$0.9$）。尽管如此，其奖励预测误差（$0.091$）仅比较小的模型略高，且仍处于同一狭窄范围内。
  • 相关性： operator 误差与回报损失的秩相关系数为 **$-0.90$**（95$[-0.90, -0.70]$）。相比之下，贝尔曼残差和奖励误差与回报的相关性较弱（分别为$-0.10$和$-0.30$）。
  • 与价值等价性的分歧： operator-on-F 指标与贝尔曼残差存在分歧（Spearman 相关系数 $+0.30$）。值得注意的是，未归一化的纯价值误差与贝尔曼残差几乎完全相同（Spearman 相关系数 $+1.00$），这证实了标准的“价值等价性”检查无法区分 317M 模型的规划失败。

• 跨架构检查（TD-MPC2 vs. LeWM）：

  • 在共享的观测层级 $F$ 上，LeWM（纯 SSL）实现的 operator 误差为 $0.384 \pm 0.009$，而单任务 TD-MPC2 检查点的得分则为 $0.840$（95$[0.79, 0.89]$）。
  • 两者的 95% 置信区间是不相交的，且即使将探测器架构从岭回归更改为带有一个隐藏层的 MLP，其排序依然得以保持，这表明差距并非探测器容量造成的人工产物。

意义与主张
论文声称 operator-on-F 是对价值等价性的补充而非替代。作者并非认为价值等价性是错误的，而是指出传统的、未经归一化的奖励和价值预测检查在会降低规划性能的特定失效模式面前可能是“沉默”的。

• 诊断价值： operator 指标成功识别了那些标准指标未能检测到的具有劣质潜在动力学的模型。在 TD-MPC2 规模扫描中，它是唯一一个正确将 317M 模型标记为由于规划崩溃而表现最差的模型指标。
• 归一化效应： 论文指出，operator-on-F 与价值等价性代理之间的分辨率差距主要是一个归一化效应；当价值切片按其目标分布进行归一化时，它们能更好地追踪回报，但 operator-on-F 将这种归一化的推前比较扩展到了适用于没有奖励/价值头的模型的共享可观测基准上。
• 局限性： 作者明确指出，跨架构比较受限于单一环境（cheetah-run）和两个模型家族。规模扫描的相关性是基于 $n=5$ 的，并以相关性形式报告。317M 的回报崩溃是针对 cheetah-run 特有的（该模型在其他任务上表现胜任），这表明该指标捕捉的是特定的操作点关联，而非全局退化的检查点。

总之，本文提出将 operator-on-F 作为一种必要的额外诊断工具，用以检测奖励预测和贝尔曼残差检查无法揭示的潜在世界模型误差，特别是在涉及规划性能的语境下。