arXiv:2603.21546cs.LGcs.AI2026-03被引 5

探究强化学习世界模型如何内部表征环境状态

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

  • 用线性/非线性探针、因果干预等方法分析模型内部表示
  • 发现两种模型均能线性解码物体位置、得分等状态变量
  • 适合关注模型可解释性与环境模拟机制的研究者

世界模型通过经验学习环境动态,实现样本高效的强化学习。但这些模型内部实际表征了什么?我们采用可解释性技术——包括线性与非线性探针、因果干预和注意力分析——研究了两种架构不同的世界模型:基于离散令牌的Transformer模型IRIS和基于连续扩散UNet的DIAMOND,二者均在Atari Breakout和Pong上训练。使用线性探针发现,两种模型均能线性解码游戏状态变量(如物体位置、得分),且MLP探针仅略微提升R²,表明这些表示近似线性。通过沿探针导出方向扰动隐状态,模型预测发生相关变化,证明表示是功能性的而非仅相关。对IRIS注意力头的分析显示空间特化:特定注意力头优先关注与游戏物体重叠的令牌。多基准令牌消融实验一致表明,含物体的令牌具有显著更高的重要性。结果提供了可解释性证据,表明学习的世界模型在两个游戏和两种架构下均发展出结构化、近似线性的内部环境状态表示。

原文摘要 · Abstract (English)

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.

世界模型可解释性强化学习状态表征

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。