Paper: TD-MPC2: Scalable, Robust World Models for Continuous Control (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
**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.

Reviewer comment (Writing):
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.

Author response:
**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!

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

Author response:
**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.

Reviewer comment (Reproducibility):
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?

Author response:
**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).

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

Author response:
**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)

Reviewer comment (Novelty):
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)?

Author response:
**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.

Reviewer comment (Experiments):
"extending TD-MPC2 to discrete action spaces remains an open problem" What are the challenges here?

Author response:
To answer your question: we recognize that bridging continuous and discrete action spaces is an important problem that warrants further research. 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 an extended discussion of this topic to Appendix I of our updated manuscript, which should be useful for readers who want to extend our algorithm to discrete action spaces.

Reviewer comment (Experiments):
Have you considered training the multi-task model online?

Author response:
**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.

Reviewer comment (Experiments):
Have you explored other dataset setups, i.e. different random to expert ratios?

Author response:
**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.

| 80 tasks, 19M params | **No reg.** | **c = 0.001** | **c = 0.01** | **c = 0.1** |
|----------------------|-------------|---------------|--------------|-------------|
| Normalized score     | 56.54       | 58.14         | **62.01**        | 44.13       |


[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)

Reviewer comment (Experiments):
How about the impacts of task diversity, i.e. ablation on number of tasks in the training data?

Author response:
**Q:** How about the impacts of task diversity, i.e. ablation on number of tasks in the training data?

**A:** This is a very interesting question that unfortunately is difficult to give a conclusive answer to with the current task suites available to us. I suspect that the conclusion will be different for our current 80 tasks vs. e.g. a hypothetical future task suite consisting of 8,000 diverse tasks. However, we presently do observe similar “scaling laws” for both our 30-task dataset based on DMControl, and our 80-task dataset based on DMControl + Meta-World, when we control for model size. To give you a slightly more conclusive answer, we train another set of multi-task models (5M, 19M, 48M parameters) on a 15-task subset of our DMControl task suite. Our results are shown [here](https://i.imgur.com/dToYApX.png). We observe that performance still scales with model size on our smaller task suite, but that normalized scores are higher across the board. Notably, the 5M parameter model fares far better on the smaller 15-task dataset than with 30 tasks. This makes intuitive sense, since model capacity remains the same while there is comparably less knowledge to learn. However, we remark that (1) it is difficult to directly compare model performance across task suites since we cannot fully control for task difficulty and “diversity”, and (2) the 48M parameter model trained on 15 tasks has not yet fully converged; we expect that it will continue to improve a bit with more training. The 5M and 19M models are trained to convergence. We will be happy to provide an update once the model has converged.

Reviewer comment (Experiments):
How do the multitask models perform compared to single-task models on each individual task?

Author response:
**Q:** How do the multitask model perform compared to single-task models on each individual tasks?

**A:** To answer this question, we compute a single normalized score as the average score of final single-task checkpoints trained on each of 80 tasks included in our largest multitask dataset. Note that these single-task models are trained **online**, and are thus not limited by the offline RL setting (fixed dataset) that our multitask model is trained with. Additionally, single-task models are trained to convergence and thus represent a (soft) upper bound on the multitask performance. This procedure gives us a score of 91.8 compared to 70.6 for our largest multitask model, without any offline RL techniques. We conjecture that this gap is due to (1) limited similarity between tasks, (2) limited data diversity for offline RL, and (3) insufficient model parameters to accurately model all 80 tasks. However, most works on multi-task model-based RL consider significantly fewer tasks, so we believe that building an algorithm that scales *at all* will be incredibly useful to the community; both as is, and as a basis for future research on large-scale RL.

Reviewer comment (Experiments):
Have you considered leveraging offline RL techniques when learning from the fixed multi-task dataset? Do you think they can help?

Author response:
**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.

| 80 tasks, 19M params | **No reg.** | **c = 0.001** | **c = 0.01** | **c = 0.1** |
|----------------------|-------------|---------------|--------------|-------------|
| Normalized score     | 56.54       | 58.14         | **62.01**        | 44.13       |


[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)

Reviewer comment (Evaluation):
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?

Author response:
**Q:** 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?

**A:** Most contemporary RL algorithms rely on small models or policies (~1M parameters) and do not scale reliably with more model capacity or data. While the research community does not have a conclusive answer for why this is, we believe that the main culprit is non-stationarity and poor signal-to-noise ratio in the learning objective (Q-learning in particular), which are generally absent in supervised learning. As for the specifics of TD-MPC, we do observe that both layer weights and latent states $z$ trend towards large values when not constrained in any way, resulting in very large gradient norms (see Appendix G in our manuscript) and training instabilities. TD-MPC2 mitigates this problem through a series of careful design choices, virtually eliminating training instabilities. While some design choices are unique to the TD-MPC algorithm and MBRL literature (e.g. SimNorm, planning), we do believe that some of our design choices (e.g. LayerNorm, discrete regression) could also improve stability of other RL algorithms such as SAC in the future.

Reviewer comment (Reproducibility):
How did you do hyperparameter search?

Author response:
**Q:** How did you do hyperparameter search?

**A:** We did not do any noteworthy hyperparameter search – most hyperparameters are adopted from TD-MPC without modification. Instead, we set out to build an algorithm that is robust enough to task variations that we do not need to tune hyperparameters in the first place. This is achieved through careful design of the architecture and model objective. We expect this approach to be a more viable solution in the long term, especially as other researchers and practitioners in the community start to apply TD-MPC2 to new task domains that were not considered in our paper.

Reviewer comment (Experiments):
Does TD-MPC2 benefit from model size scaling on single-task setting?

Author response:
**Q:** Does TD-MPC2 benefit from model size scaling on single-task setting?

**A:** We have run additional ablations on the choice of model size in response to your question; our results are shown [here](https://i.imgur.com/t94RKfO.png). Ablations are conducted on the same three tasks as in previous ablations: Dog Run, Humanoid Walk (DMControl), and Pick YCB (ManiSkill2). We observe that different model sizes yield similar data-efficiency and asymptotic performance across model sizes in a single-task setting, with the 5M parameter model being marginally better on average. We attribute this to the fact that (1) small models are generally sufficient for modeling a single task, and (2) larger models may require more gradient updates to fit the data which we do not control for in this experiment. More research is needed to determine whether there is tangible benefit in training large models on small datasets. However, it is notable that the algorithmic robustness of TD-MPC2 does prevent the larger models from diverging during training, which is otherwise commonly observed in TD-MPC and other contemporary RL algorithms (as evidenced e.g. in our multi-task experiments).

Reviewer comment (Evaluation):
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?

Author response:
**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.

Reviewer comment (Novelty):
The novelty is somewhat limited as the modifications to TD-MPC seem to be all from some other existing methods.

Author response:
**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](https://openreview.net/forum?id=Oxh5CstDJU&noteId=xCoQaNp0RE). 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.

Reviewer comment (Theory):
The method is not directly applicable to domains with discrete action spaces.

Author response:
**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.

Reviewer comment (Theory):
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).

Author response:
**Q:** 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).

