Bisimulation Metric for Model Predictive Control

Authors: Yutaka Shimizu, Masayoshi Tomizuka
Primary area: reinforcement learning
Keywords: Reinforcement Learning, Model-based reinforcement learning, optimal control, MPC

Abstract:
Model-based reinforcement learning (MBRL) has shown promise for improving sample efficiency and decision-making in complex environments. However, existing methods face challenges in training stability, robustness to noise, and computational efficiency. In this paper, we propose Bisimulation Metric for Model Predictive Control (BS-MPC), a novel approach that incorporates bisimulation metric loss in its objective function to directly optimize the encoder. This optimization enables the learned encoder to extract intrinsic information from the original state space while discarding irrelevant details. BS-MPC improves training stability, robustness against input noise, and computational efficiency by reducing training time. We evaluate BS-MPC on both continuous control and image-based tasks from the DeepMind Control Suite, demonstrating superior performance and robustness compared to state-of-the-art baseline methods.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730220324964

thread_type:
official_review

<thread_object>
creator: reviewer_sest

<content>
summary:
This paper is concerned with a new, model-based reinforcement learning method, which utilizes bi-simulation metric. This formulation helps with the stability of training and helps with the robustness of the controller. General idea is that they seek to find states that "behave" similarly, and intuition behind it is that one can use similar control input for similar states, which simplifies the controller and make it more interpretable. Authors learn an encoder which maps states of the environment to another domain, in which similar states are identified, and are mapped to the same representation (roughly speaking). Then, this representation is utilized to train a controller. Novelty of this work is to add the encoder loss directly into the training procedure.

strengths:
Paper is well written, all ideas are clearly explained. Overall, paper is mathematically rigorous. Authors do a good job in walking the reader through the preliminaries, highlighting distinctions, and presenting their work.

Furthermore, paper presents more than 20 case studies, which helps immensely in comparing their performance to the state of the art methods.

weaknesses:
Improvements and contributions seem incremental, and overall not that beneficial according to the case studies (Figure 6), one only sees improvements in few case studies (such as humanoid walk, dog walk and trot), and basically identical to TD-MPC in others (such as humanoid stand, pendulum, cheetah). 
The only major contribution is adding the bi-simulation metric loss to the loss function, the other two contribution naturally follow from this addition. 

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

questions:
1- 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*.  

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


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

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

5- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730409484773

thread_type:
official_review

<thread_object>
creator: reviewer_hjrq

<content>
summary:
This paper considers model based reinforcement learning and proposes bisimulation metric to improve over temporal differential MPC method. The authors show theoretical analysis of the expected cumulative rewards in the latent space, and empirically demonstrate enhancement over TD-MPC and other baselines on several continuous control tasks.

strengths:
The paper is clearly written and well presented. The proposed bisimulation metric seems to work well on the experiments considered, compared to TD-MPC and other baselines. The supplementary sections are comprehensive.

weaknesses:
The novelty of the paper seems ambiguous. It seems that both on-policy bisimulation and TD-MPC methods are well studied for model based RL, and the authors plug bisimulation into TD-MPC. 

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

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

In Fig. 4, why do all non-MPC based methods only have results till 10M steps?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692433735

thread_type:
official_review

<thread_object>
creator: reviewer_cxnw

<content>
summary:
This paper proposes BS-MPC, a model-based reinforcement learning approach that introduces bisimulation metrics (loss) on top of TD-MPC. Compare to TD-MPC, BS-MPC has an explicit encoder loss term, the adaptation of bisimulation metric, and parallelizing the BS loss. The authors found that their approach can improve training stability, robustness against input noise, and computation efficiency, which is validated on a set of simulation environments.

strengths:
The paper is well-written and easy to follow. The overall presentation is good. The approach is sound and makes sense to the reviewer. The experimental results look promising, compared to TD-MPC.

weaknesses:
However, the major weakness is its novelty. 
1. The whole framework is based on TD-MPC. The difference is the authors introduce the Bisimulation metric and its corresponding loss design, which are from the existing literature, as stated in the paper. 
2. It is also a common way to introduce additional regularization loss terms for the encoder of model-based RL. 
3. The theoretical analysis mainly borrows from the existing work and does not have any major significant result. It would be great if the authors could provide "Under the BS loss training error, what's the performance gap between the final converged policy by their approach and the ideally optimal policy", and "Theoretically, how much performance gain could their approach improve, compared to TD-MPC."

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

With the above question, I'd like to know the latency overhead of BS loss term in the training.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730907684300

thread_type:
official_review

<thread_object>
creator: reviewer_w1co

