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

We thank the reviewer for their valuable feedback. We address your comments in the following.

Q: My main suggestion, highly recommended, is for the authors to include in the main text a list of the differences from TDMPC.

A: We fully agree. We actually provide a full list (~1 page) of differences between TD-MPC and TD-MPC2 in Appendix A, but regrettably did not reference it in the method section. We have added a reference to this list in the updated manuscript.

Q: 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.

A: We agree with the sentiment of this, and have updated the manuscript to describe the algorithm and results in a more neutral tone. Thank you for the suggestion!

Q: How does the DreamerV3 baseline work, doesn't that algorithm expect discrete action spaces (as discussed in Section 5)?

A: DreamerV3 (as well as previous iterations) can be applied to either discrete or continuous action spaces by changing the parameterization of its latent policy (Gaussian vs. categorical), similar to SAC. However, DreamerV3 had previously only been benchmarked on an easy subset of the DMControl tasks, with all other tasks being discrete. We find that DreamerV3 performs quite well on this set of tasks, but performs significantly worse when applied to other continuous control tasks – both harder DMControl tasks, as well as new domains.

Q: 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?

A: Correct – the task embeddings are simply learnable vectors. We maintain a dictionary of vectors (one for each task), and condition all components on the retrieved vector(s). They are trained end-to-end together with other components using the objective from Equation 3. Intuitively, the semantic meaning of task embeddings likely “emerges” because it is the most straightforward way for the model to distinguish different tasks and embodiments. Our only constraint (which we do find to be moderately important) is that we constrain the task embeddings to have a norm of 1. We ablate this in Appendix F, both in terms of normalized score and semantic coherence (T-SNE visualizations).

Q: "minimum of two randomly sub-sampled Q-functions" Why not just take the minimum over all 5?

A: While this design choice may seem somewhat arbitrary at first glance, there is some evidence that taking the minimum/average over a large ensemble of Q-functions can actually be detrimental to data-efficiency as it propagates biases from the target Q-function. For example, REDQ by Chen et al. [1] demonstrates this phenomenon for SAC trained on a set of continuous control tasks. Additionally, taking the minimum over the full ensemble would also be slower than when sub-sampling.

[1] Chen et al., Randomized Ensembled Double Q-Learning (https://arxiv.org/abs/2101.05982)

Q: 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)?

A: To some extent. The high-level motivation is similar: constraining the latent state in some way is beneficial for learning. DreamerV2/V3/Director factorizes the representation as multiple categoricals that function as discrete codes (originally proposed in VQ-VAE), whereas SimNorm can be viewed as a “soft” variant in the sense that its representation is factorized as multiple simplices (each sub-vector sums to 1). Intuitively, SimNorm biases the representation towards sparsity, whereas vector-of-categorical forces the representation to be sparse, in the same way that Softmax can be interpreted as a relaxation of argmax. While we do not adjust the temperature of the SimNorm Softmax in our experiments, it does provide a mechanism for interpolating between the same two extremes: infinite temperature would result in a discrete (one-hot) code, whereas zero temperature would result in constant (uniform) vectors. We have added an extended discussion on this in Appendix H of the updated manuscript.

Please do not hesitate to let us know if you have any additional comments.

We thank the reviewer for their valuable feedback. We address your comments in the following.

Q: Have you considered training the multi-task model online?

A: Yes, we considered that, but ultimately decided against it to (1) avoid the significant wall-time overhead of simulation, and (2) to eliminate additional factors of variation such as exploration that might influence our scaling results. However, given the strong single-task online RL results, we do not expect conclusions to change drastically based on this experimental design choice.

Q: Have you explored other dataset setups, i.e. different random to expert ratios? Have you considered leveraging offline RL techniques when learning from the fixed multi-task dataset? Do you think they can help?

A: We did not experiment with other data distributions, since we wanted to focus on building an RL algorithm that works at scale. Offline RL problem settings and techniques are somewhat orthogonal to realizing this goal, since offline RL techniques focus specifically on mitigating extrapolation errors in small or highly skewed datasets. For example, prior work [1] shows that standard online RL algorithms without explicit/implicit conservatism (i.e., without any of the popular offline RL techniques) outperform offline RL methods when data is abundant and covers a wider state-action distribution than what the replay buffer of an agent would typically provide. Thus, we would expect vanilla TD-MPC2 to outperform TD-MPC2 with any of the common offline RL techniques added when evaluating on maximally diverse data. Our current setting uses replay data, and therefore falls somewhere in between hard offline RL (narrow data distribution, mainly suboptimal data) and ideal data (uniform state-action distribution coverage). Whether TD-MPC2 (as well as any other RL algorithm) benefits or suffers from added conservatism thus depends entirely on where the specific dataset falls on this spectrum.

