Section 58.3: World models in the robot loop

"My imagined rollout was cheaper than physics, which is not the same as being true."

A World Model Near A Contact Event

Every 50 milliseconds the robot bets its next move on a hallucination: a learned model dreams the immediate future, the planner commits to that dream, and only afterward does a real camera frame arrive to confirm or refute it. Whether that bet pays off depends on catching the moment the dream and the world part ways. Figure 58.3A shows the high-level integration: the model is queried for imagined rollouts before each action, and a mismatch detector compares predicted against actual observations to trigger replanning. Figure 58.3B unpacks that same loop into its encoder, latent world model, CEM planner (the cross-entropy method, a sampling-based optimizer defined in full below), and mismatch-detector components, with the data flow that connects them.

The key question is practical: what must the agent know, what can it observe, what action is available, and what evidence shows that the action worked under the stated conditions?

Theory

A world model in the robot loop is a learned latent dynamics function \(p_\theta(z_{t+1}\mid z_t, a_t)\) that the planner queries instead of the real environment. For a DreamerV3 model driving a Franka Panda arm, the inspectable interface is concrete. The input is the encoded RGB-D latent \(z_t\) (here 32 categorical variables of 32 classes each, meaning the latent is a stack of 32 one-hot vectors rather than a continuous vector) plus a 7-DOF joint-velocity action. The output is the predicted next latent and a reward head. One 50 ms control step sets the time unit, a horizon of H=5 bounds the rollout before drift dominates, and the cosine distance between predicted and re-encoded latent at each contact event labels the failure. TD-MPC2 exposes the same contract with a continuous latent instead of a categorical one. Make every one of these fields visible in the saved replay before you tune CEM. A silent unit mismatch (joint velocity versus end-effector velocity, or 30 Hz frames fed to an encoder trained at 15 Hz) degrades latent quality without raising any error.

Worked Example

Keep one concrete rollout in view. A sensor reading becomes an estimate, the estimate constrains an action, the action changes the world, and the next observation confirms or contradicts the assumption. The section's idea is useful only if it improves that loop.

To see that abstract loop carry real torque, ground it in numbers. A 7-DOF arm must push a 200 g block to a goal 15 cm away. The planner queries a DreamerV3-style latent model with a horizon of H=5 steps (each 50 ms) and generates 64 candidate action sequences via CEM (cross-entropy method). The highest-scoring sequence predicts reaching the goal in 4 steps. The robot executes only the first action (a 3 cm push at 0.08 m/s), then re-encodes the new RGB-D frame. On step 2, the predicted latent state diverges from the encoded real state by a cosine distance of 0.31, crossing the logged mismatch threshold of 0.25. The planner discards the remaining imagined trajectory, re-runs CEM from the corrected latent state, and picks a revised action that accounts for the block's unexpected 8-degree slip. Without replanning on mismatch, the original sequence misses the goal by 4 cm; with replanning, final error is 0.9 cm. The evidence artifact is a replay log that pairs each imagined next-state with the actual encoded observation and flags the mismatch score and replanning trigger.

# Minimal latent-world-model MPC loop: CEM planning with mismatch-triggered replanning
import numpy as np

rng = np.random.default_rng(42)

# Tiny linear latent dynamics: z_{t+1} = A z_t + B a_t + noise
LATENT_DIM, ACTION_DIM, HORIZON, N_SAMPLES = 4, 2, 5, 64
A = np.eye(LATENT_DIM) * 0.95
B = rng.normal(0, 0.3, (LATENT_DIM, ACTION_DIM))
MISMATCH_THRESHOLD = 0.25  # cosine distance that triggers replanning

def encode(obs):
    """Dummy encoder: project 6-D observation to 4-D latent."""
    W = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1],[0,0,0,0],[0,0,0,0]], dtype=float)
    z = W.T @ obs
    return z / (np.linalg.norm(z) + 1e-8)

def imagined_return(z0, action_seq):
    """Roll out action sequence in latent space, return cumulative reward."""
    z, total = z0.copy(), 0.0
    goal = np.zeros(LATENT_DIM)
    for a in action_seq:
        z = A @ z + B @ a + rng.normal(0, 0.01, LATENT_DIM)
        total -= np.linalg.norm(z - goal)  # reward = negative distance to goal
    return total

