World Models: Predicting the World After Actions in Internal States

Planning is expensive and dangerous because if every candidate action had to be tried in the real world, the cost of trial and error could be too high to bear. A collision by a real robot could damage hardware, a failure in a game could lose long-accumulated progress, and a write operation in a production system could contaminate real data. An RL agent needs a way to roughly know what consequences an action will have before actually acting, so that obviously poor options can be filtered out first.

World models exist for exactly this purpose. Their core idea is to compress high-dimensional, noisy observations into an internal state, and then predict, based on this internal state: after taking an action, what the next state will be, how much reward will be received, and whether the current episode will terminate. Once it has this predictive ability, the RL agent can simulate the consequences of actions in its "imagination"—first mentally running through several candidate options, filtering out clearly bad choices, and finally actually executing only one candidate. In this way, the number of trial-and-error attempts in the real environment is greatly reduced, and the cost is correspondingly lowered.

It is worth emphasizing its inputs and outputs: the input is the current state and a candidate action, and the output is a prediction of the next state, reward, or termination probability. Together these three constitute all the key information the RL agent needs for decision-making: the next state determines "what step comes next", the reward determines "whether this step is worth it", and the termination probability determines "whether this path has reached its end."

One key boundary must be made clear: a world model is not a complete copy of the world, but an approximate simulator that retains the necessary variables for the decision task. It only cares about the variables that truly affect the consequences of actions, and the remaining irrelevant details are compressed or discarded. This means that being able to generate a realistic video does not equal predicting the consequences of actions correctly. The pixels may look reasonable, but the model could be completely wrong about collisions, inventory, or rewards. The basis for judging whether a world model is good or bad is whether its predicted action consequences are reliable, not whether the generated image looks visually appealing.

A world model that can be used for planning must at least be able to answer four quantities: what the current state is, what the state will become after taking a certain action, how much reward can be obtained at this step, and optionally whether the state can be reconstructed back into an observation. These correspond to several components in the model, and can be written as the following set of equations:

zₜ = E(o₍≤ₜ₎), ẑ₍ₜ₊₁₎ = f(zₜ, aₜ), r̂ₜ = g(zₜ, aₜ)

Let's break down the meaning of each step one by one.

The encoder E compresses the current and historical observations o₍≤ₜ₎ into an internal state zₜ. Observations may be high-dimensional, such as a frame of video or a stream of sensor readings; the state is a compact representation chosen internally by the model. The dynamics function f is the core transition model: given the current state zₜ and a candidate action aₜ, it predicts the next state ẑ₍ₜ₊₁₎. The “hat” here indicates that this is a predicted value, not the state that actually occurs in the real environment. The reward model g uses the same inputs zₜ and aₜ to estimate the action value r̂ₜ for this step—that is, whether this step is worth doing. Together, these three components allow the agent to advance the future in imagination: first encode the state, then repeatedly use f to predict the next state and g to evaluate the return at each step.

Optionally, the model may also include a decoder that reconstructs the state z back into an observation for auxiliary training. Its purpose is to provide a reconstruction target for comparison, helping the encoder learn more informative states, but it is not strictly necessary for decision-making.

Whether the state is “sufficient” is a key question. If the environment is partially observable, looking only at the current frame is often insufficient to assess the situation—for example, from a single frame you cannot tell whether an object is moving left or right, nor how much of an occluded inventory remains. In this case, the state cannot be determined solely by the current observation; it must integrate historical information, or explicitly maintain a belief distribution over hidden variables. This is exactly why the input to encoder E is written as o₍≤ₜ₎ rather than a single oₜ: it encodes all observations up to the current moment, thereby bringing cross-time information such as velocity and hidden inventory into the state.

From this we can give a precise meaning of “state sufficiency”: given a state z and the subsequent action sequence, predicting the future no longer requires earlier history. In other words, z compresses all past information relevant to decision-making, and past observations can be safely discarded. It is particularly important to clarify that state sufficiency is a modeling goal, not a property that automatically holds as long as the vector dimension is made large enough. A large dimension only provides space to hold information; if the encoder does not actually put the needed information in, the state can still lose key quantities no matter how long it is.

A world model is not guessed out of thin air; all of its predictive ability comes from trajectories collected in the real environment. A trajectory consists of several steps, and each step is denoted as (o, a, r, o′): given observation o, action a was taken, reward r was received, and the next observation o′ was reached. The goal of training is to have the model learn from a large number of such real samples “If I do this, what happens next?” During training, all parts of the model are optimized simultaneously: the state representation, the prediction of the next state or next observation, the reward prediction, and the termination prediction are all adjusted to be consistent with the real trajectories.

