TD-MPC and Control-Centric Representations

Picture a quadruped learning to walk in a simulator. The controller must choose an action now, but the consequence that matters may arrive well beyond the next few steps. A learned model can help compare immediate consequences; rolling that imperfect model far into the future can make the comparison unreliable. TD-MPC makes the boundary explicit: plan briefly with the model, then use a learned value estimate for the rest.

One alternative learns a policy or action-value function directly from interaction without an explicit transition model. Another learns dynamics and searches through predicted trajectories. Their relative data and compute costs depend on the task and algorithm; the important distinction here is where prediction ends and value estimation begins.

Longer open-loop rollouts may expose a planner to model error and out-of-distribution states, but error need not rise monotonically on every trajectory or metric. The useful planning horizon therefore has to be measured, not deduced from a general slogan about compounding error.

TD-MPC, short for Temporal Difference Learning for Model Predictive Control (Hansen, Su, and Wang, ICML 2022), combines short-horizon latent-space planning with a terminal action-value estimate. Its reported continuous-control benchmarks show strong sample efficiency against the baselines tested. The model and critic are learned together: training unrolls predicted latent states for reward, consistency, and value losses, while the released implementation computes TD bootstrap targets from encoded observed successor states.

Two commitments define the method, and both are worth stating plainly before any notation appears.

1. The model should be control-centric. TD-MPC learns a decoder-free latent representation for reward prediction, short-horizon dynamics, and value estimation. No observation-reconstruction loss requires the latent to preserve every visual detail. This resembles the decision-relevant representation discussed in MuZero and Value-Equivalent Models, now with continuous actions and sampling-based rather than tree search.
2. The model should be trusted for a measured horizon. The reported TD-MPC experiments use a maximum planning horizon of five latent steps across tasks, scheduled upward from one early in training. A bootstrapped action-value estimate supplies the tail, analogous to the terminal value in an $n$-step return.

These choices shape both the method and its limitations. A latent without a decoder need not preserve enough detail to render an observation. A short plan relies more heavily on the critic, while still depending on the model's predictions along its short rollout. Both dependencies reappear in the training objective and failure modes below.

If you have read The Dreamer Family, you already know a world model that trains its latent by reconstruction and then trains an actor inside imagined rollouts. Sampling-Based Planning and Model Predictive Control introduced the MPPI and cross-entropy shooting lineage. TD-MPC combines sampling-based planning, value bootstrapping, and a decoder-free representation. MPPI searches over actions without differentiating through a learned rollout; temporal-difference learning estimates returns beyond the short plan; the representation is trained for control-related predictions rather than pixel fidelity.

First we build the model and separate its weighted multistep loss from policy optimization. Then we assemble the latent-space planner and work through an exact terminal-value example. Finally we examine TD-MPC2's 80-task generalist, run a small executable diagnostic that may fail to learn under its tiny budget, and trace the remaining failure modes.

Throughout, keep three things separate in your head: what the model predicts, how the planner uses those predictions, and how the value function summarizes what lies beyond. Most confusion about this family of methods comes from blurring those three roles. A model that predicts poorly may still support good decisions if the planner only relies on it briefly and the value function covers the rest. A value function that looks accurate in isolation may still be exploited by a planner that searches the edges of its support. Keeping prediction, planning, and value estimation conceptually distinct is what lets you diagnose which one is failing when the agent underperforms.

A TD-MPC agent has five learned components that share a latent representation. The encoder, dynamics, reward, and value components are optimized with a weighted multistep objective; the policy prior has a separate objective and optimizer. The representation is shaped by control-related losses rather than a pixel-reconstruction loss.

To make this precise, let $o_t$ denote the observation (the raw state vector in a proprioceptive task, an image in a visual one), $a_t \in \mathbb{R}^{m}$ the action, and $r_t$ the scalar reward. The agent maintains a latent state $z_t \in \mathbb{R}^{d}$ and five functions of it:

