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

ARCHITECTURE

Figure 13: Single-task Meta-World results. Success rate ( \% ) as a function of environment steps. TD-MPC 2 performance is comparable to existing methods on easy tasks, while outperforming other methods on hard tasks such as Pick Place Wall and Shelf Place . DreamerV3 often fails to converge.

Figure 14: Single-task ManiSkill2 results. Success rate ( \% ) as a function of environment steps on 5 object manipulation tasks from ManiSkill2. Pick YCB is the hardest task and considers manipulation of all 74 objects from the YCB (Calli et al., 2015 ) dataset. We report results for this tasks at 14 M environment steps, and 4 M environment steps for other tasks. TD-MPC 2 achieves a >60\% success rate on the Pick YCB task, whereas other methods fail to learn within the given budget. Mean and 95\% CIs over 3 seeds.

Figure 16: Single-task MyoSuite results. Success rate ( \% ) as a function of environment steps. This task domain includes high-dimensional contact-rich musculoskeletal motor control ( \mathcal{A}\in\mathbb{R}^{39} ) with a physiologically accurate robot hand. Goals are randomized in tasks designated as “Hard”. TD-MPC 2 achieves comparable or better performance than existing methods on all tasks from this benchmark, except for Key Turn Hard in which TD-MPC succeeds early in training.

THE PROBLEM

Before TD-MPC2, continuous controlControl & PlanningControlThe method used to make the robot move the way you want. in robotics faced a scalability crisis. Prior model-based RLImitation & Reinforcement LearningReinforcement Learning (RL)Teaching a robot through trial and error using rewards. methods like the original TD-MPC required careful hyperparameter tuning for each new taskCore ConceptsTaskThe job the robot is supposed to complete, such as pick-and-place, navigation, or drawer opening. or domain—what worked for a simulated robotCore ConceptsRobotA physical system with sensors and actuators that can observe the world and take actions. dog didn't work for manipulationManipulation & TasksManipulationUsing a robot arm or hand to move or interact with objects. tasks. Model-free methods like SAC were more robust but learned inefficiently, requiring millions of real-world interactions because they didn't build a world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions. to reason about consequences. DreamerV3, the previous state-of-the-art, improved multi-task learning but still struggled with the diversity problem: trainingRobot LearningTrainingThe process of fitting a model using data or experience. a single agent on 80+ tasks across different embodiments (quadrupeds, hands, humanoids) and actionCore ConceptsActionA command the robot sends to its motors, controller, or low-level system. spaces (continuous jointMovement, Mechanics & Robot BodyJointA movable connection between robot parts. commands, muscle activations) seemed impossible. The core limitation was that existing world models either required task-specific tuning or collapsed under the complexity of learning multiple incompatible environments simultaneously.

HOW IT WORKS

1

Learning an Implicit World Model in Latent Space

Instead of learning to reconstruct images (which wastes model capacity), TD-MPC2 learns a "decoder-free" world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions.—it compresses observations into a compact latent representation and learns to predict how that latent stateCore ConceptsStateThe robot’s current condition, such as joint positions, velocity, object positions, or internal variables. evolves under actions, without ever reconstructing pixels. This is crucial for scaling: by avoiding pixel prediction, the model conserves parameters for actually understanding taskCore ConceptsTaskThe job the robot is supposed to complete, such as pick-and-place, navigation, or drawer opening. dynamicsMovement, Mechanics & Robot BodyDynamicsThe study of motion including forces, torques, mass, and inertia.. The implicit world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions. essentially asks: 'given what I observe now and the actionCore ConceptsActionA command the robot sends to its motors, controller, or low-level system. I take, what will the next latent stateCore ConceptsStateThe robot’s current condition, such as joint positions, velocity, object positions, or internal variables. be?' This compression is learned via temporal difference (TD) learning, the same principle that powers value-function learning in RLImitation & Reinforcement LearningReinforcement Learning (RL)Teaching a robot through trial and error using rewards., making the entire system theoretically cohesive.

