Source-linked AI summary
Understanding the limits of latent world models
Rodrigo Gonzalez Laiz, Steffen Schneider
TL;DR
Task reward alone offers limited evidence about whether latent world models capture environmental dynamics, motivating an evaluation of faithfulness beyond task success. The paper defines faithfulness and introduces LRC and state decoding; across the studied models and controlled environments, high reward did not imply faithful latent dynamics, particularly from pixels.
Problem
Task reward indicates overall agent success but does not reveal whether learned latent dynamics capture the environment, and prior studies described here do not evaluate action-conditioned rollouts in interacting agents.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
The results suggest that improving learned representations and latent dynamics jointly is central to building more faithful world models.
Source: §7 · p. 9The study covers three representative latent world-model families in simple deterministic continuous-control environments, leaving extension to stochastic, partially observable, discrete, or substantially higher-dimensional settings open.
Source: §6 · p. 9
Abstract
from arXiv · showhide
World models are widely proposed as core components for enabling agents to reason over long temporal and spatial horizons. Yet current evaluations often emphasize task performance while providing limited insight into whether learned world models acquire predictive latent dynamics that reflect the underlying structure of the environment. Here we introduce the concept of faithful latent world models as an equivariance condition between the environment's transition map and the learned latent dynamics. We propose latent rollout consistency (LRC) as a metric to evaluate action-conditioned latent prediction. Strikingly, we find that across the studied state of the art latent world models, none solves the task by actually faithfully modeling the dynamics of the environment.
1 INTRODUCTION
The paper asks whether successful latent world models actually capture their environments’ dynamics, since reward alone does not reveal what their latent dynamics learned. It defines faithfulness and uses LRC and state decoding to identify failures, especially under pixel observations.
1 INTRODUCTION
- The paper tests whether successful latent world models learn representations and dynamics that model their environment, studying action-conditioned control rather than passive video encoders or pixel-space prediction.The cited related studies examine passive video encoders [4] and observational video prediction [5], whereas this work studies agents interacting with and controlling an environment.Source: §1 · p. 1
- The evaluation examines DreamerV3, TD-MPC2, and LeWM [6–8] in Pendulum, CartPole, and Reacher, using simple, well-defined dynamics to isolate world-model accuracy.Failures in these controlled settings provide counterexamples and help isolate mechanisms.Source: §1 · p. 1
- Reward alone does not establish faithful latent dynamics: LRC and state decoding reveal inaccurate velocity estimates and difficulty learning through nonlinear observation mappings.The authors argue that improving both learned representations and latent dynamics is central to building more faithful models.Source: §1 · p. 1
- Figure 1 depicts faithfulness as agreement between encoding after a true transition and applying learned latent dynamics, and describes LRC as comparing predicted and realized-trajectory latents.The comparison is action-conditioned.Source: Figure 1 · p. 2
2 WHAT DOES IT TAKE TO MODEL THE WORLD?
The section defines a faithful latent world model through state preservation and agreement between true and learned dynamics. It introduces LRC and post-hoc state decoding as separate diagnostics, and illustrates distinct ways reward, prediction, and decoding can diverge.
2 WHAT DOES IT TAKE TO MODEL THE WORLD?
- A latent world model jointly infers a representation from observations and models its action-conditioned evolution; the paper’s usage excludes conceptually different video generators and cognitive or spatial maps.The environment is modeled as a partially observed controlled system with deterministic dynamics and observations.Source: §2 · p. 2
- A faithful model must make learned latent prediction agree with encoding after a true transition and preserve the full state through an injective state-to-latent map.The section distinguishes dynamics equivariance from state sufficiency.Source: §2 · pp. 2–3
- The equivariance condition can be optimized from observed transitions by matching predicted latents to encoded next observations, alongside additional signals used to prevent collapse.The listed signals include reconstruction, contrastive learning, regularization, and, for reward-based models, reward and value prediction.Source: §2 · p. 3
- LRC measures agreement between autoregressive latent rollouts and latents inferred from realized trajectories, while post-hoc state decoding tests whether the representation contains underlying state information.High LRC alone is insufficient because collapsed latents can be easy to predict; good decoding supports state sufficiency but does not guarantee injectivity.Source: §2 · p. 3
- In the illustrated CartPole cases, noisy transitions or encoder error can make reward and decoding diverge from LRC, including high reward despite collapsed LRC when the encoder is corrupted.With an identity encoding and noisy transitions, linear and nonlinear decoding remain at R2 = 1 while reward falls with LRC; with a nonlinear mix, linear decoding is poor even at zero noise.Source: §2 · p. 3
- The illustrated settings show that both representation and dynamics learning can limit task performance, motivating joint evaluation of latent prediction, state decoding, and reward.The examples include nonlinear coordinates that impair linear decoding and encoder error that accumulates over a rollout.Source: §2 · p. 3
3 RELATED WORK
Related work spans debates over internal models, latent predictors, model-based agents, probing of learned representations, and identifiable representation learning. The paper positions its analysis as an empirical test of whether action-conditioned world models recover environment dynamics beyond task success.
3 RELATED WORK
- World-model research ranges from compact latent dynamics learned from experience [1] and task-driven agents such as PlaNet [15], DreamerV3 [6], and TD-MPC2 [7] to video generators and JEPA-style predictors [19–23].Recent work questions whether these models capture environment physics, while Ying et al. [25] call for metrics beyond task success.Source: §3 · pp. 3–4
- Probing classifiers have been used to recover latent state in Othello-GPT [27, 28] and examine trained agents including DQN [29] and AlphaZero [30]; this paper probes state content in world models.A complementary line studies physical structure in foundation models [35, 36].Source: §3 · p. 4
- Identifiable representation learning asks when the data-generating process can be recovered from observations [37, 38], and the paper’s results reinforce calls for more empirically grounded identifiability theory [40].The authors suggest theory-assumption mismatches or numerical implementations that do not fully match theory as possible explanations for current failures.Source: §3 · p. 4
4 EXPERIMENTAL SETUP
The experiments compare three world-model families on simple continuous-control environments using state and pixel observations. Evaluation combines action-conditioned latent prediction, state decoding, reward, and targeted CartPole interventions.
4.1 BENCHMARKS AND TRAINING
- DreamerV3, TD-MPC2, and LeWM are evaluated on Pendulum, CartPole, and Reacher, spanning reconstruction and training-objective choices.DreamerV3 is reconstruction-based and reward-driven/task-specific; TD-MPC2 is reconstruction-free and reward-driven/task-specific; LeWM is reconstruction-free and reward-free/task-agnostic, with MPC applied post hoc.Source: §4.1 · p. 4, Figure 3 · p. 4
- The three low-dimensional environments have ground-truth states and dynamics, enabling direct evaluation; they are deliberately simple test cases, not a comprehensive benchmark.Pendulum and CartPole have nontrivial unforced dynamics, while Reacher remains stationary from rest without control.Source: §4.1 · pp. 4–5
- Models are trained on raw state vectors and rendered RGB images, with four seeds per model, modality, and environment.TD-MPC2 and DreamerV3 use 64×64-pixel images; LeWM uses 224×224-pixel images, and a state-based LeWM variant enables direct comparison.Source: §4.1 · p. 5
4.2 EVALUATION PROTOCOL
- Each environment has a separate evaluation set of 1,000 100-step episodes collected with a random Ornstein–Uhlenbeck policy for broad state-space coverage.The collected datasets’ state, action, and reward distributions are compared in Appendix A.Source: §4.2 · p. 5
- Latent rollout consistency compares autoregressive predicted latents with latents inferred from observed trajectories under the recorded actions.R2 is computed across episodes for each latent coordinate and aggregated by target-variance weighting; an autoregressive state-prediction MLP and persistence serve as state-space baselines.Source: §4.2 · p. 5
- Frozen world models are probed with separate nonlinear predictors of each state coordinate, with held-out error reported in physical units and relative to a shuffled-target control.Episodes are split 80% for probe training and 20% for held-out evaluation.Source: §4.2 · p. 5
- Reward is evaluated as summed dm_control reward over 500-step episodes, using learned controllers for TD-MPC2 and DreamerV3 and post-hoc MPC for LeWM.Source: §4.2 · p. 5
4.3 CONTROLLED INTERVENTIONS
- Two CartPole interventions test nonlinear state mixing and the role of velocity information in pixel models.The mixing experiment applies 0–3 fixed invertible layers before encoding without changing simulator dynamics; the velocity experiment adds ground-truth velocities and compares reduced temporal context for TD-MPC2 and LeWM.Source: §4.3 · p. 5
5 RESULTS
Across the tested models, reward can coexist with inaccurate latent rollouts, while pixel inputs worsen prediction and velocity information explains only part of the pixel–state gap. Under invertible nonlinear mixing, LeWM’s reward and LRC decline as mixing depth increases, and the controller solves only the raw-state and one-layer conditions.
5.1 TASK PERFORMANCE DOES NOT IMPLY ACCURATE LATENT PREDICTION
- High task reward can coexist with very low LRC, so reward alone does not establish accurate latent dynamics.This dissociation is most apparent in task-trained TD-MPC2 and DreamerV3.Source: §5.1 · pp. 5–6, Figure 4 · p. 6
5.2 PIXEL OBSERVATIONS DEGRADE LATENT PREDICTION
- Pixel-based rollouts deteriorate faster than state-based rollouts as the horizon grows, approaching the persistence baseline.In CartPole, pixel predictions also depart from realized trajectories over 30 control steps (0.6 seconds).Source: §5.2 · p. 6
- Position decoding errors are 0–7% and velocity errors 3–22%; state inputs consistently outperform pixels on velocity decoding.Position decoding is relatively similar across input modalities.Source: §5.2 · p. 6
5.3 MISSING VELOCITIES EXPLAIN PART OF THE PIXEL–STATE GAP
- Adding ground-truth velocity substantially improves TD-MPC2 and LeWM, reducing one-frame-context normalized velocity error from approximately 130% to 1–2%.The intervention improves LRC with both reduced and full temporal context, but has no effect on DreamerV3.Source: §5.3 · p. 7, Figure 6 · p. 8
- Improved velocity representations do not consistently raise task reward, and adding velocity explains only part of the pixel–state LRC gap.LeWM’s improved LRC at H = 30 remains insufficient for planning the complete swing-up motion.Source: §5.3 · p. 7
5.4 LEARNING LATENT DYNAMICS FROM NONLINEAR OBSERVATIONS
- For LeWM, reward and LRC decline with increasing invertible mixing depth, and the controller solves only raw-state and one-layer conditions.Even though the observation map preserves state information, the evaluated objectives struggle to learn representations from which physical state is recoverable.Source: §5.4 · p. 8, Figure 7 · p. 9
6 DISCUSSION AND LIMITATIONS
Task performance alone gives limited information about learned latent dynamics. The study identifies challenges with pixel observations while limiting its conclusions to controlled environments and three model families.
6 DISCUSSION AND LIMITATIONS
- Pixel-based models struggle to recover velocity, and explicitly providing velocity improves latent prediction without fully closing the gap to state-based models.Even with the underlying state fully observed, nonlinear transformations make representation and dynamics harder to learn.Source: §6 · p. 9
- The experiments cover simple deterministic continuous-control environments, leaving extension to stochastic, partially observable, discrete, and substantially higher-dimensional settings open.The study evaluates three representative latent world-model families rather than exhaustively comparing existing approaches.Source: §6 · p. 9
- The results are controlled counterexamples to the assumption that successful control implies an accurate model of environment dynamics, not evidence that all latent world models share these failure modes.Task performance alone provides limited information about learned dynamics.Source: §6 · p. 9
7 CONCLUSION
The paper tests whether successful latent world models represent and model their environments faithfully. Across controlled continuous-control environments, high reward does not imply faithful latent dynamics, particularly from pixels, and the findings point to representation and dynamics learning as joint challenges.
7 CONCLUSION
- Across three latent world models and controlled continuous-control environments, high reward did not imply faithful latent dynamics, particularly when learning from pixels.The study defines faithfulness through an equivariance condition and uses latent rollout consistency and state decoding to evaluate it independently of task reward.Source: §7 · p. 9
- The gap reflects both difficulty recovering dynamical state from observations and difficulty learning dynamics from nonlinear observations, even when relevant state information is available.The results motivate addressing representation and dynamics learning jointly.Source: §7 · p. 9
- The findings motivate closing the gap between theory of identifiable representation learning and numerically successful implementations of these algorithms.This is the paper's stated direction for improving world models.Source: §7 · p. 9
AI USE STATEMENT
The authors describe AI assistance in manuscript drafting and editing, experimental code, and plotting code, with author review and responsibility for the final work.
AI USE STATEMENT
- AI tools assisted with manuscript drafting and editing, experimental code, and figure plotting; authors reviewed the outputs and take responsibility for the final work.Experimental outputs and figures were checked for correctness.Source: AI USE STATEMENT · p. 10
A ENVIRONMENTS AND EVALUATION DATASETS
The evaluation uses three DeepMind Control Suite environments with simulator states as probing targets and image observations for pixel-based models. Latent rollout consistency compares action-conditioned predicted latents with representations inferred from observed trajectories, then aggregates coordinate-wise prediction scores.
A ENVIRONMENTS AND EVALUATION DATASETS
- Pendulum, CartPole, and Reacher provide low-dimensional simulator states as probing targets, alongside 64×64 and 224×224 image observations for pixel-based models.Their state dimensions are 3, 5, and 6, respectively.Source: A ENVIRONMENTS AND EVALUATION DATASETS · p. 14
- Each environment contributes 1,000 episodes of 100 steps, storing simulator states, actions, rewards, and rendered frames; OU actions encourage broad coverage of physical states.Figures 8 and 9 compare marginal state, action, and reward distributions under random OU and expert policies.Source: A ENVIRONMENTS AND EVALUATION DATASETS · p. 14, Figure 8 · p. 15, Figure 9 · p. 15
B LATENT ROLLOUT CONSISTENCY COMPUTATION
- For each model seed and environment, one autoregressive latent rollout begins at each evaluation episode's start and is compared with target latents encoded from realized observation histories.Episodes are pooled within each model-seed and environment condition, but representations are not pooled across seeds or environments.Source: B LATENT ROLLOUT CONSISTENCY COMPUTATION · p. 14
- LRC computes coordinate-wise R² across evaluation episodes and aggregates latent coordinates using target-variance weighting; 1 indicates exact prediction, 0 matches the across-episode mean, and negative values are worse.Coordinates with zero target variance receive zero weight and are omitted; negative scores are retained.Source: B LATENT ROLLOUT CONSISTENCY COMPUTATION · p. 16
- For DreamerV3's stochastic rollouts, the estimator averages latent predictions across 500 independently sampled trajectories at each horizon before computing LRC.The same inferred posterior anchor is copied across rollouts, and recorded actions are replayed.Source: B LATENT ROLLOUT CONSISTENCY COMPUTATION · p. 16
C.1 MODELS AND IMPLEMENTATION SOURCES
The evaluation uses TD-MPC2, DreamerV3, and LeWM, starting from the authors’ implementations and default configurations. The models’ latent features are then used for downstream analyses.
C.1 MODELS AND IMPLEMENTATION SOURCES
- Three world-model families—TD-MPC2, DreamerV3, and LeWM—are evaluated, with four independent seeds per family, observation modality, and environment.Table 3 summarizes the experimental matrix and defines latent dimension as the activation-vector size used in downstream analyses.Source: §C.1 · p. 16, Table 3 · p. 16
- TD-MPC2 uses a 512-dimensional latent, while state-input DreamerV3 uses a 1536-dimensional RSSM representation.DreamerV3’s representation combines a 1024-dimensional categorical stochastic state and a 512-dimensional deterministic state.Source: §C.1 · p. 16
- Pixel LeWM uses 224×224 inputs and three-frame context; its reproduced Reacher checkpoints reach 81.5% success versus 86% in the original study.Because the original code does not support state inputs, the authors implement a state-based variant with a Mish MLP encoder and action-conditioned residual MLP predictor.Source: §C.1 · p. 16
- Downstream analyses use each model’s encoder output, with DreamerV3 represented by its posterior feature vector rather than sampled categorical states.The DreamerV3 vector concatenates its deterministic state with the softmax probabilities of its categorical stochastic state.Source: §C.1 · p. 17
C.2 TRAINING PROTOCOL
The protocols train task-driven models to stable high-return regimes, train LeWM on random and expert trajectories, and evaluate task return, state decoding, and interventions. LeWM’s decoded-state cost yields higher return than its latent cost in the controller comparison.
C.2 TRAINING PROTOCOL
- TD-MPC2 and DreamerV3 train for up to one million environment steps without per-task hyperparameter tuning and retain stable high-return checkpoints.TD-MPC2 is evaluated deterministically, DreamerV3 samples from its learned actor distribution, and native return is measured over 100 episodes of 500 steps.Source: §C.2 · p. 17, §C.3 · p. 17
- LeWM trains for 50 epochs on 2,000 episodes split evenly between random OU and TD-MPC2 expert policies, using aligned state and pixel trajectories.The selected SigReg coefficients vary by modality and environment; prediction and SigReg losses converge within the training budget.Source: §C.2 · p. 17
C.3 TASK-RETURN EVALUATION
- LeWM’s decoded state cost attains higher return than its latent cost, so it is used across environments and observation modalities with the highest-reward tested planner.The state cost decodes predicted latents with a frozen MLP trained on aligned simulator states, then sums weighted squared goal distances over the horizon.Source: §C.3 · p. 17
D PROBING DETAILS
- State probes use episode-level activation datasets, with linear, nonlinear, and LeWM pixel-context window probes to decode physical state coordinates.The data are split by episode, with 80% for training and 20% for testing; the windowed probe uses three consecutive latent vectors.Source: D PROBING DETAILS · pp. 17–18
- Probe error is normalized to a shared shuffled-target control: 100% matches control error, while lower values indicate better decoding than control.The control uses ground-truth state inputs with randomly shuffled target correspondence; probes are evaluated on held-out episodes.Source: D PROBING DETAILS · p. 18
E.1 NONLINEAR STATE MIXING
- The nonlinear state-mixing intervention applies a fixed, frozen mixer to CartPole’s five-dimensional state before each model’s encoder.Models are trained from scratch at mixing depths zero through three, with standard pixel models included as a reference.Source: §E.1 · p. 18
E.2 VELOCITY OBSERVABILITY
- The velocity intervention supplies ground-truth CartPole velocities only during pixel observation encoding, not during open-loop latent prediction.TD-MPC2 and LeWM vary one- versus three-frame context and velocity access; DreamerV3 processes one image per timestep and maintains context recurrently.Source: §E.2 · pp. 18–19
- Figure 20 compares terminal latent and decoded-state costs across planners for the CartPole velocity and mixing checkpoints.Those comparisons use the state cost and the highest-reward planner from the task-return protocol.Source: §E.2 · p. 19
F SYNTHETIC PERTURBATIONS OF THE FAITHFUL CONDITION
The synthetic CartPole study uses an oracle planner to isolate corruption in either the transition dynamics or the observation encoder. It tracks reward alongside representation and prediction diagnostics, while the supplementary figures describe additional model and intervention evaluations.
F SYNTHETIC PERTURBATIONS OF THE FAITHFUL CONDITION
- The synthetic study uses an oracle CartPole planner and corrupts either the transition dynamics or the encoder; it does not use a learned world model.The planner uses the ground-truth transition, random-shooting MPC, and a canonical-state goal cost over 500-step episodes.Source: F SYNTHETIC PERTURBATIONS OF THE… · p. 19
- The three conditions compare canonical-state observations, an invertible state mix with conjugated dynamics, and the same mix with a noisy inverse encoder.The fixed depth-three mixer is shared across the sweep, and only one corruption channel is active per condition as relative Gaussian noise varies.Source: F SYNTHETIC PERTURBATIONS OF THE… · p. 19
- Figure 2 plots cumulative reward against linear R2, nonlinear R2, and LRC, with point color encoding corruption magnitude.The linear R2 is set to 1 for the identity representation; LRC is evaluated at horizon H = 30.Source: F SYNTHETIC PERTURBATIONS OF THE… · pp. 19–20
Supplementary evaluations
- Figures 11 and 12 describe cross-environment LRC rollouts and episode reward versus LRC at H = 30, distinguishing model family and input modality.Figure 11 compares state and pixel models with a persistence baseline; both captions report variation across model seeds.Source: Figure 11 · p. 22, Figure 12 · p. 22
- Figure 13 compares DreamerV3’s deterministic rollout, a mean over 500 stochastic rollouts, and an oracle best-matching sample.The best-matching sample is selected using the realized trajectory and is explicitly an upper bound, not a deployable estimator.Source: Figure 13 · p. 23
- Figure 15 reports 81.5% success for reproduced pixel LeWM on Reacher versus 11.5% for a random policy.The figure also marks the original study’s 86% result as a fixed reference without uncertainty.Source: Figure 15 · p. 24
- Figures 16–18 describe LeWM planner-cost comparisons and state-decoding errors, including normalized linear and nonlinear probe results.Figure 18 separates positional and velocity decoding and marks the shuffled-target control error at 50%.Source: Figure 16 · p. 25, Figure 17 · p. 26, Figure 18 · p. 27
- Figures 19 and 20 describe CartPole evaluations across nonlinear mixing depths and velocity, frame-context, and controller-cost interventions.Figure 19 tracks LRC, task reward, and normalized position and velocity errors as mixing depth increases.Source: Figure 19 · p. 27, Figure 20 · p. 28