What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators
On the map: World models → World models for language agents
Abstract (arXiv)
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.
Key facts. Cited by 4 papers in the map (data through 7 Oct 2026). Placed in World models for language agents, one of 11 directions in World models. Builds on 11 of its 12 references that are in the map. First posted on arXiv 23 Mar 2026.
Builds on
11 of 12 references.
Closest ancestors · 1
Same direction, newest first
Other references · 10
- 2018World Modelscited by 1,444
- 2016Understanding intermediate layers using linear classifier probescited by 1,279
- 2023Mastering Diverse Domains through World Modelscited by 1,193
- 2022Toy Models of Superpositioncited by 928
- 2021Probing Classifiers: Promises, Shortcomings, and Advancescited by 831
- 2023The Linear Representation Hypothesis and the Geometry of Large Language Modelscited by 679
- 2022Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Taskcited by 417
- 2021Causal Abstractions of Neural Networkscited by 401
- 2024Diffusion for World Modeling: Visual Details Matter in Ataricited by 285
- 2019Unsupervised State Representation Learning in Ataricited by 144
Built on by
Citing papers per year: 2026: 4.
4 papers.
- 2026The Intervention Gap in Latent World Modelscited by 2
- 2026Low-Rank Dynamics-Effective Latent Carriers for Counterfactual Rollout in Learned World Modelscited by 1
- 2026Event-Conditioned Diagnostics of Kinematic, Contact, and Object-Permanence Structure in Passive Object-State World Modelscited by 0
- 2026RolloutFaith: Auditing Persistent Internal Interventions in Visual World Modelcited by 0
- 2026RolloutFaith: Auditing Persistent Internal Interventions in Visual World Model29 Sep 2026
- 2026The Intervention Gap in Latent World Models30 Aug 2026
- 2026Low-Rank Dynamics-Effective Latent Carriers for Counterfactual Rollout in Learned World Models15 Aug 2026
- 2026Event-Conditioned Diagnostics of Kinematic, Contact, and Object-Permanence Structure in Passive Object-State World Models26 Jun 2026
Similar, not cited
Close in wording, no citation either way
- 2026Learning Task-Sufficient World Models by Synergizing Agentic Exploration and Structured Modeling
- 2025Assessing Adaptive World Models in Machines with Novel Games
- 2025Context and Diversity Matter: The Emergence of In-Context Learning in World Models
- 2025Better Decisions through the Right Causal World Model
- 2026Research on World Models Is Not Merely Injecting World Knowledge into Specific Tasks
- 2026MAP: A Map-then-Act Paradigm for Long-Horizon Interactive Agent Reasoning
- 2026From Generative Engines to Actionable Simulators: The Imperative of Physical Grounding in World Models
- 2025World-in-World: World Models in a Closed-Loop World