[論文レビュー] What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators

Xinyu Zhang|arXiv (Cornell University)|Mar 23, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

論文は2つの世界モデルアーキテクチャ（IRISとDIAMOND）をAtari BreakoutとPongで評価し、学習されたゲーム状態の内部表現がほぼ線形で、機能的に使用され、整然としたアテンションを示すことを、線形/MLPプローブ、因果的介入、トークンアブレーションを通じて明らかにする。

ABSTRACT

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.

研究の動機と目的

• RLにおける将来の観測を予測する際、世界モデルが学習する潜在表現を調査する。
• これらの表現が、位置やスコアなどのコアゲーム状態変数を線形にデコード可能な形でエンコードしているか評価する。
• 表現が単なる相関なのか、モデルの予測に因果的に関与しているのかを評価する。
• IRISとDIAMONDの間で、線形表現が現れる場所と方法の設計上の違いを検証する。

提案手法

• IRISとDIAMONDの各層にわたる凍結済み隠れ表現に対して線形およびMLPプロービングを適用する。
• R^2とRidgeおよびMLPプローブ（5-fold CV）を用いて、エンコードされたゲーム状態変数の線形性を評価する。
• プローブ由来の方向に沿って隠れ状態を摂動することで因果介入を実施し、次トークン予測の変化を測定する。
• IRISのアテンションヘッドを空間的特化で分析し、複数ベースラインのトークンアブレーションを実施してトークンの重要性を評価する。

Figure 1: Probe $R^{2}$ across layers (in network data-flow order) for IRIS (left) and DIAMOND (right) on Breakout (top) and Pong (bottom). Each line tracks one game-state property; shaded bands show $\pm$ 1 std over 5-fold CV. IRIS representations are flat across transformer layers, while DIAMOND s

実験結果

リサーチクエスチョン

• RQ1 世界モデルはゲーム状態変数の線形にデコード可能な表現を発展させるのか？
• RQ2 表現は予測において機能的に使用されており、単なる相関だけではないのか？
• RQ3 2つのアーキテクチャ（IRISとDIAMOND）は、線形表現がどこでどのように現れるかでどのように異なるのか？
• RQ4 ゲームオブジェクトを追跡する上で、最も重要な空間的領域（トークン）とアテンションヘッドはどれか？

主な発見

Representation	ball_x	ball_y	player_x	score
Random model	-1.21	-1.22	-1.14	-1.18
Shuffled labels	-0.51	-0.49	-0.53	-0.52
Raw pixels	-1.31	-0.48	0.9989±0.0006	0.9998±0.0001
IRIS (Linear)	0.85±0.006	0.58±0.03	0.9994±0.0001	1.0000±0.0000
IRIS (MLP)	0.91±0.005	0.59±0.03	0.9987±0.0002	0.9999±0.0000
Δ_IRIS	+0.06	+0.01	-0.0007	-0.0001
DIAMOND (Linear)	0.81±0.01	0.57±0.05	1.0000±0.0000	1.0000±0.0000
DIAMOND (MLP)	0.91±0.005	0.63±0.05	0.9994±0.0002	0.9998±0.0001
Δ_DIAMOND	+0.10	+0.06	-0.0006	-0.0002

• IRISとDIAMONDの両方が、ゲーム状態変数（例：ボール位置、パドル/スコア）に関してほぼ線形の表現を発展させ、BreakoutでのΔ ≤ 0.06、PongでのΔ ≤ 0.03という小さな選択性ギャップを示す。
• 因果介入により、隠れ状態をプローブ方向に沿って移動させると予測に相関する変化が生じ、機能的使用を示す（r ≥ 0.95）。
• IRISのアテンションヘッドは空間的特化を示し、特定のヘッドがゲームオブジェクトに重なるトークンに焦点を当てる；トークンアブレーションは一貫してオブジェクトを含むトークンを高く重要と判定（基準を超えるρ > 0.9）。
• DIAMONDのボトルネックは層をまたいで反転V型のパターンで抽象状態を符号化し、MLPプローブはデコーダ段で非線形のボール位置情報を回復する；両モデルはベースラインよりも優れており（生デピクセルは性能が低い）。
• ゲーム全体でPongはBreakoutよりR^2が高い傾向があり、これはシーンが単純である可能性による；アーキテクチャ固有のパターン（IRISは全層でフラット、DIAMONDはボトルネックでピーク）も両ゲームで持続する。

Figure 2: Causal intervention on Breakout: shifting IRIS layer-5 hidden states along probe directions produces correlated changes in predictions ( $r\geq 0.96$ for all properties, measured via KL divergence).