[논문 리뷰] What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators

Xinyu Zhang|arXiv (Cornell University)|2026. 03. 23.

Reinforcement Learning in Robotics인용 수 0

한 줄 요약

이 논문은 Atari Breakout과 Pong에서 두 가지 world-model 아키텍처(IRIS와 DIAMOND)를 비교하여 학습된 게임 상태의 내부 표현이 대략 선형적이며 기능적으로 사용되고, 선형/MLP 탐사, 인과적 개입, 토큰 절삭을 통해 조직화된 주의(attention)를 보임을 밝힌다.

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의 고정(hidden) 표현을 계층별로 선형 및 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 IRIS와 DIAMOND 두 아키텍처에서 선형 표현이 어디에서 어떻게 나타나는지 어떻게 다른가?
• RQ4 어떤 공간적 영역(토큰)과 주의 헤드가 게임 객체를 추적하는 데 가장 중요한가?

주요 결과

표현	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 across baselines).
• DIAMOND 병목은 계층 전체에서 역-V 패턴의 추상 상태를 인코딩하고, MLP 탐사를 통해 디코더 단계에서 비선형 볼 위치 정보를 회복한다; 두 모델 모두 baselines를 능가한다(원시 픽셀은 성능이 좋지 않다).
• 두 게임에 걸쳐 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).