Evaluating World Models: Physical Consistency, Controllability & Coherence

"They gave me a perfect score on beauty and a passing grade on motion. Then a physicist watched my video, saw a glass of water refill itself, and gave me a zero. We are now arguing about which of them was measuring the right thing."

A Video Model Awaiting Its Real Performance Review

A generated clip can earn a top score for realism and still show a glass of water that refills itself, because the worry running through this entire chapter is the one thing appearance metrics never test: whether a model has learned the structure of the world or only its surface. Section 36.2 noted that text-to-video models fail at counting and conservation; Section 36.6 noted that simulators drift and hallucinate; Section 36.7 argued the whole pixel objective might be misdirected. All of these are claims about quality that you cannot settle by looking, because the failures hide in dynamics and physics, not in any single frame. Evaluation is therefore not an afterthought here; it is the instrument that decides which models are real progress. And the first thing to establish is why the standard metric is insufficient.

1. Why FID and FVD Are Necessary but Not Sufficient Intermediate

The standard generative-video metric is the Fréchet Video Distance (FVD), the temporal cousin of the Fréchet Inception Distance (FID) you will meet in full in Chapter 37. FVD embeds real and generated clips with a pretrained video network (an I3D action-recognition model from Chapter 26) and measures the distance between the two distributions of features, modeled as Gaussians:

$$ \text{FVD} \;=\; \lVert \mu_r - \mu_g \rVert_2^2 \;+\; \operatorname{Tr}\!\left( \Sigma_r + \Sigma_g - 2\,(\Sigma_r \Sigma_g)^{1/2} \right), $$

where $(\mu_r, \Sigma_r)$ and $(\mu_g, \Sigma_g)$ are the mean and covariance of real and generated feature distributions. This is the same feature-statistics-comparison idea the cross-reference map traces from the humble histograms of Chapter 2. FVD is genuinely useful: it captures gross appearance and motion realism and tracks human judgments of overall quality. But it is blind to exactly the failures world models exhibit. A clip where a glass refills itself, an object teleports, or a "turn left" command makes the car go straight can have an excellent FVD, because the feature distribution still looks like real driving video on average. FVD measures distributional realism; it does not measure physical truth, action obedience, or long-range consistency. You must add metrics that probe those directly.

2. The Evaluation Triad Advanced

Three orthogonal axes, beyond appearance, define a good world model, and each needs its own probe.

Physical consistency. Does the generated world obey physical laws? The cleanest tests are controlled and counterfactual: generate scenarios with known physics (a ball rolling off a table, two objects colliding, a liquid being poured) and check whether the outcome matches the law. Kang et al. (2024), the physical-law study from Section 36.2, do exactly this, varying whether test conditions are in or out of the training distribution to ask whether the model learned the rule or merely interpolated seen motions. Practical proxies include object-permanence checks (track objects through occlusions and verify they reappear correctly), conservation checks (count objects, measure sizes over time), and trajectory-physics checks (fit the generated motion to the expected equation and measure the residual).

Controllability. For an interactive simulator (Section 36.6), does the generated future actually obey the action? This is measured by an action-following score: condition the model on an action, generate, then use a separate perception model (a detector or pose estimator from Parts II and III) to verify the action was carried out. If you command "turn left" and an independent lane-detector confirms the trajectory bent left, the model is controllable; if the video is gorgeous but the car drifts straight, it is not. Controllability is invisible to FVD entirely.

Long-horizon coherence. Does the world stay consistent over time? This targets the drift of Section 36.6 and the identity drift of Section 36.1. Probes include re-identification (an object seen early should be recognizable later), loop-closure (turn 360 degrees and the scene should match where you started), and the flow-based warping error of Section 36.1 extended over long windows. Coherence degrades gracefully in good models and catastrophically in bad ones, and only a long rollout reveals it.

Figure 36.8.1 organizes the triad over the appearance baseline. The code below implements a concrete controllability probe, the action-following score, which uses an independent perception model exactly as the figure prescribes.

# Controllability probe: command an action, generate the future, then let an
# INDEPENDENT perception model read off what actually happened and check it
# matches the intended effect. This catches the failure FVD is blind to.
import numpy as np

