Simplifying Latent Dynamics with Softly State-Invariant World Models

Authors: tankred  saanum
† peter dayan
eric  schulz
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 08:59:59

Abstract: To solve control problems via model-based reasoning or planning, an agent needs to know how its actions affect the state of the world. The actions an agent has at its disposal often change the state of the environment in systematic ways. However, existing techniques for world modelling do not guarantee that the effect of actions are represented in such systematic ways. We introduce the Parsimonious Latent Space Model (PLSM), a world model that regularizes the latent dynamics to make the effect of the agent's actions more predictable. Our approach minimizes the mutual information between latent states and the change that an action produces in the agent's latent state, in turn minimizing the dependence the state has on the dynamics. This makes the world model softly state-invariant. We combine PLSM with different model classes used for i) future latent state prediction, ii) planning, and iii) model-free reinforcement learning. We find that our regularization improves accuracy, generalization, and performance in downstream tasks, highlighting the importance of systematic treatment of actions in world models.

Main paper: 
 introduction: 
 In Reinforcement Learning (RL), the actions that an agent can use to solve tasks often have systematic and predictable effects on the state of the environment. When stepping on the gas pedal, the car tends to accelerate, when holding down the joystick in a given direction, the video game character tends to move in that direction, and so forth. These typical effects have exceptions that are predictable as well, e.g. stepping on the gas will not lead to acceleration if the car engine is turned off, and the video game character will not move if it is facing a wall. How can we learn world models that capture these systematic properties of actions? World models predict the agent's future states, given the current state and action [1, 2, 3] . Most approaches to world modelling represent high-dimensional observations (such as images) in compact, low-dimensional latent states z t , simplifying model-based prediction and control [4] . Here we explore the possibility of compressing states and dynamics jointly to learn systematic effects of actions (see Fig. 1 ). As is common practice in many dynamics model architectures [4, 5] , we consider the case where the model predicts the next latent zt+1 state by predicting the difference ∆at t , or the change, between the current and future latent state, given an action a t .zt+1 = z t + ∆at t (1)Even if z t is low-dimensional, the effects of actions might not be represented parsimoniously within the world model: Performing the same action a t in two similar states z, z ′ might produce two very different deltas ∆, ∆ ′ , due to how the observation encoder has constructed the latent state space (see Query Dynamics Encoder Figure 1 : Overview: World models are commonly used to predict latent trajectories, predict sequences of pixel observations, and perform planning. We propose an architecture together with an information bottleneck for learning simple and parsimonious world models. Our method relies on a query network that extracts a sparse representation h t for predicting latent transition dynamics. Combining our method with auxiliary loss functions for i) contrastive learning ii) planning and iii) and model-free RL, we see consistent performance improvement in all domains. Lines and bars show mean performance from three sets of RL benchmarks. Error bars represent 95% confidence interval. Fig. 2 ). We introduce the Parsimonious Latent Space Model, or PLSM for short, a world model where actions have more predictable effects on the inferred latent states of the agent. Dynamics are simplified by minimizing how much the predicted dynamics ∆at t depend on z t . This pushes the world model to represent states in a way that makes actions have coherent and predictable effects on the dynamics. We still allow the dynamics to vary depending on the state, but we penalize the extent of this dependence, resulting in dynamics that are softly state-invariant.We combine PLSM with two classes of world models: Contrastive World Models (CWM) [3] for latent state prediction, and with Self Predictive Representations (SPR) for model-free and modelbased control (TD-MPC) [6, 5, 7] . Across control experiments and prediction experiments we see improvements in planning, representation learning for control, robustness to noise, world model accuracy, and generalization. 
