TD-MPC and Control-Centric Representations › Code Implementation: A Minimal TD-MPC Agent

Michael BrenndoerferJuly 17, 202618 min read
World Models HandbookAll chapters
Part 8 of 12Decision-Centric Research Lineages
Chapter 47TD-MPC and Control-Centric Representations › Code Implementation: A Minimal TD-MPC Agent
Reading section 6 of 9 · Code Implementation: A Minimal TD-MPC Agent · 18 minSection 6 of 9 · 18 min

Book contents

Section 6 of 9 · 18 min read

TD-MPC and Control-Centric Representations › Code Implementation: A Minimal TD-MPC AgentOne section · Sections in this chapter6 / 9 ↓
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 will now build a compact TD-MPC-style cart-pole agent from scratch so the environment and controller are visible. It uses one-step replay samples rather than either paper's recurrent multistep sequence training, and borrows SimNorm from TD-MPC2. Its eight episodes are a diagnostic of code behavior, not a reproduction of either paper or a guaranteed learning demonstration.

Three diagnostics come out of this deliberately short run:

  1. A learning curve, to check whether this tiny training budget improves behavior at all.
  2. An open-loop error measurement on held-out trajectories, without assuming error must rise with horizon.
  3. Two small planning comparisons: terminal value on versus off at fixed horizon, and horizon three versus six with terminal value off. The analytic LQR example above, not this short run, establishes the terminal-value mechanism.

Start with imports and the environment. The cart-pole uses an Euler integrator and a near-upright initial state, so the task is balance rather than swing-up. The dense reward is rt=cos⁡θt−0.05 xt2r_t = \cos\theta_t - 0.05\,x_t^2; episodes end when the pole falls past twelve degrees or the cart leaves the track. Unlike a constant survival reward, this signal gives the reward head a state-dependent target to learn. The experiment can therefore test its predictions separately from the terminal value estimate.

In[11]:
Code
import copy
import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
np.random.seed(0)

GAMMA = 0.97  # effective horizon 1/(1-GAMMA) ~ 33 steps
LATENT_DIM = 64  # divisible by the SimNorm group count
HIDDEN = 128
MAX_EPISODE_STEPS = 200
In[12]:
Code
class CartPole:
    """Near-upright-start cart-pole with a dense, state-dependent reward."""

    dt = 0.02
    gravity = 9.8
    masscart = 1.0
    masspole = 0.1
    length = 0.5  # half the pole length
    force_mag = 10.0
    theta_limit = 12 * math.pi / 180
    x_limit = 2.4

    def observe(self):
        return np.asarray(self.state, dtype=np.float32)

    def reset(self, rng=None, state=None):
        if state is None:
            state = rng.uniform(-0.05, 0.05, size=4)
        self.state = np.asarray(state, dtype=np.float64).copy()
        self.steps = 0
        return self.observe()

    def step(self, action):
        u = float(np.asarray(action).reshape(-1)[0])
        force = self.force_mag * max(-1.0, min(1.0, u))
        x, xdot, theta, thetadot = self.state
        cos, sin = math.cos(theta), math.sin(theta)
        total_mass = self.masscart + self.masspole
        polemass_length = self.masspole * self.length
        temp = (force + polemass_length * thetadot**2 * sin) / total_mass
        thetaacc = (self.gravity * sin - cos * temp) / (
            self.length * (4.0 / 3.0 - self.masspole * cos**2 / total_mass)
        )
        xacc = temp - polemass_length * thetaacc * cos / total_mass
        self.state = np.asarray(
            [
                x + self.dt * xdot,
                xdot + self.dt * xacc,
                theta + self.dt * thetadot,
                thetadot + self.dt * thetaacc,
            ]
        )
        self.steps += 1
        x, xdot, theta, thetadot = self.state
        reward = math.cos(theta) - 0.05 * x**2
        terminated = bool(
            abs(x) > self.x_limit or abs(theta) > self.theta_limit
        )
        truncated = self.steps >= MAX_EPISODE_STEPS
        return self.observe(), reward, terminated, truncated

A random policy gives us a reference line and confirms the environment terminates sensibly.

