We are grateful for the reviewer’s constructive suggestions and provide our responses below.

Q1. Based on your case studies, your method does not seem to change the episode return that much, except for a few cases like dog walk, dog trot and humanoid walk. In your Appendix, you provide a rough explanation of why that may be. Looking at Figure 7 and 8, it appears that loss, and consequently the gradient, explode (in TD-MPC), however, in RL, gradient clipping is used to tackle this issue. When you compared your method to TD-MPC, did you employ gradient clipping for it or not? It does not appear to be a fair comparison if you didn't, and perhaps that is why your method did not do significantly better in other case studies; as loss did not explode.

A1. Yes, we do use gradient clipping for TD-MPC in our experiments. Specifically, the setting file (yaml file) used in our experiments includes a parameter labeled 'clip_grad_norm,' which specifies the gradient clipping threshold. In Figure 8, we present the gradient values before clipping the gradient for transparency.

Although gradient clipping is applied, we observe that the error values continue to grow and eventually explode, leading to degraded performance. This issue highlights the inherent instability in TD-MPC. Moreover, while BS-MPC achieves similar performance to TD-MPC in many cases, it demonstrates significantly superior robustness in noisy environments, as shown in Figure 5.

Finally, we want to emphasize that the primary objective of BS-MPC is to provide a more stable learning process and performance compared to TD-MPC, rather than merely outperforming it across all metrics. This is one of the open problems in TD-MPC as shown in Figure 1.

Q2. I suggest you revise the experiments' section, and run all case studies on TD-MPC2 rather than TD-MPC. I realize it is touched upon in appendix D, however since TD-MPC2 is the updated version, I suspect it would make for a more fair comparison. Moreover, adding a thorough comparison would certainly present your method better; between training time, sample complexity, number of parameters used, and hyper parameter tuning and different configurations; it will strengthen your case if you could show how it might fail. I would also like to know the rationale behind using TD-MPC in the main body and mentioning TD-MPC2 in the appendix. To the best of my knowledge, TD-MPC2 can have many parameters, since it can be used on different domains. Thus, it is not an apple to apple comparison, unless it is specifically mentioned in the paper.

A2. Thank you for the suggestion. We agree with the reviewer's idea and are working on this. Since TD-MPC2 does not provide some of the results in image-based tasks, we are currently running their official code.

There are three primary reasons why we chose TD-MPC over TD-MPC2 as our baseline:

• Parameter Count: As the reviewer highlighted, TD-MPC2 requires significantly more parameters compared to both BS-MPC and TD-MPC. While BS-MPC and TD-MPC have approximately 1M parameters, TD-MPC2 requires 5M parameters to achieve comparable performance. This increased complexity substantially raises computational costs, which is a concern given that TD-MPC already incurs high computational overhead.

• Discrete Regression: TD-MPC2 introduces discrete regression to handle variations in reward magnitude and Q-values. While this eliminates the need for certain parameter tuning, it also necessitates discretizing these values, further increasing computational costs. This trade-off made TD-MPC2 less aligned with our goal of balancing performance with computational efficiency.

• Network Complexity: TD-MPC2 employs additional layers and normalization mechanisms, such as SimNorm normalization, which biases the latent state representation towards sparsity and maintains a small ℓ2-norm. While these techniques have shown empirical success, they lack a robust theoretical foundation to justify their inclusion. In contrast, BS-MPC intentionally adopts the same simple network architecture as TD-MPC, augmenting it with a bisimulation metric to regularize the latent space projection. This simplicity aligns with our design philosophy while having preferrable theoretical supports.

Q3. Is there any theoretical results on why your method requires less parameters, and converges faster? or is it mainly based off of experiments? Since Theorem 3 only offers an upper bound for expected cumulative rewards for the optimal policy.

A3. BS-MPC has the exact same number of parameters as TD-MPC as both of them use the exact same architecture. However, since our method has a different computational structure, BS-MPC can enjoy parallel computation when computing and optimizing its cost. This is attributed to the faster computational time. We add Appendix C.2 for more details about the computational structure change.

Q4- I suggest a more rigorous approach for robustness, such as comparing the Lipschitz constant of your controller to TD-MPC's.

A4. We explored the possibility of computing the Lipschitz constant for both our approach and TD-MPC. However, due to the complexity introduced by the use of neural networks in both methods, deriving a rigorous mathematical analysis proved to be challenging.