latent dynamics: 
 We assume that sequences of states, actions and rewards arise in a Markov Decision Process (MDP). An MDP consists of a state space S, an action space A, and transition dynamics s t+1 ∼ P (s t+1 |s t , a t ) determining how the state evolves with the actions the agent performs. In RL, we additionally care about the reward function r(s t , a t ), which maps state-action pairs to a scalar reward term. Here, the goal is to learn the policy π θ (a t |s t ) that maps states to the actions with the highest possible Q-values Q π θ (s t , a t ) = E π θ T t=1 γ t r(s t , a t ) , where γ is a discount factor. In this paper, we consider latent dynamics learning both in reward-free and RL settings. The heart (left) can appear on any x, y coordinate in a two-dimensional latent space with boundaries, on which it can transition in 9 different ways (moving in eight directions and standing still, for instance when moving into a boundary). Encouraging dynamics to be parsimonious recovers these 9 different possible transitions (see right), whereas an unconstrained model (see center) does not.To predict environment dynamics, the current state s t is first transformed using an encoder into a latent state z t that compactly represents the agent's sensory observation (2) . The world model predicts the change in the latent state, ∆at t that is induced by the action a t (3) & (4).EQUATIONEQUATIONWe omit the action superscript from ∆t for simplicity. Here, e θ is the encoder network mapping states to latent states, and d θ is the dynamics network mapping z t and a t to ∆t .Consider the probability distribution of transitions P ( ∆t |a t ) given the agent's actions, marginalizing across latent states. If actions have predictable effects on the state of the environment, the entropy of P ( ∆t |a t ) will be low -knowing the latent state gives little information about the effect of a t . To make the world model handle actions more systematically, we propose to minimize the amount of information the world model needs from z t in order to predict correctly how an action changes the state of the world. This quantity is represented in the mutual information between the latent state z t and the dynamics ∆t :EQUATIONwhere H[•] denotes the Shannon entropy. If this quantity is 0, the latent dynamics ∆t only depend on the agent's action a t and not z t . In this extreme, all ∆t are predicted exclusively by the action. However, it is rarely the case that the action can capture an environment's full dynamics, and making dynamics contingent on states is often necessary to some degree.To allow only the relevant information from z t to influence the dynamics, we introduce a query network f θ which maps latent state-action pairs to a latent code h t . We modify the next-step prediction components accordinglyEQUATIONWe give the query network information about the action that the transition is conditioned on. This allows the network to attend to the relevant bits in z t to output an appropriate h t . Finally, to make h t represent only the minimal amount of information needed to predict the next state, provided that a t is known, we penalize the norm of h t [8, 9] . This type of regularization has been used to constrain representations of deterministic Autoencoders in past work [8, 9] , with [8] showing that it is equivalent to minimizing the KL divergence to a constant variance zero mean Gaussian. Other regularizers and stochastic formulations are also possible (see Appendix C). For simplicity we use the deterministic variant and leave stochastic versions for future work.The strength of the penalization is controlled through a hyperparameter β. This regularization minimizes how much h t can vary with z t and hence their mutual information (see Appendix A).We then train our encoder e θ , query network f θ , and dynamics d θ jointly to minimize the following information regularized loss function.EQUATIONOur loss function encourages that the mutual information term is kept as low as possible, while still allowing the model to predict the next latent state accurately. In contrast to information bottlenecks imposed on the latent states themselves, we apply an information bottleneck to the dynamics, making ∆t easier to predict simply given a t . Regularizing h t differs from regularizing z t in important ways.In environments where the dynamics ∆ can be predicted perfectly from the actions and independently of the state, regularizing h t will not lead to a loss in information in the latent representation z t . This is because the bottleneck on h t only constrains the model in using information from z t to predict ∆ t , and not necessarily in predicting z t+1 . See Appendix B for a comparison of our method against L 1 and L 2 norm regularization on latent states. Notably, our method can also lead to state compression, in that e θ will be encouraged to omit features from s t+1 that cannot be predicted easily with little information from z t .Unfortunately, the above loss function has a trivial solution: it can be minimized completely if d θ and e θ output a constant 0 vector for all states and state-action pairs [7] . This issue is referred to as representational collapse. Representational collapse can be remedied in various ways. To show the generality of our information bottleneck, we combine it with two different approaches for mitigating representational collapse, a self-supervised approach for model-based and model-free RL in Section 3, and a contrastive approach for future state prediction in Section 4.3 Parsimonious dynamics for Reinforcement Learning 
model-based rl: 
 We evaluated the PLSM's effect on planning algorithms' ability to learn policies in continuous control tasks. To do so, we built upon the TD-MPC algorithm [6] , an algorithm that jointly learns a latent dynamics model and performs policy search by planning in the model's latent space.TD-MPC makes use of a Task-Oriented Latent Dynamics (TOLD) model. This dynamics model is trained to predict its own future state representations from an initial state and action sequence while making sure that the controller's policy π θ and Q-value function are decodable from the latent state (hence the name task-oriented latent dynamics). TOLD falls under the category of Self Predictive Representation (SPR) models, since it uses an exponentially moving target encoder e - θ with the stop-gradient operator to learn to predict its own representations. For planning TD-MPC uses the Cross-Entropy Method [10] , searching for actions that maximize Q-values.Only minimal adjustments to the TOLD model are necessary to attain parsimonious dynamics. Instead of predicting the next latent directly from the current latent and action zt+1 = d θ (z t , a t ), we use a query network f θ , mapping latent state-action tuples to h t and then minimizeL SPR = ||sg(e θ (s t+1 )) -(z t + d θ (h t , a t ))|| 2 2 + β||h t || 2 2We evaluated the efficacy of parsimonious dynamics for control in five state-based continuous control tasks from the DeepMind Control Suite (DMC) [11] . We chose the following environments: i) acrobot-swingup, due to its challenging and chaotic dynamics. ii) finger-turn hard, which poses a challenging exploration problem that TD-MPC was found to struggle with. iii) quadruped-walk, iv) quadruped-run and v) humanoid-walk due to the high-dimensional dynam-Figure 3 : PLSM, when incorporated into either the TD-MPC algorithm (A), or RePo (B), improves planning in continuous control tasks with high-dimensional and complex dynamics, and visual distractions, respectively. Lines show the average return attained across 15 evaluation episodes, averaged over five seeds. The shaded region represents the 95% confidence interval.ics. These tasks have dynamics that appear complex in the original state-space but could potentially be simplified in an appropriate latent space by introducing a dynamics bottleneck.We trained the latent dynamics and planning models in the five tasks up to a million environment steps. Scores for TD-MPC are obtained from the original implementation provided by the authors 1 . Again we used β = 0.1 for all tasks except for humanoid-walk, where β = 0.001 was more successful.Otherwise we relied on the standard hyperparameters from [6] . We see clear performance gains due in all tasks except quadruped-run (Fig. 3A ). Our results suggest that modeling the world with simple dynamics can be beneficial for RL and trajectory optimization.Since our regularization compresses away aspects of the environment whose dynamics are unpredictable given the agent's actions, it could be beneficial in control tasks with distractors. We implemented PLSM on top of RePo [12] , relying on the authors' official implementation 2 . RePo is a model-based RL algorithm based on the Dreamer [2] architecture. Here the environment dynamics are represented through a GRU network which is updated recurrently with latent state and action variables. To make the dynamics parsimonious, we update the GRU using a compressed query representation h t subject to L 2 regularization instead of the full state representation z t , resulting in recurrent dynamics that are softy state-invariant. We then trained RePo with and without our regularization on the Distracting Control Suite [13] , a challenging visual control benchmark based on DMC, where the background is replaced with a random, distracting video from the 2017 Davis video dataset. These videos are independent of the agent's actions, irrelevant for rewards, and change from episode to episode.We trained RePo with and without the PLSM objective in five Distracting Control Suite tasks for one million environment steps across five seeds. We used a regularization coefficient of β = 1e-7 for all tasks. This produced a considerable improvement over RePo in more challenging tasks like hopper-stand, walker-run and finger-spin (see Fig. 3B ). Our results suggest that encouraging the dynamics model to represent the effects of actions more consistently can improve its generalization ability in environments with distractors.Next we tested whether the learned latent space of PLSM could provide useful for model-free learning. Several methods rely on latent dynamics learning as an auxiliary objective for model-free RL [5, 7, 14, 15] . Since PLSM arranges the latent space in way that makes state transitions more predictable, it may discover useful state features and ignore aspects of the environments that would make the dynamics unpredictable otherwise. We build upon the SPR implementation for Atari3 [16] , which uses a latent dynamics model for next latent state prediction. We alter the architecture of this latent dynamics model in the same way we did for TD-MPC, and add the h t norm to the loss function.Setting β = 5 and leaving all other hyperparameters as per the standard implementation, we train the PLSM augmented SPR algorithm on 100k environment steps across 5 seeds on all 26 Atari games. We use the SPR scores reported in [17] as our baseline.Across several games we see substantial improvements to human normalized score (see Fig. 4 ). Averaging over all games, using PLSM dynamics improves human normalized scores by 5.6 percentage points (61.5 % for SPR vs 67.1% with PLSM). While we see improvements in games such as UpNDown and Kangaroo, there are other games where the regularization impacts performance negatively. Performance could potentially improve by fine-tuning the regularization strength for these domains. See Appendix I for the full score table. 