In[13]:
Code
_rand_rng = np.random.default_rng(0)
_env = CartPole()
random_returns, random_lengths = [], []
for _ in range(20):
    _env.reset(_rand_rng)
    total = 0.0
    for _ in range(MAX_EPISODE_STEPS):
        _, reward, terminated, truncated = _env.step(
            _rand_rng.uniform(-1, 1, size=1)
        )
        total += reward
        if terminated or truncated:
            break
    random_returns.append(total)
    random_lengths.append(_env.steps)

random_return_baseline = float(np.mean(random_returns))
random_length_baseline = float(np.mean(random_lengths))
Out[14]:
Console
Random policy: mean return 24.2, mean episode length 24.3 steps

Now the five components. This compact implementation combines dynamics and reward prediction in one module and borrows SimNorm from TD-MPC2. It is an executable teaching example, not a reproduction of either paper's full architecture. A combined module shares computation; no benchmark here establishes a specific speedup.

In[15]:
Code
class SimNorm(nn.Module):
    """Softmax within groups of latent units (Simplex Normalization)."""

    def __init__(self, dim, groups=8):
        super().__init__()
        self.groups = groups
        self.group_size = dim // groups

    def forward(self, x):
        shape = x.shape
        x = x.view(*shape[:-1], self.groups, self.group_size)
        x = F.softmax(x, dim=-1)
        return x.view(*shape)


class Encoder(nn.Module):
    def __init__(self, obs_dim, latent_dim=LATENT_DIM, hidden=HIDDEN):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden),
            nn.Mish(),
            nn.Linear(hidden, hidden),
            nn.Mish(),
            nn.Linear(hidden, latent_dim),
            SimNorm(latent_dim),
        )

    def forward(self, obs):
        return self.net(obs)


class LatentDynamics(nn.Module):
    """Predicts the next latent and the reward for one transition."""

    def __init__(self, latent_dim=LATENT_DIM, action_dim=1, hidden=HIDDEN):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(latent_dim + action_dim, hidden),
            nn.Mish(),
            nn.Linear(hidden, hidden),
            nn.Mish(),
        )
        self.next_latent = nn.Sequential(
            nn.Linear(hidden, latent_dim), SimNorm(latent_dim)
        )
        self.reward_head = nn.Linear(hidden, 1)

    def forward(self, z, a):
        h = self.trunk(torch.cat([z, a], dim=-1))
        return self.next_latent(h), self.reward_head(h).squeeze(-1)


class QFunction(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM, action_dim=1, hidden=HIDDEN):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim + action_dim, hidden),
            nn.Mish(),
            nn.Linear(hidden, hidden),
            nn.Mish(),
            nn.Linear(hidden, 1),
        )

    def forward(self, z, a):
        return self.net(torch.cat([z, a], dim=-1)).squeeze(-1)


class PolicyPrior(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM, action_dim=1, hidden=HIDDEN):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, hidden),
            nn.Mish(),
            nn.Linear(hidden, hidden),
            nn.Mish(),
            nn.Linear(hidden, action_dim),
            nn.Tanh(),
        )

    def forward(self, z):
        return self.net(z)


class TDMPCAgent(nn.Module):
    def __init__(
        self, obs_dim=4, action_dim=1, latent_dim=LATENT_DIM, hidden=HIDDEN
    ):
        super().__init__()
        self.action_dim = action_dim
        self.encoder = Encoder(obs_dim, latent_dim, hidden)
        self.dynamics = LatentDynamics(latent_dim, action_dim, hidden)
        self.q = QFunction(latent_dim, action_dim, hidden)
        self.prior = PolicyPrior(latent_dim, action_dim, hidden)

    def encode(self, obs):
        return self.encoder(obs)

Here is a small MPPI-style planner. The policy contributes noisy proposal sequences; Gaussian samples fill out the population; elites are weighted by their raw return scores after the numerically stable softmax shift. The executed action is sampled from the refitted Gaussian and clipped. Unlike the released planner, this toy starts each decision with a zero-mean, fixed-width Gaussian and has no previous-plan warm start or scheduled proposal variance.

