TD-MPC and Control-Centric Representations › Code Implementation: A Minimal TD-MPC Agent
World Models HandbookAll chapters 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:
- A learning curve, to check whether this tiny training budget improves behavior at all.
- An open-loop error measurement on held-out trajectories, without assuming error must rise with horizon.
- 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 ; 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.
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 = 200class 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, truncatedA random policy gives us a reference line and confirms the environment terminates sensibly.
_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))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.
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.
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, stdThe 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.
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.
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.
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()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.
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()
}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.
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])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.0001With 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.

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.


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.

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.