<content>
summary:
This paper presents a new method for model-based reinforcement learning (MBRL) called BS-MPC. The key innovation lies in incorporating a bisimulation metric loss into the objective function to improve encoder stability, robustness to noise, and computational efficiency. By using the bisimulation metric, BS-MPC aims to ensure behavioral equivalence in the latent space, maintaining key characteristics of the original state space. The method is benchmarked against the Temporal Difference Model Predictive Control (TD-MPC) and other model-free and model-based methods on various tasks, showing superior stability and resilience to noise.

strengths:
The paper provides a new perspective by integrating the bisimulation metric to address known challenges in MBRL, particularly around stability and robustness to noise. The experimental results demonstrate how BS-MPC performs well in both state-based and image-based tasks, showing increased resilience to noise and achieving faster training times due to parallel computation. The theoretical analysis adds depth by bounding cumulative rewards in the learned latent space, suggesting that BS-MPC retains meaningful state information effectively.

weaknesses:
While the theoretical foundations are thorough, certain explanations, particularly on encoder stability and noise resilience, could be made clearer to broaden accessibility. The parameters require extensive tuning, which may be impractical for real-world applications lacking automated parameter selection. Additionally, the approach to introducing perturbations, particularly with visual distractions, doesn’t seem entirely effective. It would be beneficial to test perturbations that are more representative of realistic environmental changes, which could better showcase BS-MPC’s resilience.

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

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

Are there additional computational costs associated with bisimulation metric loss, especially in high-dimensional latent spaces?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731812292327

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Genenral Response

comment:
First and foremost, I would like to sincerely thank all reviewers for their insightful comments and constructive feedback. They help me critically analyze and enhance my approach. I have replied to each reviewer about their questions and concerns. Here I would like to summarize the changes I made in my manuscript and respond to common questions and concerns raised by multiple reviewers.

# Summary of Paper Revisions
- A new subsection has been added to Appendix A, explaining why BS-MPC outperforms TD-MPC, particularly in noisy environments (as suggested by [seSt](https://openreview.net/forum?id=F07ic7huE3&noteId=N3v9RuT4mp).
- A new subsection has been added to Appendix C (C.2), detailing the computational structure differences between BS-MPC and TD-MPC using pseudo-code. It gives detailed explanations why BS-MPC has faster computational time than TD-MPC.
- Sentences highlighted by [hjrQ](https://openreview.net/forum?id=F07ic7huE3&noteId=21oVCAFWxF) have been corrected and revised to improve clarity and precision.

# Responses to Common Questions
Q1. BS-MPC appears sensitive to hyperparameters. Is there a way to mitigate this issue?

A1. Indeed, tuning the new weight term $c_4$ is essential for achieving optimal performance. 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. The contribution appears incremental, as BS-MPC mainly integrates the bisimulation metric into TD-MPC.

A2. We agree that BS-MPC builds upon TD-MPC, and inherits many components from TD-MPC. However, we believe that incremental yet meaningful advancements are essential to driving progress in reinforcement learning and beyond. Many significant contributions in the field, such as TD-MPC2, Dreamer, and Rainbow, have advanced the state of the art by refining and extending existing methods. In this context, we view BS-MPC as a step forward, offering two key contributions:

- Integration of the Bisimulation Metric: While the bisimulation metric has been utilized in some model-based reinforcement learning approaches, this is the first work to incorporate it into a planning-based method. This simple addition provides a novel perspective on how to improve planning robustness and performance. BS-MPC is also the first paper that provides theoretical support for the estimated value.

- Modification of the Computational Flow: BS-MPC revisits and modifies the computational flow used in TD-MPC for calculating costs. Although the values being computed are similar, our changes significantly improve computational efficiency and stabilize the training process compared to TD-MPC.

While these changes may appear incremental, we believe that minimal modifications to existing simple algorithms can play a crucial role in advancing the field and addressing important challenges within the community. Maintaining the simplicity of TD-MPC while addressing three of its open problems demonstrates the importance of targeted refinements to simple yet powerful algorithms. We hope this contribution will inspire further exploration in the community.

---

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732083514335

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Paper Update

comment:
As suggested by the reviewer [seSt](https://openreview.net/forum?id=F07ic7huE3&noteId=N3v9RuT4mp), we have conducted the additional experiments for TD-MPC2 and updated the PDF with the new results. In the updated PDF, 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.

We have finished responding to all of the comments we got from the reviewers, and we are looking forward to having discussions with each reviewer.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732527982159

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Paper Update2

comment:
We added a new experimental result in Carla simulator suggested by reviewer [w1co](https://openreview.net/forum?id=F07ic7huE3&noteId=ggQrsRBRSf). 

## Summary of the new experimental result
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.

## Where to find the results
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."
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents Bisimulation Metric for Model Predictive Control (BS-MPC), a technique for model-based reinforcement learning that uses a bi-simulation metric to improve encoder stability, robustness to noise, and computational efficiency. The benefits of BS-MPC are demonstrated via experiments on both continuous control and image-based tasks from the DeepMind Control Suite.