In[16]:
Code
def plan(
    agent,
    z,
    horizon=3,
    num_samples=128,
    num_elites=16,
    iterations=3,
    num_prior_trajs=8,
    sigma=0.5,
    temperature=0.5,
    use_terminal_value=True,
    generator=None,
    action_low=-1.0,
    action_high=1.0,
):
    """MPPI in latent space with an optional terminal value."""
    with torch.no_grad():
        action_dim = agent.action_dim
        mean = torch.zeros(horizon, action_dim)
        std = torch.full((horizon, action_dim), sigma)

        # Distinct noisy proposal trajectories from the learned policy prior.
        z_prior = z.expand(num_prior_trajs, -1)
        prior_steps = []
        for _ in range(horizon):
            a = (
                agent.prior(z_prior)
                + 0.1
                * torch.randn(num_prior_trajs, action_dim, generator=generator)
            ).clamp(action_low, action_high)
            prior_steps.append(a)
            z_prior, _ = agent.dynamics(z_prior, a)
        prior_seq = torch.stack(prior_steps, dim=1)  # (P, H, A)

        for _ in range(iterations):
            eps = torch.randn(
                num_samples, horizon, action_dim, generator=generator
            )
            sampled = (mean + std * eps).clamp(action_low, action_high)
            candidates = torch.cat([prior_seq, sampled], dim=0)

            zs = z.expand(candidates.shape[0], -1)
            returns = torch.zeros(candidates.shape[0])
            for t in range(horizon):
                zs, r = agent.dynamics(zs, candidates[:, t])
                returns = returns + (GAMMA**t) * r

            if use_terminal_value:
                returns = returns + (GAMMA**horizon) * agent.q(
                    zs, agent.prior(zs)
                )

            # Softmax subtracts the elite maximum for numerical stability.
            elite_idx = torch.topk(returns, num_elites).indices
            elites = candidates[elite_idx]
            weights = torch.softmax(temperature * returns[elite_idx], dim=0)

            mean = (weights[:, None, None] * elites).sum(dim=0)
            std = (
                torch.sqrt(
                    (weights[:, None, None] * (elites - mean) ** 2).sum(dim=0)
                )
                + 1e-6
            )

        # Sample the executed action for exploration.
        action = mean[0] + temperature * std[0] * torch.randn(
            action_dim, generator=generator
        )
        action = action.clamp(action_low, action_high)
    return action, mean, std

The toy optimizes the model/critic losses and then the policy loss in separate steps. It uses one-step replay samples rather than the papers' recurrent multi-step training and adds a small action-magnitude penalty to the toy policy objective. The penalty is not entropy and is not part of the original TD-MPC objective.

In[17]:
Code
class ReplayBuffer:
    def __init__(self, capacity=50_000):
        self.capacity = capacity
        self.storage = []
        self.next_index = 0

    def add(self, obs, action, reward, next_obs, terminated):
        item = (
            np.asarray(obs, dtype=np.float32),
            np.asarray(action, dtype=np.float32).reshape(-1),
            np.float32(reward),
            np.asarray(next_obs, dtype=np.float32),
            np.float32(terminated),
        )
        if len(self.storage) < self.capacity:
            self.storage.append(item)
        else:
            self.storage[self.next_index] = item
        self.next_index = (self.next_index + 1) % self.capacity

    def __len__(self):
        return len(self.storage)

    def sample(self, batch_size, generator):
        idx = torch.randint(
            len(self.storage), (batch_size,), generator=generator
        ).tolist()
        obs, act, rew, nobs, term = zip(*(self.storage[i] for i in idx))
        return (
            torch.as_tensor(np.stack(obs)),
            torch.as_tensor(np.stack(act)),
            torch.as_tensor(np.stack(rew)),
            torch.as_tensor(np.stack(nobs)),
            torch.as_tensor(np.stack(term)),
        )


