Simplifying Latent Dynamics with Softly State-Invariant World Models

We would like to thank the reviewer for the thoughtful review and we are glad that the reviewer enjoyed reading our paper. The reviewer raised several important points regarding the effect of our regularization and how it compares to other regularization methods in RL. In response. we show that our regularization indeed has a principled effect (see Supporting Figure D), that PLSM shows considerable improvements over another robust baseline, RePO, in the Distracting DeepMind Control Suite (see Supporting Figure A), and lastly that regularizing a separate query representation $h_t$ provides a substantial improvement over regularizing $z_t$ (see Supporting Figure F). We believe that addressing the concerns raised by the reviewer has helped us improve our paper substantially. We provide detailed responses to each question below:

I'm not sure how simply minimizing the norm $h_t$ of leads to minimal representations

We thank the reviewer (along with reviewer 2) for bringing this criticism to our attention. We investigated whether PLSM learns higher norm weight matrices to compensate for the low norm $h_t$. First, we found that the norm of PLSM’s dynamics weights were not substantially larger than those of the baseline model, whereas the norm of $h_t$ often was several orders of magnitude smaller than $z_t$ (see Supporting Figure D). Additionally, we show in Figure 5 in the paper that the query representations carry less information about the generative factors than the latent representations. Lastly, even though our model does not compensate with large weights, we propose two methods for preventing this. i) by penalizing the L2 norm of the dynamics weights, and ii) by only using the top-k largest features in the query representation. We trained these two approaches on a subset of the datasets and found that they perform comparably, and sometimes better than the original PLSM formulation (see Supporting Figure F).

I'm also not fully convinced that we need an additional variable $h_t$. Why can't we just impose a bottleneck on $z_t$?

To address this point we evaluated a version that imposes an information bottleneck directly on $z_t$, instead of the query representation $h_t$. We observed a general reduction in performance on the three datasets tested on (See Supporting Figure F), suggesting that $h_t$ is necessary. The advantage of regularizing the query representation is that we only encourage the model to compress away aspects of the environment whose dynamics are unpredictable. In an environment with easily predictable dynamics, our method won’t necessarily lead to any loss of information, whereas imposing a bottleneck on z will invariably do that.

one of the emphasized contributions seems to be the fact that PLSM makes difference prediction ($z_{t+1} - z_t$) instead of latent prediction ($z_{t+1}$), I believe this is an extremely common trick in model-based RL

We agree with the reviewer that this is common practice in model-based RL. We did not intend to make this sound like a key contribution in our paper, but we emphasized this modeling choice because our regularization does not make much sense without it. We have updated the introduction to reflect this:

"As is common practice in many dynamics model architectures [1, 2], we consider the case where the model predicts the next latent $\hat{z}_{t+1}$ state by predicting the difference $\Delta$, or the change, between the current and future latent state, given an action $a_t$."

it's unclear how PLSM compares to other regularization methods (e.g., RePo, other regularization techniques, etc.)

We followed the reviewer’s suggestion and we evaluated PLSM against another robust model-based RL algorithm, RePO [3], in five Distracting DMC environments [4]. In these environments, the background was replaced with a distracting video from the DAVIS 2017 dataset. RePO uses a GRU-based dynamics model, akin to Dreamer, but without pixel reconstruction. We incorporated our regularization into RePO by learning a compressed query representation, and passed this alongside the action to the GRU, making the GRU dynamics softly state-invariant. We observe considerable performance gains in the more challenging environments like hopper-stand, finger-spin and walker-run, suggesting that our regularization can be beneficial together with methods designed for robustness (see Supporting Figure A, scores are averaged across 5 seeds, with standard error shown as shaded region).

References:

[1] PILCO: A Model-Based and Data-Efficient Approach to Policy Search, Deisenroth et al., ICML 2011

[2] Data efficient Reinforcement Learning With Self-Predictive Representations, Schwarzer et al., ICLR 2021

[3] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023

[4] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021

We thank the reviewer for bringing these points to our attention. We believe that addressing them has substantially improved our paper, and we are grateful that the reviewer raised the score accordingly. We will also revise Figure 1 to clarify that the dynamics model predicts $\Delta$. Thank you again!

