TD-MPC2: Scalable, Robust World Models for Continuous Control

Motivation

• TD-MPC2 is designed for multitask RL on large uncurated datasets spanning domains, embodiments, action spaces, and behavior-policy skill levels.It avoids requiring hyperparameter tuning across tasks.Source: §1 · p. 2

Approach

• TD-MPC2 performs local trajectory optimization in the latent space of a learned implicit, decoder-free world model.The method builds upon TD-MPC while revisiting core design choices for improved robustness and scalability.Source: §1 · p. 2

Evaluation

• 104 diverse continuous-control tasks span DMControl, Meta-World, ManiSkill2, and MyoSuite, including high-dimensional actions, images, sparse rewards, manipulation, locomotion, and musculoskeletal control.The evaluated tasks include action spaces up to A ∈R39.Source: §1 · p. 2

Results

• TD-MPC2 consistently outperforms existing model-based and model-free methods using the same hyperparameters across all tasks.The results include especially challenging locomotion and multi-object manipulation subsets.Source: §1 · p. 2

Scaling

• A single 317M-parameter world model performs 80 tasks across multiple task domains, embodiments, and action spaces.The paper also reports that agent capabilities increase with model and data size.Source: §1 · p. 2

Reinforcement Learning

• Continuous-control RL seeks a policy that maximizes the expected discounted sum of rewards in an infinite-horizon MDP.The MDP is represented by states, actions, transitions, rewards, and a discount factor.Source: §2 · pp. 2–3

Model-Based RL

• TD-MPC2 obtains a control policy by learning a world model and selecting actions through planning with that model.The world model represents the environment for model-based action selection.Source: §2 · p. 3

Model Predictive Control

• Model Predictive Control optimizes finite-length action sequences over a horizon H to maximize return or minimize cost.The resulting policy is obtained from a finite-horizon optimization problem.Source: §2 · p. 3

Planning Limitation

• Candidate trajectory returns are estimated by simulating them with the learned model.This makes the resulting policy temporally locally optimal rather than guaranteed to solve the general RL problem.Source: §2 · p. 3

TD-MPC2 Extension

• TD-MPC2 addresses finite-horizon local optimization by bootstrapping return estimates beyond H with a learned terminal value function.This extends planning information beyond the explicit trajectory horizon.Source: §2 · p. 3

Overview

• TD-MPC2 builds on TD-MPC with improvements that enable the same hyperparameters across tasks and scale the world model to 300× more parameters.The paper presents TD-MPC2 as a practical algorithm for massively multitask world models.Source: §3 · p. 3

Implicit world model

• The model learns a decoder-free, control-centric world model by combining joint-embedding prediction, reward prediction, and TD-learning.It predicts task-relevant outcomes rather than reconstructing future observations, then supports local trajectory optimization in latent space.Source: §3.1 · p. 3
• The architecture encodes observations into normalized latents and recurrently predicts actions, rewards, and terminal values without decoding future observations.Its five components are an encoder, latent dynamics model, reward predictor, terminal-value predictor, and policy prior, with learnable task embeddings for multitask models.Source: Figure 3 · p. 3, §3.1 · p. 3

Learning objectives

• TD-MPC2 jointly optimizes latent representations, dynamics, reward prediction, and value prediction, while discrete log-space reward and value regression accommodates differing reward magnitudes.The method uses cross-entropy objectives and an EMA ensemble of Q-functions to reduce bias in TD targets.Source: §3.1 · p. 4

Model predictive control

• MPPI plans over sampled latent action trajectories, bootstrapping returns beyond horizon H with a learned terminal value and executing only the first action.The planner estimates Gaussian action-sequence parameters, can use the policy prior for some samples, and warm-starts from the previous decision step.Source: §3.2 · pp. 4–5

Evaluation and multitask scaling

• Across 104 tasks in four domains, TD-MPC2 achieves higher data-efficiency and asymptotic performance than existing methods using one hyperparameter set.The multitask design addresses task semantics and heterogeneous observation and action spaces through learnable task embeddings and action masking.Source: Figure 4 · p. 5, §3.3 · p. 5

• 104 tasks span four domains, including high-dimensional actions up to A ∈R39, image observations, sparse rewards, manipulation, musculoskeletal control, and complex locomotion.Source: §4 · p. 5
• TD-MPC2 outperforms SAC, DreamerV3, and TD-MPC across all task domains, using the same hyperparameters across tasks.It achieves especially large margins on high-dimensional locomotion and multi-object manipulation.Source: §4.1 · p. 6, §4 · p. 6
• TD-MPC2 remains stable where TD-MPC can diverge from exploding gradients, while DreamerV3 shows numerical instabilities on Dog and struggles with fine-grained manipulation.Source: §4.1 · pp. 6–7

