一个透镜，众多世界：面向世界模型可解释性的能力类型化接口

# 一镜览万象：面向世界模型可解释性的能力类型化接口
来源：https://arxiv.org/html/2606.09936
Bhavith Chandra Challagundla1, Sanskar Pandey2, Param Thakkar3, Rishikesh Mallagundla4, Yugandhar Reddy Gogireddy5, Wenhao Lu6, Hindol Roy Choudhury2, Shravani Challagundla7, Mohamed Deraz Nasr8, and Spursh Deshpande2 1纽约大学2独立研究者3Veermata Jijabai Technological Institute 4Mercity5南加州大学6独立研究者，MIT 7独立研究者，GITAM8佐治亚理工学院

###### 摘要

世界模型如今构建在截然不同的计算基座上。诸如 PlaNet 和 Dreamer 家族的潜在递归状态空间模型将观测压缩为递归状态；基于 token 的模型（如 IRIS）将观测量化到学习到的码本中，并使用 transformer 进行自回归预测；而联合嵌入预测架构（如 I-JEPA）则在学习的潜在空间中预测，无需像素解码器。应用于这些模型的可解释性方法，包括探针、激活修补、稀疏自编码器和惊奇分析，都共享一组基本原语，但每种架构都需要从头重新实现，因为现有的钩子-缓存工具假设是一个没有动作、环境步骤或想象展开概念的 transformer 语言模型。我们认为这种碎片化是工具而非模型本身的问题，而世界模型的共享结构可以通过一个类型化的小型接口来捕获。我们提出了 WorldModelLens，一个围绕能力类型化适配器组织的开源可解释性基座：每个模型实现四个必需方法（编码、转移、初始状态、采样），并通过显式的能力描述符声明一组可选头（解码、奖励、续存、策略、价值），使得强化学习和自监督世界模型成为一等公民，而无需模仿对方。一个单一的钩子和缓存层在此接口上暴露时间索引的激活、想象展开和干预重放，使得每个分析只需编写一次。我们在 I-JEPA 上端到端地展示了完整的分析套件，提供了跨越潜在递归、transformer-token 和联合嵌入系列的适配器，以证明该接口通过构造具有泛化性，并展示了钩子层在非激活状态下大约增加 12% 的每步开销，这使得始终开启的可解释性遥测在训练和控制循环中变得实用。

## 1 引言