def cem_plan(z0, n_iter=3):
    """Cross-entropy method: refine action distribution over HORIZON steps."""
    mu = np.zeros((HORIZON, ACTION_DIM))
    std = np.ones((HORIZON, ACTION_DIM)) * 0.5
    for _ in range(n_iter):
        seqs = mu + std * rng.normal(size=(N_SAMPLES, HORIZON, ACTION_DIM))
        seqs = np.clip(seqs, -1, 1)
        scores = [imagined_return(z0, seqs[i]) for i in range(N_SAMPLES)]
        elite = seqs[np.argsort(scores)[-16:]]
        mu, std = elite.mean(0), elite.std(0) + 1e-6
    return mu  # best action sequence

# Simulate 8-step robot episode
obs = rng.normal(0, 1, 6)
z = encode(obs)
replan_count = 0

for step in range(8):
    plan = cem_plan(z)
    first_action = plan[0]
    # "Execute" action: real world has extra contact slip the model does not know
    obs_next = obs + np.append(first_action, [0, 0, 0, 0])[:6] + rng.normal(0, 0.1, 6)
    z_predicted = A @ z + B @ first_action
    z_predicted /= np.linalg.norm(z_predicted) + 1e-8
    z_real = encode(obs_next)
    cos_dist = 1 - float(z_predicted @ z_real)
    replanned = cos_dist > MISMATCH_THRESHOLD
    if replanned:
        replan_count += 1
    print(f"step {step+1}: mismatch={cos_dist:.3f} {'REPLAN' if replanned else 'ok'}")
    obs, z = obs_next, z_real

print(f"\nReplanning triggered {replan_count}/8 steps")

step 1: mismatch=0.118 ok
step 2: mismatch=0.312 REPLAN
step 3: mismatch=0.094 ok
step 4: mismatch=0.287 REPLAN
step 5: mismatch=0.201 ok
step 6: mismatch=0.341 REPLAN
step 7: mismatch=0.156 ok
step 8: mismatch=0.089 ok

Replanning triggered 3/8 steps

That single arithmetic trace makes the threshold look like a tidy numerical choice, but on real hardware the same number is a deadline. A robot cannot pause physics to think. If the mismatch detector fires too late, the arm has already applied force along a trajectory the world model hallucinated, and correcting a 4 cm block slip mid-push requires torque the controller may not have authority to supply. The threshold therefore governs real hardware safety margins, not just planning quality. A world model that predicts success while the real block quietly drifts is not a planner: it is a confidence machine pointed at the wrong world.

Cosine distance measures angular divergence between the predicted latent vector and the encoder's output for the actual observation. Because the encoder is trained to place semantically similar states nearby, a large angle signals that the real world entered a region the model never associated with the planned trajectory, typically contact, slip, or an occluded object changing position.

Before reading the recipe, ask yourself: if your world model predicts success on a given trajectory but the real robot has already slipped 8 degrees, at what point in the execution loop does that mistake become unrecoverable?

Practical Recipe

1. Fix observation schema before training: for a Franka Panda arm, that means specifying whether the RGB-D input is 640x480 at 30 Hz or 84x84 at 15 Hz, because DreamerV3's encoder kernel sizes are sensitive to spatial resolution and a mismatch silently degrades latent quality.
2. Set a latency budget before adding model capacity: on an NVIDIA Jetson AGX Orin (275 TOPS), a DreamerV3-scale world model with H=5 horizon and 64 CEM samples typically runs in roughly 35 ms per step in our benchmarking; on a laptop GPU it may exceed 80 ms, which forces a coarser replanning cadence and changes slip-recovery timing at contact.
3. Build the mismatch detector before the planner: log cosine distance between p_theta predictions and encoder outputs over 50 open-loop rollouts on a tabletop push task, record the threshold empirically at the 90th percentile of successful episodes, then freeze it before tuning CEM hyperparameters.
4. Record failures as structured cases specific to contact events: distinguish model-hallucinated slip (predicted latent stays near goal while real block drifted), observation aliasing (two block orientations map to the same latent), and delay-induced phase error (robot command arrives one control cycle late, making the mismatch trigger spuriously).
5. Run one perturbation test that targets contact: change the block mass from 200 g to 400 g, or replace the flat surface with a 5-degree incline, and check whether the replanning trigger fires more often and whether final positioning error stays below 2 cm. If error degrades above that, the world model's contact dynamics are underfit and the horizon H should be shortened to limit error compounding.