Scaling

• TD-MPC2 capabilities consistently increase with model size on 80-task and 30-task datasets, without apparent saturation at 317M parameters.The normalized score appears to scale linearly with the log of model parameters, although the authors refrain from formulating a scaling law.Source: §4.1 · p. 7, Figure 7 · p. 7

Few-shot learning

• A 19M-parameter agent finetuned on 70 tasks improves 2× over learning from scratch on 10 held-out tasks in the 20k-step low-data regime.Finetuning uses an initially empty replay buffer and unchanged hyperparameters.Source: §4.1 · pp. 7–8

Ablations

• Ablations show that the proposed design choices contribute meaningfully to robustness and performance in both single-task and multitask RL.Normalized task embeddings are moderately important for stable multitask training and more meaningful task relations.Source: §4.1 · p. 8

Visual RL

• With a shallow convolutional encoder, TD-MPC2 performs comparably to DrQ-v2 and DreamerV3 on 10 image-based DMControl tasks without hyperparameter changes.Source: Figure 10 · p. 8, §4.1 · p. 9

Lessons

• TD-MPC2 addresses RL sensitivity to architecture, hyperparameters, task characteristics, and random seed, which makes successful deep-RL application resource-intensive.Source: §5 · p. 9

Opportunities

• Massively multitask world models could support zero-shot skills on seen embodiments, rapid finetuning to new embodiments, and higher-level tasks combined with vision-language models.The paper presents these as future directions rather than demonstrated capabilities.Source: §5 · p. 9
• Generalist behavior will likely require several orders of magnitude more tasks than are currently available.Source: §5 · p. 9

Risks

• Misspecified rewards, unconstrained physical-robot autonomy, and expensive data requirements pose risks of unintended outcomes, catastrophic failures, and concentrated power.The authors call for new research innovations to mitigate these challenges.Source: §5 · p. 9

• Prior robustness methods target hyperparameters, architecture, training stability, exploration, or representation quality through techniques such as Q-learning variants and data augmentation.Source: §6 · p. 9
• The cited robustness methods focus strictly on data-efficient and robust single-task online RL.Source: §6 · p. 10
• Scaling literature for decision-making commonly assumes large datasets of near-expert demonstrations for behavior cloning.Source: §6 · p. 10
• Other RL-scaling approaches include behavior-cloning pretraining followed by RL, or population-based training without pre-existing data.Source: §6 · p. 10

Architectural design

• TD-MPC2 replaces the original architecture with LayerNorm and Mish MLPs, SimNorm latent normalization, five Q-functions, dropout, and randomized double-Q targets.Source: A SUMMARY OF IMPROVEMENTS · p. 16

Task coverage

• Task visualizations cover 104 tasks whose objectives, embodiments, and action spaces vary substantially.Source: Figure 11 · p. 17

Supplementary analyses

• Additional materials provide task details and videos, while the paper directs readers to appendices for further task curves and ablations.Source: Figure 11 · p. 17

Domains

• 104 tasks span four domains: DMControl, Meta-World, ManiSkill2, and MyoSuite.The domains include 39, 50, 5, and 10 tasks, respectively.Source: C TASK DOMAINS · p. 18, Table 2 · p. 18, Table 3 · p. 19

Task diversity

• The tasks include action spaces up to A ∈R39, image observations, sparse rewards, multi-object manipulation, complex locomotion, and physiologically accurate musculoskeletal control.Source: C TASK DOMAINS · p. 18
• MyoSuite includes high-dimensional dexterous-hand motor-control tasks with either fixed goals or randomized goals labeled Easy and Hard.Source: Table 5 · p. 19

Evaluation

• All four domains are infinite-horizon continuous-control environments evaluated with fixed episode lengths, action repeats, environment-step budgets, and domain-specific performance metrics.Source: C TASK DOMAINS · p. 20, Table 6 · p. 20
• An episode counts as successful only when the final step is successful, so picking up and later dropping an object does not count.Source: C TASK DOMAINS · p. 20

DMControl

• TD-MPC2 remains stable on Walker Stand and Walker Walk, whereas TD-MPC diverges on these tasks.Source: Figure 12 · p. 21

Meta-World

• TD-MPC2 is comparable to existing methods on easy Meta-World tasks and outperforms them on hard tasks such as Pick Place Wall and Shelf Place.Source: Figure 13 · p. 22

ManiSkill2

• TD-MPC2 exceeds 60% success on ManiSkill2 Pick YCB, while other methods fail to learn within the given budget.Pick YCB involves manipulation of all 74 YCB objects.Source: Figure 14 · p. 23