def action_following_score(simulator, perception_model, action_tests):
    """Controllability: does the generated future obey the commanded action?
    perception_model independently measures what the generated world actually did."""
    correct = 0
    for init_frame, action, expected_effect in action_tests:
        simulator.reset(init_frame)
        rollout = [simulator.step(action) for _ in range(16)]   # generate the future
        # an INDEPENDENT model reads off what actually happened (no peeking at the action)
        measured = perception_model.measure_effect(rollout)     # e.g. ego-trajectory bend
        # score 1 if the measured effect matches what the action should have caused
        correct += int(np.sign(measured) == np.sign(expected_effect)
                        and abs(measured) > 0.2 * abs(expected_effect))
    return correct / len(action_tests)

# A model with high FVD but low action-following looks real but ignores control.
# This single number catches the failure FVD cannot see.
print(f"action-following: {action_following_score(sim, detector, tests):.2f}")

3. The Benchmarks and the Open Problem Intermediate

A 2024 to 2025 wave of benchmarks operationalizes the triad. VBench and its successors decompose video-generation quality into many interpretable dimensions (subject consistency, motion smoothness, temporal flicker, and increasingly physics and action dimensions) rather than one FVD number. Physics-focused suites (Physion-style intuitive-physics benchmarks adapted to generation, and the controlled physical-law probes of Kang et al., 2024) test prediction of collisions, support, and containment. Action-controllable benchmarks score generated environments on whether the commanded action produced the right effect. The common thread is decomposition: replace a single opaque score with a panel of targeted probes, each measuring one property, so a model cannot hide a dynamics failure behind appearance realism.

# VBench: a decomposed video-generation benchmark (schematic; see the repo's
# README for the exact constructor and evaluate() arguments, which evolve).
# pip install vbench ; then:
from vbench import VBench
bench = VBench(device="cuda", full_info_path="vbench_info.json")
results = bench.evaluate(
    videos_path="generated_clips/",
    dimension_list=["subject_consistency", "motion_smoothness",
                    "temporal_flickering", "dynamic_degree"])
print(results)   # a per-dimension score panel, not one opaque number

4. What "Good" Means, and Where the Chapter Lands Beginner

Pulling the chapter together: a world model is good to the extent that it predicts the consequences of being in and acting on a world, measured by physical consistency, controllability, and long-horizon coherence, with appearance realism (FVD) as a necessary floor rather than the goal. This reframing answers the worry that opened the chapter. Video models, 3D generators, and world simulators are all attempts to generate something that must stay consistent with itself, and the evaluation triad is simply the operational definition of that self-consistency along the three axes the chapter added, time, space, and action.

The arc that ran from a single denoised pixel in Chapter 33 to a controllable simulator of reality is now complete, and one capstone lab remains to make that arc tangible: the hands-on build at the end of this section assembles the chapter's single thread, a latent space with dynamics, into a working agent you can train and score. After that lab, the chapter ends on an evaluation question that is bigger than world models. How do we measure any generative model, manage its risks, and turn it from a curiosity into a trustworthy tool and a data engine? Section 36.8's FVD is one instance of the distribution metrics that Chapter 37: Evaluation, Safety & Generative Data Engines treats in full, and the controllable-simulator capability this chapter celebrated is exactly the capability whose safety and provenance Chapter 37 must confront. The instrument you built here generalizes there.

5. A Concrete Evaluation Toolkit: Named Benchmarks for Each Axis Intermediate

The triad of Section 2 is a principle; a practitioner setting up a leaderboard needs named, citable instruments. The problem with stopping at the principle is that "measure physical consistency" is not runnable: you cannot file a pull request against an abstraction. Each axis has, as of 2025, at least one peer-reviewed or widely adopted benchmark that operationalizes it on real data with a published protocol, and knowing which instrument measures which property (and whether it has been peer-reviewed) is what separates a defensible evaluation from a hand-waved one. This section names the instruments; the comparison table at the end lets you pick the right one for each axis at a glance.

5.1 Physical consistency: Physics-IQ and MORPHEUS