future state prediction: 
 We found that our regularization improved model-free and model-based performance across several environments. Next we evaluated whether PLSM also generally improves world models' long-horizon prediction accuracy in latent space. Using the evaluation framework and environments from [3] (see Fig. 13 for example observations), we generated datasets of image, action, next-image triplets from two Atari games (Pong and Space Invaders), and grid worlds with moving 3D cubes and 2D shapes, where each action corresponds to moving an object in one of four cardinal directions. To make the learning tasks more challenging, we increased the number of movable objects from 5 to 9. Additionally we created an environment based on the dSprite dataset [18] , with four sprites traversing latent generative factors on a random walk. The sprites had 6 generative factors: Spatial x, y coordinates, scale, rotation, color, and shape. The sprites could vary in coordinates, scale, and rotation within episodes, and additionally vary in color across episodes. Lastly, we evaluate PLSM on a dynamic object interaction dataset with realistic textures and physics without actions, MOVi-E [19] , to see if our method can be beneficial in action-free settings.Following [3] , we pair PLSM with a contrastive loss function to mitigate representational collapse:The contrastive loss encourages that different states are distinguishable in the latent space. Given a latent state and action, the model should minimize the distance between the predicted and true future latent state, while maximizing the distance between the predicted future latent state z t+1 and all other latent states in the training batch z -, up to a margin λ.EQUATIONTo 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 fitted the regularization coefficient β with a grid search and found a value of 0.1 to work the best. As a baseline we used the unregularized contrastive model from [3] , referred to as CWM (for Contrastive World Model), and its dynamics are defined through Equation ( 2), ( 4) and (9) . We also combine PLSM with the slot-based version of this model, called C-SWM (see Appendix H for results).The models were scored based on their ability to correctly predict future states (e.g. latent prediction accuracy). The models were trained and evaluated using the same parameters and metrics as in [3] : Given a state s t , a sequence of N actions a t , ..., a t+N -1 , and the resulting state s t+N , we make the model predict its latent representation of s t+N from the initial state and action sequence. The models were evaluated at several prediction horizons. Finally, we report the Hits at Rank 1 accuracy for transitions in the test set, a common test metric for contrastive models [3, 20, 21] .PLSM could better predict its own representations further into the future in five out of the six datasets (Fig. 5 ). We see the greatest gains in the cubes and shapes environments. Here, all 9 objects can collide with each other and the grid boundaries depending on their position. Always considering all possible interactions makes it challenging for a model to generalize to novel transitions. A more parsimonious solution is simply to represent whether or not the object in question would collide or not if moved in the direction specified by a t . Our results suggest that our regularization can help learn such representations, affording better generalization to left-out transitions. In one environment, Pong, we do not see an advantage in encouraging parsimonious dynamics. Here, various components are outside of the agent's sphere of influence, for instance, the movement of the opponent's paddle. This makes it challenging for the PLSM to capture all aspects of the environment state in its dynamics. In environments with non-controllable dynamics, we offer a remedy by only enforcing half of the latent space to be governed by parsimonious dynamics, and allowing the other half to be unconstrained. This hybrid model in turn shows the strongest performance in the Atari environments. See Appendix G for details.Figure 7 : Latent states z t carry decodable information about the data generating factors, whereas query states h t do not. When conditioning on an action a t , query states carry more information about the object that the action changes.Next, we evaluated the generalization and robustness properties of PLSM. For the cubes and shapes environments, we generated novel datasets where the number of moving objects was lower than in the original training data. In these datasets, the number of moving objects varied from 1 to 7. We also probed the models' robustness to noise. We corrupted test data from the dSprite environment with noise sampled from Gaussian distributions. Models were tested both in a low noise (σ = 0.1) and a high noise (σ = 0.2) condition. See Supplementary Fig. 14 for example data.When tested on scenes with fewer objects than the training data, we see a general decrease in accuracy because of the domain shift. Still, with our information-theoretic dynamics bottleneck the model generalizes significantly better to the out-of-distribution transition data (Fig. 6 ). Since PLSM seeks to predict the next state using as little information from its latent representations as possible, it is more likely to learn that the blocks move in a way that is generally invariant to the number of other blocks in the scene.Lastly, PLSM proved more robust to Gaussian noise than the unconstrained dynamics models:In the high noise condition, PLSM dynamics still accurately predict almost 60% of the transitions after 10 steps, whereas the unconstrained dynamics models, more susceptible to use noisy information in the state, predict less than 50% accurately.We also investigated the representations learned by the query network (see Fig. 7 ). To verify that PLSM learns objects' attributes (such as position and orientation), we attempted to decode ground truth object positions, scales and orientations in the dSprite dataset from PLSM latent states. We find that one can decode these ground truth factors with high accuracy from the latents. However, trying to decode these variables from the query states yielded substantially lower accuracy, indicating that they contain less information about the generative factors of the environment. Interestingly, when we condition PLSM on an action that affects only one object, we can decode attributes of that object more accurately from the query state than average. We observe this because the query state is designed to only encode information that is relevant to predict the effect of individual actions.In sum, parsimonious dynamics regularization improves both generalization and robustness properties in the three environments tested. The mutual information bottleneck on ∆t therefore not only improves prediction accuracy for data in the training distribution, but may also allow the model to generalize better to out-of-distribution data, and improve its robustness to noisy observations. 
