TD-MPC2

BetaWorks, but under active development — APIs and results may shift. Works; MPPI planning is intricate and the world model is sensitive.

Temporal Difference Model Predictive Control learns a latent world model and plans with it at every decision, sampling action sequences and optimizing them online.

from noeira.deep_agents.tdmpc2 import TDMPC2Agent

Papers

• TD-MPC2: Scalable, Robust World Models for Continuous ControlHansen, Su, Wang · 2024 · ICLR 2024 · arXiv:2310.16828
• Temporal Difference Learning for Model Predictive ControlHansen, Wang, Su · 2022 · ICML 2022 · arXiv:2203.04955TD-MPC, the predecessor.

Planning, not a policy

This is the distinguishing property. SAC has a policy network that maps observation to action. TD-MPC2 has a model, a value function and a planner, and each action comes from a fresh optimization:

1. Sample many action sequences over a short horizon
2. Roll each through the learned latent dynamics
3. Score them: predicted rewards plus the terminal value estimate
4. Refit the sampling distribution toward the best (MPPI)
5. Repeat a few iterations, then execute the first action of the winner

real observation / chosen actionsampled rolloutsscoringdistribution refit

A policy prior guides the sampling so the search stays tractable — without it, random sequences in a high-dimensional action space almost never find anything good.

Two parts of that loop carry more weight than their line count suggests.

The terminal value is what makes a 3-step horizon sufficient. Scoring a sequence purely on its predicted rewards would need a horizon long enough to reach the consequences of an action — impossible for a learned model, since error compounds. Appending γᴴ Q(z_H) hands the rest of the future to the value function, which is trained by TD learning on real data. Planning covers the short horizon where the model is accurate; the critic covers everything after. That split is the “TD” in the name, and it is why horizon = 3 is a reasonable setting rather than a compromise.

Refitting is weighted, not just filtered. MPPI takes a hard top-MPC_ELITES cut like CEM does, and then — unlike CEM — weights the survivors by exp(score / temperature) rather than averaging them equally. So a clearly better sequence pulls the next mean further than a marginal one, instead of both counting once. That gradient inside the elite set is the whole difference from CEM, and it is what makes the refit harder for a single lucky sample to capture.

The planner itself lives in noeira.planners and knows nothing about this agent — TD-MPC2 supplies a rollout callback and the planner supplies the search.

The advantage over MBPO is that planning can correct model error at decision time, rather than baking it into a policy through training. The cost is that every action requires an optimization, which makes inference much more expensive.

The world model

Trained end to end on:

• Latent dynamics — (z, a) → z′
• Reward prediction
• Value prediction (distributional)
• Policy prior for sampling

Like MuZero, it never reconstructs observations. Only the quantities planning needs are predicted, which is a far easier learning problem than pixel reconstruction.

Latents are normalized with SimNorm — simplicial normalization, projecting the latent onto a set of simplices, which keeps its scale bounded without the collapse that plain normalization can cause.

Reward and value are predicted as discrete distributions over 101 bins in log-transformed space, not as scalars — the same two-hot representation DreamerV3 uses, and for the same reason. A regression head has to commit to a single number and gets pulled around by reward-scale differences between tasks; a classification head over fixed bins does not, which is what lets one set of hyperparameters transfer across environments with wildly different reward magnitudes.

Building it

The TDMPC2[...] preset applies the reference-tuned defaults and takes the architecture and planning budget as compile-time parameters:

from noeira.deep_agents.tdmpc2.config import TDMPC2

var ag = TDMPC2[
    "gpu", OBS, ACT, BATCH, REPLAY_CAPACITY,
    ENC, LATENT, MLP,          # 256, 512, 512
    BINS, SIMNORM_DIM,         # 101, 8
    VMIN, VMAX,                # -10, 10
    HORIZON,                   # 3
    MPC_SAMPLES, MPC_PI_TRAJS, MPC_ELITES, MPC_ITERS,
](ctx=ctx, lr=3e-4, learning_starts=5_000)

The last four are the planning budget, and they are the whole inference cost:

Parameter	Reference	Repo example	Meaning
MPC_SAMPLES	512	256	sequences sampled per iteration
MPC_PI_TRAJS	24	12	of those, drawn from the policy prior
MPC_ELITES	64	32	sequences the refit is weighted over
MPC_ITERS	6	4	refit rounds per action

The example runs at roughly half the reference budget, which is the difference between a feasible first run and a very long one.

Acting is gated by a compile-time USE_MPC switch. With it off, the agent acts as a = π(encode(obs)) — a single forward pass through the policy prior, no planning at all. That is worth knowing for two reasons: it is a legitimate fast mode for evaluation, and it is the cleanest way to measure what planning is actually buying you on your task.

pixi run -e nvidia mojo run -I . examples/half_cheetah/tdmpc2_half_cheetah_gpu.mojo
pixi run -e nvidia mojo run -I . examples/hopper/tdmpc2_hopper_gpu.mojo
pixi run mojo run -I . examples/half_cheetah/tdmpc2_half_cheetah_eval.mojo

Note the separate eval script — because planning is expensive, evaluation is usually a distinct run rather than something folded into training.

Apple Silicon is slower than CPU here

TD-MPC2 is kernel-launch-bound rather than compute-bound: the planner fires many small kernels per action. On CUDA the launch overhead is low enough that the GPU path wins comfortably; on Metal it is not, and the CPU path is faster. This is the one agent in the framework where -e apple is not the right default.

Watch out for

HalfCheetah’s angle penalty was removed because of this algorithm

HalfCheetah once carried ANGLE_PENALTY_WEIGHT = 0.5, added to suppress a “running on its head” local optimum that SAC found.

That penalty actively blocked TD-MPC2 — it produced a Q-pessimism collapse. Later optimizations made it unnecessary for SAC too, so it is now 0.0, matching reference dm_control.

The general lesson: reward shaping tuned for one algorithm can break another. Planning methods are especially sensitive, because a shaped penalty distorts the value landscape the planner is searching.

Planning cost dominates. Simulations × horizon × iterations forward passes per action. Profile before assuming training is the bottleneck.

Normalization layers matter. A bug in this repo’s dynamics LayerNorm cost real time to find — the model trained, the losses descended, and the planner produced nonsense. If planning underperforms while losses look healthy, suspect the world model before the planner.

GPU view shapes. An MPC callback regression here came from a 2D view passed to a 1D kernel. Batched planning tensors are easy to reshape wrongly.

TD-MPC2, MBPO or DreamerV3?

All three learn a model. What they do with it differs completely:

	Model used for	Inference cost
MBPO	generating training data	cheap — a policy forward pass
TD-MPC2	planning at decision time	expensive — an optimization per action
DreamerV3	training a policy in imagination	cheap — a policy forward pass
MuZero	tree search at decision time	expensive — MCTS per action

TD-MPC2 and MuZero both plan; TD-MPC2 uses sampling-based optimization over continuous actions, MuZero uses tree search over discrete ones.

Next

• MuZero — planning with tree search instead of MPPI
• DreamerV3 — learning a policy in imagination
• MBPO — using a model for data instead of planning