Pretrained Visual Models and DINO-WM › Action-Conditioned Latent Rollouts

Michael BrenndoerferJuly 20, 20265 min read
World Models HandbookAll chapters
Part 9 of 12Foundation and World-Action Models
Chapter 50Pretrained Visual Models and DINO-WM › Action-Conditioned Latent Rollouts
Reading section 4 of 10 · Action-Conditioned Latent Rollouts · 5 minSection 4 of 10 · 5 min

Book contents

Section 4 of 10 · 5 min read

Pretrained Visual Models and DINO-WM › Action-Conditioned Latent RolloutsOne section · Sections in this chapter4 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Once the predictor exists, everything else is planning. A rollout is just recursive application:

z^t=ztobs,z^t+k+1=Fϕ(z^t+k,at+k),k=0,…,H−1.\hat{z}_{t} = z_t^{\text{obs}}, \qquad \hat{z}_{t+k+1} = F_\phi(\hat{z}_{t+k}, a_{t+k}), \qquad k = 0, \ldots, H-1.

At every step after the first, the input is the model's own previous output. This recursive, self-fed prediction is often called imagination. Errors can accumulate as it rolls forward, which is why the diagnostics in the code below matter.

Specifying the task with a goal image

The elegant part of the DINO-WM recipe is how it turns a rollout into a plan. There is no reward function and no value network. Instead, the task is specified by a goal observation ogo_g, and its encoded form

zg=fθ(og)z_g = f_\theta(o_g)

becomes the target. The planning objective is

at:t+H−1∗=arg⁡min⁡at:t+H−1∈AHd ⁣(z^t+H, zg),a^*_{t:t+H-1} = \arg\min_{a_{t:t+H-1}\in\mathcal A^H} d\!\left(\hat{z}_{t+H},\, z_g\right),

where dd is a distance in feature space, such as mean squared error over patch tokens, and A\mathcal A is the feasible action set. This terminal-state cost matches the DINO-WM planning objective. The toy implementation below deliberately uses a sum of intermediate-state costs to reward progress throughout its short horizon; that is a different objective and can select a different plan.

Three design consequences follow.

A test goal can be specified by one image. The planner does not need a task-specific reward model or demonstration for that goal, but the feature distance is still a chosen cost and the dynamics model needs offline observation–action trajectories. Some DINO-WM training datasets replay noisy expert trajectories; “zero-shot” here concerns a new test goal, not an absence of behavioral data.

Patch-level costs retain spatial indexing. Comparing corresponding patches can expose where images differ, unlike one pooled vector. The cost gradient need not point toward a useful action: encoder invariances, the learned transition, and occlusions can all distort it. A global embedding can still support search if it preserves goal-relevant information, but it offers less direct spatial diagnosis.

The cost is only as good as the representation and distance. If the encoder is insensitive to a task-critical difference, the cost can miss it; if it responds strongly to irrelevant variation, the planner can pursue the wrong match. DINOv2's behavior under position, texture, and lighting changes must be tested in the intended domain rather than presumed.

One practical caveat: the goal image may not show the exact configuration the agent must reach. If the agent needs to grasp an object, an image of the object at a destination may not show grip quality or contact force. Some tasks need additional process or contact information; a goal trajectory, tactile feedback, explicit constraints, or a final-approach policy are possible ways to supply it. The limitation is structural: a single frame does not specify the path or interaction process by which the configuration is reached.

Searching for actions

With the cost defined, planning optimizes an action sequence. If the predictor and cost are differentiable, one can backpropagate through a rollout; Part VII: Planning and Agency covers that option. The DINO-WM experiments use a sampling-based planner, CEM. It accommodates clipped action candidates and does not require cost gradients, but finite samples can miss good plans too. Neither family has a general advantage for every learned cost.

The standard choice is the cross-entropy method (CEM), a derivative-free optimizer:

  1. Initialize a Gaussian distribution N(μ,σ2I)\mathcal{N}(\mu, \sigma^2 I) over action sequences of length HH.
  2. Sample MM candidate sequences from the distribution.
  3. Roll out the predictor for each candidate and evaluate the cost against zgz_g.
  4. Keep the KK lowest-cost candidates (the elites).
  5. Refit μ\mu and σ\sigma to the elite set.
  6. Repeat for a fixed number of iterations. Return μ\mu.

The intuition behind CEM is that it repeatedly refits a sampling distribution to low-cost candidates, often concentrating search around promising action sequences without needing gradients. Its sampled variance need not shrink on every iteration.

Then use the result as model predictive control: execute only the first action (or the first few), observe the new frame, re-encode it, and replan. Frequent replanning limits open-loop exposure to model error, but cannot guarantee recovery from a poor representation, bad action coverage, or an unsafe intermediate action.

The horizon HH and sample count MM are important compute and accuracy controls. Longer horizons let the planner anticipate further, but later predictions may be less accurate; their quality must be measured. More samples cover more candidate sequences, with rollout cost roughly linear in MM for a fixed horizon and model. Elite count, iteration count, covariance handling, and the learned cost also matter. A longer horizon makes each candidate rollout more expensive, so horizon and sample count compete for a finite compute budget.

Receding-horizon control

Model predictive control is the loop “plan HH steps, execute one, replan.” Feedback from new observations can correct some model error. Re-encoding an observed frame also replaces the imagined latent at the next planning cycle. Neither operation guarantees task success or cancels a systematic bias within the executed action.

Does the cost being minimized correspond to the thing you want?

Sometimes, and this is the crux of evaluation for this family of models. Feature-space distance to a goal image is a proxy, and optimization can exploit a misaligned proxy. Two failure patterns recur:

  • Latent-space shortcuts. The planner may favor an action sequence whose predicted features look close to the goal but whose real outcome does not. Representation aliasing and model error are two different causes. Per-dimension whitening can amplify noisy small-variance directions, but it is not a necessary condition for failure.
  • Proxy insensitivity. The cost may barely change across task-relevant configurations. A separate final-approach controller is one possible remedy, but task-specific validation is needed.

Evaluate closed-loop task success alongside feature prediction and goal costs. The toy experiment reports both kinds of measure, but agreement on a small sample does not establish that the proxy will rank plans correctly in new states or goals.