Show more ↓

Figure 3: The TD-MPC 2 architecture. Observations \mathbf{s} are encoded into their (normalized) latent representation \mathbf{z} . The model then recurrently predicts actions \hat{\mathbf{a}} , rewards \hat{r} , and terminal values \hat{q} , without decoding future observations.

2

Planning via Latent Trajectory Optimization

Once the world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions. is trained, TD-MPC2 doesn't sample random actionCore ConceptsActionA command the robot sends to its motors, controller, or low-level system. sequences and pick the best one (like prior methods). Instead, it performs gradient-based optimization directly in the latent spaceRobot LearningLatent spaceA compressed internal representation space inside a model.: it starts with a guess about the next 8-12 actions, then iteratively improves that sequence by asking the learned world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions. 'if I tweak this actionCore ConceptsActionA command the robot sends to its motors, controller, or low-level system., how much better does my predicted rewardImitation & Reinforcement LearningRewardA score that tells the robot how well it is doing. become?' This is vastly more efficient than sampling thousands of trajectories. The optimization happens on each step as the robotCore ConceptsRobotA physical system with sensors and actuators that can observe the world and take actions. acts, creating a tight loop: observe → plan in latent spaceRobot LearningLatent spaceA compressed internal representation space inside a model. → execute → observe again. This local, online planningControl & PlanningPlanningFiguring out what the robot should do before or during movement. is what makes the system robust to model errors.

Show more ↓

3

Scaling Through Architectural Improvements and Better Data

TD-MPC2 makes several technical improvements over the original: better latent regularizationData, Distributions & Training IssuesRegularizationMethods used to reduce overfitting. (keeping the learned representation stable), improved actionCore ConceptsActionA command the robot sends to its motors, controller, or low-level system. normalizationData, Distributions & Training IssuesNormalizationRescaling inputs or features to stabilize learning., and more careful rewardImitation & Reinforcement LearningRewardA score that tells the robot how well it is doing. scaling. But the real scaling breakthrough comes from trainingRobot LearningTrainingThe process of fitting a model using data or experience. on massive, diverse datasets. The authors collected 545M environmentCore ConceptsEnvironmentThe external world the robot operates in, including objects, obstacles, people, and surfaces. transitions from 240 separate single-task agents across 80 tasks, then train one unified 317M-parameter agent on all of it. The key insight: collecting this data is cheap (done offline via parallel RLImitation & Reinforcement LearningReinforcement Learning (RL)Teaching a robot through trial and error using rewards. trainingRobot LearningTrainingThe process of fitting a model using data or experience.), but trainingRobot LearningTrainingThe process of fitting a model using data or experience. one large multi-task model generalizes better than 80 separate specialists. This mirrors how large language models work—scale up data and model size, and emergent capabilities appear.

Show more ↓

Figure 5: High-dimensional locomotion. Episode return as a function of environment steps in Humanoid ( \mathcal{A}\in\mathbb{R}^{21} ) and Dog ( \mathcal{A}\in\mathbb{R}^{38} ) locomotion tasks from DMControl. SAC and DreamerV3 are prone to numerical instabilities in Dog tasks, and are significantly less data-efficient than TD-MPC 2 in Humanoid tasks. Mean and 95\% CIs over 3 seeds. See Appendix D for more tasks.

Figure 6: Object manipulation. Success rate ( \% ) as a function of environment steps on 5 object manipulation tasks from ManiSkill2. Pick YCB considers manipulation of all 74 objects from the YCB (Calli et al., 2015 ) dataset. TD-MPC 2 excels at hard tasks. Mean and 95\% CIs over 3 seeds.