The choice of training signal determines where the model spends its capacity and also determines its blind spots. There is an obvious tension here: if only pixels are reconstructed, the model tends to spend a lot of capacity on details irrelevant to decision-making, such as background texture; no matter how well the image is reconstructed, it may not be useful for control. Conversely, if only rewards are predicted, the model may discard factors that are important for future actions but have no reward at the current step—such as an intermediate action that does not score now but paves the way for later scores. Therefore, in practice multiple objectives are usually combined, each preserving part of the information and compensating for each other's blind spots.

The objective of observation reconstruction is to preserve visual details. Its value lies in providing the encoder with a comparable reconstruction signal, but it is also the easiest to slip into the blind spot of “beautiful but irrelevant to control.” The objective of latent transition is to preserve state dynamics, that is, to make the state evolve according to the true dynamics in the latent space; its risk is representation collapse—the model may map all inputs to almost the same state, making the transition meaningless, in which case additional constraints are needed to prevent this degradation. The objective of reward and termination prediction is to align directly with the task signal; its blind spots lie in sparse rewards and proxy loopholes—when rewards rarely appear, pure reward supervision is difficult to propagate, and once the reward itself has exploitable loopholes, the model will learn to exploit flaws in the reward function rather than complete the real task. The objective of contrastive prediction is to enable states to distinguish different futures; its blind spots lie in the definition bias of negative samples and augmentation methods. If negative samples are chosen poorly or the augmentation method is inappropriate, the standard for “distinguishable” will be distorted.

These objectives must be used in combination, because no single objective can cover all necessary information at the same time. In addition, there is an easily overlooked requirement: the exploration data used for training must cover the actions that the planner may choose in the future. If the model is trained only on trajectories of a certain type of action, but the planner selects actions outside the training distribution, then the model's predictions for those actions will be unreliable, and subsequent planning will be built on unreliable predictions.

A concrete small scenario can make clear how planning actually happens inside the model. Imagine a cart that is still 3 cells away from the goal, and there is a pit in the 2nd cell directly ahead of it. Now there are two candidate routes: go straight and detour. What the world model needs to do is to roll out two steps in imagination for each of these two routes, predict the reward at each step, and then discount the rewards into total returns for comparison.

The predicted reward at each step is denoted as r̂₁ and r̂₂, representing the predicted reward at step 1 and step 2, respectively. To compare rewards at different time points, a discount factor γ is used to discount future rewards to the present; here γ = 0.9, so the two-step predicted return is computed as r̂₁ + γ × r̂₂, and the true two-step return is likewise computed with the true rewards using the same formula.

First consider going straight. At step 1 it moves forward one cell, with predicted reward +2 (closer to the goal); at step 2 it continues forward and steps right into the pit, with predicted reward −10. Therefore the predicted return is +2 + 0.9 × (−10) = −7.0, exactly matching the true two-step return of −7.0. Now consider the detour. At step 1 it goes around, with predicted reward 0; at step 2 it moves toward the goal again, with predicted reward +2. The predicted return is 0 + 0.9 × 2 = 1.8, likewise equal to the true 1.8. The planner compares these two values: −7.0 versus 1.8, and naturally chooses the detour. This is exactly the value of two-step imagination—it lets the planner see the disaster that immediately follows a short-term positive reward, thereby avoiding the short-sighted choice of going straight.

The crucial third row reveals how the model's error is amplified. Suppose the world model missed the pit information during training—it did not learn from observations that there is a pit ahead. Then for the straight path, the model will still predict +2 at step 1 and +2 at step 2, giving a predicted return of +2 + 0.9 × 2 = 3.8. This number is higher than the detour's 1.8, so the model will rank the straight path first, while the true return of going straight is actually −7.0. This is a wrongly overestimated estimate: the model thinks it is taking the safest route, but it is actually taking the most dangerous route.

The meaning of this row is worth emphasizing: the planner can only optimize predicted returns, not directly optimize true returns. When the world model misses the pit, it will rank the most dangerous route first; moreover, the stronger the planning ability and the more thorough the search, the more stably it will exploit this error—because a stronger planner will more thoroughly mine the model's internal wrong conclusion that 'going straight is beneficial.' Therefore the solution is not to 'search more,' because no amount of searching can find a pit the model does not know about. The correct direction is: collect data near the pit to let the model fill in this information, express predictive uncertainty, penalize out-of-distribution trajectories, and shorten the length of open-loop rollouts. This last point directly raises the next question: why does executing only the first step and then re-observing reduce the error caused by long rollouts?