def td_update(agent, target, model_optimizer, prior_optimizer, batch):
    obs, action, reward, next_obs, terminated = batch

    z = agent.encode(obs)
    with torch.no_grad():
        z_next = agent.encode(next_obs)
        z_next_consistency_target = target.encode(next_obs)
        q_target = reward + GAMMA * (1.0 - terminated) * target.q(
            z_next, agent.prior(z_next)
        )

    pred_next_z, pred_reward = agent.dynamics(z, action)
    consistency_loss = F.mse_loss(pred_next_z, z_next_consistency_target)
    reward_loss = F.mse_loss(pred_reward, reward)
    q_loss = F.mse_loss(agent.q(z, action), q_target)

    model_loss = consistency_loss + reward_loss + q_loss
    model_optimizer.zero_grad()
    model_loss.backward()
    # Clip only parameters updated by this optimizer. Prior gradients from the
    # previous actor step may still exist and must not scale the model update.
    model_parameters = [
        p for group in model_optimizer.param_groups for p in group["params"]
    ]
    nn.utils.clip_grad_norm_(model_parameters, 20.0)
    model_optimizer.step()

    # Policy: maximize Q through the action, not through critic parameters.
    z_prior = z.detach()
    q_flags = [p.requires_grad for p in agent.q.parameters()]
    for p in agent.q.parameters():
        p.requires_grad_(False)
    prior_optimizer.zero_grad()
    prior_action = agent.prior(z_prior)
    prior_loss = -agent.q(z_prior, prior_action).mean()
    action_magnitude_penalty = 0.1 * (prior_action**2).mean()
    (prior_loss + action_magnitude_penalty).backward()
    nn.utils.clip_grad_norm_(agent.prior.parameters(), 20.0)
    prior_optimizer.step()
    for p, was_trainable in zip(agent.q.parameters(), q_flags):
        p.requires_grad_(was_trainable)
    return q_loss.detach().item()

The optimizer boundary is worth testing: gradients left on the policy prior after an actor step must not change the next model update. This assertion runs on two identical agents and the same batch; only one starts with large stale prior gradients. SGD makes any unintended gradient rescaling visible in the model parameters.

In[18]:
Code
def assert_model_update_ignores_stale_prior_grads():
    with torch.random.fork_rng():
        torch.manual_seed(47)
        clean = TDMPCAgent()
    stale = copy.deepcopy(clean)
    clean_target, stale_target = copy.deepcopy(clean), copy.deepcopy(stale)
    for target_model in (clean_target, stale_target):
        for parameter in target_model.parameters():
            parameter.requires_grad_(False)

    batch_rng = torch.Generator().manual_seed(47)
    batch = (
        torch.randn(8, 4, generator=batch_rng),
        torch.tanh(torch.randn(8, 1, generator=batch_rng)),
        torch.randn(8, generator=batch_rng),
        torch.randn(8, 4, generator=batch_rng),
        torch.zeros(8),
    )
    assert batch[2].shape == batch[4].shape == (8,)

    def optimizers(agent):
        model_parameters = (
            list(agent.encoder.parameters())
            + list(agent.dynamics.parameters())
            + list(agent.q.parameters())
        )
        return (
            torch.optim.SGD(model_parameters, lr=0.01),
            torch.optim.SGD(agent.prior.parameters(), lr=0.01),
            model_parameters,
        )

    clean_model_opt, clean_prior_opt, clean_model_params = optimizers(clean)
    stale_model_opt, stale_prior_opt, stale_model_params = optimizers(stale)
    for parameter in stale.prior.parameters():
        parameter.grad = torch.full_like(parameter, 1e5)

    td_update(clean, clean_target, clean_model_opt, clean_prior_opt, batch)
    td_update(stale, stale_target, stale_model_opt, stale_prior_opt, batch)
    for clean_parameter, stale_parameter in zip(
        clean_model_params, stale_model_params
    ):
        torch.testing.assert_close(
            clean_parameter, stale_parameter, rtol=0, atol=1e-7
        )


assert_model_update_ignores_stale_prior_grads()

The training loop alternates collection and learning. The buffer stores true termination separately from time-limit truncation, so the toy bootstraps at time limits. Two random-action episodes seed replay. Four updates every fourth environment step give roughly one update per post-warm-up step once the buffer gate opens; this is not evidence of improvement by itself.

