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

Xinyu Zhang

cs.LGMar 23, 2026

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

Authors: Xinyu Zhang

Organizations: Anyscale

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.

Figures & tables

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 shows a peaked inverted-V centered on the UNet bottleneck. Note:

y

-axis includes negative

R^{2}

values, revealing that DIAMOND’s early encoder layers are worse than a constant predictor for ball position.

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\pm.0006$	$0.9998\pm.0001$
IRIS (Linear)	$\mathbf{0.85}\pm.006$	$0.58\pm.03$	$\mathbf{0.9994}\pm.0001$	$1.0000\pm.0000$
IRIS (MLP)	$0.91\pm.005$	$0.59\pm.03$	$0.9987\pm.0002$	$0.9999\pm.0000$
$\Delta_{\text{IRIS}}$	$+0.06$	$+0.01$	$-0.0007$	$-0.0001$

Table 1: Best-layer

R^{2}

(mean

\pm

std over 5-fold CV) for Breakout. Both linear and MLP probes are shown; the small selectivity gap (

\Delta

) confirms approximately linear representations.

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

Figure 3: Three-way token ablation on Breakout ( $4\times 4$ grid). Zero, mean, and random replacement produce consistent importance rankings ( $\rho>0.92$ ), with token 0 (score/brick region) most critical.

Explore similar work

Jun 7, 2026

One Lens, Many Worlds : A Capability-Typed Interface for World-Model Interpretability

World models are now built on substantially different computational substrates. Latent recurrent state-space models such as PlaNet and the Dreamer family compress observations into recurrent states; token-based models such as IRIS quantize observations into a learned codebook and predict autoregressively with a transformer; and joint-embedding predictive architectures such as I-JEPA predict in a learned latent space with no pixel decoder. The interpretability methods applied to these models, including probing, activation patching, sparse autoencoders, and surprise analysis, share a common set of primitives, yet they are re-implemented from scratch for each architecture because existing hook-and-cache tooling assumes a transformer language model with no notion of actions, environment steps, or imagined rollouts. We argue that this fragmentation reflects the tooling rather than the models, and that the shared structure of world models is captured by a small typed interface. We present WorldModelLens, an open-source interpretability substrate organized around a capability-typed adapter: every model implements four required methods (encode, transition, initial state, sample) and declares a set of optional heads (decode, reward, continue, actor, critic) through an explicit capability descriptor, so that reinforcement-learning and self-supervised world models are first-class without either imitating the other. A single hook and cache layer exposes time-indexed activations, imagination rollouts, and intervention replay over this interface, allowing each analysis to be written once.

Bhavith Chandra Challagundla, Sanskar Pandey, Param Thakkar +7

New York University · Independent Researcher · Veermata Jijabai Technological Institute +5

World Models Interpretable ML

Ask PDF

Jul 28, 2026

Learning Implicit Causal World Models from Multi-Agent Demonstrations

In model-based reinforcement learning, world models exist as internal simulators, but their training often conflates statistical correlations with causal mechanisms. This problem is exacerbated in multi-agent systems where physical transitions are intertwined with strategic agent intents, causing world models to fail under distribution shift. We introduce Implicit Causal World Models to recover environmental dynamics from offline demonstrations without requiring pre-defined causal graphs. By incorporating policy variance, we render world models discoverable via the sequential backdoor condition. Evaluations across coordination tasks (Two-Door, Navigation, and Giveway) demonstrate that these models provide interpretable causal representations under both full and partial observability, with model accuracy scaling directly with interventional strength.

Jasorsi Ghosh

Department of Computer Science Purdue University West Lafayette, Indiana, USA

Causal Representation Learning Causal World Models World Models Causal Structure Learning Causal RL Offline MARL+1

Ask PDF

Jul 7, 2026

A Definition and Roadmap for World Models

World models -- internal simulators that learn the structure and dynamics of an environment -- have become one of the most actively debated concepts in AI. From model-based reinforcement learning and video generation to embodied robotics and ultimately, physical AI, researchers across AI subfields are building systems that they call "world models", yet there is no consensus on what a world model fundamentally is, what it should predict, or how it should be built. This perspective article provides a scientific definition of world models, discussions of their key technical aspects, and a staged roadmap for developing effective world models.

Xinyuan Chen, Haoyu Guo, Shi Guo +10

Shanghai Artificial Intelligence Laboratory

World Model Learning World Models

Ask PDF