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

• URL: http://arxiv.org/abs/2603.21546v1
• Date: Mon, 23 Mar 2026 04:00:53 GMT
• Title: What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators
• Abstract summary: We apply interpretability techniques to two architecturally distinct world models: IRIS and DIAMOND.<n>We find that both models develop linearly decodable representations of game state variables.<n>Analysis of IRIS attention heads reveals spatial specialization: specific heads attend preferentially to tokens overlapping with game objects.
• Score: 6.400424118459376
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Towards Real-world Human Behavior Simulation: Benchmarking Large Language Models on Long-horizon, Cross-scenario, Heterogeneous Behavior Traces [81.41397370235102]
  We introduce OmniBehavior, the first user simulation benchmark constructed entirely from real-world data.<n>We show that current models struggle to accurately simulate complex behaviors, with performance plateauing even as context windows expand.<n>This results in the loss of individual differences and long-tail behaviors, highlighting critical directions for future high-fidelity simulation research.
  arXiv  Detail & Related papers  (2026-04-09T15:26:21Z)
• Emergence of Linear Truth Encodings in Language Models [64.86571541830598]
  Large language models exhibit linear subspaces that separate true from false statements, yet the mechanism behind their emergence is unclear.<n>We introduce a transparent, one-layer transformer toy model that reproduces such truth subspaces end-to-end.<n>We study one simple setting in which truth encoding can emerge, encouraging the model to learn this distinction in order to lower the LM loss on future tokens.
  arXiv  Detail & Related papers  (2025-10-17T16:30:07Z)
• Linear Spatial World Models Emerge in Large Language Models [4.9185678564997355]
  We investigate whether large language models implicitly encode linear spatial world models.<n>We introduce a formal framework for spatial world models and assess whether such structure emerges in contextual embeddings.<n>Our results provide empirical evidence that LLMs encode linear spatial world models.
  arXiv  Detail & Related papers  (2025-06-03T15:31:00Z)
• ExPLAIND: Unifying Model, Data, and Training Attribution to Study Model Behavior [39.590138981646696]
  Post-hoc interpretability methods typically attribute a model's behavior to its components, data, or training trajectory in isolation.<n>We present ExPLAIND, a unified framework that integrates all these perspectives.
  arXiv  Detail & Related papers  (2025-05-26T14:53:11Z)
• SPARTAN: A Sparse Transformer Learning Local Causation [63.29645501232935]
  Causal structures play a central role in world models that flexibly adapt to changes in the environment. We present the SPARse TrANsformer World model (SPARTAN), a Transformer-based world model that learns local causal structures between entities in a scene. By applying sparsity regularisation on the attention pattern between object-factored tokens, SPARTAN identifies sparse local causal models that accurately predict future object states.
  arXiv  Detail & Related papers  (2024-11-11T11:42:48Z)
• A Theoretical Analysis of Self-Supervised Learning for Vision Transformers [66.08606211686339]
  Masked autoencoders (MAE) and contrastive learning (CL) capture different types of representations.<n>We study the training dynamics of one-layer softmax-based vision transformers (ViTs) on both MAE and CL objectives.
  arXiv  Detail & Related papers  (2024-03-04T17:24:03Z)
• Beyond DAGs: A Latent Partial Causal Model for Multimodal Learning [80.44084021062105]
  We propose a novel latent partial causal model for multimodal data, featuring two latent coupled variables, connected by an undirected edge, to represent the transfer of knowledge across modalities.<n>Under specific statistical assumptions, we establish an identifiability result, demonstrating that representations learned by multimodal contrastive learning correspond to the latent coupled variables up to a trivial transformation.<n>Experiments on a pre-trained CLIP model embodies disentangled representations, enabling few-shot learning and improving domain generalization across diverse real-world datasets.
  arXiv  Detail & Related papers  (2024-02-09T07:18:06Z)
• Intriguing Equivalence Structures of the Embedding Space of Vision Transformers [1.7418480517632609]
  Pre-trained large foundation models play a central role in the recent surge of artificial intelligence. Due to their inherent complexity, these models are not well understood. We show via analyses and systematic experiments that the representation space consists of large piecewise linear subspaces.
  arXiv  Detail & Related papers  (2024-01-28T04:59:51Z)
• Dynamic Latent Separation for Deep Learning [67.62190501599176]
  A core problem in machine learning is to learn expressive latent variables for model prediction on complex data. Here, we develop an approach that improves expressiveness, provides partial interpretation, and is not restricted to specific applications.
  arXiv  Detail & Related papers  (2022-10-07T17:56:53Z)
• Bridging the Gap to Real-World Object-Centric Learning [66.55867830853803]
  We show that reconstructing features from models trained in a self-supervised manner is a sufficient training signal for object-centric representations to arise in a fully unsupervised way. Our approach, DINOSAUR, significantly out-performs existing object-centric learning models on simulated data.
  arXiv  Detail & Related papers  (2022-09-29T15:24:47Z)
• On the Transfer of Disentangled Representations in Realistic Settings [44.367245337475445]
  We introduce a new high-resolution dataset with 1M simulated images and over 1,800 annotated real-world images. We propose new architectures in order to scale disentangled representation learning to realistic high-resolution settings.
  arXiv  Detail & Related papers  (2020-10-27T16:15:24Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.