• Encoder $h_\theta$: maps an observation to a latent state, $z_t = h_\theta(o_t)$.
• Latent dynamics $d_\theta$: predicts the next latent from the current latent and action, $\hat{z}_{t+1} = d_\theta(z_t, a_t)$.
• Reward model $R_\theta$: predicts the immediate reward, $\hat{r}_t = R_\theta(z_t, a_t)$.
• Value function $Q_\theta$: estimates the discounted return from a latent state and an action.
• Policy prior $\pi_\theta$: maps a latent state to an action and supplies noisy proposal trajectories to the planner.

Notice that the five functions divide naturally into two groups. The encoder, dynamics, and reward model describe the world: they say where you are, where you go next, and what you earn along the way. The value function and policy prior make decisions: one estimates how good a state-action pair is, and the other proposes actions worth trying. Joint training is what binds the two groups, because the value loss reaches back through the encoder and forces the world-modeling half to preserve whatever information the decision-making half needs.

The original method has deterministic latent dynamics rather than a belief distribution or stochastic transition model. Sampling candidate actions during planning does not make those latent transitions probabilistic. This distinction matters when the environment has stochastic outcomes.

Training samples sequences of observed transitions from replay, encodes the first observation, and unrolls the latent dynamics along recorded actions. Write $\hat z_0=h_\theta(o_t)$ and $\hat z_{k+1}=d_\theta(\hat z_k,a_{t+k})$. For an unroll of $K$ transitions, the original paper's model objective has three coefficient-weighted losses at each step, discounted in the loss by $\lambda^k$:

$$\mathcal{J}_{\mathrm{model}}= \mathbb{E}_{\mathcal D}\sum_{k=0}^{K-1}\lambda^k\left[ c_r\big(R_\theta(\hat z_k,a_{t+k})-r_{t+k}\big)^2 +c_q\big(Q_\theta(\hat z_k,a_{t+k})-y_{t+k}\big)^2 +c_z\left\|\hat z_{k+1}-\operatorname{sg}\!\left[h_{\bar\theta}(o_{t+k+1})\right]\right\|_2^2 \right].$$

Here $c_r,c_q,c_z$ balance reward prediction, TD value fitting, and latent consistency; they are not equal by definition. $\bar\theta$ denotes slow-moving target parameters. The consistency target encodes the observed next state, while its prediction is reached by rolling forward from the first encoded state. A collapsed constant latent can make consistency easy, but reward and value prediction generally require distinctions between states. Their losses supply task-relevant pressure; none individually guarantees a useful representation.

For example, a constant encoder and matching dynamics achieve zero consistency error, yet cannot distinguish two observations whose rewards or returns differ under the same action. In tasks where those targets vary, minimizing the other losses disfavors that collapse. This is a task-dependent argument, not a proof of convergence or a guarantee against every degenerate local minimum.

The released TD-MPC implementation computes the one-step TD target using the observed successor encoded by the online encoder. With its twin target critics and a noisy action sampled from the current policy, the target is

$$z_{t+1}^{\mathrm{obs}}=h_\theta(o_{t+1}),\qquad \tilde a_{t+1}\sim\pi_\theta(\cdot\mid z_{t+1}^{\mathrm{obs}}),\qquad y_t=r_t+\gamma\min_{j\in\{1,2\}}Q_{\bar\theta,j}(z_{t+1}^{\mathrm{obs}},\tilde a_{t+1}).$$

The current policy samples around its mean with truncated noise for this target; the paper's single-$Q$ notation abstracts the released twin critics. The value loss evaluates both $Q_{\theta,j}(\hat z_k,a_{t+k})$ on predicted latent states during the multistep unroll. The bootstrap target above instead uses the next real observation. The original paper writes its TD target with $z_{k+1}$ in the recurrent loss; reading that notation as a claim that the released code bootstraps from the predicted successor confuses the two paths. A target network slows target drift; it does not guarantee convergence of nonlinear off-policy learning.