related work: 
 Our approach introduces a mutual information constraint between the latent states z t and the latent dynamics ∆t inferred by the model. Several methods focus on state compression for dynamics modeling [1, 22, 5, 6, 23, 24] : The Recurrent State Space Model (RSSM) [22, 2, 25] , uses a variational Auto-Encoder in combination with a recurrent model (e.g. a GRU [26] ) to infer compact latent states in partially observable environments. Latent consistency is enforced by minimizing the Kullback-Leibler (KL) divergence between latent states predicted by the model and latent states inferred from pixels. The KL term regularizes the latent state not to contain more information than can be predicted by the dynamics model [2, 27] . Unlike our approach, this information bottleneck is not applied to the dynamics ∆t themselves, but to the representations z t . Expanding on this line of work, RePo [12] discards image reconstruction from the RSSM, and simply enforces that the model can reconstruct the environment's reward function, leading to stronger compression and improved performance in tasks with unpredictable elements. Similar approaches like Denoised MDP [28] only model controllable and reward-relevant aspects of the environment. While simplified latents can make transition dynamics more tractable to model, they do not necessarily give rise to the systematic action representations that we are interested in. Lastly, Self Predictive Representation (SPR) models [7, 6, 5] learn dynamics by predicting the future representations of a target encoder. SPR models have been used both for model-free and model-based control.Mutual information minimization is used in many deep learning frameworks more generally: [29] and [30] use variational methods for minimizing the mutual information between the network's inputs X and latent representations Z while maximizing the mutual information between representations Z and outputs Y . Mutual information minimization also has links to generalization ability [31, 32, 33, 34, 35] , robustness in RL [36, 37] , and exploration [38, 39] . Our information bottleneck differs from previous approaches in that it directly constrains the effect the latent state can have on the residual term in the latent dynamics over and above the agent's actions.Closest to our regularization method is the past-future information bottleneck [40, 41] . Here the mutual information between sequences of past states and future states is minimized [42, 43] . While this method simplifies dynamics, our approach differs in important ways: Rather than representing the environment's dynamics, say, using a low number of its principal components, we treat the dynamics operator ∆t itself as a random variable, and minimize its conditional dependence on z t . Furthermore, when ∆t is fully disentangled from z t , each action can be seen as a transformation that acts on the latent state space in the same way, invariantly of z t [44, 45] . We model the dynamics as softly state-invariant, allowing us to predict future latents both accurately and parsimoniously.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We will release code with model implementation upon publication.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 We have proposed a world model that tries to represent the effect of actions parsimoniously. Our model predicts future states while minimizing the dependence between the predicted dynamics ∆t and the latent state representations z t . Optimizing this objective makes the effect of the actions on the agent's latent state more predictable. Combining our objective with different model classescontrastive world models and SPR models -we observed consistent improvements in models' ability to predict their own representations accurately, generalize to novel and noisy environments, and perform planning and model-free control in high-dimensional and pixel-based environments with complex dynamics. Overall, our results suggest that systematic action representations can offer important improvements to the generalization ability and data-efficiency of world models.Limitations: Our model, in its current formulation, assumes that actions have predictable and deterministic effects on the environment. Aspects of the environment that do not behave predictably conditioning on actions are susceptible to be ignored by the model, even if they are relevant for the downstream task. While performance in Atari was improved on average, this aspect led to reduced performance in some tasks. Similarly, the degree of regularization β needs to be tuned to achieve the right level of dynamics compression for different environments.Future directions: A promising future direction is to combine our mutual information regularization with recurrent models. In partially observable, non-Markovian environments, next-state prediction is often done with a recurrent model that uses the agent's history as input. Using our bottleneck here would correspond to the assumption that the effect of the agent's actions are softly history-invariant. Another promising avenue of research is to combine PLSM with discrete dynamics. Many successful world modelling approaches assume that the latent state space is categorical [25, 46] . Instead of modelling transitions using affine transformations ∆, transitions could be modelled by predicting transition matrices. Finally, recent recent advances in controllable video-generation like Genie [47] and UniSim [48] both incorporate actions into their video models. Genie infers latent actions that explain transitions in videos, and UniSim conditions on actions and language instructions to generate controllable videos. Modelling the effect of actions in a parsimonious way could improve the accuracy and generation capabilities of such systems.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Details are found in Experiment section and Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Error bars and confidence intervals are reported in figures.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We found no potential societal impacts worth discussing. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: All datasets, models and assets used in the paper are cited properly.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: To simplify dynamics, our algorithm minimizes the mutual information between z t and ∆t conditioned on a t (Eq. 8 & 3). With stochastic gradient descent, we search for parameters θ that optimize the following objective:EQUATIONTo do so, we introduce a new variable h t that depends on z t and a t (Eq. 6). We use this variable, together with a t , to produce ∆t , instead of z t . For a specific a t , the Markov chain behind ∆at t can be written asz t → h t → ∆at t (11)Thus, to minimize the mutual information between z t and ∆t , we can instead minimize the mutual information between z t and h t : min θ I(z t ; h t |a t ). Following [29] , we write the mutual information between z t and h t as follows:EQUATIONIn our case, p(h|z, a) is deterministic. The marginal distribution conditioned on a, p(h|a), is challenging to obtain. We instead use a variational approximation: With our variational distribution q(h|a), we can upper bound the mutual information as follows I(z t ; h t |a t ) = dz dh da p(z, a)p(h|z, a) log p(h|z, a) q(h|a)First, this upper bound lets us see that the mutual information can be minimized by minimizing the number of bits of information z t contains about h t over and above a t . This could potentially be done by introducing a parameterized prior that depends on a t , q θ (h t |a t ) and minimize the KL divergence between these two probability distributionsEQUATIONWe opt for a simpler approach, that allows us to do this without introducing a new prior and parameterization. Following [8] , we assume that q(h t |a t ) is a standard, d-dimensional, isotropic Gaussian distribution N (0, I). Assuming further that h t is another d-dimensional isotropic Gaussian, the KL divergence is a function of the mean and standard deviation of h t [8, 27] :EQUATION)Here µ(•) and σ(•) return the mean and standard deviation of the Gaussian, respectively. Importantly, we can minimize the KL divergence by minimizing the first term in the above equation, which is simply the norm of the mean of h t . In our deterministic setup, we therefore simply minimize the norm of h t itself.Figure 8 : To illustrate how our mutual information minimization impacts the latent space, we trained PLSM to learn the dynamics of a dot moving in four directions on a gridworld with a wall in the middle. We compared PLSM to contrastive models that regularize the L 1 and L 2 norms of the latent space, respectively. We observe that PLSM regularization leads to a more regular representation of the states in the gridworld, whereas no regularization leads to a warped representation. Moreover, L 1 and L 2 regularization simply shrinks the latent space. This type of shrinkage is not present in PLSM.We tested some alternative formulations and ablations of our PLSM objective on a subset of the datasets (see Fig. 9 ). One ablation removed the query representation h t and simply regularized the latent state z t . This model achieved substantially worse performance in the three datasets we tested it on. As an alternative to regularizing the L 2 norm of h t , we also tested a top-k bottleneck, which used the top 15 features from h t instead of the full representation. This model performed on par with the original PLSM. Lastly, minimizing the L 2 norm of h t could potentially be compensated for by the dynamics MLP by increasing the norm of the weight matrices. As a control, we added weight decay to the dynamics model to prevent this, and observe comparable performance.One potential issue with minimizing the L 2 norm of the query representation h t is that the model can compensate for this by increasing the magnitude of the weights in the ensuing dynamics module. While this can be prevented by adding weight decay to the dynamics model's weights, we show that PLSM does not suffer from this empirically in the datasets we evaluated it on. We calculated the average norm of h t of PLSM relative to the average norm of z t in the unregularized model. Comparing them we see indeed that h t has several orders of magnitude lower norm. However, comparing the norm of the ensuing linear layers, we see that they most often do not differ significantly (see Fig. 10 ). This suggests that shrinking h t truly minimizes the amount of information the model can extract from it to predict ∆. We report more detailed Atari results in this section, including median, IQM, and the probability of improvement using the RLiable package [17] (see Fig. 11 ) . We also investigated how the regularization strenght of PLSM impacted performance in Atari. 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. See Fig. 13 for example frames from the training datasets and Fig. 14 for example frames from the generalization tests. We propose a hybrid model for the Atari environments. The hybrid model splits the latent state in two, one parsimonious state space z 1 t and an unconstrained state space z 2 t . To predict the next state, the model uses two dynamics MLPs z2t+1 = z 2 t + d 1 θ (z 1 t , z 2 t , a t ) and z1 t+1 = z 1 t + d 2 θ (h t , a t ).The next latent prediction is then defined as zt+1 = Concatenate(z 1 t+1 , z2 t+1 ). Results are shown in Fig. 15 .We evaluated the C-SWM on our cubes 9 task, both with and without our regularization. The performance of C-SWM dropper by a small amount due to the increased difficulty of the task for longer prediction horizons (t = 10). Interestingly, using our regularization method proved beneficial for the slotted model, outperforming the unregularized C-SWM in long horizon latent prediction.Table 1 : PLSM with slots outperforms the C-SWM [3] in latent prediction at prediction horizon t = 10, averaged over five seeds, plus and minus standard error of the mean.Accuracy PLSM C-SWM Cubes 9 99.05% ±0.18 97.45%±1.3 • If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.See https://github.com/nicklashansen/tdmpcSee https://github.com/zchuning/repoSee https://github.com/mila-iqia/spr 