Figure 7: Massively multi-task world models. (Left) Normalized score as a function of model size on the two 80-task and 30-task datasets. TD-MPC 2 capabilities scale with model size. (Right) T-SNE (van der Maaten & Hinton, 2008 ) visualization of task embeddings learned by a TD-MPC 2 agent trained on 80 tasks from DMControl and Meta-World . A subset of labels are shown for clarity.

Figure 11: Task visualizations. Visualization of a random initial state for each of the \mathbf{104} tasks that we consider. Tasks vary greatly in objective, embodiment, and action space. Visit https://tdmpc2.com for videos of TD-MPC 2 performing each task. See Appendix C for task details.

tdmpc2 humanoid spin

tdmpc2 humanoid run

tdmpc2 humanoid stand

4

Single Hyperparameter Set Across All Domains

A hidden win in TD-MPC2 is robustnessModern Robot LearningRobustnessHow well a robot keeps working despite noise, disturbances, or variation. to hyperparameter choices. The method achieves strong performance on DMControl (simulated humanoids and quadrupeds), Meta-World (dexterous hand manipulationManipulation & TasksManipulationUsing a robot arm or hand to move or interact with objects.), ManiSkill2 (object assemblyManipulation & TasksAssemblyPutting components together in a structured way.), and MyoSuite (muscle-actuated locomotionNavigation & LocomotionLocomotionMovement of the robot body through space, like walking, rolling, or running.)—all with one learning rate, batch size, and planningControl & PlanningPlanningFiguring out what the robot should do before or during movement. horizon. This is not a technical detail: in prior work, you'd need different hyperparameters for simulationSimulation & Sim-to-RealSimulationA virtual environment where robots can be trained or tested. vs. muscle controlControl & PlanningControlThe method used to make the robot move the way you want., or for vision-based tasks vs. proprioceptive ones. The unified performance suggests TD-MPC2 found a sweet spot in the algorithm design that generalizes across this landscape.

Show more ↓

Figure 1: Overview. TD-MPC 2 compares favorably to existing model-free and model-based RL methods across \mathbf{104} continuous control tasks spanning multiple domains, with a single set of hyperparameters ( right ). We further demonstrate the scalability of TD-MPC 2 by training a single 317 M parameter agent to perform \mathbf{80} tasks across multiple domains, embodiments, and action spaces ( left ).

Figure 4: Single-task RL. Episode return (DMControl) and success rate (others) as a function of environment steps across \mathbf{104} continuous control tasks spanning 4 diverse task domains. TD-MPC 2 achieves higher data-efficiency and asymptotic performance than existing methods, while using the same hyperparameters across all tasks. Mean and 95\% CIs over 3 seeds.

Figure 9: Ablations. (Curves) Normalized score as a function of environment steps, averaged across three of the most difficult tasks: Dog Run , Humanoid Walk (DMControl), and Pick YCB (ManiSkill2). Mean and 95\% CIs over 3 random seeds. (Bars) Normalized score of 19 M parameter multitask ( 80 tasks) TD-MPC 2 agents. Our ablations highlight the relative importance of each design choice; red is the default formulation of TD-MPC 2 . See Appendix D for more ablations.

Figure 12: Single-task DMControl results. Episode return as a function of environment steps. The first 4 M environment steps are shown for each task, although the Humanoid and Dog tasks are run for 14 M environment steps; we provide those curves in Figure 15 as part of the “Locomotion” benchmark. Note that TD-MPC diverges on tasks like Walker Stand and Walker Walk whereas TD-MPC 2 remains stable. We visualize gradients on these tasks in Appendix G . Mean and 95\% CIs over 3 seeds.

tdmpc2 assembly

tdmpc2 dog run

tdmpc2 turn faucet 1

tdmpc2 cheetah run

MORE DEMONSTRATIONS

FIGURES

Figure 8: Finetuning. Score of a 19 M parameter TD-MPC 2 agent trained on 70 tasks and finetuned online to each of 10 held-out tasks for 20 k environment steps. 3 seeds.