The bootstrap action comes from the learned policy rather than an explicit maximization over a continuous action space. This is a tractable actor-critic approximation, and its errors can affect the critic. It does not imply a strict upper bound on critic quality or that an action omitted by the prior is never represented in replay.

The released implementation optimizes the policy separately on the initial and $K$ predicted recurrent latents. It detaches these $K+1$ latents, freezes both critics' parameters for the actor step, and maximizes the smaller critic estimate:

$$\mathcal{J}_{\pi}=-\mathbb{E}_{\mathcal D}\sum_{k=0}^{K}\lambda^k \min_{j\in\{1,2\}}Q_{\theta,j}\big(\operatorname{sg}[\hat z_k],\pi_\theta(\operatorname{sg}[\hat z_k])\big).$$

This writes the released update at the level of its mean objective and omits the action sampling noise. The detached latents prevent this actor step from updating the encoder or dynamics, while the frozen critics still pass gradients through the policy's action. Neither an entropy bonus nor an action-magnitude penalty appears in the original objective. The small code example below adds an explicit action-magnitude penalty as a toy stabilizer and labels it accordingly.

The toy code below includes an episodic termination mask, $(1-\mathrm{terminated}_t)$, but not a time-limit truncation mask. That is a sensible convention for its finite episodes, not a term shown in the original paper's displayed TD equation. Whether a timeout should bootstrap depends on the environment's observation and episode semantics; confusing timeout with a true terminal state can bias fitted values.

Dropping the decoder is not an oversight. Three arguments motivate it, and they are worth stating because they apply well beyond this one algorithm.

Sufficiency beats fidelity. If the robot's arm color carries no information about the task, dynamics, or reward, a control policy need not preserve it. It does need quantities that change the return: joint angles, velocities, contacts, and distances. Color could matter, however, if it identifies a task-relevant object or instruction. Reconstruction asks the latent to retain visual detail even when that detail is irrelevant to control. A decoder-free objective does not impose that demand. This is the same argument that motivates the predictive-embedding family covered in Predictive Embeddings, JEPA, and Energy-Based Models.

The gradient budget goes to task-related predictions. The encoder is trained through reward, latent consistency, and value losses rather than a pixel decoder. This avoids allocating a direct reconstruction loss to visual detail irrelevant to the reward. It does not prove that every retained feature matters for control, nor that reconstruction would dominate every alternative model.

It changes what you monitor. Low reconstruction error alone cannot certify good decisions. Removing the decoder also removes that potentially misleading metric, but a falling value loss is no guarantee of good control either: the critic can fit biased targets or be exploited by the planner. Closed-loop evaluation remains necessary.

Reconstruction can nevertheless provide a useful auxiliary signal when reward and value targets are weak or sparse. The original paper evaluates decoder-free learning on continuous-control tasks, including sparse-reward cases; it does not establish a universal dense-versus-sparse boundary. Whether observation reconstruction helps is an empirical question for the task and data regime.

The original TD-MPC model and TD-MPC2 should not be conflated here. The following normalization choices are associated with TD-MPC2, not the original paper's encoder and dynamics.

Layer normalization is used throughout TD-MPC2's MLP components. The released original TD-MPC implementation does not place it in the dynamics and reward networks as previously claimed. Normalization can improve optimization stability, but its exact effect cannot be inferred from one layer in isolation.

Simplicial normalization (SimNorm), introduced in TD-MPC2, splits the latent vector into fixed-size groups and applies a softmax within each group. Each coordinate lies in $[0,1]$ and each group sums to one. This bounds the coordinates of a normalized latent; it does not guarantee accurate or stable multistep predictions. Our toy below borrows SimNorm from TD-MPC2, so it is an illustrative hybrid rather than a faithful reproduction of the 2022 architecture.

