TD-MPC and Control-Centric Representations › Temporal-Difference Model Predictive Control

We now have a model that predicts rewards and latents, and a value function that estimates returns. The remaining question is how to convert them into an action at every control step. TD-MPC uses a receding-horizon sampling planner, run entirely inside latent space, with the value function supplying the tail of the objective.

Recall from Sampling-Based Planning and Model Predictive Control that the MPPI-style planner samples action sequences, evaluates them under the model, and refits a sampling distribution toward high-scoring candidates. It does not differentiate through the planned rollout. That avoids relying on rollout gradients, but sampling can still exploit errors in the learned model or critic.

Let $z_t$ be the current latent, and let $\mathbf{a} = (a_t, \dots, a_{t+H-1})$ be a candidate action sequence over a horizon of $H$ steps. Roll the latent dynamics forward deterministically, $\hat{z}_{t+k+1} = d_\theta(\hat{z}_{t+k}, a_{t+k})$. The planner scores each candidate by

$$J(\mathbf{a}) \;=\; \sum_{k=0}^{H-1} \gamma^{k} R_\theta\!\left(\hat{z}_{t+k}, a_{t+k}\right) \;+\; \gamma^{H} \widehat V_\theta\!\left(\hat{z}_{t+H}\right).$$

The first term estimates rewards along a candidate model rollout; the second evaluates the tail after that rollout. The paper's Equation (3) writes the terminal term as $Q_\theta(z_H,a_H)$, leaving the terminal action abstract in that notation. The released 2022 implementation instantiates it as

$$\widehat V_\theta(z) = \min_{j\in\{1,2\}} Q_{\theta,j}(z,\tilde a), \qquad \tilde a \sim \pi_\theta(\cdot\mid z;\sigma_{\min}).$$

The terminal action is sampled from the learned policy using the configured minimum standard deviation because the candidate sequence ends at step $H-1$. The two online critics are combined by a minimum; this choice alone is not a guarantee against optimistic scoring. The planner needs only a numeric score, while the policy's separate training objective uses gradients through its proposed action. The score resembles an $n$-step bootstrapped return, but planning optimizes actions using model predictions, whereas TD learning fits a value function to targets. They are related operations, not identical algorithms.

Two error sources enter the score. Errors in reward and latent predictions can change the scores of candidate sequences; for a fixed sampled terminal action, error in the twin-critic terminal estimate contributes $\gamma^H\epsilon_Q(\hat z_{t+H},\tilde a)$, where $\epsilon_Q$ is the estimation error of $\min_j Q_{\theta,j}$ at that terminal state and action. Policy sampling adds another source of score variation. These effects depend on which states the optimizer visits. A shorter horizon reduces exposure to model rollouts but puts more weight on the critic. No fixed ranking of the errors follows without measuring the model, critic, and candidate distribution.

MPPI maintains a time-varying Gaussian $(\mu \in \mathbb{R}^{H \times m}, \sigma \in \mathbb{R}^{H \times m}_{+})$ over action sequences and improves it over a few iterations. Each iteration does four things.

1. Sample. Draw $N$ action sequences from $\mathcal{N}(\mu, \mathrm{diag}(\sigma^2))$ and clip each action to the environment's bounds. Add policy-guided trajectories; the original implementation injects noise into policy actions so these proposals are not duplicates.
2. Score. Roll every candidate through $d_\theta$ and $R_\theta$, accumulate the discounted predicted reward, and add the terminal value.
3. Fit. Keep the top $K$ candidates (the elites) and refit:

$$w_i = \frac{\exp(\kappa\, (J_i-J_{\max}))}{\sum_{j} \exp(\kappa\, (J_j-J_{\max}))}, \qquad \mu' = \sum_i w_i\, \mathbf{a}_i, \qquad \sigma'^2 = \sum_i w_i\,(\mathbf{a}_i - \mu')^2 ,$$

where $J_{\max}$ is the largest elite score and $\kappa$ controls weight sharpness. Subtracting the maximum is a numerically stable softmax shift; it does not normalize score variance or make the weighting invariant to reward scale. The original planner uses this max shift, not the z-score normalization illustrated separately below. 4. Repeat for a fixed number of iterations.

In the released TD-MPC planner, the first action comes from a trajectory sampled among the final elites according to their weights. During training, exploration noise is then added using the final elite spread; evaluation omits that extra noise. The code does not clamp the action again after adding this final noise. This is distinct from the toy planner below, which samples from its refitted Gaussian and clips the result. Neither action-selection rule guarantees useful coverage when the model and critic are poor.

The next illustration asks a counterfactual question: what would z-score normalization do to weights under a reward-scale change? It is useful for understanding softmax sensitivity, but it is not a diagram of the original TD-MPC planner. The weights use deterministic return vectors, with no sampling.

This is the point where TD-MPC most sharply diverges from Dreamer, so it is worth being precise.

The model is used in training and at decision time. Training samples real trajectories from replay, then repeatedly applies learned latent dynamics along their recorded actions. Reward and value predictions on these predicted latents receive gradients through the unroll. At decision time, the planner searches over new candidate action sequences in the same learned model.

The TD bootstrap in released code uses observed successors. This is narrower than saying model predictions never enter learning: the value prediction is fitted on model-unrolled latent states, and its gradient shapes the model. Unlike Dreamer's policy-learning rollouts, these unrolls follow replayed actions rather than a newly imagined policy trajectory.

This distinction matters for diagnosing failure. A model error can affect both planning scores and value predictions during recurrent training, even when the TD target begins from an encoded observed successor. Off-policy replay permits reuse of past data but does not make arbitrary stale or poorly covered trajectories harmless.

The planner and critic remain coupled through the terminal score and policy proposals. If the critic is optimistic for candidate actions or terminal states outside replay coverage, optimization can exploit that error. The method limits rollout length but offers no general divergence guarantee.

The complete online loop is short enough to state in full.

1. Reset the environment, encode the observation to $z_0$.
2. At each step, run MPPI on the current latent with the model and value function, take the sampled action, and step the environment.
3. Store $(o_t, a_t, r_t, o_{t+1}, \text{terminated}_t)$ in the replay buffer.
4. Sample replay sequences, optimize the weighted model/critic losses and the separate policy objective, then update target parameters by exponential moving average.
5. Repeat for a fixed number of episodes or environment steps, using a short random-action warm-up at the start to seed the buffer.

Warm-up supplies initial replay data before relying on an untrained planner. The released loop normally runs one optimization update per collected decision transition, not per underlying simulator frame when actions repeat; our short toy groups updates differently. Neither setting establishes that the toy has learned a useful controller.

The original implementation plans at each control step, uses a short horizon and sampled action sequences, and normally takes one training update per collected decision transition. The horizon is scheduled upward during early training. TD-MPC2 uses a short horizon too, but differences in architecture, task mix, and evaluation prevent attributing its gains solely to a stronger critic or a shorter plan.