However, to address your question experimentally, we conduct an additional experiment where we enforce conservatism in a 19M parameter TD-MPC2 agent trained on our 80-task dataset. It is worth noting that most existing offline RL techniques (1) regularize predictions at training-time, and (2) require task-specific tuning of the coefficient that balances return maximization and conservatism / uncertainty, which combined makes it difficult to apply in a multitask setting like ours. Instead, we choose to regularize the planning of TD-MPC2 at test-time (i.e., we apply it to our existing checkpoint without any additional model updates) using the technique proposed by Feng et al. [2], which discourages trajectories with large uncertainty in the return estimation. Like most other offline RL methods, however, this method was proposed for a single-task setting and requires per-task tuning of the regularizer. 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 a given latent state. We provide a more precise mathematical definition in Appendix J of the updated manuscript. Our results indicate that additional regularization (and using our heuristic for automatic tuning) can indeed improve the average model performance for some values of our task-agnostic coefficient $c$, but we view this finding as orthogonal to the main contributions of our paper. Similar to what one would expect in a single-task setting, we find that large values of c (strong regularization) decrease performance, whereas small (but $>0$) values of c tend to improve performance compared to TD-MPC2 with no regularization. Given that our heuristic with $c = 0.01$ leads to meaningful improvements across 80 tasks, we expect it to also work reasonably well for other datasets. We have added this experiment + discussion to Appendix J, and thank the reviewer for the suggestion.