The coordinate bound can be checked on a single deterministic vector. This demonstrates the operation of SimNorm, not a claim about rollout errors.

These normalization choices affect representation scale in TD-MPC2. The planner does not score candidates by latent distance: it sums predicted rewards and a terminal value. Scale can still affect optimization, but a fixed latent distance need not have a fixed semantic meaning during training.

Finally, note what is not in the original model: no observation decoder, stochastic latent transition, or dynamics ensemble. Those are architectural choices, not proofs that uncertainty is unimportant.

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.

When can a short plan plus a terminal value outperform a longer model rollout? A good terminal value can replace many predicted steps, but its own errors then matter more. The choice is a measured tradeoff rather than a theorem that three steps beats thirty.

Over $H$ steps, an imperfect dynamics model feeds its predicted state into the next prediction. Errors may accumulate, cancel, or vary with the visited states; neither mean-squared prediction error nor return variance is guaranteed to increase at every horizon. Longer rollouts can also search more reward opportunities. Their net effect has to be evaluated on held-out trajectories and in closed-loop control.

If the critic's error at the reached terminal state-action pair is $\epsilon_Q(\hat z_H,a_H)$, its direct contribution to the planning score is $\gamma^H\epsilon_Q(\hat z_H,a_H)$. The discount attenuates a fixed error, but the terminal state itself changes with $H$ and may enter a region where critic error is larger. One critic evaluation is computationally cheap; that does not make its error constant in horizon.

The direct critic term shrinks with $\gamma^H$ only under the fixed-error comparison; model error may grow with $H$. Those tendencies can create an interior optimum, but they can also produce a boundary optimum or a nonmonotone curve. If the critic is weak, planning longer helps only if the model remains accurate enough and the added search cost is worthwhile. Horizon selection belongs to validation, not a rule of thumb based on critic quality alone.

The next plot deliberately assumes one increasing model-error component and one fixed critic-error magnitude discounted by $\gamma^H$. Under these chosen functions, their sum has an interior minimum. It is a hypothetical error budget, not a measured TD-MPC curve or a general upper bound on planning error.

There is also a classical control-theoretic reading. With exact dynamics and the exact optimal cost-to-go as terminal cost, optimizing the first action in a one-step plan recovers the optimal action by the Bellman equation. This requires the optimal terminal state value $V^*(s)=\max_a Q^*(s,a)$, equivalently exact $Q^*$ evaluated at an optimal terminal action; exact dynamics and $Q^*$ with a suboptimal terminal policy do not suffice. Stability claims for approximate MPC require additional assumptions; an arbitrary terminal cost that upper-bounds cost-to-go is not by itself a sufficient blanket condition. The worked example below isolates the oracle-value case before the neural approximation complicates it.

Make the argument concrete with a problem we can solve by hand. Take a one-dimensional linear system with quadratic cost,

$$s_{t+1} = A s_t + B a_t, \qquad A = 1.0,\; B = 0.5, \qquad r(s, a) = -(s^2 + a^2),$$

with discount factor $\gamma = 0.9$. For this linear-quadratic problem the optimal value function is exactly quadratic, $V(s) = -P s^2$, where $P$ solves the discounted Riccati equation

$$P = 1 + \gamma A^2 P - \frac{\gamma^2 A^2 B^2 P^2}{1 + \gamma B^2 P}.$$

The Riccati equation is the fixed point of the Bellman equation for this problem: substituting a quadratic value function into the Bellman optimality condition and matching coefficients yields exactly this algebraic equation for $P$. Substituting the numbers gives $P = 1 + 0.9P - \frac{0.2025 P^2}{1 + 0.225P}$. Multiplying both sides by $(1 + 0.225P)$ and collecting terms yields

$$0.225 P^2 - 0.125 P - 1 = 0 \quad\Longrightarrow\quad P = \frac{0.125 + \sqrt{0.125^2 + 0.9}}{0.45} \approx 2.404 .$$

