mlsanityData to 7 Oct 2026

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

arXiv 2603.21546 ↗First posted 23 Mar 2026cs.LG cs.AIcited by 4

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

  1. 2023Emergent Linear Representations in World Models of Self-Supervised Sequence Models2 Sep

Other references · 10

  1. 2018World Modelscited by 1,444
  2. 2016Understanding intermediate layers using linear classifier probescited by 1,279
  3. 2023Mastering Diverse Domains through World Modelscited by 1,193
  4. 2022Toy Models of Superpositioncited by 928
  5. 2021Probing Classifiers: Promises, Shortcomings, and Advancescited by 831
  6. 2023The Linear Representation Hypothesis and the Geometry of Large Language Modelscited by 679
  7. 2022Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Taskcited by 417
  8. 2021Causal Abstractions of Neural Networkscited by 401
  9. 2024Diffusion for World Modeling: Visual Details Matter in Ataricited by 285
  10. 2019Unsupervised State Representation Learning in Ataricited by 144

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 4.

4 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026Learning Task-Sufficient World Models by Synergizing Agentic Exploration and Structured Modeling
  2. 2025Assessing Adaptive World Models in Machines with Novel Games
  3. 2025Context and Diversity Matter: The Emergence of In-Context Learning in World Models
  4. 2025Better Decisions through the Right Causal World Model
  5. 2026Research on World Models Is Not Merely Injecting World Knowledge into Specific Tasks
  6. 2026MAP: A Map-then-Act Paradigm for Long-Horizon Interactive Agent Reasoning
  7. 2026From Generative Engines to Actionable Simulators: The Imperative of Physical Grounding in World Models
  8. 2025World-in-World: World Models in a Closed-Loop World