The sharpest instrument for physical-principle understanding is Physics-IQ (Motamed et al., 2025), a benchmark built with DeepMind. The protocol is a clean prediction test: the model sees a short conditioning segment of a real video and must predict the next five seconds, after which the prediction is compared against the true continuation. The benchmark contains 396 real videos spanning five physics domains, fluid dynamics, optics, solid mechanics, magnetism, and thermodynamics, chosen so that the correct continuation is dictated by a physical law rather than by stylistic plausibility. A model that has learned the law (water settles, a magnet attracts, a candle shortens) predicts the right continuation; a model that has only learned appearance produces a plausible-looking but physically wrong one.

The result that should reshape how you read any video-generation leaderboard is Physics-IQ's headline finding: across leading models (Sora, Runway, Pika, Lumiere, VideoPoet, and others), physical understanding is severely limited and is uncorrelated with visual realism. A model can top the appearance metrics and sit near the floor on physics. This is the empirical hammer behind the entire "FVD is not sufficient" argument of Section 1: it is no longer an intuition that pretty video can be physically wrong, it is a measured, near-zero correlation across the strongest models of the day.

MORPHEUS (Zhang et al., 2025) complements Physics-IQ from the conservation-law angle. It evaluates physical reasoning on 80 real videos of physical experiments (pendulums, collisions, and similar setups) whose outcomes are pinned by conservation laws, then scores a model with physics-informed metrics derived from the conserved quantity for each setting (energy, momentum). Because the conservation law is exact for the chosen experiments, MORPHEUS can ask not merely "does this look like physics?" but "is the conserved quantity actually conserved in the generated rollout?" Its finding echoes Physics-IQ: current models generate aesthetically pleasing video yet struggle to encode the underlying physical principles, even with video conditioning and advanced prompting.

5.2 Controllability: action-conditioned prediction error

Controllability is measured by action-conditioned prediction error: condition the model on a known action sequence, generate the next frame or latent, and measure the error of the prediction against the true action-induced outcome. This is the quantity DIAMOND (Alonso et al., 2024) and the Genie line use in their controllability analyses, and it is exactly the action-following idea of Code Fragment 1 expressed as a regression error rather than a binary check. For decoder-free predictors (Section 36.7), the analogous score is planning success: V-JEPA 2-AC reports the success rate of a planner that uses the action-conditioned predictor to reach goals, measuring controllability by whether the model is good enough to plan with, not by any pixel.

5.3 Long-horizon coherence: VBench and rollout consistency

Long-horizon coherence is the drift, object permanence, and scene-persistence dimension that Genie 3 explicitly targets and that only a long rollout exposes. The most widely adopted instrument here is VBench (Huang et al., 2023; CVPR 2024 Highlight), which decomposes video-generation quality into 16 disentangled, human-preference-validated dimensions (subject consistency, background consistency, motion smoothness, temporal flickering, spatial relationship, dynamic degree, and more) rather than one opaque number. Three of those dimensions, subject consistency, background consistency, and temporal flickering, directly operationalize parts of long-horizon coherence: a model whose subject identity or background drifts across frames scores poorly on exactly the consistency dimensions that a renderer-not-simulator would fail. VBench++ (Huang et al., 2024) extends the suite to image-to-video evaluation and adds trustworthiness dimensions. Beyond VBench, the axis is also measured directly as rollout consistency: accumulated open-loop prediction error as a function of horizon, the drift curve the capstone lab plots in Step 6.

6. The Gold Standard: Downstream-Task Value Advanced

Every metric above scores the world model in isolation, but for an agent world model there is a more honest question that sidesteps pixels entirely: is the simulator good enough to be useful? The trouble with appearance and even physics scores is that they measure the model against a notion of correctness the designer chose, not against the task the model exists to serve. A simulator earns its keep only if an agent trained inside it performs well in the real environment. This is the gold standard, evaluate the simulator by the utility of the policy it produces, not by how its frames look.

