Core concept

Physical robots do not receive low-dimensional state vectors from an oracle. They observe camera pixels, depth maps, and motor currents. Building a world model directly in pixel space by autoregressively generating future video frames is computationally prohibitive and semantically misaligned. Most pixels in a camera frame describe static background, ambient illumination, or texture variations that have no causal influence on manipulation dynamics.

A latent world model resolves this issue by splitting predictive dynamics into three distinct components:

1. An encoder $z_t = \text{enc}(o_t)$ that compresses high-dimensional observations $o_t \in \mathbb{R}^{H \times W \times C}$ into a compact latent state $z_t \in \mathbb{R}^d$.
2. An action-conditioned latent transition model $z_{t+1} = z_t + g_\theta(z_t, a_t)$ that evolves the latent representation forward in time under applied control inputs.
3. An observation decoder $\hat{o}_t = \text{dec}(z_t)$ that reconstructs images from latent codes, or downstream probes that predict task variables such as end-effector and object coordinates.

This lesson constructs a latent world model on the tabletop scene, compares pixel-only training against position-probed auxiliary supervision, and demonstrates why high pixel reconstruction fidelity diverges from control utility.

Visual intuition

The figure displays open-loop multi-step visual rollouts up to 2.0 s (20 steps) into the future from a held-out initial state. The top row shows ground-truth frames rendered by MuJoCo. The subsequent rows show rollouts produced by latent models trained under two distinct objectives: pure pixel reconstruction and pixel reconstruction augmented with position probes.

At short horizons (0.5 s), both models maintain plausible tabletop geometry. As the prediction horizon extends to 1.0 s and 2.0 s, the pure pixel model blurs contact boundaries. Because the puck occupies fewer than 3% of the image pixels, the pixel loss permits the puck to dissolve or drift while still minimizing aggregate mean squared error across the table surface.

Mathematics

Latent dynamics architecture

Observations $o_t$ are $64 \times 64 \times 3$ RGB images captured from the overhead camera. The encoder consists of three convolutional layers with stride 2 and ReLU activations, followed by a linear projection to a $d = 32$ latent vector $z_t \in \mathbb{R}^{32}$.

The latent transition model computes an action-conditioned residual update:

where $g_\theta$ is a three-layer multi-layer perceptron with 256 units and SiLU activations taking the concatenated vector $[z_t, a_t] \in \mathbb{R}^{34}$ as input.

The decoder mirrors the encoder using transposed convolutions with stride 2 to map $z_t$ back to $\hat{o}_t \in [0, 1]^{64 \times 64 \times 3}$ via a sigmoid output layer:

An optional linear probe predicts the planar coordinates of the pusher and all three pucks directly from the latent embedding:

Training objectives

Given an observed trajectory sequence $(o_t, a_t, o_{t+1}, \dots, o_{t+K})$, the models are trained across unrolled horizons of $K = 5$ steps. Under the pure pixel objective, the loss sums reconstruction error across the unrolled sequence:

where $\hat{z}t = \text{enc}(o_t)$ and $\hat{z}{t+k+1} = \hat{z}{t+k} + g\theta(\hat{z}{t+k}, a{t+k})$.

Under the position-augmented objective, privileged coordinate supervision penalizes geometric drift in latent space:

with scaling coefficient $\lambda = 100$.

Why pixel loss diverges from control utility

Let an image $o \in \mathbb{R}^{N}$ contain $N = 64 \times 64 \times 3 = 12288$ subpixel values. Suppose a puck has radius $r = 30$ mm on a table of width $0.6$ m, covering approximately $M \approx 300$ pixels.

If a predicted puck is displaced by distance $\Delta x$ such that it does not overlap with the true puck at all, the penalty incurred in the total loss is bounded by:

Conversely, a slight shift in lighting, camera vignette, or table reflection that changes every pixel by merely $5\%$ produces:

A gradient step optimizing $\mathcal{L}_{\text{pixel}}$ dedicates its capacity to matching the $97\%$ of background pixels. For downstream model-predictive control, a 5 mm error in puck position causes a planned push to miss contact completely, converting a task success into a failure, while background illumination errors are irrelevant. Evaluating world models by pixel MSE alone obscures this failure mode.

Implementation