We thank the reviewer for the constructive feedback and for the helpful suggestions. To address the concerns raised in the review, we have added another set of experiments showing that PLSM improves robustness in visual control tasks with distracting background videos (see Supporting Figure A), as well as improving our results on the Atari control suite by tuning the regularization strength parameter (see Supporting Figure C). We agree with the reviewer that mutual information (MI) based regularizers themselves are not novel in RL. However, our particular MI method is novel, and aims at regularizing a part of latent dynamics models that to our knowledge has not been explored before. We extend the discussion on other MI based methods in RL in the related works section, including the papers cited above. In sum, we believe that addressing the points the reviewer outlined has made our paper better. We give detailed answers to the reviewer’s questions below:

The contribution is limited since the MI regularization for RL representation has been adopted in several previous works, including [1-3].

We have extended the related works section with the following paragraph:

"Mutual information minimization methods have been effectively applied for representation learning in RL too. Previous works have focused on using information bottlenecks to improve generalization [1, 2, 3], robustness [4] and exploration [5, 6]. Our information bottleneck differs from these by constraining the effect the latent state can have on the residual term in the latent dynamics."

Since the MI between representation and state is constrained, should the algorithm exhibit stronger robustness in various domains?

We evaluated the robustness properties of PLSM across five Distracting DeepMind Control Suite environments (Distracting DMC). Distracting DMC consists of challenging visual control tasks where the background has been replaced with videos that are irrelevant for solving the task. We used RePO [7] as a baseline, and implemented the PLSM regularization by minimizing how much the latent state influences the dynamics of the GRU in RePOs world model, while still conditioning on actions. We observe considerable improvements in more challenging tasks (hopper-stand, finger-turn and walker-run) and smaller improvements in walker-walk (see Supporting Figure A). This suggests that PLSM really shines in cases where the agent needs to learn to compress away irrelevant information to generalize.

The empirical results presented in this paper do not show significant performance improvements.

We have extended our evaluation of PLSM in Distracting DMC, and the standard DeepMind Control Suite. We also optimized our regularization strength in Atari and found that $\beta=5$ gave the best performance overall, improving human normalized score in Atari from 61.5 % for SPR, to 67% with PLSM (see Supporting Figure C). Note that SPR and DrQ interquartile ranges are lower since they are based on 100 seeds. In DMC, we see improvements in many domains, especially in challenging tasks like humanoid-walk (see Supporting Figure B). In the Distracting DMC tasks, we see larger performance gains as well.

References:

[1] Generalization and Regularization in DQN, Farebrother et al., ArXiv 2018

[2] DeepMDP: Learning Continuous Latent Space Models for Representation Learning

[3] Learning Robust Representations via Multi-View Information Bottleneck, Federici et al., ICLR 2020

[4] Robust Predictable Control, Eysenbach et al., NeurIPS 2022

[5] Dynamic Bottleneck for Robust Self-Supervised Exploration, Bai et al., NeurIPS 2021

[6] InfoBot: Structured Exploration in ReinforcementLearning Using Information Bottleneck, Goyal et al., ICLR 2019

[7] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023

We thank the reviewer for their constructive feedback and helpful suggestions. We have taken steps to address the concern that performance improvements are relatively modest: We evaluated PLSM in environments with visual distractions and show that parsimonious dynamics can offer considerable performance gains (see Supporting Figure A). We have also performed an analysis in Atari where we investigate the sensitivity to regularization strength (see Supporting Figure F). In sum, we believe that adding these clarifications and showing that there are other settings where PLSM offers a greater advantage has considerably improved our paper, and we thank the reviewer again for suggesting these. We give more detailed answers to the reviewer’s questions below:

While PLSM exhibits strong performance on specific prediction tasks, the improvements in RL tasks are relatively modest.

We evaluated PLSM in tasks where our regularization might yield a larger improvement. We chose the Distracting DeepMind Control Suite [1] for this, a challenging visual control environment where static backgrounds are replaced with distracting videos. Since PLSM seeks to compress away unpredictable aspects of the environment, we hypothesized that performance should improve in this setting. We evaluated PLSM against another robust model-based RL algorithm, RePO [2], in five tasks with distracting background videos. We incorporated our regularization into RePO by learning a compressed query representation, and passed this alongside the action to the GRU which RePO uses to model the latent dynamics. The GRU dynamics were therefore softly invariant of the latent state. We observe considerable performance gains in the more challenging environments like hopper-stand, finger-spin and walker-run, suggesting that our regularization is beneficial when the agent needs to learn to ignore irrelevant distractors (see Supporting Figure A, scores averaged across 5 seeds with standard error shown as shaded regions).