[1] Yarats et al., Don’t Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning, arXiv 2022 (https://arxiv.org/abs/2201.13425)

[2] Feng et al., Finetuning Offline World Models in the Real World, CoRL 2023 (https://arxiv.org/abs/2310.16029)

We thank the reviewer for their valuable feedback. We address your comments in the following.

Q: Why did you decide to use the Mish activation function?

A: We observe that Mish activations lead to slightly more consistent results compared to ReLU (common in RL) and ELU (used in TD-MPC) activations. However, TD-MPC2 works well with either of them. We have run additional ablations on the choice of activation function in response to your question; links to results are shown here. Ablations are conducted on the same three tasks as in previous ablations: Dog Run, Humanoid Walk (DMControl), and Pick YCB (ManiSkill2). Results indicate that performance (return / success rate) is similar across activations, but that Mish leads to smaller and more stable gradients compared to the alternatives. Computational overhead of Mish vs. ReLU/ELU is negligible, so we decided to use Mish activations. We have added these additional ablations to our updated manuscript.

Q: 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?

A: Short answer: yes. While TD-MPC2 does not use a Transformer architecture per se, we indeed find that TD-MPC2 benefits from many of the same design choices that have been instrumental for scaling Transformers to large datasets. Examples of such design choices include MLPs with LayerNorm, Dropout, Softmax (key to our proposed SimNorm representation), and multiple prediction heads. The referenced paper on Dropout Q-functions [2] shows that Q-learning can benefit from one of such components (Dropout), whereas we demonstrate that all of these components complement each other when carefully applied to the TD-MPC algorithm (along with our other proposed changes).

Q: “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?

A: While it is difficult to give a conclusive answer to why SAC and DreamerV3 experience numerical instabilities for this set of tasks, we can attempt to extrapolate our knowledge about the tasks and instabilities in TD-MPC to these two algorithms. Firstly, it is worth noting that the Dog tasks have a very large action space compared to other locomotion tasks (38 continuous dimensions vs. e.g. 21 for Humanoid and 6 for Walker). We observe that SAC and DreamerV3 also fail to learn in MyoSuite, which has a similarly large action space (39 dimensions). Based on our ablations in Figure 9 (actor) + the fairly strong performance of TD-MPC1 on these tasks, we believe that planning plays a key role in tasks with large action spaces. Secondly, we observe that rewards increase for SAC during its initial stages of training, but then diverge without recovery. This is qualitatively similar to what we observe in TD-MPC, for which gradients grow extremely large before ultimately diverging. TD-MPC2 improves training stability through a series of changes to the architecture and objective, which result in consistently small gradients through training.

Q: I suggest elaborating what are the challenges involved in applying TD-MPC2 on discrete action spaces.

A: Good suggestion. We recognize that bridging continuous and discrete action spaces is an important problem. The challenge mainly lies in the choice of planning algorithm. TD-MPC2 relies on the Model Predictive Control (MPC) framework 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 (used in MuZero), although that is not the focus of this paper. It is also possible that there exists a way to apply MPC to discrete action spaces that is yet to be discovered (to the best of our knowledge), similar to how researchers have discovered ways to apply MCTS to continuous action spaces through adaptive discretization and clustering methods. We have added this discussion to Appendix I of our updated manuscript.

(1/2)

I thank the authors for their detailed responses to my questions and commentary. I have no further concerns, and hence I am maintaining my decision of acceptance of the paper.

Thank you for acknowledging our response! We are glad to hear that our response has addressed your primary concerns.

Regarding alternatives to discrete regression:

Firstly, we'd like to clarify that while the problem of inconsistent reward ranges between tasks may be particularly relevant for multi-task learning, it is also problematic in single-task learning when one wants to use fixed hyper-parameters, since magnitude of reward/value directly affects the magnitude of gradients and individual loss terms. You will find that many open-source implementations of popular algorithms scale rewards manually, and that the scale depends both on the specific task and algorithm used. For example, the minimalRL (2.6k stars) implementations of DQN and SAC scale rewards by $0.01$ and $0.1$, respectively, for the same task (Cartpole). While we cannot possibly know how the code authors decided on these values, it is likely that they were discovered simply by exhaustively trying different values. It is thus very desirable to build an algorithm that eliminates the need for such tuning altogether.

Either way, discrete regression is one of the most straightforward ways to address this problem. There are other options, but they each have their own drawbacks:

• Sparse rewards: Binarizing the reward signal based on either success criteria or reward thresholding addresses the initial problem of inconsistent reward ranges, but would drastically reduce the available learning signal for tasks in which rewards are highly shaped (common in locomotion), and also makes the exploration problem significantly harder (no learning signal unless a success state / reward threshold is reached). Lastly, there exists many tasks for which defining a binary success flag is difficult; while it may work for e.g. pick-and-place style manipulation tasks where success is tied to concrete states (is the object at the target location?), this may not work for other tasks such as learning locomotion behaviors (abstract goals extended in time, such as "run indefinitely" or "make a backflip") for which defining a reward threshold would require knowledge about (1) the exact definition of the reward function and its range of possible values for each individual task, and (2) which reward threshold constitutes "successful" behavior, neither of which are feasible to define at scale.
• Manual normalization: Normalizing rewards manually is certainly an option, but requires privileged information about the reward function in question and its possible values, and is not feasible for a human to do beyond maybe a dozen tasks. We kind of addressed this approach in our first paragraph, so we'll keep this one brief.
• Automatic normalization: Rather than specifying reward scaling coefficients manually, it is also possible to automatically scale rewards based on running statistics. However, this approach has seen limited use in existing literature since running statistics introduce non-stationarity in the prediction problem, and is sensitive to outliers (infrequent very small/large rewards).

One last point that also favors our discrete regression approach is that -- even with normalized single-step rewards -- value targets may still be on different scales for different tasks. This is true for our considered task domains, since they differ both in episode length, distribution of rewards within an episode (rewards are not evenly distributed within successful episodes), discount factor (prediction horizon), and sparsity (some of our tasks use binary rewards whereas others use shaped rewards).

We hope that our responses have addressed the reviewer's concerns. Please let us know if not! Otherwise, we kindly ask the reviewer to consider raising their score accordingly.

We thank the reviewer for their valuable feedback. We address your comments in the following.

Q: 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?

A: Correct; TD-MPC performs well on many of the tasks that we consider. TD-MPC2 performs comparably on most tasks, and improves significantly over TD-MPC on difficult tasks (Dog, Humanoid, Pick YCB, Pick Place Wall, MyoSuite, etc.). However, TD-MPC2 does so without task-specific hyperparameters, whereas TD-MPC requires hyperparameter tuning for each embodiment / domain. We believe that developing an algorithm that can learn diverse tasks without any tuning is extremely valuable. We would also like to reiterate that the success of our multi-task scaling experiments (which is unprecedented in literature) are a direct result of the design choices that also lead to our tuning-free single-task results.

Q: The novelty is somewhat limited as the modifications to TD-MPC seem to be all from some other existing methods.

A: Respectfully, this can be said about many of the RL papers that have had a profound impact on the field; Rainbow, EfficientZero, RAD, CURL, DrQ, Dreamer all come to mind, as well as the DroQ algorithm brought up by reviewer BqZp. We believe that there is value in pushing the boundaries of what can be done with existing algorithms, through a series of carefully made design choices informed by large-scale experimental results backed by existing literature. Additionally, we would like to reiterate that TD-MPC2 is the first RL algorithm to (1) learn continuous control tasks from diverse domains without any tuning, and (2) scale to large models trained on massively multi-task datasets, neither of which are trivial to achieve.

Please do not hesitate to let us know if you have any additional comments.