PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nicklas Hansen, Hao Su, Xiaolong Wang

arXiv:2310.16828v2cs.LGcs.AIcs.CVcs.RO

TL;DR

Generalist embodied agents need RL methods that handle diverse tasks and uncurated data without task-specific hyperparameter tuning. TD-MPC2 improves latent-space planning with an implicit world model, outperforming baselines across 104 tasks while scaling to a 317M-parameter agent that performs 80 tasks.

  • Problem

    Existing RL algorithms are commonly designed for single-task learning with per-task hyperparameters, while generalist embodied control across diverse embodiments and action spaces remains challenging.

  • Method

    TD-MPC2 performs local trajectory optimization in a latent, decoder-free world model and learns task-conditioned components for multitask control.

  • Results

    TD-MPC2 outperforms prior methods across 104 tasks in 4 domains using the same hyperparameters, while a 317M-parameter model performs 80 tasks.

  • Takeaways & Limitations

    Model and data scaling consistently increase agent capabilities, and finetuning improves performance 2× over learning from scratch in the low-data regime on new tasks.

  • Takeaways & Limitations

    Generalist world models face risks from misspecified rewards, unconstrained physical-robot autonomy, and the cost of obtaining data at the required scale.

Abstract

from arXiv · show

TD-MPC is a model-based reinforcement learning (RL) algorithm that performs local trajectory optimization in the latent space of a learned implicit (decoder-free) world model. In this work, we present TD-MPC2: a series of improvements upon the TD-MPC algorithm. We demonstrate that TD-MPC2 improves significantly over baselines across 104 online RL tasks spanning 4 diverse task domains, achieving consistently strong results with a single set of hyperparameters. We further show that agent capabilities increase with model and data size, and successfully train a single 317M parameter agent to perform 80 tasks across multiple task domains, embodiments, and action spaces. We conclude with an account of lessons, opportunities, and risks associated with large TD-MPC2 agents. Explore videos, models, data, code, and more at https://tdmpc2.com

1 INTRODUCTION

TD-MPC2 targets generalist embodied control by learning across diverse tasks and task conditions without per-task hyperparameter tuning. It combines improved robustness and scalable architecture design to achieve broad continuous-control performance.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

2 BACKGROUND

The background formulates continuous-control RL as maximizing discounted rewards in an infinite-horizon MDP. Model-based control learns an environment model, plans over finite action sequences, and estimates candidate returns through simulation.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

3 TD-MPC2

TD-MPC2 improves TD-MPC for robust, scalable multitask world modeling by learning an implicit latent model and planning locally with it. Its design combines decoder-free outcome prediction, task-aware handling of varied spaces, and MPC enhancements for broad continuous-control performance.

  • 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.
  • 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.
  • Implicit world model: 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.
  • 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.
  • 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.
  • 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.

4 EXPERIMENTS

TD-MPC2 is evaluated on 104 diverse continuous-control tasks and compared with strong model-free and model-based baselines. Results cover robustness, scaling, visual control, few-shot finetuning, and ablations of design choices.

  • 104 tasks span four domains, including high-dimensional actions up to A ∈R39, image observations, sparse rewards, manipulation, musculoskeletal control, and complex locomotion.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

5 LESSONS, OPPORTUNITIES, AND RISKS

The paper frames TD-MPC2 as a step toward robust, generalist world models while emphasizing unresolved requirements and risks. It highlights scaling opportunities alongside limits involving data, reward specification, autonomy, and power concentration.

  • Lessons: TD-MPC2 addresses RL sensitivity to architecture, hyperparameters, task characteristics, and random seed, which makes successful deep-RL application resource-intensive.
  • 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.
  • Opportunities: Generalist behavior will likely require several orders of magnitude more tasks than are currently available.
  • 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.

6 RELATED WORK

Related work improves RL robustness, scaling, or generalization through algorithmic changes, data augmentation, demonstrations, or population-based methods. These approaches generally differ from TD-MPC2 in scope, data assumptions, or multitask emphasis.

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

APPENDICES

The appendices and supplementary material describe TD-MPC2’s design differences, task coverage, visualizations, and additional analyses. The listed materials emphasize architectural normalization, multitask task diversity, and further experimental detail.

  • 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.
  • Task coverage: Task visualizations cover 104 tasks whose objectives, embodiments, and action spaces vary substantially.
  • Supplementary analyses: Additional materials provide task details and videos, while the paper directs readers to appendices for further task curves and ablations.

C TASK DOMAINS

The benchmark covers 104 continuous-control tasks across four domains, spanning diverse embodiments, observations, action spaces, task difficulties, and manipulation settings. Evaluation uses fixed episode lengths and a strict final-step success criterion.

  • Domains: 104 tasks span four domains: DMControl, Meta-World, ManiSkill2, and MyoSuite.The domains include 39, 50, 5, and 10 tasks, respectively.
  • 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.
  • Task diversity: MyoSuite includes high-dimensional dexterous-hand motor-control tasks with either fixed goals or randomized goals labeled Easy and Hard.
  • 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.
  • Evaluation: An episode counts as successful only when the final step is successful, so picking up and later dropping an object does not count.

D SINGLE-TASK EXPERIMENTAL RESULTS

Across single-task benchmarks, TD-MPC2 is stable and competitive across diverse control settings, with especially strong results on difficult manipulation tasks. It remains stable where TD-MPC can diverge and succeeds on tasks where other methods fail to learn within the budget.

  • DMControl: TD-MPC2 remains stable on Walker Stand and Walker Walk, whereas TD-MPC diverges on these tasks.
  • 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.
  • 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.
  • Locomotion: The high-dimensional locomotion benchmark includes Humanoid and Dog embodiments with action dimensions A ∈R21 and A ∈R38.
  • 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.

E FEW-SHOT EXPERIMENTAL RESULTS

Few-shot experiments finetune a multitask-pretrained TD-MPC2 agent on held-out tasks and find that finetuning efficacy varies by task. The authors identify task similarity and source-task performance as plausible influences, but leave their relative roles unresolved.

  • Setup: A 19M-parameter TD-MPC2 agent pretrained on 70 tasks is finetuned to each of 10 held-out tasks.
  • 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.
  • Interpretation: The authors conjecture that task similarity and pretrained source-task performance both influence finetuning results, requiring further research to distinguish them.
  • 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.

F ADDITIONAL ABLATIONS

Ablations examine task-embedding normalization, activation functions, and gradient behavior. Normalized embeddings improve multitask performance and semantic structure, while Mish produces smoother gradients without changing asymptotic performance or data efficiency.

  • Task embeddings: Normalizing task embeddings to a maximum ℓ2-norm of 1 improves multitask performance.
  • Task embeddings: Normalized task embeddings are more semantically meaningful, placing related tasks such as Door Open and Door Close closer together.
  • 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.
  • Activation function: Mish leads to smoother gradients overall despite comparable performance and data efficiency.
  • Gradient stability: TD-MPC2 gradients remain stable throughout training, whereas TD-MPC is prone to exploding gradients that can cause divergence.

H IMPLEMENTATION DETAILS

TD-MPC2 uses a modular latent-world-model architecture with normalized representations, task embeddings, and shared hyperparameters across tasks. Its experiments vary model scale, modalities, and action-space settings while documenting scope boundaries for discrete actions and offline RL.

  • 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.
  • Architecture: 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.
  • 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).
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
Loading 2310.16828v2…