**A:** Agreed! We have added a new subsection on the intuition behind SimNorm in Appendix H of the updated manuscript. Intuitively, SimNorm can be thought of as a "soft" variant of the vector-of-categoricals approach to representation learning proposed in VQ-VAE. Whereas VQ-VAE represents latent codes using a 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 $\tau \in [0,\infty)$ of the softmax used in SimNorm in our experiments, it is useful to note that it provides a mechanism for interpolating between two extremes. For example, $\tau \rightarrow \infty$ would force all probability mass onto single categories, resulting in the discrete codes (one-hot encodings) of VQ-VAE. The alternative of $\tau = 0$ would result in trivial codes (constant vectors; uniform probability mass) and prohibit propagation of information. SimNorm thus biases representations towards sparsity without enforcing discrete codes or other hard constraints.

Reviewer comment (Experiments):
Why did you decide to use the Mish activation function? Did it provide significant improvements compared to other commonly used activations, e.g., ReLU?

Author response:
**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](https://i.imgur.com/BG5novj.png). 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.

Reviewer comment (Experiments):
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.

Author response:
**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).

Reviewer comment (Theory):
“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?

Author response:
**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.

Reviewer comment (Theory):
Below Eq. (3), it is missing a $\gamma$ in the TD-target $q_t = r_t + \gamma Q$.

Author response:
We have also addressed your other minor comments, such as the missing $\gamma$ below Eq. 3. Thank you for pointing this out! Please do not hesitate to let us know if you have any additional comments.