In[19]:
Code
def collect_and_train(num_episodes=8, warmup_episodes=2, seed=1):
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    gen = torch.Generator().manual_seed(seed)

    agent = TDMPCAgent()
    target = copy.deepcopy(agent)
    for p in target.parameters():
        p.requires_grad_(False)
    model_params = (
        list(agent.encoder.parameters())
        + list(agent.dynamics.parameters())
        + list(agent.q.parameters())
    )
    model_optimizer = torch.optim.Adam(model_params, lr=3e-4)
    prior_optimizer = torch.optim.Adam(agent.prior.parameters(), lr=3e-4)

    env, buffer = CartPole(), ReplayBuffer()
    log = {
        "episode_return": [],
        "episode_length": [],
        "env_steps": [],
        "updates": [],
    }
    total_steps = 0
    total_updates = 0

    for episode in range(num_episodes):
        obs = env.reset(rng)
        ep_return = 0.0
        for t in range(MAX_EPISODE_STEPS):
            if episode < warmup_episodes:
                action = rng.uniform(-1.0, 1.0, size=1).astype(np.float32)
            else:
                with torch.no_grad():
                    z = agent.encode(torch.as_tensor(obs).unsqueeze(0))
                    action, _, _ = plan(agent, z, generator=gen)
                action = action.numpy()

            next_obs, reward, terminated, truncated = env.step(action)
            buffer.add(obs, action, reward, next_obs, float(terminated))
            obs = next_obs
            ep_return += reward
            total_steps += 1

            if episode >= warmup_episodes and len(buffer) > 64 and t % 4 == 0:
                for _ in range(4):
                    td_update(
                        agent,
                        target,
                        model_optimizer,
                        prior_optimizer,
                        buffer.sample(128, gen),
                    )
                    total_updates += 1
                with torch.no_grad():
                    for p, tp in zip(agent.parameters(), target.parameters()):
                        tp.mul_(0.99).add_(p, alpha=0.01)

            if terminated or truncated:
                break

        log["episode_return"].append(ep_return)
        log["episode_length"].append(env.steps)
        log["env_steps"].append(total_steps)
        log["updates"].append(total_updates)

    return agent, log


agent, log = collect_and_train()
Out[20]:
Console
Environment steps collected: 127
Model/critic and policy update pairs: 72
Episode returns:  39.8   30.9    8.9    8.9    8.9    8.9    9.9    9.9
Episode lengths:    40     31      9      9      9      9     10     10
Random-policy reference return: 24.2

The executed notebook collected 127 environment steps and performed 72 model/critic-plus-policy update pairs. Its two random-action warm-up returns were 39.8 and 30.9; the six planner-driven returns were 8.9, 8.9, 8.9, 8.9, 9.9, and 9.9, below the independently measured random-policy mean of 24.2. Updates happened, but this run did not demonstrate improved control. The sharp drop is precisely why an executable learning curve matters more than a plausible training loop.

Now a small controlled diagnostic with the same trained weights and initial evaluation states. Comparing the first two settings isolates terminal-value use at horizon three. Comparing the second and third holds terminal-value use off and changes only the horizon from three to six. Three starts are too few to infer a reliable terminal-value benefit or choose a generally useful horizon.

In[21]:
Code
def evaluate_planning(
    agent, planning_kwargs, eval_states, max_steps=120, seed=0
):
    gen = torch.Generator().manual_seed(seed)
    env = CartPole()
    returns, lengths = [], []
    for state in eval_states:
        obs = env.reset(state=state)
        total = 0.0
        for t in range(max_steps):
            with torch.no_grad():
                z = agent.encode(torch.as_tensor(obs).unsqueeze(0))
                action, _, _ = plan(agent, z, generator=gen, **planning_kwargs)
            obs, reward, terminated, truncated = env.step(action.numpy())
            total += reward
            if terminated or truncated:
                break
        returns.append(total)
        lengths.append(t + 1)
    return {
        "return": float(np.mean(returns)),
        "length": float(np.mean(lengths)),
    }


_eval_rng = np.random.default_rng(7)
eval_states = [_eval_rng.uniform(-0.05, 0.05, size=4) for _ in range(3)]

planning_configs = {
    "3 steps + terminal value": dict(horizon=3, use_terminal_value=True),
    "3 steps, no terminal value": dict(horizon=3, use_terminal_value=False),
    "6 steps, no terminal value": dict(horizon=6, use_terminal_value=False),
}
eval_results = {
    name: evaluate_planning(agent, cfg, eval_states)
    for name, cfg in planning_configs.items()
}
Out[22]:
Console
planning configuration        mean return  mean length
3 steps + terminal value           9.9158         10.0
3 steps, no terminal value         9.9156         10.0
6 steps, no terminal value         9.9337         10.0