Decision: Learning latent state representation is crucial for scaling reinforcement learning to real-world problems as these often operate in this limit. In this work, the authors proposed a new and practical approach for learning latent dynamics in high-dimensional spaces. The new method, PLSM, learns the latent state representation by learning the action conditional mutual information between the state  and residual changes in the representation.  The authors conducted thorough empirical investigations and managed to show their method improves performance of model free and model based algorithms in standard reinforcement learning benchmarks. 

In summary, this work advances our understanding of latent state representation learning and explores a novel algorithmic idea. Reviewers and I agree that it is valuable to the community, and I support accepting it.

Reviews: review: paper_summary: paper_summary: The authors propose an approach to learn latent dynamics from high dimensional observations. Their method seeks to minimize the mutual information between the current latent state and the change in the latent state conditioned on the next action, which they argue minimizes the dependence that the state representation has on the dynamics. They evaluate their method on predicting future states, and on model-based and model-free reinforcement learning. They show that their method improves accuracy and downstream task success in difficult continuous control environments, and also helps in some Atari game settings.
strength_weakness: strength_weakness: The biggest glaring weakness is in the continuous control experiments. The graphs cutoff at 10^5 environment steps which is much shorted than the original TD-MPC paper, and not till convergence for most tasks. This begs the question what the method's asymptotic performance is compared to the original method. 