The cleanest instance is DIAMOND (Alonso et al., 2024): a reinforcement-learning agent is trained entirely inside the learned diffusion world model, never touching the real environment during policy learning, and is then scored in the real environment. DIAMOND reports a mean human-normalized score of 1.46 on the Atari-100k benchmark, the best at publication for an agent trained purely within a world model. That single number evaluates the simulator more stringently than any pixel metric could: a world model with beautiful frames but wrong dynamics would train a policy that fails on the real game, dragging the score down. Downstream-task value is the metric that cannot be gamed by appearance, because the policy, not the human eye, is the judge, and the policy only cares whether the dreamed dynamics matched reality. This is the evaluation analogue of the capstone lab's philosophy: the dream is only as good as what you can do with it.

7. Where the Generative Metrics Fit Intermediate

Chapter 37 develops a full toolkit of generative-model metrics, and it is worth recapping them here to fix exactly why they are necessary but insufficient for world models. The sample-based metrics, FID and FVD (Fréchet distances between feature distributions) and KID (the kernel analogue, an unbiased MMD that behaves better on small samples), measure how close the generated distribution is to the real one in a perceptual feature space. Precision and recall for generative models split that into two numbers: precision asks what fraction of generated samples are realistic, recall asks what fraction of the real distribution the model can cover, separating a model that produces a few perfect samples from one that captures the full variety. CLIPScore measures text-conditioning fidelity: the cosine similarity between a CLIP embedding of the prompt and of the generated frame, answering "did the model generate what was asked?" There are also likelihood-based measures (bits-per-dimension for models with a tractable density), which score how well the model assigns probability to held-out real data.

Every one of these is necessary: a world model that scores terribly on FVD or CLIPScore is generating unrealistic or off-prompt video and has failed a floor test. But every one is insufficient for the same structural reason FVD is. They all compare distributions or static prompt-frame agreement; none of them conditions on an action and checks the consequence, probes a physical law counterfactually, or follows identity across a long rollout. A self-refilling glass passes FID, FVD, KID, precision-recall, and CLIPScore without complaint, because each frame is realistic and on-prompt and the feature distribution matches real footage. The generative metrics rule out the gross failures; the triad and the downstream task catch the failures that hide in dynamics. You need both layers, never the appearance layer alone.

7.1 Comparison: which instrument measures which axis

This is the place to assemble the whole chapter into one running system. The capstone lab below builds a tiny latent-dynamics world model, the recurrent state-space recipe of Section 36.5 stripped to its essentials, trains a policy entirely in imagination, and then scores it on the consistency triad this section defined. It is the chapter's single thread, a latent space with dynamics, made into something you can run and measure.

pip install torch numpy matplotlib

import torch, torch.nn as nn, torch.nn.functional as F, numpy as np
torch.manual_seed(0); np.random.seed(0)
ACTIONS = torch.tensor([-1.0, 0.0, 1.0])   # left, stay, right

def step_env(pos, action_idx):
    pos = np.clip(pos + 0.1 * ACTIONS[action_idx].item() + 0.01 * np.random.randn(), -1, 1)
    obs = pos + 0.02 * np.random.randn()    # noisy observation of true position
    reward = 1.0 - abs(pos)                  # reward peaks at the center (pos = 0)
    return pos, np.array([obs], dtype=np.float32), reward

def log_episode(T=30):
    pos, obs_seq, act_seq, rew_seq = np.random.uniform(-1, 1), [], [], []
    for _ in range(T):
        a = np.random.randint(3)
        pos, obs, r = step_env(pos, a)
        # TODO: append obs, a one-hot of the action (length 3), and r to the three lists.
        # Hint: one_hot = np.eye(3, dtype=np.float32)[a]
        ...
    return np.stack(obs_seq), np.stack(act_seq), np.array(rew_seq, dtype=np.float32)

episodes = [log_episode() for _ in range(200)]   # the replay buffer
print(episodes[0][0].shape, episodes[0][1].shape)  # (30, 1) (30, 3)