For open-loop diagnostics, collect trajectories with the learned prior plus noise, fit a linear state probe on half the segments, and evaluate both probe and rollout errors on the other half. These trajectories are a diagnostic distribution; we have not shown that they match MPPI deployment. Latent MSE, reward MSE, and state-probe L2 error have different units and answer different questions. The probe's held-out no-rollout error is a reference for interpreting its rolled-out error, not a measure of distance to the latent manifold.

In[23]:
Code
def collect_validation_segments(
    agent, num_segments=8, max_seg_steps=40, noise=0.3, seed=11
):
    rng = np.random.default_rng(seed)
    env = CartPole()
    segments = []
    for _ in range(num_segments):
        obs = env.reset(rng)
        obs_list, act_list, rew_list = [obs.copy()], [], []
        for _ in range(max_seg_steps):
            with torch.no_grad():
                a = (
                    agent.prior(agent.encode(torch.as_tensor(obs).unsqueeze(0)))
                    .squeeze(0)
                    .numpy()
                )
            a = np.clip(a + noise * rng.normal(size=1), -1.0, 1.0).astype(
                np.float32
            )
            next_obs, reward, terminated, truncated = env.step(a)
            obs_list.append(next_obs.copy())
            act_list.append(a)
            rew_list.append(np.float32(reward))
            obs = next_obs
            if terminated or truncated:
                break
        if len(act_list) >= 2:
            segments.append(
                (np.stack(obs_list), np.stack(act_list), np.stack(rew_list))
            )
    return segments


def open_loop_diagnostics(agent, segments, max_horizon=8):
    # Fit a linear probe on one split; evaluate all metrics on held-out segments.
    assert len(segments) >= 2
    split = len(segments) // 2
    fit_segments, heldout_segments = segments[:split], segments[split:]
    latents_all = []
    states_all = []
    for obs_seq, _, _ in fit_segments:
        with torch.no_grad():
            latents_all.append(agent.encode(torch.as_tensor(obs_seq)).numpy())
        states_all.append(obs_seq)
    X = np.concatenate(latents_all, axis=0)
    Y = np.concatenate(states_all, axis=0)
    design = np.hstack([X, np.ones((len(X), 1))])
    probe, *_ = np.linalg.lstsq(design, Y, rcond=None)
    heldout_obs = np.concatenate(
        [obs_seq for obs_seq, _, _ in heldout_segments]
    )
    with torch.no_grad():
        heldout_latents = agent.encode(torch.as_tensor(heldout_obs)).numpy()
    heldout_design = np.hstack(
        [heldout_latents, np.ones((len(heldout_latents), 1))]
    )
    probe_baseline = float(
        np.mean(np.linalg.norm(heldout_design @ probe - heldout_obs, axis=1))
    )

    latent_err = {h: [] for h in range(1, max_horizon + 1)}
    reward_err = {h: [] for h in range(1, max_horizon + 1)}
    state_err = {h: [] for h in range(1, max_horizon + 1)}

    for obs_seq, act_seq, rew_seq in heldout_segments:
        with torch.no_grad():
            latents = agent.encode(torch.as_tensor(obs_seq))
        T = len(act_seq)
        for i in range(T):
            z = latents[i : i + 1]
            for h in range(1, max_horizon + 1):
                j = i + h - 1
                if j >= T:
                    break
                with torch.no_grad():
                    z, r_hat = agent.dynamics(
                        z, torch.as_tensor(act_seq[j : j + 1])
                    )
                latent_err[h].append(
                    float(F.mse_loss(z, latents[j + 1 : j + 2]))
                )
                reward_err[h].append(float((r_hat - float(rew_seq[j])) ** 2))
                decoded = np.hstack([z.numpy(), np.ones((1, 1))]) @ probe
                state_err[h].append(
                    float(np.linalg.norm(decoded - obs_seq[j + 1]))
                )
    return latent_err, reward_err, state_err, probe_baseline


segments = collect_validation_segments(agent)
latent_err, reward_err, state_err, probe_baseline = open_loop_diagnostics(
    agent, segments
)