To address this limitation, we have added an intuitive explanation in Appendix A.2 to provide insight into why BS-MPC demonstrates improved robustness compared to TD-MPC. This discussion complements the response provided in Q5 and offers a qualitative understanding of the underlying factors.

Q5. I believe if you can theoretically confirm in which case studies your method is going to perform better than TD-MPC, it will strengthen your paper significantly.

A5. While we have not yet established a rigorous theoretical framework to precisely predict such cases, we have included an additional analysis in Appendix A.2. This section explores the factors that contribute to BS-MPC's improved performance in certain scenarios.

We kindly invite the reviewer to review this new addition. In summary, BS-MPC demonstrates superior performance when TD-MPC struggles to train its encoder to effectively capture the essential information in the original state $s$. This provides a partial explanation for BS-MPC's higher scores in noisy environments (Figure 5), where the TD-MPC encoder is more likely to fail to learn an accurate representation.

Q6. As the authors have mentioned, the hyperparameters play a huge role, and one wonders how much time is needed to tune these parameters.

A6. In our experiment, we use the grid search for hyperparameter optimization. Since there are only 6 candidates, it does not take much time to adjust the parameter. We also find smaller values are suitable for more complex environments (e.g. dogs and humanoids) from our experimental results, and this insight helps us to narrow down the range for searching parameters.

Q7. Contribution weakness

A7. We wrote our thoughts in general reposense.

We have conducted the additional experiments for TD-MPC2 and updated the PDF with the new results. Following the reviewer’s suggestion, we have moved Appendix D to Section 5 to better highlight the experimental results, specifically by including the TD-MPC2 results in the main body of the paper.

Thank you for adding this theorem.

I believe this theorem was added in haste, I originally asked the authors that if you are claiming robustness, you should do Lipschitz constant analysis. To which authors responded that it is intractable/computationally expensive, without knowing the Lipschitz constants, how can you claim you are more robust?

Moreover, equation 20 is misleading, Lipschitz continuity also applies to your encoding, and you also suffer from noise, if you employ global Lipschitz continuity. I personally think this theorem would hurt your paper rather than helping it.

We are grateful for the reviewer’s constructive suggestions and provide our responses below.

Q1: “We assume that the learned policy in BS-MPC continuously improves throughout training and eventually converges to the optimal policy π∗, which supports Theorem 1.” This seems to be a very strong assumption. For example, by looking at the training curve, the return does not improve monotonically, and we have no information about if the learned policy is converging to the optimal policy. How do you explain such a strong assumption? Is it possible to remove it for the theoretical results?

A1: Yes! We believe we can loosen this assumption by adapting the theory from the Robust Bisimulation Metric[1]. In this paper, they enable the bisimulation metric to have the same guarantee with some more realistic assumptions. In our paper, we choose the optimal policy assumption \pi^* because we would like to make the algorithm and proof simpler and straightforward. However, as the reviewer suggested, we believe our method can be more adaptable with realistic assumptions by using their robust bisimulation metric.

[1]. Kemertas, M., & Aumentado-Armstrong, T. (2021). "Towards Robust Bisimulation Metric Learning." https://arxiv.org/pdf/2006.10742

Q2: In Fig. 4, why do all non-MPC based methods only have results till 10M steps?

A2: In image-based experiments (fig 4), we cite our baseline results from their official papers. I assume the computational cost is very expensive and authors just run their code until 1M. We run BS-MPC and TD-MPC until 3M to show the stability of the proposed methods. We also choose the computational steps based on TD-MPC2 paper [2].

[2]Hansen, N., Su, H., & Wang, X. (2024). TD-MPC2: Scalable, Robust World Models for Continuous Control. https://arxiv.org/abs/2310.16828

Q3. There are several typos in the paper. “In BS-MPC, the latent dynamics are modeled using an MLP. We also model the latent dynamics model with an MLP” I believe BS-MPC should be TD-MPC. “we sample M action sets from Gaussian distribution N (μ0, σ0) based on the initial meanμ0 and standard deviation σ0” Missing spacing between mean and \mu_0

A3. We sincerely thank the reviewer for identifying these typos and providing detailed feedback. We have carefully revised the manuscript to address these issues. We corrected the sentence and simplified the phrasing of the affected sentences to enhance clarity. A revised version of the paper has been uploaded for your review.

Q4. Contribution weakness