class RSSM(nn.Module):
    def __init__(self, stoch=8, deter=32, act=3, obs=1, hidden=64):
        super().__init__()
        self.gru = nn.GRUCell(stoch + act, deter)
        self.prior = nn.Sequential(nn.Linear(deter, hidden), nn.SiLU(), nn.Linear(hidden, 2 * stoch))
        self.post  = nn.Sequential(nn.Linear(deter + obs, hidden), nn.SiLU(), nn.Linear(hidden, 2 * stoch))
        self.dec   = nn.Sequential(nn.Linear(deter + stoch, hidden), nn.SiLU(), nn.Linear(hidden, obs + 1))
        self.stoch = stoch

    def sample(self, params):
        mean, logstd = params.chunk(2, -1)
        std = torch.exp(logstd.clamp(-5, 2))
        return mean + std * torch.randn_like(std), mean, std

    def forward(self, z, a, h, obs=None):
        h = self.gru(torch.cat([z, a], -1), h)        # advance deterministic memory
        prior = self.prior(h)
        if obs is None:                                # IMAGINATION: dream the next state
            z, _, _ = self.sample(prior)
            return z, h, prior, None
        # TODO: compute the posterior from cat([h, obs]), sample z from it, and return
        #       (z, h, prior, post) so the KL(post || prior) loss can be formed in Step 3.
        ...

def kl(prior, post):
    pm, pls = prior.chunk(2, -1); qm, qls = post.chunk(2, -1)
    pv, qv = torch.exp(pls.clamp(-5, 2)) ** 2, torch.exp(qls.clamp(-5, 2)) ** 2
    return (0.5 * ((qm - pm) ** 2 / pv + qv / pv - 1 + (pv.log() - qv.log()))).sum(-1).mean()

model = RSSM(); opt = torch.optim.Adam(model.parameters(), lr=2e-3)
for epoch in range(40):
    obs, act, rew = [torch.tensor(np.stack(x)) for x in zip(*episodes)]  # [B, T, .]
    B, Tn = obs.shape[0], obs.shape[1]
    z, h = torch.zeros(B, model.stoch), torch.zeros(B, 32)
    rec_loss, kl_loss = 0.0, 0.0
    for t in range(Tn):
        z, h, prior, post = model(z, act[:, t], h, obs=obs[:, t])
        pred = model.dec(torch.cat([h, z], -1))
        # TODO: add MSE(pred[:, :1], obs[:, t]) + MSE(pred[:, 1:], rew[:, t:t+1]) to rec_loss,
        #       and add kl(prior, post) to kl_loss.
        ...
    loss = rec_loss / Tn + 1.0 * kl_loss / Tn
    opt.zero_grad(); loss.backward(); opt.step()
    if epoch % 10 == 0: print(f"epoch {epoch}  rec {rec_loss.item()/Tn:.4f}  kl {kl_loss.item()/Tn:.4f}")

policy = nn.Sequential(nn.Linear(32 + 8, 64), nn.SiLU(), nn.Linear(64, 3))
popt = torch.optim.Adam(policy.parameters(), lr=3e-3)
for it in range(300):
    z, h = torch.zeros(64, 8), torch.zeros(64, 32)   # 64 parallel dreams
    total_reward = 0.0
    for _ in range(15):                               # 15-step imagined horizon
        logits = policy(torch.cat([h, z], -1))
        a = F.softmax(logits, -1)                     # soft action keeps it differentiable
        z, h, _, _ = model(z, a, h, obs=None)         # dream forward with the prior
        # TODO: decode the reward (last column of model.dec(cat([h, z]))) and add its
        #       mean to total_reward.
        ...
    loss = -total_reward                              # maximize dreamed return
    popt.zero_grad(); loss.backward(); popt.step()
    if it % 100 == 0: print(f"iter {it}  dreamed return {total_reward.item():.3f}")

@torch.no_grad()
def open_loop_error(horizon):
    obs, act, _ = [torch.tensor(np.stack(x)) for x in zip(*[log_episode() for _ in range(64)])]
    z, h = torch.zeros(64, 8), torch.zeros(64, 32)
    z, h, _, _ = model(z, act[:, 0], h, obs=obs[:, 0])   # one observed warm-up step
    err = []
    for t in range(1, horizon):
        z, h, _, _ = model(z, act[:, t], h, obs=None)    # dream, no observation
        pred_obs = model.dec(torch.cat([h, z], -1))[:, :1]
        err.append(F.mse_loss(pred_obs, obs[:, t]).item())
    return err

