RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability

Summary

This paper presents a model-based reinforcement learning (RL) method which learns a task-relevant latent-space dynamic model. It learns a latent representation containing task-relevant features by 1) maximising the mutual information between latent states and rewards whilst 2) minimising the mutual information between the latent states and observations. This contrasts approaches which reconstruct the image observations and thus learn task-irrelevant features. They highlight that their method fails to adapt to environment changes which results in a distribution shift to the observations, e.g. changes in lighting. To this end, they present a method for adapting to distribution shifts at test time. They test their method in modified versions of the DeepMind Control Suite and Maniskill with (distracted) video backgrounds.

Strengths

The paper addresses an important problem which is needed for deploying model-based RL algorithms in real-world robotics scenarios. I am not familiar with test-time adaptation for visual model-based RL, however, I do think this is an imprtant problem and the approach seems sensible. The experiments seem to be reproducible from the implementation details in the appendix. However, I have not ran the code to check.

Weaknesses

The paper's biggest weakness is its positioning against prior work. There exist visual model-based RL methods that learn task-relevant latent representations which are not cited in this paper. For example: ``` - Denoised MDPs @inproceedings{tongzhouw2022denoisedmdps, title={Denoised MDPs: Learning World Models Better Than The World Itself}, author={Wang, Tongzhou and Du, Simon S. and Torralba, Antonio and Isola, Phillip and Zhang, Amy and Tian, Yuandong}, booktitle={International Conference on Machine Learning}, organization={PMLR}, year={2022} } - TD-MPC @inproceedings{Hansen2022tdmpc, title={Temporal Difference Learning for Model Predictive Control}, author={Nicklas Hansen and Xiaolong Wang and Hao Su}, booktitle={ICML}, year={2022} } - ALM @inproceedings{ghugare2023simplifying, title={Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective}, booktitle = {Advances in Neural Information Processing Systems}, author={Raj Ghugare and Homanga Bharadhwaj and Benjamin Eysenbach and Sergey Levine and Ruslan Salakhutdinov}, year={2022}, } ``` Given the similarity of the work, I think it is important to not only discuss them in the related work but to also compare to them in the experiments section. Especially as it is not obvious to me that this is a better approach to learning a task-relevant latent space. I am interested to see how this method compares but I do not think this can be published without such comparisons. The training curves should show the converged values with dashes lines. In many figures it is not possible to see where algorithms converge. For example, dreamer in Walker Stand/Walker Walk/Walker Run in Figure 5. There are a lot of minor corrections and typos. Minor corrections and typos: - How many seeds were used for each experiment? - Do the shaded regions on figures show variance or std or something else? - Line 22 - operator should be operate - Line 22 - the sentence doesn't read properly - Perhaps you could replace "and learning to operator the same arm amidst of plentiful dynamic distractors" with "and learning to operate the same arm amidst dynamic distractors". - Eq 4.1 - $I_{p,h}(z_{1:T};x_{1:T} \mid a_{1:T})$ should be $I_{p,h}(z_{1:T}; o_{1:T} \mid a_{1:T})$ - Figures 4/7/8/9 captions should have full stops. - Figures 1/2/3 are never referenced. - Figure 1 is not very informative so I am unsure why it is included? - Line 113 - brackets in expectation - Line 123/132 - What does the $h$ subscript in $s_{h}$ represent? You introduce $s_{t}$ where the subscript is time $t$ but then it looks like you're indexing at the encoder? - Line 84-87 - This sentence doesn't read properly. - Line 221 - $\bar{P} : \mathcal{X} \times \mathcal{A} \rightarrow \Delta(\mathcal{X})$ should have $\times$ not $\rightarrow$? - Line 324 - "Does representations" should be "Do representations" - Line 360 - Incorrect use of semicolon - Line 366 - "Can we apply our algorithm could be used in a continual learning setup?" doesn't make sense. - Eq. 4.1/4.5/4.6 are missing full stops. - Be consistent with Eq. vs Equation, vs Eqn. vs Eqs. - Line 197 - Eqn should be Eq. for consistency - Line 197 - Eqn should be Eq. for consistency - Line 177 - (4.3) should be Eq. 4.3? - Line 167 - (4.1) should be Eq. 4.1? - Line 261 - (4.5) should be Eq. 4.5? - Line 265 - (4.5) should be Eq. 4.5? - Similarly, with Fig vs Fig. vs Figure. - Parsimonious/eschew seem unnecessarily complicated words. I had to look them up... I recommend simpler words to make the paper more accessible. I asked around and 60% of people would need to look up both words.

Questions

1. Are the authors aware of the literature mentioned above? Do you think you need to compare to these methods? 2. $q$ is usually used for approximate posteriors and $p$ is used for priors. Is there a reason you have swapped these around?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors mention that their method is not task agnostic and I think this is sufficient.

Summary

This paper presents an approach that aims to enhance visual model-based RL by learning state representations that filter out irrelevant information while retaining task-specific information. The proposed method achieves this by maximizing the mutual information (MI) between the current representation and future rewards, while simultaneously minimizing the MI between the representation and observations. In order to facilitate the transfer of the learned representation to new environments with distribution shifts, this paper introduces a support constraint regularization technique. This technique ensures that a specific set of "calibration" states, paired across training and testing domains, share the same encoding. The empirical results obtained demonstrate the effectiveness of the proposed method.

Strengths

1. This paper addresses an important open problem in RL and is very well motivated. 2. The logic flow of the paper is good, and the presentation is clear and coherent making it a pleasant read. 3. Test-time adaptation is a reasonable and interesting topic in this field. 4. Empirical studies show the performance gain of the proposed method.

