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

arXiv: https://arxiv.org/abs/2603.21546 · HTML: https://mlsanity.com/paper/2603.21546/

## Key facts

- Cited by 5 papers in the map (data through 14 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.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 11 of 11)

- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Causal Abstractions of Neural Networks](https://mlsanity.com/paper/2106.02997/) (2021)
- [Emergent Linear Representations in World Models of Self-Supervised Sequence Models](https://mlsanity.com/paper/2309.00941/) (2023)
- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)

## Built on by (most cited first, 5 of 5)

- [The Intervention Gap in Latent World Models](https://mlsanity.com/paper/2608.29998/) (2026)
- [Low-Rank Dynamics-Effective Latent Carriers for Counterfactual Rollout in Learned World Models](https://mlsanity.com/paper/2608.15156/) (2026)
- [Event-Conditioned Diagnostics of Kinematic, Contact, and Object-Permanence Structure in Passive Object-State World Models](https://mlsanity.com/paper/2606.28455/) (2026)
- [RolloutFaith: Auditing Persistent Internal Interventions in Visual World Model](https://mlsanity.com/paper/2609.36843/) (2026)
- [When Generative Models Act as Simulators](https://mlsanity.com/paper/2610.12847/) (2026)

## 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.