In addition to the continuous control experiments, it would have been nice to see some explanation or hypothesis as to why some of the Atari  games performance degraded so much compared to not using their method.
questions: questions: Have you run the TD-MPC experiments longer? This is a super important experiment to run and should be included in the paper
limitations: limitations: Yes
paper_summary: paper_summary: The paper addresses learning a world model with state-invariant dynamics. To this end, it proposes Parsimonious Latent Space Model (PLSM), which introduces an information bottleneck to the additive dynamics residual. The influence of the state on the dynamics is summarized in the bottleneck variable, whose norm is regularized in the loss function. The proposed PLSM is combined with contrastive learning (Contrastive World Model) to demonstrate its advantage in accuracy, generalization, and robustness for future state prediction. The paper also conducts experiments on PLSM with Self-Predictive Representation and TD-MPC to show the benefits of PLSM to downstream control.
strength_weakness: strength_weakness: - There seem to be important flaws in the proposed PLSM.
  - First, the math derivation in Appendix A contains several loose ends. Why replacing $p(h|a)$ by $q(h|a)$ in (13) results in a upper bound? Also, why is (13) an equality? Why is it reasonable to assume $q(h_t|a_t)$ is the standard Gaussian distribution and $h_t$ is another isotropic Gaussian? Why does $(15)$, the $KL_D[p(h_t |z_t , a_t ) || q_{\theta} (h_t |a_t )]$ have no dependence on $z_t$ at all? And even if we accept the Gaussian assumption, why immediately after (15), does $h_t$ become a deterministic quantity in Line 440?
  - Second, regardless of the math derivations, I have concerns about the fundamental validity of the regularization on $||h_t||$. It just seems to me that the norm of $h_t$ can be made arbitrarily small without changing the dynamics. For example, let $h_t = 0.0001 f_{\theta}(z_t, a_t)$ and let $\Delta_t = d_{\theta}(10000 h_t, a_t)$. Then, $||h_t||$ becomes 10000 times smaller, but the dynamics remains the same. The constants 0.0001 and 10000 can be absorbed in the parameter $\theta$. This flaw undermines my confidence in the proposed approach.
- The experiments do not appear to show that the proposed PLSM has a significant advantage, especially in the control setting. Actually, Figure 7 shows with PLSM, the performance drops in more environments than the performance improves.
- The writing has much room for improvement. 
  - First, the experiments is organized in three sections, namely Sections 3-5. It is unclear to me why Section 3 is a standalone section (and if so, is it better to call it "Experiment setup", rather than "Experiments"). I could be wrong, but it actually seems to me that Section 4 is in parallel to Section 4.1, both of which may be better organized as subsections of Section 3.
  - Some figures are dangling in the paper without being referenced, e.g., Figure 1.
  - Equations are better referenced in parentheses like (1)(2); in LaTeX the command is `\eqref`.
  - Some other minor problems are detailed in the "Questions" section below.
questions: questions: - Line 9: "minimizing the dependence the state has on dynamics". Logically, maybe "the dependence the dynamics has on the state", or "the impact the state has on the dynamics"? E.g., in Line 34, it also says "the degree to which the predicted dynamics depend on $z_t$".
- In Equation (5), should there be a comma between $z_t$ and $a_t$?
- Line 111: Could the authors elaborate a bit on how (9) is paired with (8)?
- In Figure 5, in the x labels, it is suggested that the math symbols are added for better clarity, e.g., "Latent state $z_t$".
- Line 185: "We evaluated the PLSM's effect": remove "the" to be consistent with other places in the paper?
limitations: limitations: Limitations are commented above. No societal impact needs to be particularly addressed.
paper_summary: paper_summary: This paper introduces a method to enforce parsimonious latent dynamic models. The key idea is that if we can minimise the influence of states on the dynamic, i.e. the conditional mutual information $I(z_t, \Delta_t | a_t)$, the dynamic can generalise better to unseen states during prediction. In order to achieve this goal, the author proposes PLSM based on information bottleneck. Experiments show that the proposed method can be widely applied to prediction, and both model-based and model-free RL. On specific tasks that fit well with the parsimonious assumption, PLSM significantly improve the result.
strength_weakness: strength_weakness: - While PLSM exhibits strong performance on specific prediction tasks, the improvements in RL tasks are relatively modest.
- As is common with information bottleneck-based methods, PLSM's performance may be highly sensitive to the hyperparameter of the regularizer for the bottleneck.
questions: questions: - Could you provide Interquartile Mean (IQM) results for the Atari experiments, as suggested by Agarwal et al. (2021) in "Deep Reinforcement Learning at the Edge of the Statistical Precipice"? This would offer a more robust measure of aggregate performance across multiple tasks.
- Since the regularisor weights $\beta$ is very important to the performance, I wonder how robust is PLSM to different $\beta$? Could you also share some insights about tuning the parameter? This can be a valuable information to the community.
limitations: limitations: Yes, the limitations are well-discussed by the authors.
paper_summary: paper_summary: This paper presents an information bottleneck principle to regularize the latent dynamics, which makes the effect of the agent’s actions more predictable. This approach minimizes the mutual information between latent states and the change that action produces in the agent’s latent state, in turn minimizing the dependence the state has on the dynamics. The proposed world model can be combined with model-based planning and model-free RL.
strength_weakness: strength_weakness: - The contribution is limited since the MI regularization for RL representation has been adopted in several previous works, including [1-3].

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