Since the planner has already rolled out a ten-step optimal action sequence in the model, why not execute the whole sequence and instead only execute the first step? The answer is: the model's predictions accumulate error with rollout length; the farther it rolls, the less trustworthy the later steps become. The approach of model predictive control is: each round, roll forward H steps within the model, select the action sequence with the highest predicted return, but actually execute only the first action in the sequence. After execution, the environment returns a real observation, the agent uses this real observation to update its state, and then restarts a round of rollout and optimization. And so on: plan H steps, execute 1 step, observe again, plan again.

The core benefit of doing this is frequent correction. Because each round only takes a small step and then stops to correct with real observations, prediction errors have no chance to keep accumulating over a long chain; at the same time, the agent can also respond to the behavior of other entities in the environment or sudden changes—if an opponent or the environment changes the situation midway, the next round of planning will immediately incorporate this new information.

The cost is equally clear: each step requires replanning once, and the computational burden increases significantly. Therefore, a trade-off must be made among the rollout length H, the number of candidate actions, and real-time latency. The larger H is, the farther a single round of planning can see, but the more single-step prediction errors accumulate, and the heavier the computation for a single round becomes; the more candidates there are, the more thorough the search, but the more it slows down the response. In real systems, these three are mutually squeezing budgets.

One more point must be made clear: the “checkpoints” used for correction must come from the real environment. Only the positions actually returned by actuators, inventory quantities, or test results are trustworthy bases for correction; a “now we have arrived at such-and-such location” that the model itself continues to write out must never be used as a checkpoint. If the intermediate states written out by the model are treated as real states for correction, that is equivalent to letting the error closed loop prove itself, losing the meaning of using real observations to interrupt error accumulation. The entire value of model predictive control lies in its constantly returning to the real environment to demand the truth of the next step.

Intuitively, one might ask: if the model only errs by a little on average at each step, why does rolling out the plan over a long horizon cause it to completely deviate from the true trajectory? The answer lies in how errors accumulate over time. Let the upper bound on the state error at each step be ε, and let the sensitivity of the dynamics to the state be approximately L—this measures how much a small perturbation in the previous state is amplified or reduced by the next transition. Then the rough recurrence for the error can be written as:

e₍ₜ₊₁₎ ≤ L · eₜ + ε

Reading this expression from left to right: the error at the next step is at most equal to the previous step's error amplified by the dynamics by a factor of L, plus the new error ε introduced at the current step. The two terms combine, so the error builds up step by step.

The value of L determines the overall trend of the error. When L < 1, the previous round's error is reduced at each step, and the error may be contained within a controlled range; when L ≈ 1, the previous step's error is passed down almost unchanged, and with the new ε added at each step, the error accumulates approximately linearly; when L > 1, the previous round's error is not only not attenuated but amplified, and the error expands rapidly, so a long rollout will soon deviate completely from the true trajectory. In real high-dimensional dynamics, L is often not less than 1, which is exactly the mathematical source of danger in long-horizon planning.

There is another factor that makes the situation worse: the actions chosen by the planner are not the randomly sampled actions found in the training data, but actions that specifically seek out high-return regions predicted by the model. This means the planner actively and with bias pushes the trajectory into places the model “considers advantageous”—and these places may be precisely the regions where the model's predictions are least reliable, because the training data may not cover them. As a result, the distribution of errors itself shifts: a model that looks good on a random validation set can be wildly wrong on trajectories induced by the planner.

This leads to a direct evaluation conclusion: one should not report only the one-step mean squared error (MSE) on a random validation set. One-step MSE reflects only the average single-step error under a random action distribution and completely misses the consequences of error accumulation along planning trajectories. A meaningful evaluation should also measure multi-step prediction error, reward error, calibration of predictions, and the final true return on trajectories induced by the planning policy. Only these metrics can reveal whether the error will spiral out of control when the planner actually uses this model to roll out the future.

A tricky situation is when two frames look exactly the same but require completely opposite actions. Suppose the cart is moving, but from a single frame you cannot tell whether it is moving left or right. If the encoder maps these two different true states to the same z, state aliasing occurs. The consequence of aliasing is direct: the same state z and the same action a correspond to two different futures in the real world. At this point, if the model uses a deterministic mean to predict the next state, it is likely to give an 'intermediate result' that does not actually exist in reality—for example, an average position of the cart that is both left and right, and this position does not correspond to any real physical evolution.

The way to handle aliasing is to explicitly carry hidden information in the state and pass uncertainty clearly to the planner. Specific approaches include: using historical recursion or memory to encode the differences between the past few frames into the state; using probability distributions or particle beliefs to represent the hidden state instead of a single point estimate. In this way, the model no longer compresses 'left' and 'right' into the same point, but retains their respective possibilities and weights.

