Section 38.1: Why predict in latent space

"A world model earns its latent only when the latent captures what the robot can change."

A Compact Latent State

Problem First

A robot can observe a high-dimensional image stream while the task depends on a small hidden state, such as object pose behind occlusion, wheel slip, or whether a drawer is already latched. Predicting every pixel exactly is expensive and often unnecessary; predicting too little causes aliasing, where two physically different situations look the same to the controller. The middle ground is a compact state that is small enough to roll forward quickly yet rich enough to choose safe actions. The deeper reason this is hard is partial observability: the agent never sees the full world state directly. A belief state (a probability distribution over what the true hidden state might be, given everything observed so far) is the standard tool for reasoning under this kind of uncertainty, and it is what the latent state below is built to approximate.

Core Model

Model-based control under partial observability is naturally written as a belief-state problem. The latent state plays the role of a learned belief: $$z_t \sim q_\phi(z_t \mid h_t, o_t), \qquad h_t = f_\theta(h_{t-1}, z_{t-1}, a_{t-1}).$$ The deterministic memory \(h_t\) carries long-range context, while the stochastic variable \(z_t\) captures the uncertainty that remains after seeing the new observation. Concretely: \(q_\phi\) is the encoder network (subscript \(\phi\) denotes its trainable weights) that turns a new observation into a latent sample, and \(f_\theta\) is the dynamics network (weights \(\theta\)) that advances memory given the previous latent and action; the two networks are trained jointly so that rolling \(h_t\) forward stays consistent with what \(q_\phi\) would have inferred from a real observation.

Prediction matters because planning and policy learning happen over future latent states: $$\hat z_{t+k+1} \sim p_\theta(z_{t+k+1} \mid h_{t+k+1}), \qquad J(\pi) = \mathbb{E}\Big[\sum_{k=0}^{H-1} \gamma^k r(\hat z_{t+k}, a_{t+k})\Big].$$ Here \(p_\theta\) is the same dynamics model used to imagine forward without waiting for a real observation, \(\gamma\) is the discount factor that down-weights distant rewards, and \(J(\pi)\) is the expected return the policy \(\pi\) is trying to maximize over a planning horizon of length \(H\). The important question is not whether \(z_t\) reconstructs pretty images. The question is whether the rollout preserves the reward-relevant and safety-relevant variables well enough for the action chosen at time \(t\) to still look sensible at time \(t+H\). This is also why value estimation, a critic network scoring how good a state is, needs the same decision-sufficient latent as the planner: a critic trained on \(z_t\) can only rank two states correctly if \(z_t\) still distinguishes them, so a value function built on an aliased latent will misjudge exactly the states where the correct action matters most.

For that reason, control-relevant abstraction is stricter than compression alone. A latent variable that discards background texture is useful; a latent variable that discards contact mode, object identity, or actor intent is dangerous because the planner will optimize the wrong future. Figure 38.1B contrasts the two rollout paths and shows why the compact latent route is the one the planner can afford to search inside a tight control loop.

Minimal Probe

The probe below shows the basic economic argument for latent planning. It compares the cost of rolling out pixel states versus compact latent states, then checks whether the compressed state still tracks the task variable the planner needs.

# Compare rollout cost in pixel space and latent space.
# Then verify that the latent variable still tracks task progress.
import numpy as np

pixel_dim = 84 * 84 * 3
latent_dim = 64
horizon = 15
transition_cost = np.array([pixel_dim * horizon, latent_dim * horizon])
task_progress = np.array([0.15, 0.33, 0.49, 0.71])
latent_proxy = np.array([0.12, 0.31, 0.52, 0.69])
tracking_error = np.abs(task_progress - latent_proxy).mean()
print(
    {
        "pixel_rollout_scalars": int(transition_cost[0]),
        "latent_rollout_scalars": int(transition_cost[1]),
        "mean_progress_error": round(float(tracking_error), 3),
    }
)

Expected behavior: The latent rollout is dramatically cheaper to evaluate, yet the average task-progress error stays small. If the compression ratio improved while the progress error exploded, the latent state would be too lossy for control.

Practical Recipe

1. Write down the hidden variable the task actually depends on, such as contact mode, object pose, or progress-to-goal.
2. Define how the latent state should expose that variable to the planner or critic.
3. Measure whether rollout cost drops faster than decision quality degrades.
4. Stress the model with occlusion, delay, or an unseen distractor, then inspect which latent coordinate or prediction head fails first.

Latent-space prediction wins in embodied systems for three reasons. First, planning cost scales with state dimension and horizon, so compression makes search or imagination feasible. A 15-step rollout in pixel space touches roughly 330,000 scalar operations per step. The same rollout in a 64-dimensional latent touches fewer than 1,000. That ratio determines whether online planning fits inside a 50 ms control loop at all. Second, the latent can align with hidden variables, such as contact mode or intent, that are easier to reason about than raw pixels. The sample-efficiency gap this produces is striking. A model-free policy learning entirely from real robot interaction typically needs 40,000 to 80,000 episodes to solve a tabletop manipulation task. A latent world model that plans inside imagination reaches the same success rate in roughly 300 real episodes, because the robot rehearses thousands of imagined futures for every single physical attempt. Third, rollout error is often more benign in latent coordinates because the model preserves task structure instead of every texture and shadow.

The tradeoff is aliasing. If two states look similar in the learned representation but require different actions, this is called the decision-aliasing failure mode, and the controller can become overconfident. That is why long-horizon visual plausibility is not enough. The deployment question is whether the latent state remains decision-sufficient under the disturbances that matter for the robot, vehicle, or interactive world being built. Why that question carries such weight becomes clear once we notice that aliasing on hardware is not merely a scoring error. On a physical robot, aliasing has irreversible consequences. If the latent conflates "drawer slightly open" with "drawer fully closed," the planner may command a pull force large enough to tear the hardware. Unlike a simulation reset, physical damage cannot be undone, and the robot cannot retry from a corrected state. No reset restores the torn drawer. Embodied systems therefore require the latent to discriminate precisely the states where wrong actions cause harm, not just the states where wrong actions reduce reward. Aliasing arises from how encoders allocate capacity. Reconstruction loss over the full observation rewards capacity in proportion to visual area, so a large uniform floor dominates the gradient while a small fingertip-contact indicator, worth only a handful of pixels, gets underweighted. The encoder then collapses physically distinct contact states into one latent cluster whenever they look alike, producing aliasing exactly where the task is most sensitive. If capacity allocation is what plants the flaw, the next thing to pin down is which operating condition finally trips it.

Before reading the next paragraph, guess where this breaks first: a mobile robot trained indoors deploys outdoors on a wet day. Does it fail because rollout cost spikes, because the planner runs out of horizon, or because the encoder maps a novel observation to the wrong latent region?

When it breaks in the wild

Consider a specific case. A mobile robot trained indoors with DreamerV3 achieves a 64-dimensional latent that tracks pose and obstacle distance well enough to plan 10-step paths with under 3 cm error. Tested outdoors, sunlight reflections on wet pavement produce observation patterns outside the training distribution. The encoder maps these novel inputs to latent coordinates near "clear floor," the planner imagines safe forward motion, and the robot drives into a puddle. The rollout math holds; the failure is a decision-sufficiency failure caused by distribution shift. This is when latent prediction breaks: not when rollout cost rises, but when the encoder can no longer reliably place novel observations in the part of latent space that the planner was trained to interpret correctly. A world model that predicts the wrong future with complete confidence is more dangerous than one that admits it does not know.