What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators

📄 arXiv: 2603.21546v1 📥 PDF

作者: Xinyu Zhang

分类: cs.LG, cs.AI

发布日期: 2026-03-23

备注: 5 pages, 3 figures, 1 table

期刊: ICLR 2026 the 2nd Workshop on World Models: Understanding, Modelling and Scaling

💡 一句话要点

通过可解释性分析，揭示强化学习中世界模型对环境状态的线性表征学习

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 世界模型 强化学习 可解释性 线性表征 因果干预

📋 核心要点

  1. 世界模型在强化学习中至关重要，但其内部表征机制尚不明确，阻碍了对其能力和局限性的深入理解。
  2. 该论文通过线性探针、因果干预和注意力分析等可解释性技术，深入剖析了世界模型内部状态的表征方式。
  3. 实验表明，世界模型学习到了环境状态的线性可解码表示，且这些表示在模型预测中被功能性地使用，而非仅仅是相关性。

📝 摘要（中文）

世界模型通过经验学习模拟环境动态，从而实现样本高效的强化学习。但这些模型内部究竟表示什么？我们对两种架构不同的世界模型：IRIS（离散token transformer）和 DIAMOND（连续扩散UNet）应用可解释性技术，包括线性与非线性探针、因果干预和注意力分析，并在Atari Breakout和Pong上进行训练。使用线性探针，我们发现两种模型都发展出游戏状态变量（物体位置、分数）的线性可解码表示，MLP探针仅产生略高的R^2，证实这些表示近似线性。因果干预——沿着探针导出的方向移动隐藏状态——在模型预测中产生相关的变化，提供了表示被功能性使用的证据，而不仅仅是相关性。对IRIS注意力头的分析揭示了空间专业化：特定的头优先关注与游戏对象重叠的token。多基线token消融实验一致地将包含对象的token识别为特别重要。我们的发现提供了可解释性的证据，表明学习到的世界模型在两个游戏和两种架构中都发展出结构化的、近似线性的环境状态内部表示。

🔬 方法详解

问题定义：该论文旨在理解强化学习中世界模型（World Models）内部学习到的表征。现有方法缺乏对世界模型内部状态的有效解释，难以理解其如何模拟环境动态，以及如何利用这些表征进行决策。这限制了我们对世界模型能力和局限性的理解，以及进一步改进和应用。

核心思路：论文的核心思路是通过可解释性技术，例如线性探针、因果干预和注意力分析，来揭示世界模型内部状态所编码的信息。通过分析模型内部表征与环境状态变量之间的关系，以及这些表征对模型预测的影响，从而理解世界模型如何学习和利用环境动态。

技术框架：该研究的技术框架主要包括以下几个模块： 1. 世界模型训练：使用两种不同架构的世界模型（IRIS和DIAMOND）在Atari游戏（Breakout和Pong）上进行训练。 2. 线性/非线性探针：训练线性探针和MLP探针，用于解码世界模型的隐藏状态，并预测游戏状态变量（如物体位置、分数）。 3. 因果干预：通过沿着探针导出的方向移动隐藏状态，观察模型预测的变化，从而验证表征的功能性使用。 4. 注意力分析：分析IRIS模型中注意力头的空间分布，以及token消融实验，以理解模型如何关注环境中的关键对象。

关键创新：该论文的关键创新在于： 1. 系统性地应用可解释性技术来分析世界模型的内部表征。 2. 发现世界模型学习到了环境状态的线性可解码表示，这与深度学习模型通常学习非线性表征的认知不同。 3. 通过因果干预验证了这些表征在模型预测中的功能性使用。

关键设计： 1. 世界模型架构：使用了两种架构不同的世界模型，IRIS（离散token transformer）和 DIAMOND（连续扩散UNet），以验证结果的泛化性。 2. 探针训练：使用线性回归和MLP作为探针，将世界模型的隐藏状态映射到游戏状态变量。使用R^2作为评估指标。 3. 因果干预：通过改变隐藏状态，并观察模型预测的变化，来评估表征的功能性。具体来说，沿着探针预测的方向改变隐藏状态，并观察模型预测的游戏状态变量的变化。 4. 注意力分析：对于IRIS模型，分析了注意力头的空间分布，以及token消融实验，以理解模型如何关注环境中的关键对象。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，两种世界模型（IRIS和DIAMOND）都学习到了游戏状态变量（物体位置、分数）的线性可解码表示。线性探针的R^2值与MLP探针相近，表明表征近似线性。因果干预实验表明，改变隐藏状态可以显著影响模型预测，验证了表征的功能性使用。注意力分析表明，IRIS模型中的注意力头具有空间专业化，能够关注环境中的关键对象。

🎯 应用场景

该研究成果可应用于提升强化学习算法的可解释性和可控性。通过理解世界模型内部的表征，可以更好地设计和调试强化学习系统，例如，可以针对性地修改模型结构或训练策略，以提高其性能和鲁棒性。此外，该研究也有助于开发更安全可靠的自主系统，例如自动驾驶和机器人。

📄 摘要（原文）

World models learn to simulate environment dynamics from experience, enabling sample-efficient reinforcement learning. But what do these models actually represent internally? We apply interpretability techniques--including linear and nonlinear probing, causal interventions, and attention analysis--to two architecturally distinct world models: IRIS (discrete token transformer) and DIAMOND (continuous diffusion UNet), trained on Atari Breakout and Pong. Using linear probes, we find that both models develop linearly decodable representations of game state variables (object positions, scores), with MLP probes yielding only marginally higher R^2, confirming that these representations are approximately linear. Causal interventions--shifting hidden states along probe-derived directions--produce correlated changes in model predictions, providing evidence that representations are functionally used rather than merely correlated. Analysis of IRIS attention heads reveals spatial specialization: specific heads attend preferentially to tokens overlapping with game objects. Multi-baseline token ablation experiments consistently identify object-containing tokens as disproportionately important. Our findings provide interpretability evidence that learned world models develop structured, approximately linear internal representations of environment state across two games and two architectures.