Could you provide Interquartile Mean (IQM) results for the Atari experiments, as suggested by Agarwal et al. (2021)?

We optimized $\beta$ and found that $\beta = 5$ gave the strongest performance, improving mean human normalized score in Atari from 61.5 % for SPR, to 67% with PLSM. We now also report Interquartile Mean, Median and Mean Human normalized performance for this model. In all three metrics PLSM is advantageous (see Supporting Figure C). Note that SPR and DrQ interquartile ranges are lower since they contain scores from 100 seeds.

Since the regularisor weights $\beta$ is very important to the performance, I wonder how robust is PLSM to different $\beta$?

After performing a more extensive hyperparameter search for Atari, we see that there are certain games that are more sensitive to the regularization strength than others. For instance, in tasks where there are important features that are unpredictable, like in Boxing and Hero, regularizing too much can be bad for performance (see Supporting Figure F). Still, we find two hyperparameter configurations that both beat the original SPR in terms of mean performance ($\beta = 5$ and $\beta = 0.1$). These two configurations only differ in mean performance by roughly 2 percentage points, indicating that performance is relatively stable on average with respect to $\beta$.

References:

[1] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021

[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023

We thank the reviewer for their thorough and detailed review of our paper. Based on the reviewer’s suggestions, we have clarified the details of the regularization in the appendix, and performed several analyses to show that our regularization works in a principled way. In short, we find that our regularization does not increase weight norms in order to compensate for our sparsity term (see Supporting Figure D). We also show that PLSM improves considerably upon existing methods in visual control tasks with distracting backgrounds (see supporting Figure A), suggesting that our model offers more gains in environments where the agent needs to learn to ignore irrelevant cues. We would like to thank the reviewer again for raising these great points - we believe that addressing them has improved our paper considerably. We answer the reviewer’s questions in more detail below.

First, the math derivation in Appendix A contains several loose ends

We thank the reviewer for making us aware of a mistake in equation 13 which should say that the mutual information is upper bounded (instead of equivalent to) the variational approximation. For the variational distribution we can choose the variational family ourselves, and we choose the Gaussian distribution. In equation 14, the KL divergence does depend on z, but the prior of h is invariant to z. Lastly, instead of making our model stochastic, we implement a deterministic variant following [1]. Here, L2 regularization is also applied to the dynamics weights, which we have added to prevent unbounded optimization (see Supporting Figure E).

I have concerns about the fundamental validity of the regularization on $||h_t||$

The reviewer raises a good point: the query representation can be made arbitrarily small by increasing the magnitude of the dynamics weights. We offer several analyses to address this point. First, we show that the norm of the dynamics weights of PLSM does not increase substantially across several datasets, whereas the norm of the query representation is often several orders of magnitude smaller (see Supporting Figure D). Secondly, while we do not see this compensation in our experiments, it can be prevented by penalizing the L2 norm of the dynamics weights. We also evaluate top-k sparsification as an alternative sparsification method. These methods perform comparably and sometimes better in the three datasets we tested them on (see Supporting Figure D).

The experiments do not appear to show that the proposed PLSM has a significant advantage

We compared PLSM against another robust RL algorithm, RePO [2], on five challenging visual control tasks with distracting backgrounds (distracting DMC [3]). In this setting, where the agent needs to learn to ignore distracting cues, we see considerable improvements by adding our regularization to RePO (see Supporting Figure A, scores are averaged across 5 seeds, with shaded regions representing standard errors). This makes sense, as PLSM compresses away aspects of the environment whose dynamics are unpredictable given the agent’s actions. We further optimized the regularization parameter $\beta$ in Atari, and find that $\beta = 5$ works the best, improving performance in Atari from 61.5 % human normalized score for SPR, to 67% with PLSM (see Supporting Figure C).

It is unclear to me why Section 3 is a standalone section

We thank the reviewer for raising this point. We have decided to follow the reviewer’s advice and combine section 3 and 4 into one section, as they belong together. We leave the RL experiments in their own section.

Some figures are dangling in the paper without being referenced

Thank you for pointing this out, we now reference the first figure in the following sentence in the introduction.

"Here we explore the possibility of compressing states and dynamics jointly to learn systematic effects of actions (Fig 1.)"

Line 111: Could the authors elaborate a bit on how (9) is paired with (8)?

To clarify the workings of our model, we expounded on the connection between (8) and (9) in the paper.

"To apply our regularization on the contrastive dynamics model, we simply add the norm of the query representation to the contrastive loss, similarly to equation (8)."

We also make minor adjustments to the text as per the reviewer’s recommendation, like referring to PLSM in a consistent way throughout the paper, refer to equations using \eqref, using math symbols in Fig. 5, add a comma in equation (5), and change our formulation in line 9 to “the impact the state has on the dynamics”.

References:

[1] From Variational to Deterministic Autoencoders, Ghosh et al., ICLR 2020

[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023

[3] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021

We thank the reviewer for the thoughtful review, and pointing out ways in which we can improve the paper. To address the reviewer’s concern about the low number of training steps, we extended the training runs for the DMC tasks where performance had not converged, and see that PLSM retains its advantage (see Supporting Figure B). We also tuned the regularization strength for Atari and find that $\beta = 5$ yields the best performance overall, increasing mean normalized scores from 61.5 % for SPR, to 67% with PLSM (see Supporting Figure C). We also offer an explanation for why PLSM yields worse performance in some Atari games, which we have added to the Appendix. We thank the reviewer again for their feedback, as we think addressing these criticisms has helped us clarify the contribution of our model, and improved our paper. We answer the reviewer’s questions in more detail below:

Have you run the TD-MPC experiments longer?

We extended training runs by another 500k steps for acrobot-swingup and finger-turn hard. We extended the runs for humanoid-walk by another 250k steps since the algorithms were already close to convergence. Results are shown in Support PDF Figure B, averaged over 5 seeds, with shaded regions representing standard error of the mean. We see that PLSM retains the advantage, with higher converged performance for acrobot-swingup, and faster convergence for finger-turn hard. After an additional 250k steps, PLSM is still superior to TDMPC in humanoid-walk.

We sought to strengthen the continuous control experiments further by testing PLSM in the Distracting DeepMind Control Suite environments [1], where the static background in DMC is replaced with a video that is irrelevant for solving the task. We built upon RePO [2] as a baseline, using the official implementation. RePO uses a GRU for latent dynamics prediction, and we incorporated our regularization into it by learning a compressed query representation $h_t$, and passed this alongside the action to the GRU, making it softly state-invariant. Using PLSM, we see important performance gains in several tasks, suggesting that our regularization can be beneficial in environments where the agent needs to learn to ignore distracting cues (see Supporting Figure A).

It would have been nice to see some explanation or hypothesis as to why some of the Atari games performance degraded

We ran PLSM on a subset of the Atari games with varying degrees of regularization strengths to see if over regularization could explain why PLSM degrades performance in some games. See Support PDF Figure F, error bars representing standard error of the mean. We see that in games where important aspects of the environment are not controllable by the agent, weaker regularization is beneficial. For instance, in Boxing, the movement of the opponent is hard to predict, and here having a lower regularization strength is advantageous. We have added the following paragraph in the appendix:

"We see that in games where important features of the environment are not controllable by the agent, weaker regularization is beneficial. For instance, in Boxing, the movement of the opponent is hard to predict, and here having a lower regularization strength is advantageous"

References:

[1] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021

[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023

We would like to express our gratitude to all the reviewers and AC for engaging with our work and providing valuable feedback. The reviewer’s comments were indispensable for improving our paper. The reviewers made important suggestions regarding the validity and impact of our regularization method. To address these we first showed that PLSM works in a principled way. We also extended our experimental evaluation to include challenging visual control tasks with distractions and another baseline that benefits from our regularization. Now all five reviewers recommend accepting our paper, with an average score of 5.8. In sum, we think implementing the changes and extra evaluations suggested by the reviewers made our paper much stronger, and we are glad that this is reflected in the increased scores. Thank you all again!