TD-MPC2: Scalable, Robust World Models for Continuous Control
Authors: Nicklas Hansen
Hao Su
Xiaolong Wang
Venue: ICLR 2024 spotlight (2024)
Keywords: reinforcement learning
model-based reinforcement learning
world models
Primary area: reinforcement learning
TL;DR: TD-MPC2 is a scalable, robust model-based RL algorithm that can be applied to diverse single-task and multi-task continuous control domains with a single set of hyperparameters.

Abstract: 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

Published as a conference paper at ICLR 2024
TD-MPC2: SCALABLE, ROBUST WORLD MODELS FOR CONTINUOUS CONTROL
Nicklas Hansen', Hao Su Xiaolong Wang"1
University of California San Diego, "Bqual advising (nihansen, haosu, xiw012 eucdd.ddl edu
Multi-task
Single-task
DMControl & Dsta-World
DMControl
Meta-Word
ManiSkill2
80 tasks
tasks
50 tasks
tasks
100
100
100
80
4BM
317M
Dreamer
Dre mreamerv3
75-
DreamerV3
Normalized score
60
5M
19M
50
SAC
TD-MPC
TD-MPC2
50-
SAC
TD-MPC
TD-MPC2
SAC
TD-MPL
TD-MPC2
TD-MPC2
40
TD-MPC
Locomotion
MyoSuite
Pick YCB
7 tasks
10 tasks
task
100
20
75
Dre preamerv3
Preamerva
Dreamerva
50-
TD.MPC
TD-MPC2
TD-MPC
TD-MPC2
TD-MPC
TD-MPC2
0-
1M
10M
100M
1B
25
SAC
SAC
SAC
Model parameters
Figure J. Overview. TD-MPC2 compares favorably to existing model-free and model-based RL methods across 104 continuous control tasks spanning multiple domains, with single set of hyper- parameters (righf). We further demonstrate the scalability of TD-MPC2 by training single 317M parameter agent to perform 80 tasks across multiple domains, embodiments, and action spaces (left).
ABSTRACT
TD-MPC is a model-based reinforcement learning (RL) algorithm that performs local trajectory optimization in the latent space of learned implicit (decoder- free) world model. In this work, we present TD-MPC2: series of improvements upon the TD-MPC algorithm. We demonstrate that TD-MPC2 improves signifi- cantly over basclines 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 suc- cessfully train single 317M parameter agent to perform 80 tasks across multiple task domains, embodiments, and action spaces. We conclude with account of lessons, opportunities, and risks associated with large TD-MPC2 agents.
Explore videos, models, data, code, and more at sattps://tdmgge httppcdd.....//////// com com
INTRODUCTION
Training large models on internet-scale datasets has led to generalist models that perform wide variety of language and vision tasks (Brown et al., 2020; He et al., 2022; Kirillov et al., 2023). The success of these models can largely be attributed the availability of enormous datasets, and care- fully designed architectures that reliably scale with model and data size. While researchers have recently extended this paradigm to robotics (Reed et al., 2022; Brohan al., 2023), a generalist embodied agent that learns to perform diverse control tasks via low-level actions, across multiple embodiments, from large uncurated (i.e., mixed-quality) datasets remains an elusive goal. We argue that current approaches generalist embodied agents suffer from (a) the assumption of near-expert trajectories for behavior cloning which severely limits the amount available data (Reed et al., 2022: Lee al., 2022: Kumar et al., 2022; Schubert al.. 2023; Driess et al., 2023; Brohan et al., 2023), and (b) a lack of scalable continuous control algorithms that are able to consume large uncu- rated datasets.

---