Uncertainty has more than one source besides insufficient observation; you can examine the corresponding strategies by source. Environmental randomness means that the same action itself has random outcomes in the real world—for example, with the same push, different ground slipperiness leads to different sliding distances for the cart. For this randomness, the correct approach is to predict a distribution and optimize risk-sensitive returns during planning, rather than focusing only on the expected value. Insufficient observation means that a key quantity is completely invisible in the observation, such as speed or hidden inventory; the countermeasure is to integrate history or maintain a belief state. Insufficient data means that some regions were not covered during training, such as actions near a pit edge that have never been seen; the countermeasure is to use model ensembles and confidence bounds to measure uncertainty, and return to the real environment to supplement exploration. Environmental change means that rules or the behavior of other agents changed over time, and the patterns the model learned have become outdated; the countermeasure is to continuously perform drift detection and online correction.

The common thread among these different sources is to treat 'I don't know' as a first-class citizen. A world model not only predicts the most likely outcome, but also tells the planner how unreliable that prediction is. Only after the planner receives the uncertainty can it make a reasonable choice between 'advancing boldly when the model is very certain and being cautious and conservative when the model is very unsure,' rather than betting everything on false certainty.

There is a failure mode worth examining separately: a planner obtains extremely high returns in the model, but performance in the real environment actually gets worse. This is model exploitation. Its mechanism is: during search, the policy finds a loophole in the model, and this loophole does not exist in the real world. A typical example is that the model incorrectly predicts that some unrealistic motive will repeatedly score, so the planner repeatedly selects that action, racking up increasingly high returns in imagination, while in the real environment it gains nothing or even suffers damage.

A direct way to detect model exploitation is to compare two returns: the return predicted by the model and the return obtained from real replay. The difference between them is called the model exploitation gap. The larger this gap, the less grounded in reality the returns that the policy “earns” in imagination are. To locate where the gap comes from, you can slice by two dimensions: action novelty (the further an action deviates from the training distribution, the more likely the model is to give wildly inaccurate predictions) and rollout length (the longer the rollout, the more likely errors are to accumulate and amplify).

There are several ways to curb model exploitation, and they can be used in combination: keep the policy from deviating too far from the training data distribution; impose penalties on states where the model is uncertain, so the planner avoids regions it is not confident about; use disagreement among model ensembles to measure uncertainty, and trust less wherever disagreement is large; and periodically put candidate actions back into the real environment for verification, and add the counterexamples found during verification to the training data so that the model can patch these loopholes in the next training run.

One final boundary must be made clear: evaluation in the real environment must be isolated and permission-limited. The reason is that “for the sake of validating the model” cannot justify authorizing unlimited trial-and-error with high-risk actions. Putting candidate actions back into the real environment for verification has costs; it may cause real collisions or real data contamination. Therefore, verification can only be carried out within a controlled and restricted scope, and the agent must never be allowed to freely engage in trial-and-error in the real environment just to close the model exploitation gap. The goal of discovering and fixing model exploitation is to realign imagined returns with real returns, not to use the real environment as an unlimited safety net.

Putting the previous steps together, we can see a complete causal chain of the World Model from problem to verifiable practice.

The starting point is: learn state representations and transition dynamics from real trajectories. The RL agent first collects trajectories in the real environment, training the encoder, dynamics, reward, and termination predictors. This is the foundation of all imaginative ability—what the model can predict is determined entirely by what is learned at this step.

The second step is: roll out candidate actions inside the model. Given the current state, the planner enumerates several candidate actions, repeatedly using the dynamics function to advance the future in imagination, producing a series of virtual trajectories.

The third step is: score candidates by predicted return and risk. Only looking at predicted return is not enough; uncertainty and model exploitation risk must also be factored in, choosing a truly reliable plan rather than the one the model thinks is most profitable.

The fourth step is: execute only near-term actions. Model Predictive Control executes only the first action after each round of rollouts, keeping the error risk of long-horizon prediction within a single step.

The fifth step is: correct the state with real observations. After execution, retrieve the next observation from the real environment and update the state, so the next planning is built on new facts, not on the model's own continuation.

The final step is: feed counterexamples of model exploitation back into training. All counterexamples with high return in imagination but low return in real replay are collected and added to the training data, so the model can patch these holes in the next iteration.

These six steps form a closed loop: learning transitions → in-model rollout → risk scoring → executing only near-term actions → real correction → counterexample feedback. Each step provides more reliable input for the next step, while each step also constrains the error of the previous step. Whether the World Model is ultimately usable depends on whether this closed loop can maintain continuous alignment between 'prediction' and 'reality': the more accurate the model's predictions, the more valuable planning becomes; and each return to the real environment is both a correction of the model and a reinvestment in planning capability.