Research map / Reinforcement learning / Policy gradient and actor-critic
TD-MPC2: Scalable, Robust World Models for Continuous Control
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TD-MPC2, a scalable and robust implicit (decoder-free) model-based RL algorithm that uses one hyperparameter set across 104 continuous control tasks and trains a 317M-parameter multitask world model on 80 tasks.
Problem
Model-based RL for continuous control is fragile across tasks (requires per-task hyperparameters, suffers instability such as exploding gradients) and does not scale to large multitask, multi-embodiment datasets of uncurated data.
Approach
Learns a decoder-free latent world model trained with joint-embedding prediction, reward and value prediction (soft cross-entropy in log-transformed space), plans with MPPI in latent space with a maximum-entropy policy prior and a learned terminal value; adds SimNorm latent normalization, LayerNorm/Mish MLPs, Q-ensembles with dropout, and learnable normalized task embeddings with action masking fo
Key results
- Across 104 online RL tasks in 4 domains, TD-MPC2 with a single hyperparameter set outperforms SAC, DreamerV3, and TD-MPC, including large margins on Humanoid/Dog locomotion and multi-object manipulation.
- A 317M-parameter multitask TD-MPC2 agent trained on 80 tasks reaches a normalized score of 70.6, up from 16.0 for a 1M model, with score rising roughly linearly in log parameters.
- In few-shot finetuning of a 19M multitask agent to 10 held-out tasks, TD-MPC2 improves about 2x over learning from scratch in the low-data regime (20k steps).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase model size from 1M to 317M parameters | TD-MPC2 · 80-task dataset | normalized score: 16.0 -> 70.6 | Helped |
| Pretrain on 70 tasks, then finetune online to a held-out task | 19M TD-MPC2 · 10 held-out tasks; low-data regime at 20k environment steps | score: 2× | Helped |
| Use TD-MPC2 instead of prior methods | Continuous control across 104 tasks · DMControl, Meta-World, ManiSkill2, and MyoSuite | task performance | Helped |
| Normalize task embeddings to maximum L2 norm 1 (ablation) | 19M TD-MPC2 multitask agent · 80-task dataset | normalized score | Helped |
| Use Mish instead of the alternative activation function (ablation) | TD-MPC2 · Dog Run, Humanoid Walk, and Pick YCB | normalized score | No clear effect |
| Remove momentum between MPPI planning iterations (ablation) | TD-MPC2 · Continuous control tasks | task performance | No clear effect |
| Replace prioritized replay sampling with uniform sampling (ablation) | TD-MPC2 · Continuous control tasks | task performance | No clear effect |
| Use a maximum-entropy policy prior instead of a deterministic policy prior with Gaussian action noise (ablation) | TD-MPC2 policy prior · Continuous control tasks | task performance | No clear effect |
| Use Mish activations (ablation) | TD-MPC2 · Dog Run, Humanoid Walk, and Pick YCB | gradient smoothness | Helped |
| Apply code-level planning optimizations | TD-MPC2 planner · Planning | throughput: approx. 2× | Helped |
Surprising findings
- Naively increasing TD-MPC model size does not consistently improve performance, whereas TD-MPC2 scales with model size.
- DreamerV3 and SAC are prone to numerical instabilities on Dog tasks, and TD-MPC diverges on some tasks due to exploding gradients.
Limitations
- Extending TD-MPC2 to discrete action spaces is an open problem.
- Few-shot finetuning results are task-dependent and exploratory; no scaling law is formulated.
- Rewards are required for task learning; leveraging non-reward signals for large-scale pretraining is unsolved.
Experimental setup
Datasets: DMControl Meta-World ManiSkill2 MyoSuite YCB
Benchmarks: DMControl Meta-World ManiSkill2 MyoSuite
Metrics: episode return success rate normalized score
Baselines: SAC DreamerV3 TD-MPC CURL DrQ-v2
Builds on: TD-MPC MPPI DreamerV3 Soft Actor-Critic
Closest papers by meaning
- Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.89)
- PWM: Policy Learning with Multi-Task World Models (ICLR 2025 · similarity 0.88)
- Discrete Codebook World Models for Continuous Control (ICLR 2025 · similarity 0.87)
- Langevin Soft Actor-Critic: Efficient Exploration through Uncertainty-Driven Critic Learning (ICLR 2025 · similarity 0.87)
- Fast TRAC: A Parameter-Free Optimizer for Lifelong Reinforcement Learning (NeurIPS 2024 · similarity 0.87)
- World Models as Reference Trajectories for Rapid Motor Adaptation (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.