We take the positive root because $P$ must be positive for the value function to be a cost (negative reward) that increases in magnitude with the state. The optimal controller is linear, $a^*(s) = -0.702\,s$, and the optimal return from $s = 1$ is $-P = -2.404$.

Now run a one-step planner from $s = 1$ with $H = 1$. Here we replace TD-MPC's learned-policy, twin-critic terminal score with the oracle $V^*(s)=Q^*(s,a^*(s))$; this is the exact optimal continuation, not a claim that any terminal policy action would work:

$$J(a) = -(1 + a^2) + \gamma V(A \cdot 1 + B a) = -1 - a^2 - 0.9 \cdot 2.404\,(1 + 0.5a)^2.$$

Expanding the square and collecting terms gives $J(a) = -3.1636 - 2.1636\,a - 1.5409\,a^2$. Setting the derivative to zero yields $a = -0.702$, and substituting back gives $J = -2.404$. A one-step plan with the exact terminal value reproduces the infinite-horizon optimum exactly, both the action and the value.

Remove the terminal value and the one-step planner maximizes $-1-a^2$ at $a=0$: do nothing now. If it follows that single myopic action and then switches to the optimal controller, its total return is $-1+\gamma V(1)\approx-3.164$, a cost magnitude about 31.6% above the optimum $2.404$. If it repeatedly chooses the myopic action, the state stays at $s=1$ and the return is $-\sum_{t=0}^\infty 0.9^t=-10$. These are different continuation assumptions; the 31.6% number belongs only to the first.

The same argument becomes a picture if we plot the planning objective over actions. The terminal value moves the maximum away from the tempting "do nothing" action and onto the optimum. All quantities are computed from the closed-form worked example, so the curves are analytic and deterministic.

Three lessons fall out of this exercise, and all three survive into the neural setting:

• With exact dynamics and the exact optimal terminal value (equivalently, exact $Q^*$ at an optimal terminal action), one-step maximization gives the optimal action. A suboptimal terminal action, approximate dynamics, or approximate values can break that equivalence.
• Myopic planning is not merely suboptimal in magnitude; it systematically prefers actions that look good now because it cannot see the consequence. The wrongness is structured, not random.
• The terminal value can change the maximizing action. For an approximate critic, its score error enters with a factor $\gamma^H$ and can alter action ranking; the size of that effect depends on the competing score gaps.

The exact example explains why terminal estimates can matter, not how accurately a learned critic will perform. A short horizon shifts responsibility to the critic; whether that helps depends on the quality of both components near the candidate trajectories.

The horizon is a hyperparameter with a real tradeoff. Use open-loop prediction checks to identify model weaknesses, then test candidate horizons in closed-loop control. A terminal-value ablation answers a different question.

Open-loop error by horizon. Roll the learned model forward with recorded actions from held-out trajectories and compare predictions with observed rewards and encoded successor states. Flat average error alone does not certify that a longer planner is safe: the planner may select different actions or exploit rare optimistic errors.

Controlled horizon sweep. With trained weights, terminal-value setting, initial states, seeds, action-repeat setting, and evaluation budget fixed, compare closed-loop returns across several horizons. The best horizon may differ from the one with lowest recorded-action prediction error because the planner selects its own actions. Include inference latency when the controller has a deadline.

Terminal-value ablation. At a fixed horizon, remove the terminal value while holding weights and evaluation conditions fixed. This isolates its contribution for the trained agent; it does not select a horizon. A small gap can mean the value estimate is unnecessary or that it is not yet useful.

Open-loop prediction quality, horizon-dependent closed-loop return, and terminal-value contribution are distinct measurements. The toy below prints the first and two small closed-loop comparisons, but its three evaluation starts are far too few for a reliable horizon choice.

