Section 38.5: TD-MPC2: latent MPC at scale

"A planner that searches in latent space trades visual completeness for the only thing control actually needs: accurate reward and value."

A Planner That Searches In Compressed State

Problem First

Dreamer uses the world model to train a policy in imagination. TD-MPC2 takes a different path: keep planning online, but make the planning problem small by searching in latent space. This matters when the action must adapt to local scene structure now, not only through a policy learned offline.

Core Model

TD-MPC2 evaluates candidate action sequences by rolling a latent model forward and scoring cumulative reward plus terminal value: $$J(a_{t:t+H-1}) = \sum_{k=0}^{H-1} \hat r(\hat z_{t+k}, a_{t+k}) + \hat V(\hat z_{t+H}).$$ The planner searches directly in latent space, often with a sampling-based optimizer such as CEM or MPPI.

The practical insight is decoder-free sufficiency. If the latent can predict rewards, values, and next latent states accurately enough for search, reconstructing pixels at every step is unnecessary overhead. That is why TD-MPC2 can stay fast even while scaling to many continuous-control tasks and multitask settings.

Its success therefore depends on two linked assumptions: the latent dynamics must stay locally smooth enough for trajectory optimization to make progress, and the terminal value must rescue the planner from short finite-horizon myopia.

Minimal Probe

The code below implements a tiny CEM-style search in latent space. It is not the full algorithm, but it exposes the planning primitive that makes TD-MPC2 different from imagined actor learning.

# Sample action sequences, score them in latent space, and keep elites.
# This mirrors the inner loop of a short-horizon latent MPC planner.
import numpy as np

rng = np.random.default_rng(3)
action_sequences = rng.normal(0.0, 0.4, size=(6, 3))
reward_weights = np.array([1.0, -0.3, 0.5])
scores = action_sequences @ reward_weights
elite_ids = np.argsort(scores)[-2:]
elite_mean = action_sequences[elite_ids].mean(axis=0)
print(
    {
        "best_score": round(float(scores[elite_ids[-1]]), 3),
        "elite_mean": np.round(elite_mean, 3).tolist(),
    }
)

Expected behavior: The elite mean summarizes which local action direction the planner should prefer next. If the elite set changes wildly under tiny observation perturbations, the latent model or reward head is too unstable for MPC to trust.

Practical Recipe

1. Keep planning horizon and wall-clock budget in the same table, because latent MPC wins only if it is fast enough to matter.
2. Warm-start the action proposal distribution from the previous planning step; this often matters as much as model accuracy.
3. Audit terminal value bias by truncating horizon and checking whether the chosen action changes drastically.
4. When scaling across tasks, inspect whether one shared latent still preserves task-specific geometry and actuation constraints.

TD-MPC2 is a reminder that world models are not one family. Some are useful because they let you train a policy cheaply in imagination; others are useful because they let you optimize the next action online. The architecture, loss, and evaluation protocol should therefore be chosen around the intended control interface, not around visual elegance.

Its broader importance is scale. The paper argues that the same core design can cover many continuous-control tasks and multi-embodiment settings. For embodied AI builders, that matters because it suggests a practical path between narrow task-specific MPC and fully general policy models.