Published as a conference paper at ICLR 2024
Figure 2. Tasks. TD-MPC2 performs 104 diverse tasks from (left to right) DMControl (Tassa et al., 2018), Meta-World (Yu et al., 2019), ManiSkil12 (Gu et al., 2023), and MyoSuite (Caggiano et al., 2022), with a single set of hyperparameters. See Appendix B for visualization of all tasks.
Reinforcement Learning (RL) is an ideal framework for extracting expert behavior from uncurated datasets. However, most existing RL algorithms (Lillicrap et al., 2016; Haarnoja et al., 2018) 8) are designed for single-task learning and rely on per-task hyperparameters, with no principled method for selecting those hyperparameters (Zhang et al., 2021). An algorithm that can consume large multi- task datasets will invariably need to robust to variation between different tasks (e.g., action space dimensionality, difficulty of exploration, and reward distribution). In this work, we present TD- MPC2: a significant step towards achieving this goal. TD-MPC2 is a model-based RL algorithm designed for learning generalist world models on large uncurated datasets composed of multiple task domains, embodiments, and action spaces, with data sourced from behavior policies that cover wide range of skill levels, and without the need for hyperparameter-tunnig tuming..
Our algorithm, which builds upon TD-MPC (Hansen et al., 2022), performs local trajectory opti- mization in the latent space of learned implicit (decoder-free) world model. While the TD-MPC family of algorithms has demonstrated strong empirical performance in prior work (Hansen et al., 2022; 2023; Yuan et al., 2022; Yang et al., 2023; Feng et al., 2023; Chitnis et al., 2023; Zhu al., 2023; Lancaster et al., 2023), most successes have been limited to single-task learning with little em- phasis on scaling. As shown in Figure 1, naively increasing model and data size of TD-MPC often leads to net decrease in agent performance, as is commonly observed in RL literature (Kumar et al., 2023). contrast, scaling TD-MPC2 leads to consistently improved capabilities. Our algorithmic contributions, which have been key to achieving this milestone, are two-fold: (!) improved algo- rithmic robustness by revisiting core design choices, and (2) careful design of an architecture that can accommodate datasets with multiple embodiments and action spaces without relying domain knowledge. The resulting algorithm, TD-MPC2, is scalable, robust, and can be applied to a variety of single-task and multi-task continuous control problems using single set of hyperparameters.
We evaluate TD-MPC2 across a total of 104 diverse continuous control tasks spanning 4 task do- mains: DMControl (Tassa et al., 2018), Werad (Yu et al., 2019), ManiSkil12 (Gu et al., 2023), and MyoSuite (Caggiano et al., 2022). We summarize our results in Figure 1, and visualize task domains in Figure 2. Tasks include high-dimensional state and action spaces (up to A E R339,, sparse rewards, multi-object manipulation, physiologically accurate musculoskeletal motor control, complex locomotion (e.g. Dog and Humanoid embodiments), and cover a wide range of task dif- ficulties. Our results demonstrate that TD-MPC2 consistently outperforms existing model-based and model-free methods, using the same hyperparameters across all tasks (Figure 1, right). Here, "Locomotion" and "Pick YCB" are particularly challenging subsets of DMControl and ManiSkill2, respectively. We further show that agent capabilities increase with model and data size, and success- fully train single 317M parameter world model to perform 80 tasks across multiple task domains, embodiments, and action spaces (Figure 1, left). In support of open-source science, we publicly release 300+ model checkpoints, datasets, and code for training and evaluating TD-MPC2 agents, which is available at https: : / /ttppc22.. com. We conclude the paper with an account of lessons, opportunities, and risks associated with large TD-MPC2 agents.
2 BACKGROUND
Reinforcement Learning (RL) aims to learn policy from interaction with an environment, for- mulated as a Markov Decision Process (MDP) (Bellman, 1957). We focus on infinite-horizon MDPs with continuous action spaces, which can be formalized as tuple (S, A. T, R. 7) where s E S are states, a E A are actions, T: S x A + S is the transition function, R: S x A + R is reward function associated with particular task, and 7 is discount factor. The goal is to
2

---

Published as a conference paper at ICLR 2024
derive control policy S A such that the expected discounted sum of rewards (return) En 1>==
1= y'rel, 7t = R(s,, क (s:)) is maximized. this work, we obtain и by learning a world model (model of the environment) and then select actions by planning with the learned model.
Model Predictive Control (MPC) is general framework for model-based control that optimizes action sequences a of finite length such that return is maximized (or cost is minimized) over the time horizon H, which corresponds to solving the following optimization problem:
H
(S:) = arg max E
Σ
++i R(S:+1,a
(1)
Best+M
1==
The return of a candidate trajectory is estimated by simulating it with the learned model (Negenborn al., 2005). Thus, a policy obtained by Equation will invariably be (temporally) locally optimal policy and is not guaranteed (nor likely) to be a solution to the general reinforcement learning problem outlined above. As we discuss in the following, TD-MPC2 addresses this shortcoming of local trajectory optimization by bootstrapping return estimates beyond horizon H with learned terminal value function.
TD-MPC2
Our work builds upon TD-MPC (Hansen al., 2022), model-based RL algorithm that performs local trajectory optimization (planning) in the latent space of a learned implicit world model. TD- MPC2 is practical algorithm for training massively multitask world models. Specifically, we propose series of improvements to the TD-MPC algorithm, which have been key to achieving strong algorithmic robustness (can use the same hyperparameters across all tasks) and scaling its world model to 300x more parameters than previously. We introduce the TD-MPC2 algorithm in the following, and provide a full list of algorithmic improvements in Appendix A.
3.1 LEARNING AN IMPLICIT WORLD MODEL.
Learning generative model of the environment using a reconstruction (decoder) objective is tempting due to its rich learning signal. However, accurately predict- ing raw future observations (e.g., images or propriocep- tive features) over long time horizons is difficult prob- lem, and does not necessarily lead to effective control (Lambert al., 2020). Rather than explicitly model- ing dynamics using reconstruction. TD-MPC2 aims to learn maximally useful model: model that accurately predicts outcomes (returns) conditioned on a sequence of actions. Specifically, TD-MPC2 learns an implicit, control-centric world model from environment interac- tion using a combination of joint-embedding prediction (Grill et al., 2020), reward prediction, and TD-learning (Sutton, 1998). without decoding observations. We ar- gue that this alternative formulation of model-based RL is key to modeling large datasets with modest model sizes. The world model subsequently be used for decision-making by performing local trajectory opti- mization (planning) following the MPC framework.
ig
a
42
åg
qo
a,
82
Z2
Z3
enc
enc
enc
S2
S3
Figure The TD-MPC2 architecture. Observations s are encoded into their (normalized) latent representation z. The model then recurrently predicts actions å, rewards i, and terminal values without decoding future observations.
Components. The TD-MPC2 architecture is shown in Figure and consists of five components:
Encoder
z = h(s, c)
D> Maps observations to their latent representations
Latent dynamics
z' d(z, a. e)
Models (latent) forward dynamics
Reward
R(z, a. e)
b Predicts reward r of a transition
(2)
Terminal value
Q e)
Predicts discounted sum of rewards (return)
Policy prior
à = =p(z,er
b Predicts action a" that maximizes Q
where s and a are states and actions, is the latent representation, and e is learnable task embed- ding for in multitask world models. For visual clarity, will omit in the following unless it is
3

---

Published as a conference paper at ICLR 2024
particularly relevant. The policy prior p serves to guide the sample-based trajectory optimizer (plan- ner), and to reduce the computational cost of TD-learning. During online interaction, TD-MPC2 maintains replay buffer B with trajectories, and iteratively (i) updates the world model using data sampled from B, and (ii) collects new environment data by planning with the learned model.
Model objective. The h, d, R. Q components are jointly optimized to minimize the objective
C (0)
E (S, our B
Σ x
2/4 - sg sg(((s())
CE(
+ CE qt
(3)
Joint-embedding prediction
Reward prediction
Value prediction
where sg is the stop grad operator, (zi, (t) are defined in Equation 2, Q: - Tz + Q p(z()) is the TD-target at step t, À E (0, is a constant coefficient that weighs temporally farther time steps less, and CE is the cross-entropy. Q used to compute the TD- "D-target is exponential moving average (EMA) of Q (Lillicrap et al., 2016). As the magnitude of rewards may differ drastically between tasks, TD-MPC2 formulates reward and value prediction class classification) problem a log-transformed space, which is optimized by minimizing cross- entropy with Tt: gr as soft targets (Bellemare al., 2017; Kumar al., 2023; Hafner et al., 2023).
discrete regression (multi-
Policy objective. The policy prior is a stochastic maximum entropy (Ziebart et al., 2008; Haarnoja al., 2018) policy that learns to maximize the objective
H X² XaQ((zz BH Ze+1 = d(zz. at), Zo = h(o)), (4)
(0)
=
E
(5.a)o
where H. is the entropy p which can be computed closed form. Gradients of (p are taken wrt. only. As magnitude of the value estimate Q(zz.((z:)) and entropy H. can vary greatly between datasets and different stages of training, is necessary to balance the two losses to prevent premature entropy collapse (Yarats et al., 2021). A common choice for automatically tuning at , is to keep one of them constant, and adjusting the other based on an entropy target (Haarnoja et al., 2018) or moving statistics (Hafner et al., 2023), In practice, opt for tuning via moving statisties, but empirically did not observe any significant difference in results between these two options.
Architecture. All components of TD-MPC2 are implemented MLPs with intermediate linear layers followed by LayerNorm (Ba et al., 2016) and Mish (Misra, 2019) activations. To mitigate ex- ploding gradients, simplices using a softmax operation (Lavoie et al., 2022). A key benefit of embedding z as simplices (as opposed e-g. discrete representation or squashing) is that naturally biases the representation towards sparsity without enforcing hard constraints (see Appendix H for motivation and implemen- tation). We dub this normalization scheme SimNorm. Let V be the dimensionality of each simplex constructed from L partitions (groups) of z. SimNorm then applies the following transformation:
normalize the latent representation by projecting into L fixeddimeensionll
/ /T
g: gL], g: E euucr//T
(5)
where x² is the simplicial embedding of z. [-] denotes concatenation, and T > 0 is temperature parameter that modulates the "sparsity" of the representation. As we will demonstrate in our ex- periments, SimNorm is essential the training stability of TD-MPC2. Finally, reduce bias in TD-targets generated by Q, we learn an ensemble of Q-functions using the objective from Equa- tion 3 and maintain Q as an EMA of each Q-function. We use Q-functions in practice. Targets are then computed as the minimum of two randomly sub-sampled Q-funciooss (Chen et al., 2021).
3.2 MODEL PREDICTIVE CONTROL WITH A POLICY PRIOR
TD-MPC2 derives its closed-loop control policy by planning with the learned world model. Specif- ically, our approach leverages the MPC framework for local trajectory optimization using Model Predictive Path Integral (MPPI) (Williams et al., 2015) as a derivative-free optimizer with sampled action sequences (a,, I.. at++)) of length H evaluated by rolling out latent trajectories with the model. At each decision step, we estimate parameters μ²' o² of a time-dependent multivariate Gaussian with diagonal covariance such that expected return is maximized, i.e.
H-1
41" 0² = arg max
E
> H Q(zi++ (. a++8) +
Σ h=t
2 R(z,, an)
(6)
(H.0) (a, Bet (2)
4

---

Published as a conference paper at ICLR 2024
DMControl 39 tasks
Meta-World 50 tasks
ManiSklll2 5 tasks
MyoSuite 10 tasks
1000
100
100
100
500
50
50-
50-
1M
1M
SAC
DreamerV3
TD-MPC
TD-MPC2
Figure 4. Single-task RL. Episode return (DMControl) and success rate (others) as function of environment steps across 104 continuous control tasks spanning 4 diverse task domains. TD-MPC2 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.
where n. R A E I Equation 6 solved by iteratively sampling action sequences from N(H, )²), evaluating their expected return, and updating μr, based on a weighted average. Notably, Equation 6 estimates the full RL. objective introduced in Section 2 by bootstrapping with the learned terminal value function beyond horizon H. TD-MPC2 repeats this iterative planning process for fixed number of iterations and executes the first action a, (ut of) in the environment. To accelerate convergence of planning, a fraction of action sequences originate from the policy prior p. and we warm-stan start planning by initializing (μ, o as the solution to the previous decision step shifted by 1. Refer to Hansen et al. (2022) for more details about the planning procedure.
3.3 TRAINING GENERALIST TD-MPC2 AGENTS
The success of TD-MPC2 in diverse single-task problems can be attributed to the algorithm outlined above. However, learning a large generalist TD-MPC2 agent that performs a variety of tasks across multiple task domains, embodiments, and action spaces poses several unique challenges: (i) how to learn and represent task semantics? (ii) how to accommodate multiple observation and action spaces without specific domain knowledge? (ii) how to leverage the learned model for few-shot learning of new tasks? We describe our approach to multitask model learning in the following.
Learnable task embeddings. To succeed in multitask setting, an agent needs learn common representation that takes advantage of task similarities, while still retaining the ability to differentiate between tasks test-time. When task or domain knowledge is available, e.g. in the form of natural language instructions, the task embedding e from Equation 2 may encode such information. How- ever, in the general case where domain knowledge cannot be assumed, we may instead choose to learn the task embeddings (and, implicitly, task relations) from data. TD-MPC2 conditions all of its five components with a learnable, siee-dimensiioal task embedding c, which is jointly trained to- gether with other components of the model. To improve training stability, we constrain the 2n-mrmm of e to < 1; this also leads to more semantically coherent task embeddings our experiments. When finetuning multitask TD-MPC2 agent to new task, the embedding of semantically similar task, or simply as a random vector.
choose to either initialize e as
Action masking. TD-MPC2 learns to perform tasks with variety of observation and action spaces, without any domain knowledge. To do so, we zero-pad all model inputs and outputs to their largest respective dimensions, and mask out invalid action dimensions in predictions made by the policy prior p during both training and inference. This ensures that prediction errors in invalid dimensions do not influence TD-target estimation, and prevents p from falsely inflating its entropy for tasks with small action spaces. We similarly only sample actions along valid dimensions during planning.
4 EXPERIMENTS
We evaluate TD-MPC2 across total of 104 diverse continuous control tasks spanning 4 task do- mains: DMControl (Tassa et al., 3018), Meta-World (Yu et al., 2019), ManiSkil12 (Gu et al., 2023), and MyoSuite (Caggiano et al., 2022). Tasks include high-dimensional state and action spaces (up to A E R39)), sparse rewards, multi-object manipulation, physiologically accurate musculoskeletal mo- tor control, complex locomotion (e.g. Dog and Humanoid embodiments), and cover a wide range of task difficulties. In support of open-source science, publicly release 300+ model check- points, datasets, and code for training and evaluating TD-MPC2 agents, which is available at https: / htttas///dmmmp//eeeee h/tdmpcc..co
5

---

Published as a conference paper at ICLR 2024
Run
Dog Trot
Dog Walk
Humanoid
Humanoid Walk
1000
500
0
414
8M
12M
4M
8M
12M
4M
8M
12M
4M
8M
12M
4M
BM
12M
SAC
Dreamerv3
TD-MPC
TD-MPC2
Figure 5. High-dimensional locomotion. Episode return as function of environment steps in Humanoid (A E IR?2) and Dog (A E 12se) locomotion tasks from DMControl. SAC and DreamerV3 are prone to numerical instabilities in Dog tasks, and are significantly less data-efficient than TD- MPC2 in Humanoid tasks. Mean and 95% CIs over seeds. See Appendix D for more tasks.
Cube
Pick Cube
Pick YCB
Stack Cube
Turn Faucet
50-
1M
3M
4M
1M
8M
12M
1.4
2M
3M
SAC
Dreamerv3
TD-MPC
TD-MPC2
Figure 6. Object manipulation. Success rate (%) as a function of environment steps on 5 object manipulation tasks from ManiSkil12. Pick YCB considers manipulation of all 74 objects from the YCB (Calli et al., 2015) dataset. TD-MPC2 excels at hard tasks. Mean and 95% CIs over 3 seeds.
We seek answer three core research questions through experimentation:
Comparison to existing methods. How does TD-MPC2 compare to state-of-the--e model -reee (SAC) and model-bas based (DreamerV3, TD-MPC) methods for data-efficient continuous control? Scaling, Do the algorithmic innovations of TD-MPC2 lead to improved agent capabilities as model and data size increases? Can single agent learn to perform diverse skills across multiple task domains, embodiments, and action spaces?
Analysis. How do the specific design choices introduced in TD-MPC2 influence downstream task performance? How much does planning contribute to its success? Are the learned task embeddings semantically meaningful? Can large multi-task agents be adapted unseen tasks?
Baselines. Our baselines represent the state-of-the-arr in data-efficient RL, and include (7) Soft Actor-Critic (SAC) (Haarnoja et al., 201 8), model-free actor-critic algorithm based on maximum entropy RL, (2) DreamerV3 (Hafner et al., 2023), model-based method that optimizes model- free policy with rollouts from a learned generative model of the environment, and (3) the original version of TD-MPC (Hansen et al., 2022), model-based RI. algorithm that performs local trajec- tory optimization (plamning) in the latent space of learned implicit (onn-enerattve) world model. SAC and TD-MPC use task-specific hyperparameters, whereas TD-MPC2 uses the same hyperpa- rameters across all tasks. Additionally, it is worth noting that both SAC and TD-MPC use larger batch size of 512, while 256 is sufficient for stable learning with TD-MPC2. Similarly, DreamerV3 uses high update-to-dant (UTD) ratio of 512, whereas TD-MPC2 uses UTD of by default. We use a 5M parameter TD-MPC2 agent in experiments (unless stated otherwise). For reference, the DreamerV3 baseline has approx. 20M learnable parameters, See Appendix H for more details.
4.1 RESULTS
Comparison to existing methods. We first compare the data-efficiency of TD-MPC2 set of strong baselines on 104 diverse tasks in an online RL setting. Aggregate results are shown in Fig- ure 4. We find that TD-MPC2 outperforms prior methods across all task domains. The MyoSuite results are particularly noteworthy, as we did not run any TD-MPC2 experiments on this bench- mark prior to the reported results. Individual task performances on some of the most difficult tasks lhigh-dimensional locomotion and multi-object manipulation) are shown in Figure 5 and Figure 6. TD-MPC2 outperforms baselines by a large margin on these tasks, despite using the same hyper- parameters across all tasks. Notably, TD-MPC sometimes diverges due to exploding gradients, whereas TD-MPC2 remains stable. We provide per-task visualization of gradients in Appendix G. Similarly, observe that DreamerV3 experiences occasional numerical instabilities (Dog) and

---

Published as a conference paper at ICLR 2024
DMControl Meta-Worid tasks
DMControl 30 tasks
Hopper Hop
80
68.0
70.6
80
71.4
Cheetaih Run
Cheatah jump
Cartpoie Walkeg Run
Normalized Score
60-
57.
60-
54.2
Pick Place
TD-MPC2 TD-MPC
Nomalizez
Stick
Shis L. Push
40
40
20
Reach
18.9
Door Open
Regc
20-
20-
Door Close
Spin
Eatch
Finger, Spin
Model parameters
10M
18
Model parameters
100M
Finger Turn Hard
Finger Tugh Easy
function of model size on the two 80-task and 30-task datasets. TD-MPC2 capabilities scale with model size. (Right) T-SNE (van der Maaten & Hinton, 2008) visualization of task embeddings learned by a TD-MPC2 agent trained on 80 tasks from DMControl and World. A subset of labels are shown for clarity.
Figure 7. Massively multi-task world models. (Left) Normalized score
generally struggles with tasks that require fine-grained object manipulation (lift, pick, stack). See Appendix D for the full single-task RL results.
Massively multitask world models. To demonstrate that our proposed improvements facilitate scaling of world models, we evaluate the performance of 5 multitask models ranging from 1M 317M parameters on collection of 80 diverse tasks that span multiple task domains and vary greatly in objective, em- bodiment, and action space. Models are trained on a dataset of 545M transitions obtained from the replay buffers of 240 single-task TD-MPC2 agents, and thus contain a wide variety of behaviors ranging from random to expert policies. The task set consists of all 50 Meta-Worr World tasks, as well as 30 DMCon- trol tasks. The DMControl task set includes 19 original DM- Control tasks, as well as 11 new tasks. For completeness, we
Table Training cost. Ap- proximate TD-MPC2 training cost on the 80-task dataset, reported in GPU days on single NVIDIA GeForce RTX 3090 GPU. We also list the normalized score achieved by each model at end of training.
Params (M)
GPU days
Score
3.7
16.0
4.2
49.5
19
5.3
57.1
include separate set of scaling results on the 30-task DMCon- trol subset (345M transitions) as well. Due to our careful de-
48
12
68.0
317
33
70.6
sign of the TD-MPC2 algorithm, scaling up is straightforward:
to improve rate of convergence we use a 4x larger batch size (1024) compared to the single-task experiments, but make no other changes to hyperparameters.
Scaling TD-MPC2 to 317M parameters. Our scaling results are shown in Figure To summarize agent performance with a single metric, we produce normalized score that is an average of all indi- vidual task success rates (Meta- World) and episode returns normalized to the [0, 100] range (DMControl). We observe that agent capabili- ties consistently increase with model size on both task sets. Notably, performance does not appear have saturated for our largest models (317M parameters) on either dataset, and we can thus expect results to continue improving beyond our considered model sizes. We refrain from formulating a scaling law, but note that normalized score appears scale linearly with the log of model parameters (gray line in Fig- ure 7). We also report approximate training costs in Table 1. The 317M parameter model can be trained with limited computational resources. To better understand why multitask model learning is successful, we explore the task embeddings learned by TD-MPC2 (Figure 7, right). Intriguingly, tasks that are semantically similar (e.g., Door Open and Door Close) are close in the learned task embedding space. However, embedding similarity appears to align more closely with task dynamics (embodiment, objects) than objective (walk, run). This makes intuitive sense, as dynamics are tightly coupled with control.
Finetuning 10 tasks
Normallized allizizdd score
60
47.0
40-
24.0
20
From scratch Finetuned
Figure 8. Finetuning. Score of 19M parameter TD-MPC2 agent trained on 70 tasks and finetuned online each of 10 held- out tasks for 20k environ- ment steps. 3 seeds.
Few-shot learning. While our work mainly focuses the scaling and robustness of world models, we also explore the efficacy of finetuning pretrained world models for few-shot learning of unseen tasks. Specifically, we pretrain 19M parameter TD-MPC2 agent on 70 tasks from DMControl and Meta- World, and naively finetune the full model to each of 10 held-out tasks (5 from each domain) via online RI. with initially empty replay buffer and changes hyperparameters. Aggregate
7

---

Published as a conference paper at ICLR 2024
Actor tasks
Actor
Normalization tasks
Normalization Multitask (80)
Multitask (80)
Normalized score
Normalzzed score
Normaized scare
Normalzed score
53.0
46.8
20. 80-
40-
20-
4M
BM
12M
BM
32M
Policy
Policy
No Norm
No Narm
Planning
Planning
SimNorm
Simnerm
Planning policy
Planning policy
LN SimNorm
LN SimNorm
Regression
Regression
Q-functions tasks
o-functions
Multitask (80)
Multitask (80)
Normaized score
00
Nermalized score
Normalized score
Normalized score
53.5
$7.0
60-
20
49.6
80-
60-
RD-
40-
20-
12M
Continuous
Continuous
10
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% Cls over 3 random seeds. (Bars) Normalized score of 19M parameter multitask (80 tasks) TD-MPC2 agents. Our ablations highlight the relative importance of each design choice; red is the default formulation of TD-MPC2. See Appendix D for more ablations.
results are shown in Figure 8. We find that TD-MPC2 improves 2x over learning from scratch on new tasks in the low-data regime (20k environment steps') ). Although finetuning world models new tasks is very much an open research problem, our exploratory results are promising. See Appendix E for experiment details and individual task curves.
Ablations. We ablate most of our design choices for TD-MPC2, including choice of actor, various normalization techniques, regression objective, and number of Q-functions. Our main ablations, shown in Figure 9, are conducted on three of the most difficult online RL tasks, as well large- scale multitask training (80 tasks). We observe that all of our proposed improvements contribute meaningfully to the robustness and strong performance of TD-MPC2 in both single-task RL and multi-task RL. Interestingly, we find that the relative importance of each design choice is consistent across both settings. Lastly, we also ablate normalization of the learned task embeddings, shown in Appendix F. The results indicate that maintaining normalized task embedding space (62 (er-nomm of 1) is moderately important for stable multitask training, and results in more meaningful task relations.
5 LESSONS, OPPORTUNITIES, AND RISKS
Lessons. Historically, RL algorithms have been notoriously sensitive to architecture, hyperparame- ters, characteristics of the task, and even random seed (Henderson al., 2018), with no principled method for tuning the algorithms. As result, successful application of deep RL often requires large teams of experts with significant computational resources (Berner et al., 2019; Schrittwieser al., 2020; Ouyang et al., 2022). TD-MPC2 along with several other contemporary RL methods (Yarats et al., 2021; Ye et al., 2021; Hafner et al., 2023) - seek to democratize use of RL (i.e., lowering the barrier of entry for smaller teams of academics, practitioners, and individuals with fewer resources) by improving robustness of existing open open-source algorithms. We firmly believe that improving al- gorithmic robustness will continue to have profound impact the field. A key lesson from the development of TD-MPC2 is that the community has yet to discover an algorithm that truly masters everything out-of-the-box- box. While e.g. DreamerV3 (Hafner et al., 2023) has delivered strong results on challenging tasks with discrete action spaces (such as Atari games and Minecraft), we find that TD-MPC2 produces significantly better results difficult continuous control tasks. At the same time, extending TD-MPC2 to discrete action spaces remains an open problem.
Opportunities. Our scaling results demonstrate path for model-based RL in which massively multitask world models are leveraged as generalist world models. While multi-task world models
20k environment steps corresponds 20 episodes in DMControl and 100 episodes in Hera-World.
8

---

Published as a conference paper at ICLR 2024
remain relatively underexplored in literature, prior work suggests that the implicit world model of TD-MPC2 may be better suited than reconstruction based approaches for tasks with large visual variation (Zhu et al., 2023). We envision a future in which implicit world models are used zero-shot to perform diverse tasks on seen embodiments (Xu et al., 2023; Yang et al., 2023), finetuned to quickly perform tasks on new embodiments, and combined with existing vision-language models to perform higher-level cognitive tasks in conjunction with low-level physical interaction. Our results are promising, but such level of generalization will likely require several orders of magnitude more tasks than currently available. Lastly, we want to remark that, while TD-MPC2 relies on rewards for task learning, is useful to adopt a generalized notion of reward as simply a metric for task comple- tion. Such metrics already exist in the wild, e.g., success labels, human preferences or interventions (Ouyang et al., 2022), or the embedding distance between a current observation and goal (Eysen- bach et al., 2022; Ma et al., 2022) within a pre-existing learned representation. However, leveraging such rewards for large-scale pretraining is an open problem. To accelerate research in this area, we are releasing 300+ TD-MPC2 models, including 12 multitask models, as well as datasets and code, and we are beyond excited to see what the community will do with these resources.
Risks. While we are excited by the potential of generalist world models, several challenges main: (i) misspecification of task rewards can lead to unintended outcomes (Clark & Amodei, 2016) that may be difficult to anticipate, (ii) handing over unconstrained autonomy of physical robots to a learned model can result in catastrophic failures if no additional safety checks are in place LLan- caster et al., 2023), and (iii) data for certain applications may be prohibitively expensive for small teams to obtain at the scale required for generalist behavior to emerge, leading to a concentration of power. Mitigating each of these challenges will require new research innovations, and we invite the community to join us in these efforts.
6 RELATED WORK
Multiple prior works have sought to build RL algorithms that are robust to hyperparameters, archi- tecture, well variation in tasks and data. For example, (1) Double Q-learning (Hasselt et al., 2016), RED-Q (Chen et al., 2021), SVEA (Hansen et al., 2021), and SR-SPR (D'On (D'Oro et al., 2023) each improve the stability of Q-learning algorithms by adjusting the biri-vari trade-off in TD- target estimation,
(2) C51 (Bellemare et al., 2017) and DreamerV3 (Hafner et al., 2023) improve robustness to the magnitude of rewards by performing discrete regression in a transformed space, and (3) model-free algorithms DrQ (Kostrikov et al., 2020) and DrQ-v2 (Yarats et al., 2021) improve training stability and exploration, respectively, through use of data augmentation and several other minor but important implementation details. However, all of the aforementioned works strictly focus on improving data-efficiency and robustness in single-task online RL.
Existing literature that studies scaling of neural architectures for decision makinng typically assume access to large datasets of pear-expert demonstrations for behavior cloning (Reed et al., 2022; Lee et al., 2022; Kumar et al., 2022; Schubert et al., 2023; Driess et al., 2023; Brohan et al., 2023). Gato (Reed et al., 2022) learns to perform tasks across multiple domains by training a large Transformer- based sequence model (Vaswani et al., 2017)
enormous dataset of expert demonstrations, and RT-1 (Brohan et al., 2023) similarly learns a sequence model for object manipulation on a single (real) robot embodiment by training on large dataset collected by human teleoperation. While the empirical results of this line of work are impressive, the assumption of large demonstration datasets is impractical. Additionally, current sequence models rely discretization of the action space (tokenization), which makes scaling to high-dimensional continuous control tasks difficult.
Most recently, researchers have explored scaling of RL algorithms as solution to the aforemen- tioned challenges (Baker et al., 2022; Jia et al., 2022; Xu et al., 2023; Kumar et al., 2023; Hafner et al.. 2023). For example, VPT (Baker al., 2022) learns to play Minecraft by first pretraining a behavior cloning policy on a large human play dataset, and then finetuning the policy with RL. GSL (Jia et al., 2022) requires no pre-existing data. Instead, GSL iteratively trains a population of "specialist" agents on individual task variations, distills them into a "generalist" policy via behavior cloning, and then uses the generalist as initialization for the next population of specialists. However, this work considers strictly single-task RL and assumes full control over the initial state in each episode. Lastly, D3eamerV3 (Hafner et al., 2023) successfully scales its world model in terms of parameters and shows that larger models generally are more data-cfficient in online RL setting, but does not consider multitask RL.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This project was supported, in part, by grants from NSF CAREER Award (2240160), NSF TILOS AI Institute (21 (2112)65), NSF CCF-2112666 (TILOS), NSF 1730158 CI-New: Cognitive Hardware and Software Ecosystem Community Infrastructure (CHASE-CI), NSF ACI-1541349 CC*DNI Pacific Research Platform, and gifts from Qualcomm.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Advances in Neural Information Processing Systems, 2016.
Bowen Baker, 1lge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24654, 2022.
Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforcement learning. In International Conference on Machine Learning, pp. 449-458. PMLR, 2017.
Richard Bellman. A markovian decision process. Indiana Univ. Math. J., 6:679-684, 1957. ISSN
0022-2518.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv: 1912.666)) 2019.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choroman- ski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language-aciion models transfer web knowledge robotic control. arXiv preprint arXiv:2307.15818 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
Vittorio Caggiano, Huawei Wang, Guillaume Durandau, Massimo Sartori, and Vikash Kumar. Myosuite contact-rich simulation suite for musculoskeletal motor control, 2022.
Berk Calli, Arjun Singh, Aaron Walsman, Siddhartha Srinivasa, Pieter Abbeel, and Aaron M. Dollar. The ycb object and model set: Towards common benchmarks for manipulation research. In 2015 International Conference on Advanced Robotics, pp. 510-517, 2015. doi: 10. 10911///..2005.. 7251504.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q-learning: Learning fast without model. International Conference Learning Representations, 2021.
Rohan Chitnis, Yingchen Xu, Bobak Hashemi, Lucas Lehnert, Urun Dogan, Zheqing Zhu, and Olivier Delalleau. Iql-td-mpe: Implicit q-learning for hierarchical model predictive control. arXiv preprint ar Xrxi::2306. X1v:22306.00867, 2023.
Jack Clark and Dario Amodei. Faulty reward functions in the wild. OpenAl Blog, 2016.
Djork-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv: 1,70229, 2015.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In H. Jaap van den Herik, Paolo Ciancarini, and H. H. L. M. (Jeroen) Donkers (eds.), Computers and Games, pp. 72-83, Berlin, Heidelberg, 2007. Springer Berlin Heidelberg.
Pierluca D' D'Oro, Max Schwarzer. Evgenii Nikishin, Pierre-Luc Bacon, Mare G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In The Eleventh International Conference Learning Representations, 2023.
10

---

Published as a conference paper at ICLR 2024
Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multi- modal language model. arXiv preprint
2023.
Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learn- ing as condiconnd reinforcement learning, Advances in Neural Information Processing Sys- lems, 35:35603-35620 2022.
Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, and Xiaolong Wang. Fine- tuning offline world models in the real world. Conference on Robot Learning, 2023.
Jean-Bastien Grill, Florian Strub, Florent Altch'e Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning. Advances in Neural Information Processing Systems, 2020.
Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiaing Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie, Zhiao Huang, Rui Chen, and Hao Su. Maniskil12: A unified benchmark for generalizable manipulation skills. In International Conference on Learning Representations, 2023.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, G. Tucker, Sehoon Ha, Jie Tan, Vikash Ku- mar, Henry Zhu, Abhishek Gupta, Abbeel, and Sergey Levine. Soft actor-critic algorithms and applications. ArXiv, abs/1812.05900 12.05905 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04100 040044 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Annual Conference on Neural Information Processing Systems, 2021.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In ICML, 2022.
Nicklas Hansen, Yixin Lin, Hao Su, Xiaolong Wang, Vikash Kumar, and Aravind Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations. 2023.
H. V. Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In Aaai, 2016.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. Proceedings of the Thirty-Second AAAI Confer- ence Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence, AAAI'18/AAT 1ATTT/AAATT//AAATTTTT AAAI Press, 2018. ISBN 978-1-57735-800-
Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces, 2021.
Zhiwei Jia, Xuanlin Li, Zhan Ling, Shuang Liu, Yiran Wu, and Hao Su. Improving Policy Opti- mization with Seneralist-Specialte Learning. In International Conference Machine Learning, 2022.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yer Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything. 3046432304.006633, 2023.
11

---

Published as a conference paper at ICLR 2024
Iyy Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. International Conference on Learning Representations, 2020.
Aviral Kumar, Joey Hong. Anikait Singh, and Sergey Levine. Should i run offline reinforcement learning behavioral cloning? International Conference Learning Representations, 2022.
Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline q- learning on diverse multi-task data both scales and generalizes. International Conference on Learning Representations, 2023.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. Conference on Learning for Decision and Control, 2020.
Patrick Lancaster, Nicklas Hansen, Aravind Rajeswaran, and Vikash Kumar. Modem-v2: Visuo- motor world models for real-world robot manipulation. arXiv preprint, 2023.
Samuel Lavoie, Christos Tsirigotis, Max Schwarzer, Ankit Vani, Michael Noukhovitch, Kenji Kawaguchi, and Aaron Courville. Simplicial embeddings in self-supervised learning and down- stream classification. arXiv preprint arXiv:2204.0065 2022.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-gama decision trans- formers. Advances in Neural Information Processing Systems, 35:27921-27958 2022.
Lillicrap, J. Hunt, A. Pritzel, N. Heess, Erez, Y. Tassa, D. Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. CoRR, abs/1509.02971 2016.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value-implieit pre-training. arXiv preprint arXiv:22 10.00030 2022.
Diganta Misra. Mish: A self regularized non monooootooicc neural activation function. arXiv preprint arXiv:/908, 08681, 2019.
Mitsuhiko Nakamoto, Yuexiang Zhai, Anikait Singh, Max Sobol Mark, Yi Ma. Chelsea Finn, Aviral Kumar, and Sergey Levine. Cal-ql: Calibrated offline rl pre-training for efficient online fine- tuning, 2023.
Rudy R. Negenborn, Bart De Schutter, Marco Wiering, and Hans Hellendoorn. Learning-basse model predictive control for markov decision processes. IFAC Proceedings Volumes, 38(1 38(1):354 359, 2005. 16th IFAC World Congress.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing, arXiv preprint arXiv:1711.00931 2017.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35; 27730-27744, 2022.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arX1v:2205.06177 2022.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609 2020.
Ingmar Schubert, Jingwei Zhang, Jake Bruce, Sarah Bechtle, Emilio Parisotto, Martin Riedmiller, Jost Tobias Springenberg. Arunkumar Byravan, Leonard Hasenclever, and Nicolas Heess. A generalist dynamics model for control. arXiv preprint arXiv:2305. /09/: 2023.
12

---

Published as a conference paper at ICLR 2024
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, 1lya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):1929-1958, 2014.
R. Sutton. Learning predict by the methods of temporal differences. Machine Learning, 3:9-44, 1998.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, al. Deepmind control suite. Technical report, DeepMind, 2018.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Machine Learning Research, 9:2579- 2605, 2008.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Grady Williams, Andrew Aldrich, and Evangelos A. Theodorou. Model predictive path integral control using covariance variable importance sampling. ArXiv, abs/1509.01 1199001140 149, 2015.
Yifan Xu, Nicklas Hansen, Zirui Wang, Yung-Chich Chan, Hao Su, and Zhuowen Tu. On the feasibility of cross-task transfer with model-base -bsee reinforcement learning. 2023.
Sizhe Yang, Yanjie Ze, and Huazhe Xu. Movie: Visual model-based policy adaptation for view generalization. Advances in Neural Information Processing Systems, 2023.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con-
trol: Improved lata-augmented reinforcement learning. International Conference on Learning Representations, 2021.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 14:25476-25488, 2021.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning. 2019.
Yifu Yuan, Jianye Hao, Fei Ni, Yao Mu, Yan Zheng, Yujing Hu, Jinyi Liu, Yingfeng Chen, and Changjie Fan. Euclid: Towards efficient unsupervised reinforcement learning with multi-choice dynamics model. arXiv preprint 10.0022200.004998 2022.
Baohe Zhang, Raghu Rajan, Luis Pineda, Nathan Lambert, André Biedenkapp. Kurtland Chua, Frank Hutter, and Roberto Calandra. On the importance of hyperparameter optimization for model-based reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 4015-4023. PMLR, 2021.
Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta, Repo: Resilient model-basee reinforcement learning by regularizing posterior predictability. arXiv preprint arXiv:2309 2000000000000000 2023.
Brian D Ziebart, Andrew Maas, Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In Proceedings of the 23rd National Conference on Artificial Intelligence, volume 3, 2008.
13

---

Published as a conference paper at ICLR 2024
APPENDICES
A Summary of Improvements
15
Task Visualizations
16
C Task Domains
17
D Single-task Experimental Results
20
Few-shot Experimental Results
23
Additional Ablations
24
G Gradient Norm and Training Stability
25
H Implementation Details
25
Extending TD-MPC2 to Discrete Action Spaces
28
J Test-Time Regularization for Offline RL
29
K Additional Multi-task Results
30
14

---

Published as a conference paper at ICLR 2024
A SUMMARY OF IMPROVEMENTS
We summarize the main differences between TD-MPC and TD-MPC2 as follows:
Architectural design. All components of TD-MPC2 are MLPs with LayerNorm (Ba et al,, 2016) and Mish (Misra, 2019) activations after each layer. We apply SimNorm normaliza- tion to the latent state z which biases the representation towards sparsity and maintaining a small fy-norm. We train an ensemble of Q-functions (5 by default) and additionally apply 1% Dropout (Srivastava et al., 2014) after the first linear layer in each Q-function. TD-targets are computed the mininum of two randomly subsampled Q- p-functions (Chen et al., 2021). )). In contrast, TD-MPC is implemented as MLPs without LayerNorm, and in- stead uses ELU (Clevert et al., 2015) activations. TD-MPC does not constrain the latent state at all, which in some instances leads to exploding gradients (see Appendix G for ex- perimental results). Lastly, TD-MPC learns only Q-functions and does not use Dropout. The architectural differences in TD-MPC2 result in a 4M net increase in learnable param- eters (5M total) for our default single-task model size compared to the 1M parameters of TD-MPC. However, as shown in Figure/ 7, naïvely increasing the model size of TD-MPC does not lead to consistently better performance, whereas does for TD-MPC2.
Policy prior. The policy prior of TD-MPC2 is trained with maximum entropy RL (Ziebart et al., 2008: Haarnoja et al., 2018), whereas the policy prior of TD-MPC is trained as a deterministic policy with Gaussian noise applied to actions. We find that carefully tuned Gaussian noise schedule is comparable to a policy prior trained with maximum entropy. However, maximum entropy RI. can more easily be applied with aask-agnostic hyperparam- eters. We only compute policy entropy over valid action dimensions in multi-task learning with multiple action spaces.
Planning. The planning procedure of TD-MPC2 closely follows that of TD-MPC. How- ever, we simplify planning marginally by not leveraging momentum between iteration, as we find it to produce comparable results. We also improve the throughput of planning by approx. 2x through series of code-level optimizations.
Model objective. We revisit the training objective of TD-MPC and improve its robustness to variation in tasks, such as the magnitude of rewards. TD-MPC2 uses discrete regression (soft (ross-nntoppy) of rewards and values in a log-transformed space, which makes the magnitude of the two loss terms independent of the magnitude of the task rewards. TD- MPC uses continuous regression which leads to training instabilities in tasks where rewards are large. While this issue can be alleviated by, e.g., normalizing task rewards based on moving statistics, in the single-task case, it is difficult design robust reward normalization schemes for multi-task learning. TD-MPC2 retains the continuous regression term for joint- embedding prediction the latent representation is already normalized by SimNorm, and discrete regression is computationally expensive for high-dimensional spaces (requires N bins for each dimension of z).
Multi-task model. TD-MPC2 introduces a framework for learning multi-task world mod- els across multiple domains, embodiments, and action spaces. We introduce a normalized learnable task embedding space which all components of TD-MPC are conditioned on, and we accommodate multiple observation and action spaces by applying zero padding and action masking during both training and inference. We train multi-task models large number of tasks, and finetune the model to held-out tasks (across embodiments) us- ing online RL. TD-MPC only considers multi-task learning on a small number of tasks with shared observation and action space, and does not consider finetuning of the learned multi-task model.
Simplified algorithm and implementation. TD-MPC2 removes momentum in MPPI (Williams et al., 2015), and replaces prioritized experience replay sampling from the re- play buffer with uniform sampling, both of which simplify the implementation with no significant change in experimental results. Finally, we also use faster replay buffer imple- mentation that uses multiple workers for sampling, and we increase training and planning throughput through code-level optimizations such as Q-function ensemble vectorization, which makes the wall- time of TD-MPC2 comparable to that of TD-MPC despite a larger architecture (5M vs. 1M).
15

---

Published as a conference paper at ICLR 2024
B TASK VISUALIZATIONS
Figure 10. Task visualizations. Visualization of a random initial state for each of the 104 tasks that we consider. Tasks vary greatly in objective, embodiment, and action space. Visit https: // ttmpc22 com for videos of TD-MPC2 performing each task. See Appendix C for task details.
16

---

Published as a conference paper at ICLR 2024
C TASK DOMAINS
We consider a total 104 continuous control tasks from task domains: DMControl (Tassa et al., 2018), Meta-World (Yu et al., 2019), ManiSkil12 (Gu al., 2023), and MyoSuite (Caggiano al., 2022). This section provides an exhaustive list of all tasks considered, as well as their observation and action dimensions. Environment details are listed at the end of the section. We provide (static) task visualizations in Appendix B and videos of TD-MPC2 agents performing each task at https: //www. Edmpo2... com.
We consider total of 39 continuous control tasks in the DMControl domain, including 19 original DMControl tasks and 11 new (custom) tasks created specifically for TD-MPC2 benchmarking and multitask training. We list all considered DMControl tasks below. The Locomo- fion task set shown in Figure corresponds to the Humanoid and Dog embodiments of DMControl, with performance reported at 14M environment steps.
Table DMControl.
Task
Observation dim
Action dim
Sparse?
New?
Acrobot Swingup
6
N
N
Cartpole Balance
5
1
N
N
Cartpole Balance Sparse
5
1
Y
N
Cartpole Swingup
5
I
N
N
Cartpole Swingup Sparse
5
Y
N
Cheetah Jump
17
6
N
Y
Cheetah Run
17
6
N
N
Y
Cheetah Run Back
17
6
N
Y
Cheetah Run Backwards
17
6
N
Cheetah Run Front
17
6
N
Y
Cup Catch
8
2
Y
Cup Spin
8
2
N
Dog Run
223
38
N
Dog Trot
223
38
N
Dog Stand
223
38
Dog Walk
223
38
I
Finger Spin
9
2
Y
Finger Turn Easy
12
2
Y
Finger Turn Hard
12
2
Y
Fish Swim
24
5
N
Hopper Hop
15
4
N
N
Hopper Hop Backwards
15
4
Hopper Stand
15
4
N
Humanoid Run
67
24
N
Humanoid Stand
67
24
N
Humanoid Walk
67
24
N
Pendulum Spin
3
N
Pendulum Swingup
3
1
N
Quadruped Run
78
12
Quadruped Walk
78
12
Reacher Easy
6
2
Y
Y
Y
Reacher Hard
6
2
Reacher Three Easy
8
3
Reacher Three Hard
8
3
Walker Run
24
6
Walker Run Backwards
24
6
Walker Stand
24
6
Walker Walk
24
6
Walker Walk Backwards
24
6
N
Y
17

---

Published as a conference paper at ICLR 2024
Table 3. Meta-World. We consider total of 50 continuous control tasks from the Meta- Woold domain. The Meta- World benchmark designed for multitask and meta-learning research and all tasks thus share embodiment, observation space, and action space.
Task
Observation dim
Action dim
Assembly
39
4
Basketball
39
4
Bin Picking
39
4
Window Open
39
4
Table 4. ManiSkill2. We consider total of 5 continuous control tasks from the ManiSkill2 domain. The ManiSkill2 benchmark is designed for large-scale robot learning and contains a high degree of randomization and task variations. The Pick YCB task shown in Figure 1 corresponds to the ManiSkil12 task of the same name, with performance reported at 14M environment steps.
Task
Observation dim
Action dim
Lift Cube
42
Pick Cube
51
Pick YCB
51
Stack Cube
55
Turn Faucet
40
Table 5. MyoSuite. We consider total of 10 continuous control tasks from the MyoSuite domain. The MyoSuite benchmark is designed for high-dimensional physiologically accurate muscoloskele- tal motor control and involves particularly complex object manipulation with dexterous hand. The MyoSuite domain consists of tasks with and without goal randomization. We consider both settings, and refer to them as Easy (fixed goal) and Hard (random goal), respectively.
Task
Observation dim
Action dim
Reach Easy
115
39
Reach Hard
115
39
Pose Easy
108
39
Pose Hard
108
39
Pen Twirl Easy
83
39
Pen Twirl Hard
83
39
Object Hold Easy
91
39
Object Hold Hard
91
39
Key Turn Easy
93
39
Key Turn Hard
93
39
18

---

Published as a conference paper at ICLR 2024
Environment details. We benchmark algorithms DMControl, Meta-World, ManiSkill2, and MyoSuite without modification. All four domains are ments for which use fixed episode length and termination conditions. We list episode lengths, action repeats, total number of environment steps, and the performance metric used for each domain in Table 6. In all experiments, of an episode is successful. This is a stricter definition of success than used in some of the related literature, which e.g. may consider an episode successful if success is achieved at any step within given episode. In tasks that require manipulation of objects, such as picking up object, our definition of success ensures that an episode which an object picked up but then dropped again is not considered successful.
continuous control environ-
only consider episode successful if the final step
Table 6. Environment details. We list the episode length and action repeat used for each task domain, as well as the total number of environment steps and performance metrics that we use for benchmarking methods. All methods use the same values for all tasks.
DMControl
Meta-World World
Meta-
ManiSkil12
MyoSuite
Episode length Action repeat Effective length
1, 1,000
200
200
100
2
2
2
1
500 4M 14M
100
100
100
Total env. steps
2M
4M - 14M
2M
Performance metric
Reward
Success
Success
Success
Appendices continue on next page
19

---

Published as a conference paper at ICLR 2024
D SINGLE-TASK EXPERIMENTAL RESULTS
Acrobot Swingup
Cartp Balance
Cartp Balance Sparse
Cartp Swingup
Cartp Swingup Sparse
1000
500
0
Cheetah Jump
Cheetah Run
Cheetah Run Back
Cheetah Run Bkwd
Cheetah Run Front
500
Cup Catch
Cup Spin
Dog Run
Dog Stand
Dog Trot
500
0-
Dog Walk
Finger Spin
Finger Tum Easy
Finger Turn Hard
Fish Swim
500
Hopper Hop
Hopper Hop Bkwd
Hopper Stand
Humanoid Run
Humanoid Stand
500
Humanoid Walk
Pendulum Spin
Pendulum Swingup
Quadruped Run
Quadruped Walk
500
Reacher Easy
Reacher
Reacher Easy
Reacher
Walker Run
500
Walker Run Bkwd
Walker Stand
Walker Walk
Walker Walk Bkwd
500
W
SAC
DreamerV3
TD-MPC
TD-MPC2
Figure 11. Single-task DMControl results. Episode return as a function environment steps. The first 4M environment steps are shown for each task, although the Humanoid and Dog tasks run for 14M environment steps; we provide those curves in Figure 14 as part of the "Locomotion" benchmark. Note that TD-MPC diverges on tasks like Walker Stand and Walker Walk whereas TD- MPC2 remains stable, We visualize gradients on these tasks in Appendix G. Mean and 95% Cls over seeds.
20

---

Published as a conference paper at ICLR 2024
Assembly
Basketball
Picking
Box Close
Button Press
Button Press
Button Press Td
Button Press Wall
Coffee Button
Coffee
100
Coffee Push
Dial Turn
Disassemble
Close
Lock
100
50-
Door Open
Door Unlock
Drawer Close
Drawer Open
Faucet Close
100
50
Faucet Open
Hammer
Hand Insert
Handle Press
Handle Press Side
100
50-
0-
Handle Pull
Handle Pull Side
Lever Pull
Peg Insert Side
Peg Unplug Side
100
Pick Out Of Hole
Pick Place
Pick Place Wall
Plate Slide
Plate Slide Back
100
Plate Slide Back Side
Plate Slide Side
Push
Push Back
Push Wall
Reach
Wall
Shelf Place
Soccer
Stick
Push
Sweep
Sweep
Window Close
Window Open
100
1M
2M
1M
2M
2M
1M
2M
SAC
DreamerV3
TD-MPC
TD-MPC2
Figure 12. Single-task Meta-World results. Success rate (%) as function of environment steps. TD-MPC2 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.
21

---

Published as a conference paper at ICLR 2024
Lift Cube
Pick Cube
Pick YCB
Stack Cube
Turn Faucet
100
4M 1M
8M 12M
1M
2M
3M
SAC
Dreamerv3
TD-MPC
TD-MPC2
Figure 13. Single-task ManiSkill2 results. Success rate (%) as function of environment steps on object manipulation tasks from ManiSkil12, 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 1 14M environment steps, and 4M environment steps for other tasks. TD-MPC2 achieves a > 60% success rate the Pick YCB task, whereas other methods fail to learn within the given budget. Mean and 95% CIs over seeds.
Dog Run
Dog Stand
Dog
Dog Walk
Humanoid Run
500
Humanoid Stand
Humanoid Walk
1000
500
0-
8M
12M
4M
8M
SAC
DreamerV3
TD-MPC
TD-MPC2
Figure 14. Single-task high-dimensional locomotion results. Episode return function of environment steps on all 7 "Locomotion" benchmark tasks. This domain includes high-dimensional Humanoid (AA R221) and Dog (A R) embodiments. Mean and 95% CIs over seeds.
Key Turn
Key Turn Hard
Obj Hold
Obj Hold Hard
Pen Twirl
100
50-
0-
Pen Twirl Hard
Pose
Pose Hard
Reach
Reach Hard
100
MINN
M
50
0
3M
2M
0
1M
214
1M
2M
1M
2M
1M
2M
DreamerV3
TD-MPC
Figure 15. Single-task MyoSuite results. Success rate (%) as a function environment steps. This task domain includes high-dimensional contact-rich musculoskeletal motor control (A R339 with a physiologically accurate robot hand. Goals are randomized in tasks designated as "Hard". TD-MPC2 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.
22

---

Published as a conference paper at ICLR 2024
E FEW-SHOT EXPERIMENTAL RESULTS
We finetune a 19M parameter TD-MPC2 agent trained on 70 tasks to each 10 held-out tasks. Individual task curves are shown in Figure 16. We compare data-efficiency of the finetuned model to a baseline agent of similar model capacity trained from scratch. However, we find that performance of our 19M parameter baselines trained from scratch are comparable to our 5M parameter agents also trained from scratch. Our few-shot finetuning results suggest that the efficacy of finetuning is somewhat task aes-ddepnddent However, more research is needed to conclude whether this is due task similarity (or rather lack thereof) or due to subpar task performance of the pretrained agent on the source task. We conjecture that both likely influence results.
When finetuning to an unseen task, we initialize the learnable task embedding for the new task as the embedding of a semantically similar task from the pretraining dataset. We list the source task embedding used as initialization for each experiment in Table 7. We did not experiment with other initialization schemes, nor other task pairings.
Cheetah Run
Hopper Hop
Bin Picking
Box Close
Door Lock
100
50
Door Unlock
Hand Insert
Pendulum Swingup
Reacher Hard
Walker Run
100
50-
20k
40k
20k
20k
40k
From scratch
Finetuned
Figure 16. Few-shot learning. Normalized episode return (DMControl) and success rate (Meta- World) as a function of environment steps while finetuning 19M parameter TD-MPC2 agent trained on 70 tasks each of 10 held-out tasks. 40k steps corresponds to 40 episodes in DM- Control and 200 World. Mean and 95% CIs over seeds.
Table 7. Initialization of task embeddings for few-shot learning, We list the task embeddings used as initialization when finetuning our 19M parameter TD-MPC2 agent to held-out tasks. We did not experiment with other initialization schemes, other task pairings.
Target task
Source task
Walker Run
Walker Walk
Cheetah Run
Cheetah Run Backwards
Hopper Hop
Hopper Stand
Pendulum Swingup
Pendulum Spin
Reacher Hard
Reacher Easy
Bin Picking
Pick Place
Box Close
Assembly
Door Lock
Door Open
Door Unlock
Door Open
Hand Insert
Sweep Into
23

---

Published as a conference paper at ICLR 2024
F ADDITIONAL ABLATIONS
Task embedding Multitask (80)
Normalized score
54.2
50
40-
30
Unnormalized
Normalized
Figure 17. Normalized task embeddings. Normalized score of 19M parameter multitask (80 tasks) TD-MPC2 agents, with and without normalized task embeddings e as described in Section 3.1. We find that normalizing e to have maximum fo-n of improves multitask performance.
Stand Hoppar Hop
Walker Walk Walker Run
Cartpola.Swingun
Cheetah Jump Cheetah Run
Cartpole Balance Cartpole Swingup
Walker Walk Walker
Finger Turn Easy Finger Spin
Cartpole Balance
Open
Place Wall
Pick Wha Push
Door Close
Door-Close
Puic
Hopper Stand Hopper
Dial Turn
Catch
Diaj Turn
Door Open
Reach Wall
Cup Spin
Door Close
Cup Spin
Cup Catch
Finger Spin
Finger Turn Hard
Che Jump
Finger Tum Easy
Figure 18. T-SNE of task embeddings with and without normalization. T-SNE (van der Maaten & Hinton, 2008) visualizations of task embeddings learned by TD-MPC2 agent trained on 80 tasks from DMControl and Meta-World. (Left) with normalization. (Right) without normalization. A subset of labels are shown for clarity. We observe that task embeddings are more semantically meaningful when normalized during training, e.g., "Door Open" and "Door Close" are close in embedding space on left, but far apart on the right.
Activation
Activation
tasks
tasks
0.3
Gradient norm
Normalized
0.2
0.1
20
0
12M
ELU
ELU
ReLU
ReLU
Mish
Mish
Normalized score function of environment steps, averaged across three of the most difficult tasks: Dog Run, Humanoid Walk (DMControl), and Pick YCB (ManiSkill2). Mean and 95% Cls over 3 random seeds. We find that TD-MPC2 achieves comparable asymptotic performance and effaicnfic with either activation function, but that Mish (Misra, 2019) leads smoother gradients overall.
Figure 19. Activation function.
24

---

Published as a conference paper at ICLR 2024
G GRADIENT NORM AND TRAINING STABILITY
Dog Trot
Hopper Hop
Quadruped Run
Reacher Hard
Walker Stand
10
10a
10
10a
10
12M
4M 1M
TD-MPC
TD-MPC2
Figure Gradient norm during training. We compare the gradient norm (log-scale) of TD-MPC and TD-MPC2 as function of environment steps five tasks from DMControl. TD-MPC is prone to exploding gradients, which can cause learning to diverge on some tasks (e.g., Walker Stand in Figure 11). In comparison, the gradients of TD-MPC2 remain stable throughout training. We only display seed per task for visual clarity.
H IMPLEMENTATION DETAILS
Architectural details. All components of TD-MPC2 are implemented MLPs. The encoder h contains a variable number of layers (2 -2 depending on the architecture size; all other components 3-layer MLPs. Intermediate layers consist of linear layer followed by LayerNorm and Mish activation function. The latent representation is normalized as a simplicial embedding. Q-functions additionally use Dropout. We summarize the TD-MPC2 architecture for the 5M parameter base (default for online RI) model size using PyTorchhllike notation:
Encoder pacametera: 157,
Dynamics parametars: 843, 254
Reward parameters: 397
Policy parameters: 582, 668
Q perameters: 3,156,
Task porametera: 680
Total parameters: 5, 389, 930
Architecture: TD-MPC2 base 5M
(Lask_embedding)' Embedding 96, max_norm=1) (encoder): ModuleDict!
(state): Sequential;
(G) Normedt.Ineal near "eaturessst,, out out_ foaturress,,,,, Mct-Mish)
(1): NormedLinear (i___eetrttee fet_tueuue, act-SinNorm)
(dynamics): Sequential
(0): Hormeduinear fet_fffattress5222:1 out feefffarures511, acz=xish)
(1): Normedlineas (in, fin_foatures=512, out fet_fuccurs--51,, acteMish)
(2): Normedtinnar (in_teatures=312, out feacuuresss,,, act=SimNorm!
(reward): Sequential(
(3): Normeduinear (in_ fonturee: out f5aturos=11, act=Kish)
HormedLineaz fe2, fus_fescures--12, sct-Mish)
(2): Lineer (inorrrri1.o out_faatures=100..
(pi): Sequentiall
(0): NormedLinear (in_ fe2tfuuettees1115,, out festuresss,, act-Mish)
NormedLinear (is_features=512, f02, act=Mish)
(2) Linear (1norrrrnnot,,,,, out aut_featuras=2A, biagsTrus)
(Q5) Vectorizad Koduletist
(3-4): 5 Sequent
(6) NormedLineat Se1 out__eaaures=512, dropout=0.01, act-Mishi
(1): NermedLinear (in Siaierateee out out_features-512, act-Mish)
(2) Linear (in_frtin_
[nt_ffaturesst01, biaseTrue)
where S is the input dimensionality, is the number of tasks, and is the action space. We exclude the task embedding T from single-task experiments. The exact parameter counts listed above are for S= 39, T= 80, and A= 6.
25

---

Published as a conference paper at ICLR 2024
Hyperparameters. We use the same hyperparameters across all tasks. Our hyperparameters listed in Table 8. We use the same hyperparameters for TD-MPC and SAC as Hansen et al. (2022). DreamerV3 (Hafner et al., 2023) uses a fixed set of hyperparameters.
Table 8. TD-MPC2 hyperparameters. We use the same hyperparameters across all tasks. Certain hyperparameters are set automatically using heuristics.
Hyperparameter
Value
Planning
Horizon (H)
3
Iterations
6
(+2 if ||||| 2 20)
Population size
512
Policy prior samples
24
Number of elites
64
Minimum std.
0.05
Maximum std. Temperature Momentum
2
0.5
No
Policy prior
Log std. min. Log std. max.
-10
2
Replay buffer Capacity Sampling
00000,,00 000
Uniform
Architecture (5M)
Encoder dim
256
MLP dim
512
Latent state dim
512
Task embedding dim
96
Task embedding norm
1
Activation
LayerNorm Mish
Q-function dropout rate
1%
Number of Q-functions
5
Number of seward/value bins
101
SimNorm dim (V)
8
SimNorm temperature (T)
1
Optimization
Update-to-data ratio
1
Batch size
256
Joint-embedding coef.
20
Reward prediction coef.
0.1
Value prediction coef.
0.1
Temporal coef. (A)
0.5
Q-fn. momentum coef. Policy prior entropy coef. Policy prior loss norm.
0.99
1 10-4
Moving (5%, 95%) percentiles
Optimizer
Adam
Learning rate
3 10-1
Encoder learning rate
1 x 10-4
Gradient clip norm Discount factor
20 Heuristic Heuristic
Seed steps
26

---

Published as a conference paper at ICLR 2024
We set the discount factor y for task using the heuristic
T 1 7 = clip(
[0.95, 0.995))
T 5
where T is the expected length of an episode after applying action repeat, and clip constrains the discount factor the interval [0.95, 0.995]. Using this heuristic, we obtain ? 0.99 for DMControl (T=50) which is the most widely used discount factor for this task domain. Tasks with shorter episodes are assigned lower discount factor, whereas tasks with longer episodes are assigned higher discount factor. All of the tasks that we consider are infinite-horizon MDPs with fixed episode lengths. We use individual discount factors (set using the above heuristic) for each task our multitask experiments. For tasks with variable or unknown episode lengths, we suggest using an empirical mean length, a qualified guess, or simply y = 0.99. While this heuristic is introduced in TD-MPC2, we apply the same discount factor for the TD-MPC and SAC baselines to ensure that comparison is fair across all task domains.
We set the seed steps S (number of environment steps before any gradient updates) for task using the heuristic
S = max 1000)
where 7 again is the expected episode length of the task after applying action repeat. We did not experiment with other heuristies nor constant values, but conjecture that Equation 8 will ensure that the replay buffer B has sufficient data for model learning regardless of episode lengths.
Model configurations. Our multitask experiments consider TD-MPC2 agents with model sizes ranging from 1M parameters to 317M parameters. Table lists the exact specifications for each of model sizes. We scale the model size by varying dimensions of fully- connecte layers, the latent state dimension z. the number of encoder layers, and the number of Q- Q-functions. We make no other modifications to the architecture nor hyperparameters across model sizes.
Table 9. Model configurations. We list the specifications for each model configuration (size) of our multitask experiments. Encoder dim is the dimensionality of fully connected layers in the encoder h. MLP dim is the dimensionality of layers in all other components, Latent state dim the dimensionality of the latent representation z, # encoder layers is the number of layers in the encoder h, Q-functions is the number of learned Q-functions, and Task embedding dim is the dimensionality of e from Equation 2. TD-targets are always computed by randomly subsampling two Q-functions, regardless of the number of Q-functions in the ensemble. We did not experiment with other model configurations. *hhe default (base) configuration used in our single-task RI experiments has 5M parameters.
1M
5M* *
19M
48M
317M
Encoder dim MLP dim
256
256
1024
1792
4096
384
512
1024
1792
4096
Latent state dim
128
512
768
768
1376
# encoder layers
2
2
3
5
# Q-functions
2
5
5
5
8
Task embedding dim
96
96
96
96
96
Simplicial Normalization (SimNorm). SimNorm is simple method for normalizing the latent rep- resentation z by projecting into mineddddmeniional simplices using a softmax operation (Lavoie et al., 2022). A key benefit of embedding z as simplices (as opposed to e.g. a discrete representation or squashing) is that it naturally biases the representation towards sparsity without enforcing hard constraints. Intuitively, SimNorm can be thought of as a "soft" variant of the rector-of-categorica approach to representation learning proposed by Oord et al. (2017) (VQ-VAE). Whereas VQ-VAE represents latent codes using set of discrete codes (L vector partitions each consisting of a one-hot encoding), SimNorm partitions the latent state into L vector partitions of continuous values that each sum to 1 due to the softmax operator. This relaxation of the latent representation is akin to softmax being a relaxation of the arg max operator. While we do not adjust the temperature 7 E [0, oo) of the softmax used in Nim in our experiments, is useful to note that it provides mechanism
27

---

Published as a conference paper at ICLR 2024
for interpolating between two extremes. For example, - would force all probability mass onto single categories, resulting in the discrete codes (one-hot encodings) of VQ- VAE. The alternative of 0 would result in trivial codes (constant vectors; uniform probability mass) and prohibit prop- agation of information. SimNorm thus biases representations towards sparsity without enforcing discrete codes or other hard constraints. We implement the SimNorm normalization layer (Lavoie et al., 2022) using Py PyTorch-like notation as follows:
def simnorm (self, V-8):
shope z.shape z.view (sviee(-s (t-appel:111 -1,
softmex (z, dims-1)
return 7. v.view
Here, is the latent representation z, and is the dimensionality of each simplex. The number of simplices L can be inferred from V and the dimensionality of z. We apply a softmax (optionally modulated by a temperature 7) to each of L partitions of form simplices, and then reshape to the original shape of z.
TD-MPC baseline implementation. We benchmark against the official implementation of TD- MPC available at https: : / /g1 thub. com ni wwwsshannnnddmppp The default TD-MPC world model has approx. 1M trainable parameters, and uses per-task hyperparameters. We use the suggested hyperparameters where available (DMControl and Meta- Seta-Worldd. For example, TD-MPC requires tuning of the number of planning iterations, latent state dimensionality, batch size, and learning rate in order to solve the challenging Dog and Humanoid tasks. Refer to their paper for complete list of hyperparameters.
V3eamerV3 baseline implementation. We benchmark against the official reimplementation of DreamerV3 V3 available at https: //github. com/ dani jar /erre thors' suggested hyperparameters for proprioceptive control (DMControl) and use the model size (20M parameters), as well as an update-to-data (UTD) ratio of 512. We use this model size and UTD for all tasks. Refer to their paper for a complete list of hyperparameters.
We follow the au-
SAC baseline implementation. We follow the TD-MPC (Hansen et al., 2022) paper in their decision to benchmark against the SAC implementation from https: / thub. com/ wttpissarrss//yttoreessss
and we use the hyperparameters suggested by the authors (when available). For example, this includes tuning the latent dimension, learning rate, and batch size for the Dog and Humanoid tasks. Refer to their paper for a complete list of hyperparameters.
I EXTENDING TD-MPC2 TO DISCRETE ACTION SPACES
It is desirable to develop a single algorithm that excels at tasks with continuous and discrete action spaces alike. However, the community has yet to discover such an algorithm. While e.g. Dream- erV3 (Hafner et al., 2023) has delivered strong results on challenging tasks with discrete action spaces (such as Atari and Minecraft), we find that TD-MPC2 produces significantly better results on difficult continuous control tasks. At the same time, extending TD-MPC2 to discrete action spaces remains an open problem. While we do not consider discrete action spaces in this work, we ac- knowledge the value of such extension. At present, the main challenge in applying TD-MPC2 to discrete actions lies in the choice of planning algorithm. TD-MPC2 relies on the MPC frame- work for planning, which is designed for continuous action spaces. We believe that MPC could be replaced with a planning algorithm designed for discrete action spaces, such as MCTS (Coulom, 2007) as used in MuZero (Schrittwieser et al., 2020). It is also possible that there exists way to apply MPC to discrete action spaces that is yet to be discovered (to the best of our knowledge), similar how recent work (Hubert et al., 2021) has discovered ways to apply MCTS to continuous action spaces through sampling.
28

---

Published as a conference paper at ICLR 2024
J TEST- TIME REGULARIZATION FOR OFFLINE RL
Our multi-task experiments revolve around training massively multi-task world models on fixed datasets that consist of variety of behaviors, which is offline RL. problem. We do not con- sider any special treatment of the offline RI. problem in the main paper, and simply train TD-MPC2 agents without any additional regularization nor hyperparameter-tuningg However, we recognize that models may benefit from such regularization (conservative estimations) due to extrapolation errors when the dataset has limited acte-action coverage and/or is highly skewed. Current offline RL algorithms are ill-suited for our problem setting, given that we aim to develop an algorithm that can seamlessly transition from massively multi-task offline pretraining to single-task online fine- tuning, without any changes in hyperparameters. Current offlin RL techniques rely on (1) explicit or implicit conservatism in Q-value estimation which requires modifications to the training objec- tive and empirically hampers online RL performance (Nakamoto et al., 2023), and (2) relies task-specific coefficient that balances value estimation and the regularizer. Instead, we propose to regularize the planning procedure of TD-MPC2 which can be done at test-time without any addi- tional model updates. Concretely, we apply the test-time regularizer proposed by Feng et al. (2023), which penalizes trajectories with large uncertainty (as estimated by the variance in Q-value predic- tions) during planning. While this approach eliminates the need for training-time regularization, it still requires users specify coefficient that weighs estimated value relative uncertainty for trajectory, which is infeasible in a multi-task scenario where estimated values may differ drastically between tasks. To circumvent this issue, we propose a simple heuristic for automatically scaling the regularization strength at each timestep based on the (magnitude of) mean value predictions for given latent state. Specifically, we estimate the uncertainty penalty latent state Ze of a sampled trajectory as
U = c evgg(( ......, @Nl) std(( .2.....qNI. .Nl)
where án a value prediction from Q-function n an ensemble of N Q-functions, and e is now task-agnostic coefficient that balances return maximization and uncertainty minimization. The planning objective in Equation is then redefined as
p²,0" 0² = arg max
E
(10)
(0,00 (a.. ace (
Q(L4H a14H) U6+H
(11)
(ri,, and un) -an))
(12)
using the definition of task-agnostic uncertainty in Equation 9. We conduct an experiment in which we apply our proposed test-time regularization to 19M parameter TD-MPC2 agent trained on the 80-task dataset, varying the regularization strength c. Results are shown in Table 10. Our results indicate that additional regularization (using our heuristic for automatic tuning) can indeed improve the average model performance for some values of c. Similar to what one would expect in single-task setting, we find that large values of e (strong regularization) decrease performance, whereas small (but > 0) values of tend to improve performance compared to TD-MPC2 without regularization. Given that our heuristic with e = 0.01 leads to meaningful improvements across 80 tasks, we expect it to work reasonably well for other datasets as well, but leave this for future work.
Table 10. Test-time regularization. Normalized score of 19M parameter TD-MPC2 agent trained on the 80-task dataset, varying the regularization strength of our proposed test-time regularizer. We do not apply this regularization in any of our other experiments, and only include these results to inspire future research directions.
No reg.
c == 0.001
c= 0.01
c= 0.1
Normalized score
56.54
58.14
62.01
44.13
29

---

Published as a conference paper at ICLR 2024
K ADDITIONAL MULTI-TASK RESULTS
To provide further insights into the effect of data size and task diversity on TD-MPC2 performance in multi-task setting, we provide additional experiments on a 15-task subset of DMControl, selected at random. Results for TD-MPC2 agents trained on 15 tasks, 30 tasks, and 80 tasks are shown in Figure 21. We observe that performance scales with model size across all three task suites, but numbers are higher across the board on the smallest dataset compared to similar capacity models trained on larger datasets. This makes intuitive sense, since model capacity remains the same while there are comparably fewer tasks to learn.
DMControl 15 tasks
DMControl tasks
DMControl & Meta-Worke
80 tasks
100
100
100
Normalized score
78.0
50.5
Normalized score
71.4
Normalized score
70.6
54.2
57,
49.5
TD-MPC2 TD-MPC
40
28
40
20
20
20
100M
10M
Model parameters
Model parameters
Model parameters
Figure 21. Additional results on massively multi-task world models. Normalized score as function model size on the 15-task, 30-task, and 80-task datasets. TD-MPC2 capabilities scale with model size.
30


Reviews:
[{"review_id": "xCoQaNp0RE", "reviewer": "Reviewer_BqZp", "summary": "This paper introduces TD-MPC2, an extension of the TD-MPC algorithm that introduces a set of improvements. TD-MPC is a model-based algorithm that performs Model-Predictive Control (MPC) in the latent space of a learned world model, and bootstraps the value of the final state on the planning horizon with a learned value function obtained via a model-free RL algorithm.\nThe introduced improvements include novel architectural designs (e.g., Layer Norm), using maximum-entropy RL as the policy prior, and, most importantly, a multi-task world model capable of learning the dynamics of several tasks with a single set of parameters. The experiments evaluate TD-MPC2 with TD-MPC, as well as with model-based and model-free SOTA algorithms, on several multi-task settings.", "questions": "Below, I have a few question and constructive feedback to the authors:\n\n1) Why did you decide to use the Mish activation function? Did it provide significant improvements compared to other commonly used activations, e.g., ReLU?\n\n2) I am not sure that it is possible to claim that MPC is a closed-loop policy, even when it considers the terminal value function at time step $H$. For instance, if Eq. (6) was maximized over a sequence of policies $(\\pi_t,...,\\pi_{t+h})$ instead of a sequence of actions $(a_{t},...,a_{t+h})$, it could result in a higher value. See Eq. (2) of [1] as an example. I believe this could lead to sub-optimal behavior in highly stochastic or multi-modal environments.\n\n3) Recently, [2] also showed performance gains via Layer Norm and Dropout. Is this related to how Layer Norm and Dropout were used on TD-MPC2? I suggest discussing these techniques in more depth.\n\n4) “SAC and DreamerV3 are prone to numerical instabilities in Dog tasks”. Could you elaborate on which numerical instabilities, and why TD-MPC2 can avoid them?\n\n5) “At the same time, extending TD-MPC2 to discrete action spaces remains an open problem.”\nI suggest elaborating what are the challenges involved in applying TD-MPC2 on discrete action spaces.\n\nMinor:\n- Below Eq. (3), it is missing a $\\gamma$ in the TD-target $q_t = r_t + \\gamma Q$.\n\n[1] Online Planning with Lookahead Policies. Yonathan Efroni et al. NeurIPS 2020.\n\n[2] Dropout Q-Functions for Doubly Efficient Reinforcement Learning. Hiraoka et al. ICLR 2022.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* The paper is well-written and organized, and includes a thorough discussion of the relevant related works.\n\n* The experimental results showcase impressive performance gains over the state-of-the-art on a wide range of robotics tasks (e.g., DM-Control, Meta-World) and settings (e.g., multi-task, few-shot learning).\n\n* The authors provide model checkpoints, datasets, and code for training and evaluating their proposed method. This significantly facilitates reproducing the experimental results and extending the introduced method.", "weaknesses": "* The method is not directly applicable to domains with discrete action spaces.\n\n* Because the proposed method relies on MPC, it inherits a few of its drawbacks (e.g., decision-time computational overhead, difficulty handling multi-modal transitions).\n\n* A few algorithmic decisions of the method could be better motivated, e.g., by giving an intuitive explanation of why they are expected to bring benefits. For instance, the discussion on why using SimNorm and how it biases the representation towards sparsity could be improved with an intuitive explanation of Eq. (5).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KbBPJXIdCr", "reviewer": "Reviewer_cggF", "summary": "The paper presents TD-MPC2, which is an upgrade for the model-based RL algorithm TD-MPC. The method is evaluated across a large number of RL tasks including multi-task RL benchmarks. The results show a significant boost of performance compared to other model-based RL baselines especially for Multi-task RL.", "questions": "See Weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method is evaluated on a large number of tasks in different domains. It's nice to see the huge boost of performance in multi-task scenarios.\n\n2.  The model checkpoints and code are provided. All the hyperparameters are also provided in the appendix.\n\n3. TD-MPC is one of the state-of-the-art model-based RL algorithms so it's great to see an upgrade to it.", "weaknesses": "From the results in appendix, it seems that in single task cases, TD-MPC2 is better than the original TD-MPC in terms of only a few more complicated tasks, is that true?\n\nThe novelty is somewhat limited as the modifications to TD-MPC seem to be all from some other existing methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iXsUsifOSA", "reviewer": "Reviewer_8R3o", "summary": "This paper introduces TD-MPC2, which incorporates improvements over a strong model-based RL baseline called TD-MPC. The improvements focus on generalization across tasks and action spaces. Experiment results show that TD-MPC2 is a viable method for continuous control tasks and is less sensitive to hyperparameters, and that it benefits from scaling both model and data sizes.", "questions": "On multi-task setting:\n1. Have you considered training the multi-task model online?\n2. Have you explored other dataset setups, i.e. different random to expert ratios?\n3. How about the impacts of task diversity, i.e. ablation on number of tasks in the training data?\n4. How do the multitask model perform compared to single-task models on each individual tasks?\n5. Have you considered leveraging offline RL techniques when learning from the fixed multi-task dataset? Do you think they can help?\n\nOther questions:\n1. The observation that TD-MPC's performance decreases as model and data sizes increase is very interesting. Do you have any hypothesis of why that's the case?\n2. How did you do hyperparameter search?\n3. Does TD-MPC2 benefit from model size scaling on single-task setting?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. TD-MPC2 is a strong model-based RL method that doesn't require significant task-specific hyperparameters tuning.\n2. The paper is very well-written.\n3. Extensive experiment results.", "weaknesses": "The study on multi-task learning is a bit unsatisfying, mainly due to the limited (and somewhat artificial) setup of the training data which was obtained from the replay buffers of the single-task TD-MPC2 agents. It would be great if the authors could consider extend their studies to more realistic training data distribution, e.g. mixture of few expert demonstrations and massive suboptimal data as prevalent in robotics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sWtNpoK6yU", "reviewer": "Reviewer_jMX1", "summary": "This paper proposes TDMPC2, which is a series of improvements over TDMPC for model-based RL in continuous domains. As far as I can tell, the following are the key differences:\n\n** Changes to make the algorithm multi-task **\n- Adding e (learned task embedding) to all components of the TOLD model.\n- Zero-padding invalid dimensions for observations and actions.\n- Handling differing reward magnitude by using cross-entropy instead of squared error for reward and value prediction losses.\n\n** Architectural changes **\n- SimNorm\n- Using an ensemble of Qs\n\n** Other changes **\n- Including an entropy bonus in policy optimization objective, following maximum entropy literature.\n\nAs far as I can tell, the planning algorithm remains unchanged.\n\n(If I have missed anything important, please let me know.)\n\nThe proposed modifications lead to large improvements over TDMPC when scaled up to larger model sizes, and enable multi-task learning.", "questions": "- How does the DreamerV3 baseline work, doesn't that algorithm expect discrete action spaces (as discussed in Section 5)?\n- I'd like some clarifications on the task embeddings e. Are they just trainable vectors that get passed as input to all components of the TOLD model? What forces them to be meaningfully diverse from each other in a way that is consistent with the semantics of the tasks? Is there some extra term in the loss function for them?\n- \"minimum of two randomly sub-sampled Q-functions\" Why not just take the minimum over all 5?\n- \"reward and value prediction as a discrete regression\" Are there limitations arising from this due to the resulting lack of expressivity? Or is the discretization fine enough that a wide enough range of value functions can be effectively learned? What are the tradeoffs of increasing or decreasing the granularity of this discretization?\n- The SimNorm contribution is interesting. Is it fair to think of it as a different method of achieving the same effect as the vector-of-categoricals approach in Appendix G of the Director paper (https://arxiv.org/abs/2206.04114)?\n- \"extending TD-MPC2 to discrete action spaces remains an open problem\" What are the challenges here?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The introductory figure on page 1 is a great idea, it works very well at describing the results at-a-glance and gets the reader interested to see more.\n- The paper is organized well and provides clear descriptions of the approach.\n- Experiments are thorough and convincing. Results are extremely favorable compared to other SOTA approaches like DreamerV3 and TDMPC. Ablations are also provided to delineate the impact of each modification.\n- Figure 7 (right) is very interesting, thank you for providing this visualization of the task embeddings.\n- Code, data, models are provided open-source. Training GPU costs being provided in Table 1 is also a plus.", "weaknesses": "I believe this paper should be accepted on the merit of the strong results and presentation. However, there are some major concerns I have that would be great to address in the next revision.\n\n- My main suggestion, *highly recommended*, is for the authors to include in the main text a list of the differences from TDMPC. There will be many readers (like myself) who are familiar with TDMPC, and we'd like to be able to understand the differences at a glance. The way the paper is currently written, the authors have described the full algorithm from first principles, but this mixes together their contributions versus TDMPC's contributions. At best, this causes confusion, while at worst, a reader may incorrectly attribute credit for an idea to this paper rather than the original TDMPC paper. A concrete example is the section on the planning algorithm; rather than a  vague sentence at the end saying \"Refer to [the TDMPC paper] for more details about the planning procedure\" (which hides whether or not there are any differences), it would be good to start off with a direct statement like \"The planning procedure is unchanged from [the TDMPC paper]; we re-describe it here so that this paper can be self-contained.\" Or, if there are minor changes, you can describe them candidly.\n\n- Another smaller suggestion: There are some parts of the paper that feel like too much overselling of what is ultimately a couple of modifications from TDMPC. The sentence which got me feeling this way was \"TDMPC2 marks the beginning of a new era for model-based RL\". It's great if the authors feel that way about their own work, but in a scientific paper I would prefer the language to remain neutral, and to allow the reader to form their own judgment about these matters.\n\n- I have included several smaller questions below that should be addressed in the next revision as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)