Figure 10: Visual RL. Episode return as a function of environment steps on 10 image-based DMControl tasks. Mean and 95\% CIs over 3 random seeds. TD-MPC 2 is comparable to state-of-the-art.

KEY RESULTS

PERFORMANCE COMPARISON

WHY DEVELOPERS SHOULD CARE

LIMITATIONS

TD-MPC2 does not solve several critical problems for real robots. First, all experiments are in simulationSimulation & Sim-to-RealSimulationA virtual environment where robots can be trained or tested. or with pre-recorded offline datasets—there's no demonstrationImitation & Reinforcement LearningDemonstrationAn example of a task being done correctly, often by a human. of the 317M agent actually controlling a physical robotCore ConceptsRobotA physical system with sensors and actuators that can observe the world and take actions. in real time. Sim-to-realSimulation & Sim-to-RealSim-to-real (sim2real)Transferring a policy trained in simulation to a real robot. transfer is hard, and latent dynamicsMovement, Mechanics & Robot BodyDynamicsThe study of motion including forces, torques, mass, and inertia. learned from simulationSimulation & Sim-to-RealSimulationA virtual environment where robots can be trained or tested. may not capture real-world frictionMovement, Mechanics & Robot BodyFrictionResistance between contacting surfaces that affects sliding and grasping., wear, and unexpected contacts. Second, the method requires learning a separate world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions., which adds trainingRobot LearningTrainingThe process of fitting a model using data or experience. complexity and latencySimulation & Sim-to-RealLatencyDelay between input, computation, and action. compared to pure model-free methods. If you have unlimited real-world data and don't care about sample efficiencyRobot LearningSample efficiencyHow quickly a method learns from each example or interaction., SAC might be simpler. Third, the paper doesn't thoroughly explore failure modes: what happens when the world modelModern Robot LearningWorld modelA model that predicts how the world will change after actions. is confidently wrong? PlanningControl & PlanningPlanningFiguring out what the robot should do before or during movement. in latent spaceRobot LearningLatent spaceA compressed internal representation space inside a model. can amplify model errors. Finally, scaling to 317M parameters helps, but there's no clear evidence this approach scales to robotCore ConceptsRobotA physical system with sensors and actuators that can observe the world and take actions. tasks requiring high-dimensional visual reasoning (like navigating cluttered scenes) or long-horizon reasoning (100+ steps into the future).

WHAT COMES NEXT

The natural next step is closing the sim-to-realSimulation & Sim-to-RealSim-to-real (sim2real)Transferring a policy trained in simulation to a real robot. gap: taking the 317M agent trained in simulationSimulation & Sim-to-RealSimulationA virtual environment where robots can be trained or tested. and successfully deploying it to physical robots. This requires either better domain randomizationData, Distributions & Training IssuesDomain randomizationChanging simulator visuals or physics during training so policies transfer better to reality. during trainingRobot LearningTrainingThe process of fitting a model using data or experience. or fine-tuningModern Robot LearningFine-tuningTaking a pretrained model and adapting it to a specific robot or task. with real-world data. Beyond that, the architecture suggests a path toward more general robotics AI: scale further (500M+ parameters), train on more diverse tasks (500+ instead of 80), and potentially add vision as input rather than relying on proprioceptive stateCore ConceptsStateThe robot’s current condition, such as joint positions, velocity, object positions, or internal variables.. The paper hints at scaling laws—if they hold, a 1B-parameter agent trained on 1000+ tasks across real and simulated domains might exhibit surprising generalizationModern Robot LearningGeneralizationThe robot’s ability to work in new situations it has not seen before.. Another frontier is interpretability: what does the 317M agent's latent spaceRobot LearningLatent spaceA compressed internal representation space inside a model. actually represent? Understanding this could unlock better transfer learningModern Robot LearningTransfer learningUsing knowledge from one task, domain, or robot to help with another. and faster adaptation to new tasks.

RELATED PAPERS