[1] Learning Robust Representations via Multi-View Information Bottleneck. ICLR 2020

[2] Dynamic Bottleneck for Robust Self-Supervised Exploration. NeurIPS 2021

[3] InfoBot: Structured Exploration in ReinforcementLearning Using Information Bottleneck. ICLR 2019
questions: questions: - Since the MI between representation and state is constrained, should the algorithm exhibit stronger robustness in various domains? Can you show some empirical results?
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes a Parsimonious Latent Space Model (PLSM) as a latent world model method. The main idea of PLSM is twofold: (1) the use of an additional hidden variable $h_t$, which is bottlenecked to have parsimonious information (for better latent predictability) and (2) the use of difference prediction instead of latent prediction. The authors combine PLSM with different world model/RL techniques (CWM, TD-MPC, and SPR), and show that it leads to better long-term prediction and improves performance on RL control and game benchmarks.
strength_weakness: strength_weakness: * I'm not sure how simply minimizing the norm of $h_t$ leads to minimal representations. If $h_t$ were a stochastic random variable with a fixed variance, this would make sense (since in this case, this becomes equivalent to KL divergence minimization between the posterior and the standard Gaussian prior). However, $h_t$ seems deterministic in this work, and I don't think just minimizing $\\\|h_t\\\|$ has a principled regularization effect, given that it is possible to contain the same amount of information with a very small scale (e.g., multiplying $h_t$ by $10^{-9}$ only makes the norm smaller while maintaining the same amount of information).
* I'm also not fully convinced that we *need* an additional $h_t$ variable. Why can't we just impose a bottleneck on $z_t$? Is there a unique benefit of having an additional latent $h_t$ variable that is not achievable by having $z_t$ alone?
* While 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, which I even consider to be the "default" design choice in model-based RL.
* The performance gain over previous RL methods (TD-MPC, SPR) seems marginal, despite the added complexity of PLSM. Moreover, the empirical comparisons are weak in the sense that the comparison is only made with a single method in each setting, and it's unclear how PLSM compares to other regularization methods (e.g., RePo, other regularization techniques, etc.).

In summary, I feel this is a borderline paper. The proposed method (as well as its writing) is clean and conceptually intriguing, but its novelty and performance gain seem marginal. I'd recommend a borderline reject at this point, but I'm willing to adjust the score if the authors address the concerns above (especially the second and fourth points).
questions: questions: Please see the weaknesses section above.
limitations: limitations: The authors sufficiently address the limitations in Section 7.

Comments: title: Summary of discussion period
Thank you!
Thanks for the rebuttal
Please address the author's rebuttal
Thank you
Rebuttal Response
Thanks for the additional comments
Response
Thanks for your rebuttal, I decide to maintain my score.
Thanks for reviewing our work
Thanks for the response
Thanks for the response
comment: 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!
We thank the reviewer for helping us clarify our approach, and to show the validity of our regularization, we agree that our paper is better as a consequence. We thank the reviewer for raising the score accordingly.
I appreciate the authors' response, which indeed addresses several of my concerns, especially my biggest concern about the validity of regularizing $\|h_t\|$ without considering $d_{\theta}$. I do feel the authors need to discuss this point and their solution in the paper. I am also more confident about the mathematical derivation. Earlier on I misunderstood (15) as not depending on $z_t$, but realizing $h_t$ is actually $f_{\theta}(z_t, a_t)$, I now see how (15) indeed depends on $z_t$ and does follow from (14). I suggest the authors remind the reader this point in the paper to avoid confusion. It may be helpful to also clarify that "the deterministic setup" (Line 440) refers to $h_t | z_t, a_t$ is deterministic, and $h_t$ remains to be a random variable and is not deterministic. Overall, I am happy that the authors agree to incorporate many of my suggestions, and I believe an updated version of the paper should be in a much better shape. As a result, I raise my score from 3 to 5.
Dear gEcM, please try to respond to the author's rebuttal and elaborate if it made you reassess your opinion.

Thank you,
We thank the reviewer for their helpful review, which we are sure has allowed us to make important improvements to it, and are grateful for the increase in score.
I would like to thank the authors for taking the time to respond to my concerns and answer my questions. Given the additional experimental evaluations I am increasing my score to accept.
We thank the reviewer for engaging with our paper. We believe that addressing the reviewer's concerns has helped us improve our paper, and we are grateful the reviewer raised the score as a consequence.
We thank the reviewer for raising these additional points. We agree that our contribution could be additionally strengthened by addressing these. As such, we ran an additional experiment where we evaluated Dreamer with and without PLSM on walker-walk in Distracting DMC. We first note that Dreamer does considerably worse than RePO, due to the pixel-reconstruction objective. Nevertheless, we see that adding PLSM strengthens the performance of Dreamer on the distracting DMC task. Dreamer attains a mean score of 460 ($ \pm 66$ CI), whereas Dreamer + PLSM attains a mean score of 530 ( $\pm 57$ CI), averaged over 5 seeds after 1 million environment steps.

We agree with the reviewer that it’s not easy to draw conclusions about performance across all Atari games based on five seeds. However, in several individual games we see that PLSM brings significantly better scores just based on 5 seeds. For instance, in UpNDown, Pong, Krull, Freeway and Assault we observe significantly better human normalized mean scores with just five seeds. We will therefore tone down our claims in the section discussing the Atari results.
The additional experiments strengthen the contribution and I encourage you to perform complement experiments with multiple seeds in the next version.
Thank you for your rebuttal and the additional experimental results provided. While I appreciate the effort, I still have some concerns:

