🤖 AI

Correcting a learned physical invariant improves world-model rollouts

This paper demonstrates that while world models like DreamerV3 can learn to encode physical invariants such as energy from video data, these invariants drift during autonomous rollouts, and correcting the latent state to satisfy the learned constraint significantly improves prediction accuracy.

Original authors: Richard Bao

Published 2026-08-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Richard Bao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the realm of artificial intelligence, a specific type of software known as a "world model" has emerged as a powerful tool for prediction. These systems are trained to watch video footage and learn how the world moves, allowing them to imagine what will happen next. If a robot learns to predict the path of a falling ball or the swing of a pendulum, it can use that foresight to plan its own actions. However, a critical distinction exists between a model that simply predicts pixels correctly and one that truly understands the physics governing those pixels. A model might learn to guess the next frame of a video with high accuracy by spotting patterns, without ever grasping the underlying laws of motion, such as gravity or the conservation of energy. This gap matters deeply when we ask these models to imagine scenarios far beyond what they have seen during training; if the model does not respect the fundamental rules of the universe, its long-term predictions will eventually drift into nonsense.

Researchers recently set out to test whether a sophisticated world model, trained only on video of a swinging pendulum, had actually learned the physical law of energy conservation. They used a system called DreamerV3, which had been shown to predict video frames effectively, but they stripped away any external guidance. The model was fed 64-by-64 pixel videos of a pendulum swinging in a vacuum, with no labels telling it what energy was, no physics equations, and no rewards for being correct. The goal was to see if, after training, the model had spontaneously discovered a hidden number inside its own internal state that stayed constant as the pendulum swung, just as real-world energy does.

The team first checked if the model could even be "read" to find such a number. They discovered that even before the model was trained, random guesses at its internal state could sometimes produce a number that looked like energy. This suggested that simply having the right information inside the model was not enough to prove it had learned the physics. The true test came after training. The researchers analyzed the internal state of three independently trained models and found that they had all converged on nearly the same hidden number. This number behaved exactly like the total energy of the pendulum: it remained steady as the pendulum swung back and forth, changing only slightly due to the natural limits of the simulation. Crucially, this consistency did not appear in models trained on a damped pendulum, where friction causes energy to disappear. In those cases, the search for a conserved number failed completely, proving that the models were not just finding random patterns but were specifically responding to the presence of a physical conservation law.

However, a deeper problem emerged when the researchers asked the models to imagine the future without looking at the real video. When the models were left to predict the pendulum's motion on their own, the hidden number they had learned began to drift. It slowly changed value over time, violating the very law of physics the model seemed to have learned. This drift caused the model's predictions to become increasingly inaccurate. To test if this drift was the cause of the error, the researchers applied a simple correction: at every step of the imagination, they forced the model's internal state back to the correct energy level. This intervention did not involve retraining the model or changing its code; it simply nudged the internal state back onto the correct path. The result was immediate and consistent. In all three models that had learned the physical law, this correction reduced the error in their 50-step predictions by roughly 3 percent. In contrast, applying a similar correction based on a random, meaningless number usually made the predictions worse.

The study reveals a specific failure mode in how these artificial minds operate. The models had successfully encoded a physical constraint from the pixels they watched, yet they failed to preserve that constraint when they tried to imagine forward. It is a bit like a student who can perfectly recite the rules of a game but keeps making illegal moves when playing alone. The research suggests that while world models can learn the structure of physical laws, they do not automatically respect those laws during their own internal simulations. This finding highlights a gap between learning a pattern and understanding the rule, showing that even advanced systems can hold a physical truth in their memory while simultaneously violating it in their imagination. The work does not claim to have solved this problem for all complex systems, but it provides a clear method for identifying when a model has learned a physical invariant and a concrete way to fix the errors that arise when it forgets to apply that law.

Technical Summary: Correcting a Learned Physical Invariant Improves World-Model Rollouts

Problem Statement
World models, such as DreamerV3, are capable of predicting future video observations and using those predictions for planning. However, high-quality video prediction does not necessarily imply that the model has learned the underlying physical dynamics of the observed system. A critical distinction exists between a model that merely retains information about a physical variable (decodability) and a model whose internal transition dynamics actually preserve the physical laws governing that variable. Previous methods, such as probing hidden states for physical variables, conflate these two concepts; a strong correlation between a latent state and a physical quantity (e.g., energy) can exist even in randomly initialized models or models that fail to respect conservation laws during autonomous prediction. This paper investigates whether a world model trained on pendulum video learns a scalar invariant that its own transition dynamics approximately preserve, and whether the failure to preserve this invariant contributes to prediction error during autonomous rollouts.