世界模型已成为学习环境预测性表示的核心工具[1 (https://arxiv.org/html/2606.09936#bib.bib1),2 (https://arxiv.org/html/2606.09936#bib.bib2)]，且用于构建它们的架构已出现严重分化。Dreamer 家族和 PlaNet 学习一个递归状态空间模型，其潜在变量结合了确定性递归分量与连续或离散分类随机分量[3 (https://arxiv.org/html/2606.09936#bib.bib3),4 (https://arxiv.org/html/2606.09936#bib.bib4),5 (https://arxiv.org/html/2606.09936#bib.bib5)]。IRIS 和 Decision Transformer 将观测量化到学习到的码本中，并使用 transformer 自回归地建模动态[6 (https://arxiv.org/html/2606.09936#bib.bib6),7 (https://arxiv.org/html/2606.09936#bib.bib7)]。联合嵌入预测架构（如 I-JEPA 和 TD-MPC2）直接在学习的嵌入空间中预测，可能根本没有解码器[8 (https://arxiv.org/html/2606.09936#bib.bib8),9 (https://arxiv.org/html/2606.09936#bib.bib9)]。这些模型不仅在其内部计算上不同，而且在对环境的接口上也不同：有些消耗动作并输出奖励和价值，而另一些则纯粹是自监督的，两者都没有。

相比之下，可解释性方法在精神上是一致的。线性和非线性探针读取内部表示并预测感兴趣的性质[19 (https://arxiv.org/html/2606.09936#bib.bib19),20 (https://arxiv.org/html/2606.09936#bib.bib20)]；激活修补覆盖内部激活并测量下游效应[15 (https://arxiv.org/html/2606.09936#bib.bib15),14 (https://arxiv.org/html/2606.09936#bib.bib14)]；稀疏自编码器将激活分解为可解释的特征[17 (https://arxiv.org/html/2606.09936#bib.bib17),18 (https://arxiv.org/html/2606.09936#bib.bib18)]。每种方法都归结为相同的三个原语：读取激活，可选地覆盖它，并观察效果。这种规律性正是语言模型机制可解释性工具所利用的。TransformerLens[12 (https://arxiv.org/html/2606.09936#bib.bib12)]推广了一种钩子-缓存抽象，其中命名的钩子点暴露每个内部激活以供读取或编辑，并且一系列工作在此基础上构建了电路级分析[13 (https://arxiv.org/html/2606.09936#bib.bib13),14 (https://arxiv.org/html/2606.09936#bib.bib14)]。

然而，该工具是围绕 transformer 语言模型塑造的。它假设一叠注意力块作用于 token 序列，没有动作输入的表示，没有环境步骤，没有多步展开，没有想象的未来，也没有混合连续和离散的潜在变量。世界模型违反了所有这些假设。实际效果是，可解释性代码为每个新世界模型重写，结果以模型特定的术语报告，无法迁移，而依赖于跨架构比较内部结构的问题无法用共享代码表达。通用可解释性库，如 Captum[21 (https://arxiv.org/html/2606.09936#bib.bib21)]、NNsight[22 (https://arxiv.org/html/2606.09936#bib.bib22)]和 Penzai[23 (https://arxiv.org/html/2606.09936#bib.bib23)]，提供了归因和访问原语，但没有建模世界模型的语义，如轨迹、动态展开或干预重放。

我们认为这种碎片化是工具而非模型本身的问题。一个世界模型，无论属于哪个家族，都可以通过一个编码器（将观测映射到潜在变量）、一个转移（推进潜在状态并可选择消耗一个动作）以及一组可选读出头来描述，这些头给定模型可能拥有也可能没有。我们通过一个*能力类型化适配器*使这种结构显式化，其中必需的核心必须被实现，可选能力通过描述符声明，使得强化学习代理和自监督预测器都成为一等公民，而无需任何一方伪造另一方的头。一个单一的钩子和缓存层安装在此接口之上，然后暴露时间索引的激活、想象展开和干预重放，使得每个分析都针对接口编写一次，而不是针对任何单一架构。

我们提出了 WorldModelLens，该设计的一个开源实现。我们的贡献是：

- • 一个能力类型化的世界模型接口（第3节 (https://arxiv.org/html/2606.09936#S3)），通过四个必需方法和五个可选、描述符门控的头，统一地表达强化学习和自监督世界模型，并且我们在 Dreamer、PlaNet、TD-MPC2、IRIS、Decision Transformer 和 I-JEPA 家族中实例化了该接口。
- • 该接口上的一个后端无关的钩子、缓存、展开和干预重放层，具有时间索引的激活缓存和用于长展开的设备卸载，在此基础上，探针、激活修补、稀疏自编码器和惊奇分析各实现一次（第4节 (https://arxiv.org/html/2606.09936#S4)）。
- • 在 I-JEPA 上的端到端演示（第5节 (https://arxiv.org/html/2606.09936#S5)），其中未修改的分析套件恢复了预测器中的层次分辨结构，并测量了该基座施加的开销。
- • 一个开放实现，其钩子层在非激活状态下大约增加 12% 的每步开销，使得始终开启的可解释性遥测在训练和控制循环中变得实用。

将实证分析扩展到其他家族，包括全尺度 Dreamer、V-JEPA[10 (https://arxiv.org/html/2606.09936#bib.bib10)] 和 Cosmos[11 (https://arxiv.org/html/2606.09936#bib.bib11)]，正在进行中，并构成了我们的路线图（第7节 (https://arxiv.org/html/2606.09936#S7)）。本文建立了接口、基座和深入的单一家族演示。

## 2 背景

我们将世界模型形式化地定义为一个元组

M=\(O,S,Z,A,ι,E,τ,H\),\\mathcal\{M\}=\\big\(\\mathcal\{O\},\\,\\mathcal\{S\},\\,\\mathcal\{Z\},\\,\\mathcal\{A\},\\ \\iota,\\,\\mathcal\{E\},\\,\\tau,\\ \\mathcal\{H\}\\big\),\(1\)其中O\\mathcal\{O\}是观测空间，S\\mathcal\{S\}是确定性潜在状态空间，Z\\mathcal\{Z\}是随机潜在空间，A\\mathcal\{A\}是可选的行动空间。必需的核心是三个映射和一个采样器。初始状态映射 ι:\{∗\}→S\\iota:\\\{\\ast\\\}\\to\\mathcal\{S\} 返回 s0s\_\{0\}。概率编码器

E:O×S→Δ\(Z\),zt∼qθ\(z∣ot,st−1\)\\mathcal\{E\}:\\mathcal\{O\}\\times\\mathcal\{S\}\\to\\Delta\(\\mathcal\{Z\}\),\\qquad z\_\{t\}\\sim q\_\{\\theta\}\(z\\mid o\_\{t\},s\_\{t\-1\}\)\(2\)将观测和前一状态映射到潜在编码的分布（遵循变分传统[58 (https://arxiv.org/html/2606.09936#bib.bib58)]），随后sample\_z 使用 Gumbel-softmax 松弛[31 (https://arxiv.org/html/2606.09936#bib.bib31),32 (https://arxiv.org/html/2606.09936#bib.bib32)]（当 Z\\mathcal\{Z\} 为分类时）或恒等映射（当连续时）从该分布中抽取 ztz\_\{t\}。转移

τ:S×Z×\(A∪\{∅\}\)→S,st\+1=τ\(st,zt,at\),\\tau:\\mathcal\{S\}\\times\\mathcal\{Z\}\\times\(\\mathcal\{A\}\\cup\\\{\\varnothing\\\}\)\\to\\mathcal\{S\},\\qquad s\_\{t\+1\}=\\tau\(s\_\{t\},z\_\{t\},a\_\{t\}\),\(3\)推进确定性状态，其中对于无动作模型 at=∅a\_\{t\}=\\varnothing。下游分析可见的单元是联合潜在变量 ht=\(st,zt\)h\_\{t\}=\(s\_\{t\},z\_\{t\}\)。可选头构成一个集合 H⊆\{gdec,grew,gcont,π,V\}\\mathcal\{H\}\\subseteq\\\{g\_\{\\mathrm\{dec\}\},g\_\{\\mathrm\{rew\}\},g\_\{\\mathrm\{cont\}\},\\pi,V\\\}，其签名为

gdec\\displaystyle g\_\{\\mathrm\{dec\}\}:S×Z→O,\\displaystyle:\\mathcal\{S\}\\\!\\times\\\!\\mathcal\{Z\}\\to\\mathcal\{O\},grew\\displaystyle g\_\{\\mathrm\{rew\}\}:S×Z→R,\\displaystyle:\\mathcal\{S\}\\\!\\times\\\!\\mathcal\{Z\}\\to\\mathbb\{R\},gcont\\displaystyle g\_\{\\mathrm\{cont\}\}:S×Z→\[0,1\],\\displaystyle:\\mathcal\{S\}\\\!\\times\\\!\\mathcal\{Z\}\\to\[0,1\],\(4\)π\\displaystyle\\pi:S×Z→Δ\(A\),\\displaystyle:\\mathcal\{S\}\\\!\\times\\\!\\mathcal\{Z\}\\to\\Delta\(\\mathcal\{A\}\),V\\displaystyle V:S×Z→R,\\displaystyle:\\mathcal\{S\}\\\!\\times\\\!\\mathcal\{Z\}\\to\\mathbb\{R\},分别用于解码、奖励、续存、策略和价值。强化学习世界模型[59 (https://arxiv.org/html/2606.09936#bib.bib59)]实例化了 H\\mathcal\{H\} 的大部分；自监督视频和联合嵌入模型实例化很少或没有。在观测序列 o1:To\_\{1:T\} 上的展开递归为 zt∼E\(ot,st−1\)z\_\{t\}\\sim\\mathcal\{E\}\(o\_\{t\},s\_\{t\-1\}\)，st\+1=τ\(st,zt,at\)s\_\{t\+1\}=\\tau\(s\_\{t\},z\_\{t\},a\_\{t\}\)，而*想象*是相同的递归，但将 E\\mathcal\{E\} 替换为学习的先验 pθ\(z∣st\)p\_\{\\theta\}\(z\\mid s\_\{t\}\)。每步的*惊奇*是后验与先验潜在变量之间的散度，

surpriset=DKL\(qθ\(z∣ot,st−1\)∥pθ\(z∣st\)\)\.\\mathrm\{surprise\}\_\{t\}\\;=\\;D\_\{\\mathrm\{KL\}\}\\\!\\big\(q\_\{\\theta\}\(z\\mid o\_\{t\},s\_\{t\-1\}\)\\,\\big\\\|\\,p\_\{\\theta\}\(z\\mid s\_\{t\}\)\\big\)\.\(5\)
我们针对的家族主要在潜在变量形式和转移上不同，而非在这个签名上。Dreamer 和 PlaNet 使用一个潜在递归状态空间模型，其潜在变量连接一个确定性递归部分与一个随机部分（在 V1 中连续，从 V2 起为离散分类）[3 (https://arxiv.org/html/2606.09936#bib.bib3),4 (https://arxiv.org/html/2606.09936#bib.bib4)]。IRIS 和 Decision Transformer 使用一个 transformer[56 (https://arxiv.org/html/2606.09936#bib.bib56)]作用于离散码本，因此潜在变量是一个 token 序列[6 (https://arxiv.org/html/2606.09936#bib.bib6),7 (https://arxiv.org/html/2606.09936#bib.bib7)]。I-JEPA 和 TD-MPC2 在连续嵌入空间中预测，并且在 I-JEPA 的情况下没有解码器[8 (https://arxiv.org/html/2606.09936#bib.bib8),9 (https://arxiv.org/html/2606.09936#bib.bib9)]。一个仅基于 transformer 的钩子库没有动作输入的容身之地，没有想象展开的表示，也没有同时覆盖递归潜在变量、token 序列和联合嵌入的抽象。上面的接口做到了，而 WorldModelLens 正是使这个接口具体化。

## 3 WorldModelLens 抽象

WorldModelLens 组织为三层，如图1 (https://arxiv.org/html/2606.09936#S3.F1) 所示，将模型*是什么*与它如何被*工具化*以及*测量了什么*分离开。在底层，后端适配器通过第2节 (https://arxiv.org/html/2606.09936#S2) 的必要-可选接口暴露单个世界模型，将该模型的特定内部结构转换为通用签名。在中间层，HookedWorldModel 包装器在适配器的输出上挂载命名的钩子点，并将其产生的每个激活通过单个缓存管理器路由，使得每个激活可以通过名称被读取、记录或覆盖。在顶层，一个分析库完全通过该包装器进行操作，通过名称读取和编辑激活，从不涉及特定的架构。

这种分离的价值在于责任在边界上被清晰地划分。前向传播进入适配器，适配器对每个观测调用encode，并通过transition推进潜在状态，在每个钩子点发出命名的激活。包装器记录或修改这些激活，并将它们组装成类型化的轨迹。然后分析层使用该轨迹，或安装干预措施，包装器通过适配器重放这些干预。由于只有适配器层知道给定模型如何计算，将 DreamerV3 替换为 I-JEPA 只需要改变底层，而之上的每个分析无需修改即可继续运行。这就是第1节 (https://arxiv.org/html/2606.09936#S1) 中可移植性声明的具体机制：分析针对接口编写一次，并自动继承所有当前和未来的后端。本节的剩余部分依次描述每一层。

探针激活修补SAE / 惊奇run\_with\_cacherun\_with\_hooksimagine / replayDreamer / PlaNetIRIS / DTI\-JEPA / TD\-MPC2分析层（一次编写）钩子与缓存层：HookedWorldModel适配器层：BaseModelAdapter\+\+能力激活通过名称读取和编辑必需\+\+可选模型接口图1：三层设计。适配器通过一个类型化接口暴露任何世界模型；钩子模型包装器挂载命名的钩子点和单个缓存；分析读取和编辑激活，不涉及任何架构。### 3.1 能力类型化适配器

通过继承 BaseModelAdapter（一个 torch.nn.Module）注册后端，其必需的表面是四个方法：encode(obs, h_prev)，返回后验潜在变量及其先验；transition(h, z, action)，推进确定性状态并接受可选动作；initial_state(batch_size, device)；以及 sample_z(logits, temperature)，从编码器输出中抽取潜在变量，并应用 Gumbel-softmax 松弛。