INPUT: tabletop.xml; expert-guided exploration rendered from top camera (64 x 64 RGB); standard state model of Lesson 23.2
PROCESS: train pixel-only and position-probed latent models; evaluate open-loop pixel MSE, detected puck position error, and missing object rate over horizons H in {1, 5, 10, 20}
OUTPUT: quantitative evaluation table; runs/l23_3/frames.png

```python file=examples/l23_3_latent_world_model.py import multiprocessing import os from concurrent.futures import ProcessPoolExecutor from pathlib import Path

import numpy as np import torch

from mjcourse import tabletop as tt from mjcourse.worldmodel import data, latent, standard

OUT = Path(file).resolve().parents[1] / “runs” / “l23_3” FAST = os.environ.get(“MJC_FAST”) == “1” TRAIN, TEST, STEPS = (40, 8, 300) if FAST else (600, 50, 30000) HORIZONS = (1, 5, 10, 20)

def collect(seed: int, n: int) -> list[dict]: with ProcessPoolExecutor(min(32, os.cpu_count() or 1), mp_context=multiprocessing.get_context(“spawn”)) as pool: return list(pool.map(data._collect, [(seed, i, {“render”: True}) for i in range(n)], chunksize=2))

def train(objective: str) -> str: torch.set_num_threads(max(1, int(os.environ.get(“SLURM_CPUS_PER_TASK”, 16)) // 2)) episodes = collect_cached(10, TRAIN) model = latent.fit(episodes, objective, steps=STEPS, log=lambda m: print(f” [{objective}] {m}”, flush=True)) path = OUT / f”latent_{objective}.pt” torch.save(model, path) return str(path)

def collect_cached(seed: int, n: int) -> list[dict]: import pickle

path = OUT / f"images_{seed}_{n}.pkl"
if not path.exists():
    path.write_bytes(pickle.dumps(collect(seed, n)))
return pickle.loads(path.read_bytes())

def score(episodes: list[dict], predict, starts_per_episode: int = 4) -> dict: rng = np.random.default_rng(0) rows = {h: {“pixel”: [], “puck”: [], “missing”: []} for h in HORIZONS} for e in episodes: target = tt.OBJECTS.index(e[“task”][0]) for t0 in rng.integers(0, len(e[“actions”]) - max(HORIZONS), starts_per_episode): pred = predict(e, t0, max(HORIZONS)) for h in HORIZONS: true_img, img = e[“images”][t0 + h], pred[h] rows[h][“pixel”].append(np.mean((img.astype(float) - true_img.astype(float)) ** 2) / 255.0 ** 2) found = latent.puck_positions(img)[tt.OBJECTS[target]] rows[h][“missing”].append(found is None) if found is not None: rows[h][“puck”].append(np.linalg.norm(found - e[“states”][t0 + h][4 + 4 * target:6 + 4 * target])) return {h: {“pixel”: float(np.mean(r[“pixel”])), “puck”: float(np.mean(r[“puck”])) if r[“puck”] else float(“nan”), “missing”: float(np.mean(r[“missing”]))} for h, r in rows.items()}

def state_model_errors(episodes: list[dict]) -> dict: mlp = standard.model(“one_step”) rng = np.random.default_rng(0) rows = {h: [] for h in HORIZONS} for e in episodes: target = tt.OBJECTS.index(e[“task”][0]) for t0 in rng.integers(0, len(e[“actions”]) - max(HORIZONS), 4): pred = mlp.rollout(e[“states”][t0][None], e[“actions”][None, t0:t0 + max(HORIZONS)])[0] for h in HORIZONS: i = 4 + 4 * target rows[h].append(np.linalg.norm(pred[h, i:i + 2] - e[“states”][t0 + h][i:i + 2])) return {h: float(np.mean(v)) for h, v in rows.items()}

def figure(episode: dict, models: dict) -> None: import matplotlib matplotlib.use(“Agg”) import matplotlib.pyplot as plt

shown = (0, 5, 10, 20)
rows = {"MuJoCo (true)": [episode["images"][20 + h] for h in shown]}
for name, model in models.items():
    pred = model.predict(episode["images"][20], episode["actions"][20:40])
    rows[f"latent model, {name}"] = [pred[h] for h in shown]
fig, axes = plt.subplots(len(rows), len(shown), figsize=(2.0 * len(shown), 2.1 * len(rows)), constrained_layout=True)
for r, (label, images) in enumerate(rows.items()):
    for c, img in enumerate(images):
        axes[r, c].imshow(img)
        axes[r, c].set_xticks([])
        axes[r, c].set_yticks([])
        if r == 0:
            axes[r, c].set_title(f"{shown[c] / 10:.1f} s ahead", fontsize=9)
        if c == 0:
            axes[r, c].set_ylabel(label, fontsize=8)
fig.savefig(OUT / "frames.png", dpi=150)
plt.close(fig)

if name == “main”: OUT.mkdir(parents=True, exist_ok=True) train_eps, test_eps = collect_cached(10, TRAIN), collect_cached(11, TEST) print(f”(1) {TRAIN} training and {TEST} held-out episodes with images; detector error on true images “ f”{1000 * latent.detector_error(test_eps):.2f} mm”) with ProcessPoolExecutor(2, mp_context=multiprocessing.get_context(“spawn”)) as pool: paths = dict(zip((“pixels”, “positions”), pool.map(train, (“pixels”, “positions”)))) models = {k: torch.load(p, weights_only=False) for k, p in paths.items()} print(f”(2) trained two latent models, {STEPS} steps each”) print(“(3) open-loop prediction on held-out states”) results = {“first image copied forward”: score(test_eps, lambda e, t0, h: e[“images”][t0:t0 + h + 1].repeat(h + 1, 0))} for name, model in models.items(): results[f”latent, {name}”] = score(test_eps, lambda e, t0, h, m=model: m.predict(e[“images”][t0], e[“actions”][t0:t0 + h])) print(f” {‘model’:<28}{‘quantity’:<26}” + ““.join(f”{h:>9d}” for h in HORIZONS) + “ (steps of 0.1 s)”) for name, res in results.items(): print(f” {name:<28}{‘pixel MSE (x 1000)’:<26}” + ““.join(f”{1000 * res[h][‘pixel’]:9.2f}” for h in HORIZONS)) print(f” {‘’:<28}{‘pushed puck error (mm)’:<26}” + ““.join(f”{1000 * res[h][‘puck’]:9.1f}” for h in HORIZONS)) print(f” {‘’:<28}{‘pushed puck missing’:<26}” + ““.join(f”{res[h][‘missing’]:9.2f}” for h in HORIZONS)) sm = state_model_errors(test_eps) print(f” {‘state model (23.2)’:<28}{‘pushed puck error (mm)’:<26}” + ““.join(f”{1000 * sm[h]:9.1f}” for h in HORIZONS)) figure(test_eps[2], models) print(f”wrote {OUT / ‘frames.png’}”) ```