Methodology
The study utilizes a frozen DreamerV3 architecture trained exclusively on 64×64 video frames of a Gymnasium pendulum (Pendulum-v1) without physical labels, actors, or critics. The training data consists of 204 trajectories of 120 frames each, with actions set to zero and initial states sampled uniformly.

The methodology proceeds in three stages:

Invariant Recovery: The author analyzes the deterministic recurrent states ( $h$ ) of the trained model. They project these states into a 12-dimensional latent subspace ( $z$ ) using principal component analysis. They search for a scalar function $C(z)$ , defined as a degree-4 polynomial over the latent coordinates, that is approximately conserved along observation-conditioned trajectories. The search minimizes the ratio of within-trajectory variance to total variance (a Rayleigh quotient). To distinguish between trivially conserved functions (like $E$ vs. $E^2$ ), a secondary criterion of "flow alignment" is applied, fitting an antisymmetric operator $B$ such that the model's displacement field approximates $B\nabla C$ .
Control Experiments: To ensure the recovered invariant is a result of learning dynamics rather than an artifact of the extraction pipeline, the author employs two controls:
- Untrained Baseline: Six randomly initialized DreamerV3 models are subjected to the same extraction pipeline.
- Damped Control: Three models are trained on a damped pendulum system where energy is not conserved. The damping ratio is carefully selected to ensure trajectories do not collapse trivially while removing the conservation law.
Intervention and Rollout Correction: The author tests whether the recovered invariant $C$ is violated during autonomous imagination (rollouts). They introduce a correction mechanism that projects the latent state back toward its initial level set at every step of the imagination loop: $z \leftarrow z - \alpha (C(z) - C_0) \frac{\nabla C(z)}{\|\nabla C(z)\|^2}$ . The efficacy of this correction is measured by the reduction in pixel mean squared error over 50-step rollouts. This is compared against matched random polynomial constraints and the application of the recovered constraint to the damped models.

Key Results

Reproducibility of Invariants: Three independently trained conservative models recover nearly identical energy-like invariants, with correlations to true pendulum energy ( $|\rho|_E$ ) ranging from 0.967 to 0.975. In contrast, six randomly initialized models show a wide scatter (0.17–0.908), indicating that training makes the recovery of this specific scalar reproducible, not just decodable.
Absence in Dissipative Systems: When the same pipeline is applied to models trained on damped dynamics, no comparable conserved scalar is found. The "drift" of the recovered scalar in damped models is approximately 2,400 times larger than in conservative models, and the invariance ratio is near 1.0 (indicating no conservation). This confirms the pipeline does not manufacture invariants where none exist.
Correction Improves Rollouts: During autonomous rollouts, the recovered invariant $C$ drifts. Enforcing the constraint by projecting the latent state back to its initial level set reduces 50-step rollout error by 2.9% to 3.5% across all three conservative models.
Specificity of Improvement: Matched random polynomial constraints generally increase rollout error (median increase of +53.9%). The improvement is specific to the recovered invariant on two of the three models; on the third, some random constraints performed better, though the recovered constraint still provided a consistent improvement over the baseline.
Mechanism: The correction acts preferentially along latent directions that cause high damage to rollouts (causal sensitivity), rather than simply along the highest-variance directions. The variance rank and damage rank are negatively correlated, suggesting the model's error accumulation is driven by specific, low-variance but high-sensitivity directions.

Significance and Claims
The paper claims to demonstrate a concrete failure mode in world models: a model can learn a physical constraint from pixels (encoding an energy-like invariant) yet violate that constraint when it imagines forward autonomously. The results distinguish a dynamically meaningful invariant from a merely decodable correlate.

The author emphasizes that while the improvement in rollout error is modest (2.9–3.5%) and the study is limited to a single degree of freedom and a specific architecture, the findings are significant because they show that:

Learning vs. Respecting: A world model can encode a useful dynamical constraint but fail to respect it during prediction.
Intervention Validity: Correcting the violation of a learned invariant directly improves prediction accuracy, suggesting that the drift of this invariant is a primary source of error in long-horizon rollouts.
Control Necessity: The use of damped and untrained controls is essential to prove that the recovered structure is a product of the learned dynamics, not an artifact of the representation or the extraction method.

The paper concludes by noting that whether this phenomenon persists in systems with actions, contacts, and multiple interacting objects remains an open question, and the current results are modest in effect size but clear in their implication regarding the gap between learning structure and respecting it.

Technical Summary: Correcting a Learned Physical Invariant Improves World-Model Rollouts

More like this