Locomotion

• The high-dimensional locomotion benchmark includes Humanoid and Dog embodiments with action dimensions A ∈R21 and A ∈R38.Source: Figure 15 · p. 23

MyoSuite

• TD-MPC2 achieves comparable or better performance than existing methods on all MyoSuite tasks except Key Turn Hard, where TD-MPC succeeds early in training.Source: Figure 16 · p. 23

Setup

• A 19M-parameter TD-MPC2 agent pretrained on 70 tasks is finetuned to each of 10 held-out tasks.Source: E FEW-SHOT EXPERIMENTAL RESULTS · p. 24

Results

• Few-shot finetuning efficacy is somewhat task-dependent compared with training similar-capacity agents from scratch.The 19M-parameter from-scratch baselines perform comparably to 5M-parameter from-scratch agents.Source: E FEW-SHOT EXPERIMENTAL RESULTS · p. 24

Interpretation

• The authors conjecture that task similarity and pretrained source-task performance both influence finetuning results, requiring further research to distinguish them.Source: E FEW-SHOT EXPERIMENTAL RESULTS · p. 24

Initialization

• For unseen tasks, the new task embedding is initialized from a semantically similar pretraining task, with no experiments on other initialization schemes or pairings.Source: E FEW-SHOT EXPERIMENTAL RESULTS · p. 24, Table 7 · p. 24

Task embeddings

• Normalizing task embeddings to a maximum ℓ2-norm of 1 improves multitask performance.Source: Figure 18 · p. 25
• Normalized task embeddings are more semantically meaningful, placing related tasks such as Door Open and Door Close closer together.Source: Figure 19 · p. 25

Activation function

• Mish and the alternative activation function achieve comparable asymptotic performance and data efficiency on three difficult tasks.The tasks are Dog Run, Humanoid Walk, and Pick YCB.Source: Figure 20 · p. 25
• Mish leads to smoother gradients overall despite comparable performance and data efficiency.Source: Figure 20 · p. 25

Gradient stability

• TD-MPC2 gradients remain stable throughout training, whereas TD-MPC is prone to exploding gradients that can cause divergence.Source: Figure 21 · p. 26

Architecture

• TD-MPC2 implements its encoder, dynamics, reward, policy, and Q-function components as MLPs, with architecture depth varying by model size.The encoder uses 2–5 layers; other components use three-layer MLPs with LayerNorm and Mish activations.Source: H IMPLEMENTATION DETAILS · p. 26
• The 5M-parameter base configuration totals 5,389,930 parameters across encoder, dynamics, reward, policy, Q, and task components.The listed configuration includes 7,680 task parameters and 3,156,985 Q-function parameters.Source: H IMPLEMENTATION DETAILS · p. 26

Hyperparameters

• The same hyperparameters are used across all tasks, while discount factors and seed steps are selected with task-length heuristics.The discount factor is clipped to [0.95, 0.995], and seed steps follow S = max(5T, 1000).Source: H IMPLEMENTATION DETAILS · pp. 27–28, Table 8 · p. 27

Scaling

• Multitask model sizes range from 1M to 317M parameters without other architecture or hyperparameter modifications across sizes.Scaling changes fully connected dimensions, latent-state dimension, encoder depth, and the number of Q-functions.Source: H IMPLEMENTATION DETAILS · p. 28, Table 9 · p. 28

Representation

• SimNorm partitions latent representations into fixed-dimensional simplices, applies softmax normalization, and reshapes them back to the original latent shape.The implementation uses V=8 by default; the number of simplices is inferred from V and the latent dimensionality.Source: H IMPLEMENTATION DETAILS · pp. 28–29

Visual RL

• TD-MPC2 supports visual RL by replacing its MLP encoder with a shallow four-layer convolutional encoder operating on 64 × 64 images.The visual experiments additionally use random-shift augmentation.Source: H IMPLEMENTATION DETAILS · p. 29

Scope and limitations

• The paper identifies discrete action spaces as an open extension because TD-MPC2’s MPC planning framework is designed for continuous actions.The authors suggest discrete-action planners such as MCTS as a possible replacement.Source: I EXTENDING TD-MPC2 TO DISCRETE ACTION… · p. 29

Offline RL

• Multitask training uses fixed datasets without special offline-RL regularization or hyperparameter tuning, although test-time regularization can improve performance for some strengths.The authors note that strong regularization decreases performance and that limited state-action coverage can cause extrapolation errors.Source: J TEST-TIME REGULARIZATION FOR OFFLINE… · p. 30