TD-MPC and Control-Centric Representations › TD-MPC2 and Multi-Task Control

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.