Pretrained Visual Models and DINO-WM › Code Implementation
World Models HandbookAll chapters Pretrained Visual Models and DINO-WM › Code ImplementationOne section · Sections in this chapter7 / 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
We'll build the pipeline in five stages: the world, the frozen encoder, the feature-space dynamics model, the rollout diagnostics, and the planner. Every piece is small enough to read in one sitting, and every number in the output is computed at run time. Keeping the pieces separate makes it clear which part of the system is responsible for which behavior, which is exactly the clarity a full-scale system makes hard to achieve.
Imports and configuration
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from book_plot_style import PALETTE, polish_axes, use_book_style
torch.manual_seed(0)
rng = np.random.default_rng(0)The synthetic world
The world is deliberately simple so that we can check the learned model against ground truth. A 32×32, three-channel image shows a red key disk and a blue cursor disk. The cursor moves by the commanded displacement inside a box. The frame shows approximate disk positions, but rasterization aliases nearby cursor locations and the blue disk can occlude the red key. Information loss can therefore occur in the renderer as well as the learned encoder; the diagnostics below do not isolate one cause perfectly.
IMG = 32
def render(cursor_xy, key_xy):
"""Paint a 32x32 RGB frame: red key disk plus blue cursor disk."""
yy, xx = np.mgrid[0:IMG, 0:IMG].astype(np.float32)
def disk(center, radius):
cx = float(center[0]) * (IMG - 1)
cy = float(center[1]) * (IMG - 1)
return ((xx - cx) ** 2 + (yy - cy) ** 2) <= radius**2
img = np.zeros((3, IMG, IMG), dtype=np.float32)
key_mask = disk(key_xy, 2.2)
img[0][key_mask] = 0.95
img[1][key_mask] = 0.25
img[2][key_mask] = 0.25
cursor_mask = disk(cursor_xy, 2.2)
img[0][cursor_mask] = 0.25
img[1][cursor_mask] = 0.45
img[2][cursor_mask] = 1.00
return img
def step(state, action, low=0.05, high=0.95):
"""Displacement dynamics inside a box; the rasterized frame is lossy."""
return np.clip(state + action, low, high)Before learning an encoder, check what the renderer itself can hide. With the cursor drawn over the key, two different key positions can produce exactly the same 32×32 frame:
occluding_cursor = np.array([0.50, 0.50], dtype=np.float32)
hidden_key_a = np.array([0.50, 0.50], dtype=np.float32)
hidden_key_b = np.array([0.51, 0.50], dtype=np.float32)
print(
"distinct key positions render identically:",
np.array_equal(
render(occluding_cursor, hidden_key_a),
render(occluding_cursor, hidden_key_b),
),
)distinct key positions render identically: True
No visual encoder can recover a distinction absent from identical input pixels. This is renderer occlusion and rasterization, not a failure introduced by frozen features.
Let's render one frame and look at it, to fix the scale of the problem: a disk of radius 2.2 pixels corresponds to about 7% of the image width.
frame shape: (3, 32, 32) pixel values in [0, 1]: 0.0 to 1.0 fraction of pixels that are background: 0.970703125
The reported background fraction shows why uniform pixel MSE can be a poor emphasis for this task. It does not by itself establish where another model spends its capacity.
Collecting unlabeled trajectories
We collect random-walk trajectories. Each episode gets a fresh key position and a fresh start, so the dataset covers a range of scene layouts. The first 32 episodes are the "training" split for the dynamics model; the last 8 are held out. Holding out entire episodes rather than random frames is important, because frames within an episode are highly correlated; a random-frame split would leak near-duplicate frames across the boundary and make the evaluation optimistic.
N_EPISODES = 40
EP_LEN = 100
N_TRAIN_EPISODES = 32
ACTION_SCALE = 0.06
MAX_ACTION = 0.08
frames, states, actions, keys, ep_ids = [], [], [], [], []
for ep in range(N_EPISODES):
key_xy = rng.uniform(0.15, 0.85, size=2)
state = rng.uniform(0.10, 0.90, size=2)
for _ in range(EP_LEN):
frames.append(render(state, key_xy))
states.append(state.copy())
keys.append(key_xy.copy())
ep_ids.append(ep)
action = np.clip(
rng.normal(scale=ACTION_SCALE, size=2), -MAX_ACTION, MAX_ACTION
)
actions.append(action)
state = step(state, action)
frames = np.stack(frames).astype(np.float32)
states = np.stack(states).astype(np.float32)
actions = np.stack(actions).astype(np.float32)
keys = np.stack(keys).astype(np.float32)
ep_ids = np.array(ep_ids)
train_episode_mask = ep_ids < N_TRAIN_EPISODESframes: (4000, 3, 32, 32) actions: (4000, 2) training episodes: 32, held-out episodes: 8 training frames: 3200 max displacement between consecutive frames: 0.080
Consecutive states differ by at most about 0.08 in each coordinate, roughly 2.5 pixels on this raster. That small physical step motivates trying a residual predictor. It does not prove that encoded features change smoothly at a patch boundary or that residual training beats a direct predictor; those are empirical questions for the diagnostics below.
A frozen patch encoder
This is a pedagogical substitute, not a DINOv2 reproduction. We pretrain a small convolutional encoder with a reconstruction objective and a foreground-weighted loss, then freeze it before dynamics training. DINOv2 uses a different pretraining objective and data scale; the shared feature of the two systems is only that a fixed visual encoder supplies the predictor's coordinates.
The encoder produces a 4×4 grid of 32-dimensional tokens, mirroring the patch-token structure of a ViT. The grid permits spatially indexed feature comparison for the goal-image cost; whether that cost tracks task success must be tested.
TOKEN_DIM = 32
N_TOKENS = 16
LATENT_DIM = TOKEN_DIM * N_TOKENS
class PatchEncoder(nn.Module):
"""Frozen visual backbone: image -> 4x4 grid of patch tokens."""
def __init__(self, token_dim=TOKEN_DIM):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 32, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 64, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, token_dim, 1),
)
def forward(self, x):
features = self.net(x) # (B, token_dim, 4, 4)
return features.flatten(2).transpose(1, 2) # (B, 16, token_dim)
class PatchDecoder(nn.Module):
"""Pretrains the encoder; excluded from dynamics and planning, retained for diagnostics."""
def __init__(self, token_dim=TOKEN_DIM):
super().__init__()
self.net = nn.Sequential(
nn.ConvTranspose2d(token_dim, 64, 4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, 16, 4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(16, 3, 3, padding=1),
nn.Sigmoid(),
)
def forward(self, tokens):
b, n, d = tokens.shape
side = int(round(n**0.5))
grid = tokens.transpose(1, 2).reshape(b, d, side, side)
return self.net(grid)Now we pretrain the encoder on the training episodes' frames, using no actions or goal/outcome labels. Roughly 97% of pixels are black background, so plain pixel MSE could favor a nearly black output. We threshold the rendered RGB frames to identify nonblack foreground pixels and upweight them. This is a hand-designed synthetic-domain foreground prior, not a separately supplied renderer disk mask or a DINOv2-style self-supervised objective.
imgs = torch.tensor(frames)
imgs_train = imgs[train_episode_mask]
encoder, decoder = PatchEncoder(), PatchDecoder()
ae_optimizer = torch.optim.Adam(
list(encoder.parameters()) + list(decoder.parameters()), lr=2e-3
)
AE_STEPS, AE_BATCH = 600, 64
ae_history = []
for _ in range(AE_STEPS):
batch = imgs_train[torch.randint(0, imgs_train.shape[0], (AE_BATCH,))]
reconstruction = decoder(encoder(batch))
foreground = (batch.amax(dim=1, keepdim=True) > 0.1).float()
pixel_weight = 1.0 + 20.0 * foreground
loss = (pixel_weight * (reconstruction - batch).square()).mean()
ae_optimizer.zero_grad()
loss.backward()
ae_optimizer.step()
ae_history.append(loss.item())Now freeze. From this point forward, no gradient ever flows into encoder.
encoder pretraining loss: 0.31562 -> 0.00231 frozen parameters in encoder: 58,400 tokens per frame: 16, token dim: 32, latent size: 512
The decoder remains in memory for the reconstruction diagnostic below, but is excluded from dynamics training and planning. We do not use it again after that diagnostic.
The common display scale lets us check whether the 512-number latent retains approximate disk positions, color, and intensity instead of merely inspecting a contrast-enhanced image. Even a visually plausible reconstruction would not establish that these features are sufficient for control; that is what the following rollout and planning tests examine.
Extracting and normalizing features
We encode every frame once and apply a global normalization. This changes the numerical scale of the regression loss; unlike per-dimension whitening, it does not change relative feature weights or candidate-plan rankings under Euclidean distance.
We estimate the normalization statistics from training frames only. Per-dimension whitening is an alternative, but it can amplify nearly constant noisy dimensions. With one scalar scale, high-variance directions retain their relative influence; whether that influence is useful requires a task-level check.
with torch.no_grad():
Z = encoder(imgs).numpy() # (N, 16, TOKEN_DIM)
Z_train = Z[train_episode_mask]
feature_mean = Z_train.reshape(-1, TOKEN_DIM).mean(axis=0)
feature_scale = (Z_train - feature_mean).std()
Z_norm = (Z - feature_mean) / feature_scale
Z_flat = Z_norm.reshape(len(Z_norm), LATENT_DIM)raw feature tensor: (4000, 16, 32) global feature scale: 1.3882 per-dimension std after normalization: min 0.587, median 0.885, max 1.982
The spread of per-dimension standard deviations after global scaling is a diagnostic of relative variation. Dominance in variation can affect a squared-error cost, but variance alone does not tell us whether the changing directions matter for the goal.
The cost is easier to inspect as a landscape. We render the cursor over a fixed grid with a fixed key, encode each frame, and evaluate feature distance to a goal image drawn at the key. Rasterization can make nearby positions identical, and the learned Conv–ReLU encoder is not known to be invertible or globally smooth. The observed landscape is a diagnostic for this trained toy, not a general theorem.
Building the transition dataset and learning dynamics
Transition pairs are consecutive frames within the same episode. Cross-episode pairs must be excluded, since they are not real transitions. This is the same leakage concern as the train/test split: joining the last frame of one episode to the first frame of the next would create a pair that no action could produce, and it would corrupt the learned dynamics.
same_episode = np.r_[ep_ids[1:] == ep_ids[:-1], False]
train_pairs = np.where(train_episode_mask & same_episode)[0]
Z_t = torch.tensor(Z_flat[train_pairs])
Z_next = torch.tensor(Z_flat[train_pairs + 1])
A_t = torch.tensor(actions[train_pairs])training transitions: 3,168 input dim: 514 (latent plus action), output dim: 512
The predictor is a residual MLP: it predicts a change in feature space rather than the next feature directly. The code does not zero-initialize its final layer, so the initial residual is not guaranteed to be small. We measure improvement against persistence below.
The action enters by concatenation. The physical update is action-affine only away from clipping boundaries; whether this MLP handles boundary transitions is empirical.
class LatentDynamics(nn.Module):
"""Residual predictor in normalized feature space."""
def __init__(self, latent_dim=LATENT_DIM, action_dim=2, hidden=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim + action_dim, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, latent_dim),
)
def forward(self, z, a):
return z + self.net(torch.cat([z, a], dim=-1))
dynamics = LatentDynamics()
dyn_optimizer = torch.optim.Adam(dynamics.parameters(), lr=1e-3)
DYN_STEPS, DYN_BATCH = 800, 256
dyn_history = []
for _ in range(DYN_STEPS):
idx = torch.randint(0, Z_t.shape[0], (DYN_BATCH,))
prediction = dynamics(Z_t[idx], A_t[idx])
loss = nn.functional.mse_loss(prediction, Z_next[idx])
dyn_optimizer.zero_grad()
loss.backward()
dyn_optimizer.step()
dyn_history.append(loss.item())
dynamics.eval()dynamics loss: 0.31021 -> 0.16331 predicting no change at all would give: 0.28829 full-train fitted MSE: 0.16575 paired full-train reduction over persistence: 42.5%
The paired full-training-set comparison against persistence checks whether the predictor improves on a simple baseline for this one-step metric. The preceding last-minibatch loss is a training trace, not the numerator of that comparison. Held-out transition error and closed-loop task success are needed before calling the learned dynamics useful.
Rollout diagnostics: one-step versus recursive error
Now the key measurement. We roll the predictor forward on held-out episodes and compare two error curves:
- One-step error: always start from a true feature vector and predict one step ahead. This measures the model in the regime it was trained on.
- Recursive error: start from a true feature vector and then feed predictions back in. This measures the model in the regime the planner will actually use it.
- Persistence error: compare each future observation with the fixed starting feature. This is a reference curve, not a mathematical lower bound for a learned model or a policy.
HORIZON_EVAL = 15
STRIDE = 12
start_idx = []
for ep in range(N_TRAIN_EPISODES, N_EPISODES):
for t in range(0, EP_LEN - HORIZON_EVAL, STRIDE):
start_idx.append(ep * EP_LEN + t)
start_idx = np.array(start_idx)
z0_eval = torch.tensor(Z_flat[start_idx])
src_idx = start_idx[:, None] + np.arange(HORIZON_EVAL)[None, :]
targets_eval = torch.tensor(Z_flat[src_idx + 1])
with torch.no_grad():
rolling = []
z = z0_eval
for h in range(HORIZON_EVAL):
z = dynamics(z, torch.tensor(actions[src_idx[:, h]]))
rolling.append(z)
rollout_pred = torch.stack(rolling, dim=1)
one_step_pred = dynamics(
torch.tensor(Z_flat[src_idx]),
torch.tensor(actions[src_idx]),
)
recursive_mse = (
((rollout_pred - targets_eval) ** 2).mean(dim=-1).mean(dim=0).numpy()
)
one_step_mse = (
((one_step_pred - targets_eval) ** 2).mean(dim=-1).mean(dim=0).numpy()
)
persistence_curve = (
((z0_eval[:, None, :] - targets_eval) ** 2).mean(dim=-1).mean(dim=0).numpy()
)held-out start states: 64 one-step MSE at horizon 1: 0.33500 recursive MSE at horizon 1: 0.33500 recursive MSE at horizon 15: 2.22436 persistence MSE at horizon 15: 1.14052 recursive error growth factor (h=15 compared with h=1): 6.6x
The gap between the one-step and recursive curves shows how feedback of model predictions affects this held-out set. Where recursive MSE crosses the fixed-start persistence reference, the model predicts these recorded open-loop trajectories less accurately under this feature metric. That is not a comparison of action policies, a safe-horizon certificate, or a statement about every initial state.
This diagnostic motivates testing shorter horizons, alternative models, and closed-loop control. A CEM planner optimizes goal cost rather than the random-trajectory prediction MSE plotted here, so the crossing cannot establish that doing nothing is better or that search cannot help.
Planning to a goal image with the cross-entropy method
The toy planner searches over action sequences and sums feature distance to the encoded goal at every predicted step. This favors earlier approach to the goal, but it differs from DINO-WM's terminal-cost objective stated above.
def cem_plan(
z_start, z_goal, horizon=8, n_samples=128, n_elite=12, n_iters=3, seed=0
):
"""Cross-entropy method over action sequences, scored in feature space."""
plan_rng = np.random.default_rng(seed)
mean = np.zeros((horizon, 2), dtype=np.float32)
std = np.full((horizon, 2), MAX_ACTION, dtype=np.float32)
z0 = torch.tensor(z_start[None, :])
zg = torch.tensor(z_goal[None, :])
for _ in range(n_iters):
noise = plan_rng.standard_normal((n_samples, horizon, 2)).astype(
np.float32
)
candidates = np.clip(
mean[None] + std[None] * noise, -MAX_ACTION, MAX_ACTION
)
candidate_t = torch.tensor(candidates)
with torch.no_grad():
z = z0.repeat(n_samples, 1)
cost = torch.zeros(n_samples)
for h in range(horizon):
z = dynamics(z, candidate_t[:, h])
cost = cost + ((z - zg) ** 2).mean(dim=1)
elite = candidates[torch.argsort(cost)[:n_elite].numpy()]
mean = elite.mean(axis=0)
std = elite.std(axis=0) + 1e-3
return meanThe MPC loop re-encodes a rendered observation at every real step, so the next search starts from the encoder's output rather than the previous imagined feature. That observation is still lossy because of rasterization and possible occlusion.
def encode_state(state_xy, key_xy):
"""Render, encode with the frozen backbone, and normalize."""
frame = render(state_xy, key_xy)
with torch.no_grad():
z = encoder(torch.tensor(frame[None])).numpy()[0]
return ((z - feature_mean) / feature_scale).reshape(-1).astype(np.float32)
def run_mpc(key_xy, start_xy, n_steps=20, horizon=8, seed=0):
"""Receding-horizon control. Execute one action, observe, replan."""
goal_latent = encode_state(key_xy, key_xy)
state = np.asarray(start_xy, dtype=np.float32).copy()
trajectory = [state.copy()]
for k in range(n_steps):
current_latent = encode_state(state, key_xy)
plan = cem_plan(
current_latent, goal_latent, horizon=horizon, seed=seed + k
)
state = step(state, plan[0])
trajectory.append(state.copy())
return np.array(trajectory)Evaluating closed-loop success
This is where proxy and task outcomes can be compared. We track final physical distance from cursor to key and final observed-frame feature distance to the goal. The latter is related to, but not identical to, the summed predicted feature cost CEM optimized. A random-action baseline runs on the same scenarios with the same number of steps.
N_SCENARIOS = 10
eval_rng = np.random.default_rng(1234)
eval_keys = eval_rng.uniform(0.20, 0.80, size=(N_SCENARIOS, 2)).astype(
np.float32
)
eval_starts = eval_rng.uniform(0.15, 0.85, size=(N_SCENARIOS, 2)).astype(
np.float32
)
mpc_trajectories, mpc_final_distance = [], []
for i in range(N_SCENARIOS):
traj = run_mpc(
eval_keys[i], eval_starts[i], n_steps=20, horizon=8, seed=100 * i
)
mpc_trajectories.append(traj)
mpc_final_distance.append(float(np.linalg.norm(traj[-1] - eval_keys[i])))
mpc_final_distance = np.array(mpc_final_distance)
random_final_distance = []
for i in range(N_SCENARIOS):
state = eval_starts[i].copy()
for _ in range(20):
noise = np.clip(
eval_rng.normal(scale=ACTION_SCALE, size=2), -MAX_ACTION, MAX_ACTION
)
state = step(state, noise)
random_final_distance.append(float(np.linalg.norm(state - eval_keys[i])))
random_final_distance = np.array(random_final_distance)initial distance to key, mean: 0.397 latent MPC final distance, mean: 0.378 random walk final distance, mean: 0.396 latent MPC success rate (<= 0.10): 10% random success rate (<= 0.10): 0%
The toy model was trained without goal/outcome labels or demonstrations, but its encoder pretraining used a foreground mask derived by thresholding rendered RGB pixels. It does not consistently reach the key in these held-out scenarios. It sometimes approaches the target and sometimes stalls far away. The predictor learned from random-walk transitions in 32 episodes; the test goal was supplied as an image. Reporting misses alongside successes matters because feature-space fit is not task success.
For each held-out scenario we compare final physical distance with final observed-frame feature cost. This is a descriptive cross-scenario check: each point has a different key and start, and the plotted final-state cost is not the exact summed predicted cost optimized within that scenario. Rank inversions across different goals do not prove that CEM misranks candidate plans for one goal. A within-goal candidate study would be needed for that conclusion.
A note on what just happened
The encoder in this example is not DINOv2. It is a tiny convolutional network trained by foreground-weighted reconstruction on a few thousand 32×32 RGB images. What it shares with DINOv2 is a system role: a frozen visual map supplies features for dynamics learning and goal comparison. The toy uses a one-frame MLP, global normalization, and a summed-horizon CEM cost; a real-backbone implementation must revisit those choices rather than copy them uncritically.
Scaling up changes the data distribution, patch grid, feature dimension, action dynamics, and compute budget. A pretrained backbone may improve useful visual invariance, but that must be verified for the task. A Transformer is one way to model token interactions, not a requirement imposed by dimension alone. The toy remains useful for understanding the interfaces and diagnostics, not for predicting real-robot performance.