Output and interpretation

Open-loop prediction evaluation on held-out test episodes:

Three empirical findings emerge:

1. Pixel fidelity does not predict spatial accuracy. The pixel-only latent model achieves lower pixel MSE than the position-probed model at every horizon (4.38 vs 4.51 at 1.0 s). Yet its position error on the pushed puck is more than double (38.7 mm vs 16.4 mm). The decoder minimizes pixel error by averaging puck edges, which blurs the object and shifts its apparent centroid.
2. Object disappearance under autoregression. In 22% of predictions at 20 steps, the color detector fails to find the pushed puck in images generated by the pixel-only model. The latent transition smooths away high-frequency color blobs because uniform background predictions incur lower expected $L_2$ pixel penalties than misaligned sharp shapes.
3. The privileged state model remains the upper bound. The state-space MLP from Lesson 23.2, which operates directly on ground-truth coordinates, incurs less than half the position error of the best latent model at 20 steps (20.6 mm vs 31.2 mm). Compressing through visual bottlenecks inevitably discards fine contact information.

Exercises

1. Latent bottleneck ablation. Vary the latent dimension $d \in {8, 16, 32, 64}$. Plot pixel reconstruction error and detected puck position error at $H = 10$. Identify the dimensionality where the transition model begins to overfit.
2. Decoder-free latent learning. Remove the deconvolutional decoder entirely. Train the latent transition model using a temporal contrastive loss or a Joint Embedding Predictive Architecture (JEPA) objective. Measure whether the absence of pixel reconstruction gradients improves downstream position probe accuracy.

Next step

We have evaluated dynamics prediction open-loop across both state and latent spaces. In Lesson 23.4, we close the loop: using these learned models inside model-predictive control to discover how model inaccuracies influence real robot decision making.