horizons = np.arange(1, 9)
mean_latent_err = np.array([np.mean(latent_err[int(h)]) for h in horizons])
mean_reward_err = np.array([np.mean(reward_err[int(h)]) for h in horizons])
mean_state_err = np.array([np.mean(state_err[int(h)]) for h in horizons])
Out[24]:
Console
Held-out diagnostic transitions: 44
horizon starts  latent MSE  reward MSE     state L2
      1     44     0.00085     0.00129       2.3866
      2     40     0.00079     0.00447       2.4079
      3     36     0.00072     0.00417       2.4470
      4     32     0.00065     0.00404       2.4912
      5     28     0.00058     0.00384       2.5446
      6     24     0.00052     0.00390       2.5968
      7     20     0.00046     0.00382       2.6571
      8     16     0.00042     0.00387       2.7131

Held-out no-rollout probe error (state L2): 0.0001

With all quantities computed, we can plot. The first figure makes the failed improvement test visible instead of presenting the training loop as proof of learning.

Out[25]:
Visualization
Line chart of episode return relative to a random-policy reference line.
Episode return in the executed toy run: two random-action warm-up episodes scored 39.8 and 30.9; six planner-driven episodes scored 8.9–9.9 despite 72 update pairs. The dashed random-policy mean is 24.2 across 20 episodes. This run does not show learning.

Next, the open-loop diagnostics compare predictions against held-out recorded trajectories. In the executed run, reward MSE is 0.00129 at horizon one, 0.00447 at two, and 0.00387 at eight: it rises, then falls, rather than increasing monotonically. Latent MSE falls from 0.00085 to 0.00042, while state-probe L2 error rises from 2.3866 to 2.7131. The held-out no-rollout probe reference is about 0.0001 L2. These metrics have different units. Their eligible start-state cohorts also shrink with horizon, so these descriptive means do not isolate how error changes along a fixed cohort.

Out[26]:
Visualization
Line chart of held-out reward prediction MSE by rollout horizon.
Held-out reward prediction MSE rises from 0.00129 at horizon one to 0.00447 at two, then varies nonmonotonically and is 0.00387 at eight. Recorded actions determine the rollout; the eligible start-state cohort shrinks with horizon.
Line chart of held-out state-probe L2 error by rollout horizon.
Held-out state-probe L2 error rises from 2.3866 at horizon one to 2.7131 at eight. The dashed no-rollout probe reference is about 0.0001 L2 in the same units. Eligible start states shrink with horizon, so this is not a matched-cohort error-growth curve.

Finally, the planner comparison. Rows one and two isolate terminal-value use at horizon three; rows two and three compare horizons three and six with terminal value off. Unlike the exact LQR example, a poorly trained toy critic is not expected to supply a reliable terminal estimate.

Out[27]:
Visualization
Bar chart of three similar toy evaluation returns testing terminal value at fixed horizon and horizon at fixed no-terminal setting.
Mean return from three fixed initial states with the same undertrained toy agent. All three settings round to about 9.9. Rows one versus two test terminal-value use at horizon three; rows two versus three compare horizons with terminal value off. Neither pair supports a general performance claim.

Read the printed tables before interpreting the shapes. Three questions matter:

First, is there a meaningful gap between the horizon-three rows? Both round to about 9.9, so this run does not demonstrate a terminal-value benefit. It does not refute the exact LQR mechanism. These three returns cannot tell us whether the learned critic's ranking is uninformative or whether the planner, model errors, or small evaluation sample prevent a useful ranking from changing the observed return.

Second, what happens when the no-terminal plan grows from horizon three to six? Both rows round to about 9.9. That pair changes horizon alone for this agent, but three starts and a single undertrained model cannot support a general conclusion about planning depth.

Third, the independent random-policy baseline averaged 24.2, well above these roughly 9.9 planner returns, even though 72 update pairs executed. The baseline uses different sampled starts, so it is a warning signal rather than a paired performance estimate. More episodes, seeds, and tuning would be needed before an empirical learning or ablation claim; the exact worked example above supplies the clean positive demonstration.

The open-loop table reports errors on held-out prior-plus-noise trajectories. In this run the reward MSE is nonmonotone, the latent MSE falls, and the state-probe error rises mildly from an already large one-step value. The H1 and H8 means use 44 and 16 start states respectively, so changing cohort composition may also move them. A low average error on this distribution would not certify accuracy on actions selected by MPPI, and a state-probe error includes both probe approximation and latent rollout error. Use these measurements alongside closed-loop evaluations, not as a formula for choosing the horizon.