Pretrained Visual Models and DINO-WM › Feature-Space Dynamics Prediction

Now that we have a state, we need dynamics. The predictor is a function

$$\hat{z}_{t+1} = F_\phi(z_t, a_t)$$

parametrized by learnable weights $\phi$, where $z_t$ and $z_{t+1}$ are frozen encoder outputs and $a_t$ is the action. The encoder is a fixed map from pixels to vectors. Its weights do not receive gradients from the dynamics loss; the toy predictor below is trained with Adam gradient steps. The predictor operates in the coordinate system the encoder defines.

The natural objective is plain regression on the encoder's output:

$$\mathcal{L}(\phi) = \mathbb{E}_{(z_t, a_t, z_{t+1}) \sim \mathcal{D}} \left[ \left\| F_\phi(z_t, a_t) - z_{t+1} \right\|^2 \right].$$

Notice what changed relative to a pixel-reconstruction objective. There is no decoder in this dynamics loss. The target is the next observation's frozen features, so the quality of the learned state depends on what the encoder retains. DINO-WM can train a separate, optional diagnostic decoder; the small example below also trains a decoder while pretraining its toy encoder. Neither decoder is used for latent rollout or planning.

This design offers three potential advantages:

• Dynamics training does not require pixel reconstruction. The transition predictor need not carry a pixel-generation head or optimize its loss. That does not make the encoder's pretraining free, nor does it automatically reallocate a fixed number of parameters.
• A pretrained feature target may suppress nuisance detail. Depending on the task and encoder, feature error can emphasize object-level changes more usefully than pixel error. Neither metric guarantees that it weights task-critical details correctly.
• The goal can use the same representation. Encoding a goal image gives a computable feature-distance cost without a separately trained reward model. Choosing that distance and verifying that it tracks success are still design and evaluation work.

There are two corresponding limitations:

• Feature-space distance is a proxy for task success. Two visually different futures can map to nearby features, and two visually similar futures can map to distant features. A representation trained for invariance may collapse distinctions that a controller cares about. Encoded history or actions can sometimes recover a state distinction absent from one frame. If a task-critical distinction is absent from every input available to the predictor, improving that predictor alone cannot restore it. Separately, a fixed goal cost built from aliased goal features cannot distinguish task-distinct goals; that requires more information or a different representation or cost.
• Latent predictions are not directly inspectable as images. A separately trained probe or diagnostic decoder can help, but its reconstruction quality must be checked; closed-loop behavior remains the decisive test.

A common practical refinement is to normalize features before regression. Squared differences in high-variance dimensions may dominate a raw Euclidean loss or planning distance; a large constant offset alone does not. One option is

$$\tilde{z}_t = \frac{z_t - \mu}{\sigma},$$

where $\mu$ and $\sigma$ are estimated on training data, with a nonzero floor on any per-dimension scale. Per-dimension scaling changes the relative weighting of feature errors and can amplify nearly constant noise. A single global scale changes optimization units but leaves Euclidean rankings between candidate plans unchanged; it does not balance dimensions. Whether high-variance dimensions carry task signal must be tested rather than assumed.

Three useful design choices illustrate the tradeoffs here. MLPs and transformers are possible predictor backbones; a residual or delta target is an output parameterization that can be used with either. Convolutional, recurrent, graph-based, and state-space predictors are further options. The backbone choice depends in part on spatial structure and the cost of modeling interactions between distant scene regions.

MLPs. A per-token MLP applies shared weights to each patch, optionally with the action broadcast to each token. It is cheap but cannot directly exchange information between patches. A global MLP instead flattens all patch tokens before prediction; it can mix locations, but its input size and weights are tied to a fixed grid and it lacks an explicit spatial-neighborhood bias.

Transformers over patch tokens. Self-attention can mix information across locations and frames. Some architectures support variable token counts, but changing resolution also changes positional structure and computational cost; it is not automatic resolution transfer. DINO-WM conditions a transformer on a short frame/action history, concatenating an action embedding to each patch feature. Its temporal masking prevents access to future frames while allowing within-frame interaction; it is not a rule that each token sees only earlier tokens.

Residual or delta predictors. Instead of predicting $z_{t+1}$ directly, predict the change:

$$F_\phi(z_t, a_t) = z_t + G_\phi(z_t, a_t).$$

When transitions are small in these coordinates, learning a delta can be convenient. The residual form alone does not initialize $G_\phi$ near zero: ordinary randomly initialized linear layers still produce nonzero outputs. Zero-initializing the last layer would create an identity starting point, but the toy model below does not do that. Persistence is a separate baseline to measure, not an automatic guarantee of residual training.

The action has to enter the predictor somehow, and the choice matters whenever the dynamics are strongly action-dependent. A predictor that is insensitive to the action is useless for control, and a predictor that is over-sensitive to the wrong features of the action will chase spurious correlations. Common options include:

• Concatenate the action vector to the feature vector before the first linear layer. Simple and works well when the dynamics are globally action-linear-ish.
• Broadcast or concatenate an action embedding to every patch token, as in DINO-WM, so each token receives the command.
• Insert an action token in the sequence, letting attention mix action and visual information. This is an architectural alternative, not the DINO-WM mechanism just described.
• Concatenate a short history of actions, which helps when the visible effect of an action is delayed, such as when a command starts a process that only becomes visible a few frames later.

Across robots, actions differ in shape, units, timing, and physical effect. Learning a shared action representation is one research direction discussed in Part VI: Learning World Models at Scale, but an embedding alone cannot make embodiments interchangeable; calibration and embodiment-conditioned dynamics remain necessary.

Prediction error can accumulate over long learned-model rollouts. Understanding when that happens helps you choose horizons, budgets, and training procedures; an exact model or a contracting system is a counterexample to inevitable degradation.

Suppose one-step additive errors are independent, zero-mean vectors with root-mean-square norm $\epsilon$, and the transition does not amplify perturbations. Then the root-mean-square accumulated error is at most $\epsilon\sqrt H$ under this nonexpansion assumption. Equality holds for an identity or norm-preserving transition that carries each error forward unchanged; contraction can give a smaller value. Correlation or an expanding transition Jacobian can increase error beyond this bound. Learned rollouts also feed predicted features back into the model, potentially moving off the training distribution. Neither the rate nor a sudden-failure threshold follows from one-step error alone.

Possible mitigations, with task-dependent effectiveness:

• Train on predicted states. Add multistep losses or mix model predictions with observed inputs during training; the latter requires a specified schedule rather than merely invoking “scheduled sampling.” This is distinct from Dreamer-style imagination, which trains actor/value components using a learned world model.
• Keep the rollout short and replan. Feedback can limit how long an inaccurate open-loop forecast is trusted. Later observations or encoded history may reveal state hidden in one frame, but replanning cannot restore a distinction absent from every available input or from an aliased fixed goal cost.
• Train with suitable input perturbations. Noise on inputs may help local robustness if it is paired with meaningful targets and evaluated; independent zero-mean noise on regression targets alone does not teach robustness to perturbed inputs.
• Ensembles. Predictor disagreement can expose some model uncertainty, as in PETS, but common-mode errors and representation failures can remain invisible.

An important subtlety: better one-step feature error does not always mean better closed-loop behavior. If the representation or goal cost is misspecified for the task, feature MSE can approach zero while control still fails. History or a later observation can sometimes disambiguate a one-frame state alias; fitting the same proxy more closely cannot recover information absent from all supplied inputs or distinguish goals collapsed by the fixed cost.