Weaknesses

1. The MDP setting in this paper bears strong resemblance to the exogenous distractors setting discussed in [1, 2, 3], as well as the denoised MDP setting introduced in [4]. It is important to include a discussion of these works in the related work section. Additionally, it is worth noting that the denoised MDP approach presented in [4] also offers a model-based method for learning state representations, which should be considered as a baseline method in the experiments conducted. 2. Utilizing the recurrent state-space model (RSSM) to parameterize both p and q is a reasonable choice. However, understanding Assumption A.1 is challenging. While $z_t$ can represent past latents and past actions, it raises questions about whether $z_t$ can effectively incorporate the information from $a_t$, given that the RSSM does not explicitly model their corresponding relationship. Typos: 1. Equation 4.1 is a bit strange. Seems the equation in Appendix A.1 is more reasonable. 2. Line 192 "parametrize" -> "parameterize" [1]: Yonathan Efroni, Dylan J. Foster, Dipendra Misra, Akshay Krishnamurthy, John Langford: Sample-Efficient Reinforcement Learning in the Presence of Exogenous Information. COLT 2022: 5062-5127 [2]: Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, John Langford: Provable RL with Exogenous Distractors via Multistep Inverse Dynamics. CoRR abs/2110.08847 (2021) [3]: George Trimponias, Thomas G. Dietterich: Reinforcement Learning with Exogenous States and Rewards. CoRR abs/2303.12957 (2023) [4]: Tongzhou Wang, Simon S. Du, Antonio Torralba, Phillip Isola, Amy Zhang, Yuandong Tian: Denoised MDPs: Learning World Models Better Than the World Itself. ICML 2022: 22591-22612

Questions

Most important questions are presented in above. Furthermore, what are the theoretical advantages of RePo over bisimulation-based methods? It appears that the proposed method necessitates test-time calibration to adapt to changes in distractors, whereas bisimulation-based methods can transfer their representations directly without any additional modifications. Notably, certain bisimulation approaches [1, 2] have demonstrated impressive performance not only in standard settings but also in distractor settings. [1]: Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, Mark Rowland: MICo: Improved representations via sampling-based state similarity for Markov decision processes. NeurIPS 2021: 30113-30126 [2]: Hongyu Zang, Xin Li, Mingzhong Wang: SimSR: Simple Distance-Based State Representations for Deep Reinforcement Learning. AAAI 2022: 8997-9005

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The method necessitates test-time adaptation instead of direct transfer. This requirement for additional calibration time may increase computational complexity during the evaluation stage.

Summary

The paper has two major contributions. The first one is the introduction of a model-based RL method that learns latent representations of visual RL environments that are resilient to spurious variations. The second contribution is a test-time alignment method that makes the latent representations invariant distribution shifts, thereby enabling the quick adaptation to diverse environments.

Strengths

The paper proposes a novel method to address a very important problem in reinforcement learning. The RePo method is simple, elegant, and well-motivated and the discussion around eq. 4.3 makes it easy to intuitively understand it. The presentation is something to compliment the authors for. The paper is very well written and easy to follow, despite the mathematical complexity of the underlying concepts. The relevant work section is concise, and the authors clearly relate their work with previous papers.

Weaknesses

The only minor weakness that I can note is that the authors did not discuss any subtle limitations of their method, which I am sure it does have. Adding such a discussion would make the paper more credible.

Questions

Should I assume that the method name RePo comes from Resilient Posterior? If so, please clarify that the first time you use it, e.g., line 40. Are there any constraints on the environment distribution shift? On the experiments you tested on natural video distractors. What if it was something extreme like white noise?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

I do not foresee any negative societal impact. Limitations were not discussed. The authors should add a small paragraph with the limitations.

Summary

This work proposes a modification to Dreamer [1], wherein the authors discard the reconstructions objective, and replace it with an objective that maximizes the information contained in the latent representation about the reward, while minimizing the information about the original image. This pushes the model to learn representations that are agnostic to 'spurious variations' - changes in the input that do not affect the reward. The authors compare the method to Dreamer [1], TIA [2] and bisimulation [3], and show that RePo deals with distracting background on DMC very well compared to other methods. Authors also show that the method can generalize from static training data to test-time data with distractors if an additional test-time adaptation stage is introduced. [1] Hafner D, Lillicrap T, Ba J and Norouzi M 2019 Dream to Control: Learning Behaviors by Latent Imagination _arXiv_ [2] Fu X, Yang G, Agrawal P and Jaakkola T 2021 Learning Task Informed Abstractions [3] Zhang A, McAllister R, Calandra R, Gal Y and Levine S 2021 Learning Invariant Representations for Reinforcement Learning without Reconstruction

Strengths

- The paper is clearly written - The proposed method is simple - The method performs better than baselines on DMC with distractions. - The method is able to work with distractions introduced at test-time if additional test-time adaptation stage is used

Weaknesses

- This method, as authors have mentioned, is no longer task-agnostic. - Possibly, if there's some information about the reward in the image (like a score), the model will focus on that instead

Questions

- Can the authors clarify what would happen if the image itself contains information about the score. - As a baseline, would be useful to see how Dreamer performs with distribution shfit. **typos** - 22 - learning to operator - > learning to operate - 120 - to a an -> to an - 135 - effect -> affect - 366 - could be used should be removed

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are adequately addressed in the submission.

Feedback

I appreciate the author's rebuttal, which includes some new empirical results. Yet, I still hope to see more comparisons between the proposed method and bisimulation approaches theoretically or empirically. I will suggest an acceptance but keep the score as-is.