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 section exists to define the middle ground: a compact state that is small enough to roll forward quickly yet rich enough to choose safe actions.

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.

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].$$ 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$.

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.

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.

There are three common reasons latent-space prediction wins in embodied systems. First, planning cost scales with state dimension and horizon, so compression makes search or imagination feasible. Second, the latent can align with hidden variables, such as contact mode or intent, that are easier to reason about than raw pixels. Third, rollout error is often more benign in latent coordinates because the model is asked to preserve 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, 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.