# TODO: build the controllability check: from a fixed warmed-up (z, h), dream 10 steps
#       with the all-left one-hot and again with the all-right one-hot, decode the final
#       predicted observation for each, and report the absolute difference (should be > 0).
drift = open_loop_error(horizon=25)
print("one-step error:", drift[0], " 24-step error:", drift[-1])

import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, 3, figsize=(12, 3.2))
ax[0].bar(["1-step"], [drift[0]]); ax[0].set_title("Physical consistency\n(lower is better)")
ax[1].set_title("Controllability\n(left vs right gap)")    # fill from your Step 5 check
ax[2].plot(range(1, 25), drift, marker="."); ax[2].set_title("Coherence: drift vs horizon")
ax[2].set_xlabel("imagined step"); ax[2].set_ylabel("prediction MSE")
plt.tight_layout(); plt.savefig("world_model_triad.png", dpi=130)
print("saved world_model_triad.png")

import torch, torch.nn as nn, torch.nn.functional as F, numpy as np
import matplotlib.pyplot as plt
torch.manual_seed(0); np.random.seed(0)
ACTIONS = torch.tensor([-1.0, 0.0, 1.0])

# --- Step 1: environment and logging ---
def step_env(pos, action_idx):
    pos = np.clip(pos + 0.1 * ACTIONS[action_idx].item() + 0.01 * np.random.randn(), -1, 1)
    obs = pos + 0.02 * np.random.randn()
    reward = 1.0 - abs(pos)
    return pos, np.array([obs], dtype=np.float32), reward

def log_episode(T=30):
    pos, obs_seq, act_seq, rew_seq = np.random.uniform(-1, 1), [], [], []
    for _ in range(T):
        a = np.random.randint(3)
        pos, obs, r = step_env(pos, a)
        obs_seq.append(obs); act_seq.append(np.eye(3, dtype=np.float32)[a]); rew_seq.append(r)
    return np.stack(obs_seq), np.stack(act_seq), np.array(rew_seq, dtype=np.float32)

episodes = [log_episode() for _ in range(200)]

# --- Step 2: RSSM ---
class RSSM(nn.Module):
    def __init__(self, stoch=8, deter=32, act=3, obs=1, hidden=64):
        super().__init__()
        self.gru = nn.GRUCell(stoch + act, deter)
        self.prior = nn.Sequential(nn.Linear(deter, hidden), nn.SiLU(), nn.Linear(hidden, 2 * stoch))
        self.post  = nn.Sequential(nn.Linear(deter + obs, hidden), nn.SiLU(), nn.Linear(hidden, 2 * stoch))
        self.dec   = nn.Sequential(nn.Linear(deter + stoch, hidden), nn.SiLU(), nn.Linear(hidden, obs + 1))
        self.stoch = stoch
    def sample(self, params):
        mean, logstd = params.chunk(2, -1)
        std = torch.exp(logstd.clamp(-5, 2))
        return mean + std * torch.randn_like(std), mean, std
    def forward(self, z, a, h, obs=None):
        h = self.gru(torch.cat([z, a], -1), h)
        prior = self.prior(h)
        if obs is None:
            z, _, _ = self.sample(prior)
            return z, h, prior, None
        post = self.post(torch.cat([h, obs], -1))
        z, _, _ = self.sample(post)
        return z, h, prior, post

# --- Step 3: train dynamics ---
def kl(prior, post):
    pm, pls = prior.chunk(2, -1); qm, qls = post.chunk(2, -1)
    pv, qv = torch.exp(pls.clamp(-5, 2)) ** 2, torch.exp(qls.clamp(-5, 2)) ** 2
    return (0.5 * ((qm - pm) ** 2 / pv + qv / pv - 1 + (pv.log() - qv.log()))).sum(-1).mean()

