Section 36.2: Forward/dynamics models; state vs. observation prediction

The best predictive target is the one the controller would pay to know one step earlier.

A Horizon-Aware Predictor

Prediction Targets And Control Interfaces

Ask two engineers to build a forward model for the same robot and you can get two systems that share no output variable at all. One predicts a 32-number latent vector, the other reconstructs a full camera frame, and both are legitimate. As Figure 36.2A illustrates, the same perception stack can branch two ways: it can predict a compact latent state, or it can reconstruct raw observations. Each branch serves a different control purpose. A forward model (any function that maps the current state and action to the predicted next state) can predict physical state, latent state, observation, reward, contact flags, or some mixture of them. Here latent state means a compact learned vector that summarizes the scene rather than a hand-engineered list of physical quantities. A common latent-state factorization is

$$ z_{t+1} = f_\theta(z_t, a_t), \qquad \hat o_{t+1} = g_\phi(z_{t+1}). $$

If the planner reasons directly in latent space, the decoder is optional at decision time. If the planner needs image-space occupancy, object masks, or human-interpretable diagnostics, the decoder becomes operationally important rather than decorative. The diagram below traces this split end to end: the dynamics function advances the latent state while the decoder branch, shown dashed, is invoked only when a downstream module actually needs a reconstructed observation.

This factorization matters for physical robots. The latent \(z\) stays compact enough to roll out hundreds of steps within a real-time control budget, yet it still encodes the contact forces, velocity, and object geometry that drive decisions. A raw-pixel forward pass at every planning step would exceed the compute budget of most onboard processors, so separating the dynamics from the decoder lets the planner run fast and reserves the decoder for human inspection or vision-based modules that genuinely need rendered frames.

How The Loss Shapes The Latent

Before reading further, consider this: if your planner has a one-millisecond budget per step, which prediction target do you keep and which do you discard?

Whichever target you decide to keep, the model only earns that target through how it is trained, so consider next how the loss shapes the latent. Training works by minimizing a combined loss: a reconstruction term pushes \(g_\phi(z_{t+1})\) close to the observed \(o_{t+1}\), and a consistency term keeps consecutive latents on a smooth trajectory under action \(a_t\). The reconstruction signal shapes the latent so it encodes scene content; the dynamics term shapes \(f_\theta\) so predictions stay accurate across multi-step rollouts. Neither signal alone is sufficient because a reconstruction loss without dynamics produces a good encoder but a useless predictor, while a dynamics loss without reconstruction can collapse the latent to a trivial constant.

Once both loss terms are in place and the latent trains cleanly, the remaining decision is what that latent should ultimately represent, which brings the state-versus-observation choice into sharp relief.

Real systems illustrate these trade-offs concretely. PlaNet (Hafner et al., 2019) and DreamerV3 plan entirely in latent state. The decoder trains as an auxiliary signal to shape the latent, then drops out at decision time, which keeps rollout cost low. The payoff can be substantial, though the exact magnitude is benchmark-dependent. A model-free agent solving the DeepMind Control Suite walker-walk task typically needs around 500,000 environment steps to reach expert performance. A latent-state world model reaches the same level in roughly 5,000 steps in the original PlaNet benchmark, a roughly 100x reduction on that specific task; other tasks and implementations show smaller gains, but the gap matters enormously when each real-robot step costs seconds of wear and risk. Practitioners call this predict in latent, discard the decoder. A 512-step planning rollout that reconstructs pixels at every step can take 4 to 8 seconds on a mid-range GPU (as measured on an NVIDIA V100 in 2023). The same rollout over the latent alone typically completes in under 400 milliseconds, a roughly 10x wall-clock gap, though the exact ratio depends on decoder architecture and image resolution, and gains this large should not be assumed to generalize to every control task without measurement. By contrast, Transporter Networks (Zeng et al., 2021) keep prediction in image space because the spatial layout of objects directly encodes which grasp is feasible and a compact state vector would lose that structure. TD-MPC2 (Hansen et al., 2023) takes a middle path: it predicts a learned latent but adds a reward predictor head so the planner can estimate value without reconstructing pixels at all.

Worked Probe

The code below contrasts a latent-state predictor with an observation predictor on a toy pushing task. The latent model predicts object position directly; the observation model predicts a rendered pixel coordinate and then recovers position from it.

# Compare a direct state predictor with an observation-space predictor.
# The state model predicts object position; the observation model predicts
# a pixel coordinate that must be converted back into world space.
from math import fabs

x_t = 0.40
action = 0.12
true_next = x_t + action

state_pred = x_t + 0.95 * action

pixel_scale = 320.0
predicted_pixel = pixel_scale * (x_t + 0.90 * action) + 2.0
obs_pred = predicted_pixel / pixel_scale

print(
    {
        "true_next": round(true_next, 3),
        "state_pred": round(state_pred, 3),
        "obs_pred": round(obs_pred, 3),
        "state_abs_error": round(fabs(true_next - state_pred), 4),
        "obs_abs_error": round(fabs(true_next - obs_pred), 4),
    }
)