- Regarding the Distracting DMC experiments: The comparison between RePO and RePO + PLSM is interesting, but may not fully demonstrate PLSM's effectiveness in handling distracting information. RePO already addresses this issue to some extent. A more illustrative comparison might be between Dreamer and Dreamer + PLSM, which could better showcase PLSM's specific contributions.
- On the IQM results for SPR + PLSM: The overlapping trust intervals between SPR + PLSM, SPR, and DrQ make it difficult to draw definitive conclusions about PLSM's performance improvements. More robust statistical analysis may be needed to support any claims of superiority.
- Hyperparameter sensitivity: Supporting Figure F highlights PLSM's sensitivity to the hyperparameter $\beta$ across a wide range (1e-6 to 5.0). This sensitivity could pose challenges for real-world applications where extensive hyperparameter tuning may not be feasible.

While I continue to see value in this paper and maintain my recommendation for acceptance, these remaining concerns prevent me from increasing my score. I believe addressing these points could further strengthen the paper's contributions to the field.
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!
Thanks for the response.

For the first point, yes, as I also mentioned in my initial response, if $\|h\|$ were a stochastic variable, norm regularization would be somewhat equivalent to KL minimization to a Gaussian prior. I can also imagine that it may have a similar effect even for deterministic $h$'s in practice. I believe the new clarification would help prevent this kind of confusion.

For the second point, thanks for the clarification. I got the point. Perhaps the main issue was then that Figure 1 is somewhat misleading, because it seems there's no direct flow from $z_1$ to $z_2$ (even though there is --- $z_2$ is the sum of $z_1$ and the prediction from $h_1$ and $a_1$). I would encourage the authors to revise Figure 1 to clarify this point.

In general, my main concerns have mostly been resolved to some degree. I also appreciate new preliminary results with additional baselines. Accordingly, I've adjusted my score to 6.
We thank the reviewer for the follow-up questions. The motivation for minimizing the L2 norm comes from [1], where they argue that it constrains the size of the latent space, and is therefore a suitable regularizer. They also show that minimizing the L2 norm amounts to minimizing the KL divergence if the prior is a constant variance Gaussian with zero mean. Here they also propose to regularize the L2 norm of the ensuing weights, to avoid unbounded optimization. We therefore show that our regularization does constrain the size of the latent space while the dynamics weights norms stay similar. The concern about the norm of the dynamics weights was also raised by reviewer gEcM. Note that we also show that our method works experimentally by sparsifying $h_t$ to only include the largest $k$ features (reducing from 50 to 15 in our case, see Supporting Figure E).  

Still, we agree with the reviewer that the argument presented in the paper is empirical. We have therefore added the following paragraph in the methods section to highlight this.

**”We minimize the $L_2$ norm of $h_t$ in order to constrain the size of the query representation. This type of regularization has been used to regularize representations of deterministic Autoencoders in past work [1, 2], with [1] showing that it is equivalent to minimizing the KL divergence to a constant variance zero mean Gaussian. Other regularizers are also possible. Instead of modeling $h_t$ as a deterministic variable, it can be treated stochastically. Using multivariate Gaussian distributions with learned variances, one can regularize $h_t$ by minimizing the KL divergence to a standard Gaussian [3]. For simplicity we use the deterministic variant and leave stochastic versions for future work.”**

>How can imposing an information bottleneck only on $h$ prevent information loss, given that there's no direct path from (say) $z_1$ to $z_2$ in Figure 1? (i.e., it must pass through $h_1$, so the bottleneck on $h_1$ will anyway cause information loss toward $z_2$) 

This is a good question. The reason why passing $z_t$ through the $h_t$ bottleneck does not necessarily cause information loss is because $h_1$ is not used to predict $z_2$ directly, but rather the *change* term $\Delta$. As such, the model is designed to have information loss with respect to $\Delta$, but not necessarily $z_{t+1}$. If $z$ evolves in a manner that is independent of $z$, the model can ignore it to predict $\Delta$, without having to lose any information contained in $z$. We have added the following clarification to the model section, and are happy to.

**”Regularizing $h_t$ differs from regularizing $z_t$ in important ways. In environments where the dynamics $\Delta$ can be predicted perfectly from the actions and independently of the state, regularizing $h_t$ will not lead to a loss in information in $z_t$. This is because the bottleneck on $h_t$ only constrains the model in using information from $z_t$ to predict $\Delta$, and not necessarily in predicting $z_{t+1}$.”**

References:

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

[2] Improving Sample Efficiency in Model-Free Reinforcement Learning from Images, Yarats et al., ArXiv 2020

[3] Auto-Encoding variational Bayes, Kingma et al., ICLR 2014
Thanks for the detailed response.

While I appreciate the new results about $\\\|h\\\|$, I don't think the fact that the norm of weights ($\\\|w\\\|$) remains similar supports the effect of the norm regularizer (why is it related to that?). Moreover, even if it does support it, this finding is at best purely empirical and not principled (by "principled", I mean having a theoretical or logical justification), given that it *does* have the theoretical issue that the norm of $h$ can be arbitrarily small while not losing any information. That being said, this was not a major concern for me, so I'd be fine if the authors clearly acknowledge that this is an *empirically* motivated regularizer in the paper with sufficient discussion.

I think I still didn't fully understand the necessity of $h$. In Figure 1, why can't we just merge $z_1$ and $h_1$, $z_2$ and $h_2$, and so on? How can imposing an information bottleneck only on $h$ prevent information loss, given that there's no direct path from (say) $z_1$ to $z_2$ in Figure 1? (i.e., it must pass through $h_1$, so the bottleneck on $h_1$ will anyway cause information loss toward $z_2$ --- or is this figure somewhat incorrect?) I can see that having an additional $h$ leads to *empirical* performance improvement, but I'm curious if there are any *principled* benefits of additionally having $h$, or if it is purely a practical design choice. I might have missed something, and could the authors elaborate on this point?