model = RSSM(); opt = torch.optim.Adam(model.parameters(), lr=2e-3)
for epoch in range(40):
    obs, act, rew = [torch.tensor(np.stack(x)) for x in zip(*episodes)]
    B, Tn = obs.shape[0], obs.shape[1]
    z, h = torch.zeros(B, model.stoch), torch.zeros(B, 32)
    rec_loss, kl_loss = 0.0, 0.0
    for t in range(Tn):
        z, h, prior, post = model(z, act[:, t], h, obs=obs[:, t])
        pred = model.dec(torch.cat([h, z], -1))
        rec_loss = rec_loss + F.mse_loss(pred[:, :1], obs[:, t]) + F.mse_loss(pred[:, 1:], rew[:, t:t+1])
        kl_loss = kl_loss + kl(prior, post)
    loss = rec_loss / Tn + 1.0 * kl_loss / Tn
    opt.zero_grad(); loss.backward(); opt.step()
    if epoch % 10 == 0: print(f"epoch {epoch}  rec {rec_loss.item()/Tn:.4f}  kl {kl_loss.item()/Tn:.4f}")

# --- Step 4: policy in imagination ---
policy = nn.Sequential(nn.Linear(32 + 8, 64), nn.SiLU(), nn.Linear(64, 3))
popt = torch.optim.Adam(policy.parameters(), lr=3e-3)
for it in range(300):
    z, h = torch.zeros(64, 8), torch.zeros(64, 32)
    total_reward = 0.0
    for _ in range(15):
        a = F.softmax(policy(torch.cat([h, z], -1)), -1)
        z, h, _, _ = model(z, a, h, obs=None)
        total_reward = total_reward + model.dec(torch.cat([h, z], -1))[:, 1:].mean()
    (-total_reward).backward(); popt.step(); popt.zero_grad()
    if it % 100 == 0: print(f"iter {it}  dreamed return {total_reward.item():.3f}")

# --- Step 5: triad ---
@torch.no_grad()
def open_loop_error(horizon):
    obs, act, _ = [torch.tensor(np.stack(x)) for x in zip(*[log_episode() for _ in range(64)])]
    z, h = torch.zeros(64, 8), torch.zeros(64, 32)
    z, h, _, _ = model(z, act[:, 0], h, obs=obs[:, 0])
    err = []
    for t in range(1, horizon):
        z, h, _, _ = model(z, act[:, t], h, obs=None)
        err.append(F.mse_loss(model.dec(torch.cat([h, z], -1))[:, :1], obs[:, t]).item())
    return err

@torch.no_grad()
def controllability_gap():
    obs, act, _ = [torch.tensor(np.stack(x)) for x in zip(*[log_episode() for _ in range(64)])]
    z0, h0 = torch.zeros(64, 8), torch.zeros(64, 32)
    z0, h0, _, _ = model(z0, act[:, 0], h0, obs=obs[:, 0])
    outs = []
    for idx in (0, 2):                                  # all-left, all-right
        z, h = z0.clone(), h0.clone()
        a = F.one_hot(torch.tensor(idx), 3).float().expand(64, 3)
        for _ in range(10):
            z, h, _, _ = model(z, a, h, obs=None)
        outs.append(model.dec(torch.cat([h, z], -1))[:, :1])
    return (outs[1] - outs[0]).abs().mean().item()

drift = open_loop_error(25)
gap = controllability_gap()
print("one-step error:", drift[0], " 24-step error:", drift[-1], " control gap:", gap)

# --- Step 6: plot ---
fig, ax = plt.subplots(1, 3, figsize=(12, 3.2))
ax[0].bar(["1-step"], [drift[0]]); ax[0].set_title("Physical consistency\n(lower is better)")
ax[1].bar(["gap"], [gap]); ax[1].set_title("Controllability\n(left vs right gap)")
ax[2].plot(range(1, 25), drift, marker="."); ax[2].set_title("Coherence: drift vs horizon")
ax[2].set_xlabel("imagined step"); ax[2].set_ylabel("prediction MSE")
plt.tight_layout(); plt.savefig("world_model_triad.png", dpi=130)
print("saved world_model_triad.png")

8. Further Reading: Benchmarks for World-Model Evaluation Intermediate