A4. We wrote our thoughts in general reposense.

We are grateful for the reviewer’s constructive suggestions and provide our responses below.

Q1. When you say BS-MPC improves computation efficiency, what does it mean? Is it compared to TD-MPC? It is surprising to me because BS-MPC has one additional loss term compared to TD-MPC and why is BS-MPC faster to run?

A1. Yes, BS-MPC achieves faster computational speed than TD-MPC. The primary reason for the speed-up is the structure change in the computation. In BS-MPC, we encode all states s_{t:t+H} at once (z_{t:t+H} = s_{t:t+H}), and hence eliminate the sequential computation. However, TD-MPC needs to compute its latent state z_t by using the previous step latent state z_{t-1}. Therefore, its calculation flow includes sequential computation, which becomes the bottleneck of TD-MPC. I added a new Appendix in C.2 to give more details about this explanation by using pseudo-code. Also, please refer to Figure 2 for the computational flow difference. (Adding bisimulation term does not contribute to the speed-up, but it helps BS-MPC to be robust against the noise in the original state space S)

Q2. With the above question, I'd like to know the latency overhead of BS loss term in the training.

A2. We estimate the calculation cost in Equation (8). Let d denote the dimension of latent space, the Gaussian Wasserstein Distance computational cost is O(d^3). Therefore, the computational cost of the bisimulation metric for Batch size B data becomes O(B d^3).

Q3. Contribution weakness

A3. We wrote our thoughts in general reposense.

We are grateful for the reviewer’s constructive suggestions and provide our responses below.

Q1. Could the authors expand on the sensitivity of BS-MPC to the parameter c4 and potential ways to reduce this dependency?

A1. Thank you for pointing this out. We also think this tuning term is the biggest limitation of the proposed method. In our experiments, we employed grid search for this hyperparameter, which involved six candidates ($10^{-8}$, $0.0001$, $0.001$, $0.01$, $0.1$, $0.5$) making the process computationally manageable. Our findings indicate that smaller values are suitable for high-dimensional and complex tasks, while larger values work better for simpler tasks. Additionally, we believe incorporating techniques such as reward discretization and normalization layers from TD-MPC2 could simplify or even eliminate the need for this tuning in the future.

Q2. How does BS-MPC perform in scenarios with dynamic backgrounds that align with the movement instead of pure noise?

A2. Regrettably, I did not do the experiment in an environment where scenarios with dynamic backgrounds align with the movement. We are setting up the Carla simulator (vehicle simulator) to test the proposed method to see the performance now. We will update the paper as soon as we get the results. However, we believe that BS-MPC is more likely to achieve better performance compared to other methods because the deep bisimulation metric is shown to have good performance in the Carla simulator (dynamic backgrounds that align with the movement)[1].

Q3. Are there additional computational costs associated with bisimulation metric loss, especially in high-dimensional latent spaces?

A3. We estimate the calculation cost in Equation (8). Let d denote the dimension of latent space, the Gaussian Wasserstein Distance computational cost is O(d^3). Therefore, the computational cost of the bisimulation metric for Batch size B data becomes O(B d^3). This is the additional computational cost at each time step compared with TD-MPC. However, since BS-MPC can employ parallel computations, BS-MPC still archives faster calculation time even with this additional computational cost. We add Appendix C.2 to our paper for more details about the parallel computational explanations.

[1]. Zhang, A., McAllister, R., Calandra, R., Gal, Y., & Levine, S. (2021). Learning Invariant Representations for Reinforcement Learning without Reconstruction. https://arxiv.org/abs/2006.10742

As the deadline for the rebuttal is in one week, we would like to kindly remind the reviewer to review our response to your comments.

If there are any clarifications or additional details we can provide to assist with your review, please do not hesitate to let us know.

Thank you for your time and valuable feedback.

We appreciate your feedback and would like to inform you that we have completed the experiments in scenarios with dynamic backgrounds aligned with movement. Due to the computational demands of the CARLA environment, we were able to run the experiment with only a single seed. The results indicate that BS-MPC outperforms TD-MPC, even in environments where the background image moves along with the car. Additionally, while the BS-MPC loss converges to zero, the TD-MPC loss diverges as training progresses.

We have uploaded the results to the supplemental material. The reviewer can find the relevant figures in the figures folder, specifically named "carla_reward.png" and "carla_consistency_loss.png."