One more practical consideration: horizon and action repeat interact. If the environment repeats each action for several simulator steps, a horizon of $H$ covers $H$ times as much real time, and $H = 3$ can correspond to a substantial fraction of a second. Report horizons in units of real time when comparing across papers. A horizon of three in a setup with action repeat of four is really a twelve-step horizon in simulator time, and comparing it to a paper with no action repeat would be misleading.

Most original TD-MPC benchmarks train separate agents per task, but the 2022 paper also reports a single policy trained jointly on ten Meta-World tasks (MT10). TD-MPC2 (Hansen et al., ICLR 2024) retains decoder-free latent planning with a terminal value while changing several architectural and optimization details. Its single-task online evaluations and its large offline multitask training are separate experiments; they should not be reported as one result.

The paper reports online single-task evaluations over 104 continuous-control tasks across DMControl, Meta-World, MyoSuite, and ManiSkill2 using a common hyperparameter recipe. Separately, it trains offline multitask agents on 545 million transitions from 80 tasks (50 Meta-World and 30 DMControl); the largest has 317 million parameters. On this 80-task dataset, reported normalized score increases across the tested sizes from 1 million to 317 million parameters. That observed scaling trend does not guarantee smooth gains beyond the measured models, tasks, or data regime.

Three implementation changes accompany the scale-up.

Discrete regression for reward and value. TD-MPC2 represents scalar targets on a fixed 101-bin support in a signed-log-transformed space, with soft two-hot target weights and cross-entropy loss. The transform compresses large magnitudes before binning; a categorical head is not intrinsically immune to outliers or support mismatch. The paper motivates this change with reward-scale variation across tasks and reports ablations, rather than proving universal stability.

Task conditioning. The multitask agent learns a bounded-norm embedding for each task. That embedding conditions all five components: encoder, dynamics, reward, value, and policy. Padding and action masks reconcile differing observation and action dimensions. Sharing parameters can transfer useful structure, but interference remains possible; the paper's results do not isolate a universal best embedding placement.

Visual observations. In its image-based experiments TD-MPC2 replaces the state MLP encoder with a shallow four-layer convolutional encoder on 64-by-64 images and applies random-shift augmentation. It does not use a frozen DINOv2 backbone in those reported experiments. Pretrained visual features are a different research lineage, covered later in Pretrained Visual Models and DINO-WM.

The multi-task setting is not just "more data." TD-MPC2 shares model parameters across tasks while conditioning the encoder, dynamics, reward, value, and policy on a learned task embedding. This lets the model use similarities across tasks without requiring a task-agnostic latent or a clean decomposition into shared dynamics and task-specific rewards. The reported results do not establish that such a factorization emerged.

That has three consequences worth carrying forward.

• The value function becomes task-conditional. Across tasks, reward and value targets can differ substantially in magnitude. This helps explain why TD-MPC2 uses signed-log categorical regression rather than treating the change as merely cosmetic.
• Representation sharing is conditional. The encoder shares parameters but also receives the task embedding, so two tasks need not map similar observations to identical latents. The paper's multitask and few-shot results show that this design can reuse structure in its tested settings; they do not prove a task-agnostic representation or rule out interference on different task mixes.
• Reward scale becomes a first-class engineering problem. Reward and value targets from different tasks can have different magnitudes; TD-MPC2 uses a signed-log transform and discrete targets to improve robustness. Each plan is scored within one task, so a larger reward scale on another task does not by itself make that other task's candidates enter the current elite set.

The decoder-free latent model, short-horizon planning, terminal value, and off-policy replay remain. The critic's distributional head and default ensemble size (five $Q$ heads rather than the original two), policy distribution, normalization, and task conditioning change substantially. TD-MPC2 is a continuation of the design, not the identical 2022 algorithm with more parameters.

The clearest invariant is the division of labor: learn task-relevant latent predictions, search a few steps, and use a critic for the tail. Training still unrolls learned latent dynamics along replay sequences, and neither the single-task nor multitask experiments establish that every retained component is necessary in every domain.

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 $r_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.