Simplifying Latent Dynamics with Softly State-Invariant World Models
Authors: Tankred Saanum
Peter Dayan
Eric Schulz
Venue: NeurIPS 2024 poster (2024)
Keywords: World model
latent dynamics
reinforcement learning
compression
Primary area: reinforcement_learning
TL;DR: 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.

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.

Simplifying Latent Dynamics with Softly State-Invariant World Models
Tankred Saanum
1
Peter Dayan
Eric Schulz 1,3
Max Planck Institute for Biological Cyberneties, University of Tübingen 3 Helmholtz Institute for Human-Centern AI, Helmholtz Center Munich, Neuherberg, Germany taakred. saanun@tuebingen. mpp.deppppeeee......
Abstract
To solve control problems via model-based reasoning or planning, an agent needs 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). 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 s latent state, in turn minimizing the dependence the state has on the dynamics. This makes the world model softly state-invariant. We combine PL.S with different model classes used for future latent state prediction, iř) planning, and fii) 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.
Introduction
In Reinforcement Learning (RL), the actions that an agent can use to solve tasks often have systematic and predictable effects 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 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 accelcration if the car engine is turned off, and the video game character will not move if it facing wall. How can we learn world models that capture these systematic properties of actions? World models predict the agent's '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,, 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, 5J, we consider the case where the model predicts the next latent ž.+1 state by predicting the difference A or the change, between the current and future latent state, given an action a..
2 2, + A.
Even if z, is low-dimensional, the effects of actions might not be represented parsimoniously within the world model: Performing the same action in two similar states z, z' might produce two very different deltas A, A', due how the observation encoder has constructed the latent state space (see
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

a1
a2
Z1
h1
A
2
h2
A
Z3
Distracting
Model-free RL
Encoder
DMC
Control
(Atari)
Human Norm'd Score
0:90
722
620
Avg. Return
Query
0.75
X02
400
0.80
Dynamics
12HMCC
RdPs Pa SM
Step
Step
SPR
SPR PLSM Dro
Figure Overview: World models are commonly used to predict latent trajectories, predict sequences of pixel observations, and perform planning. We propose architecture together with an information bottleneck for learning simple and parsimonious world models. Our method relies on query network that extracts sparse representation h, for predicting latent transition dynamics. Combining method with auxiliary loss functions for i) contrastive learning üi) 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 Á depend on 2,, This pushes the world model 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 dynamics that are softly invariinaaiint...
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 model- based control (TD-MPC) [6, 5, 7]. Across control experiments and prediction experiments we see improvements in planning, representation learning for control, robustness 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 state space S, action space A, and transition dynamics St+1 P(s |s, a,) determining how the state evolves with the actions the agent performs. In RL, we additionally care about the reward function r(st, at), which maps state-action pairs to scalar reward term. Here, the goal is to learn the policy To
((aa (s)) that maps states to the actions with the highest possible Q-values
Q.o (St, a²) E.g EE= (rr at) where y discount factor. In this paper, we consider latent dynamics learning both in reward-free and RL settings.

---

Unconstrained P()+)
Parsimonious P(),)
0.35
011
A 0.00
0.00
Az2
A2
Figure 2: The heart (left) can appear any x: y coordinate in two-dimensional latent space with boundaries, on which it can transition different ways (moving in eight directions and standing still, for instance when moving into 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, is first transformed using an encoder into latent state that compactly represents the agent's 's sensory observation (2). The world model predicts the change in the latent state, A. that is induced by the action a, (3) & (4).
2, = (s((s)
(2)
Å, dz at)
(3)
2141 Zt + D.
We omit the action superscript from Å, for simplicity. Here, eg is the encoder network mapping states latent states, and dg is the dynamics network mapping z, and a, to A.
2.1 Parsimonious latent dynamics
Consider the probability distribution of transitions P(Á, lar) given the agent's actions, marginalizing across latent states. If actions have predictable effects on the state of the environment, the entropy of P(A:Ja)) (a.) will low knowing the latent state gives little information about the effect of a, To make the world model handle actions information the world model needs from z, 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 2, and the dynamics Á,:
systematically,
propose to minimize the amount of
T T(::; A.la.) = H(A.la.) H1A,1z,, ai)
where H|--] denotes the Shannon entropy. If this quantity is O, the latent dynamics Ås only depend on the agent's 's action a, and not In this extreme, all Å, are predicted exclusively by the action. However, it is rarely the case that the action can capture an environment's full dynamies, and making states is often necessary to some degree.
dynamics contingent
To allow only the relevant information from z, influence the dynamics, we introduce query network fe which maps latent state-action pairs to a latent code h;- We modify the next-step prediction components accordingly
h, fa(ze, ar)
(6)
A, (bh ar)
(7)
give the query network information about the action that the transition is conditioned on. This allows the network to attend the relevant bits in z; output appropriate h,. Finally, to make

---

h1 represent only the minimal amount of information needed to predict the next state, provided that a, is known, we penalize the norm of h, [8. 9]. This type of regularization has been used to constrain representations of deterministic Autoencoders in past work [8, 9], with [8] showing that is equivalent minimizing the KI. 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 B. This regularization minimizes how much h, can vary with z, and hence their mutual information (see Appendix A). We then train our encoder eg, query network fe, and dynamics de jointly to minimize the following information regularized loss function.
L ==((o((5(+1) - (z: + (o,, a a 0/
(8)
Our loss function encourages that the mutual information term is kept as low as possible, while still allowing the model predict the next latent state accurately. In contrast information bottlenecks imposed on the latent states themselves, we apply an information bottleneck to the dynamics, making A easier to predict simply given a,- Regularizing h, differs from regularizing 2, in important ways. In environments where the dynamics A can be predicted perfectly from the actions and independently of the state, regularizing h, will not lead to a loss in information in the latent representation z,. This because the bottleneck on h, only constrains the model in using information from z, predict A,, and not necessarily in predicting 21+1- See Appendix B for comparison of our method against L1 and La norm regularization on latent states. Notably, our method can also lead to state compression, in that Ge will be encouraged to omit features from S that cannot be predicted easily with little information from 2,
Unfortunately, the above loss function has a trivial solution: it can be minimized completely de and eg output constant 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 with two different approaches for mitigating representational collapse, 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
3.1 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-Oriente 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 अy and Q-value function are decodable from the latent state (hence the name fask-oriented latent dynamics). TOLD falls under the category of Self Predictive Representation (SPR) models, since it uses an exponentially moving target encoder ep with the stop-gradiens operator to learn predict its representations. For planning TD-MPC uses the Cross- cosss-Entrpyy Method [10]. searching for actions that maximize Q 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 ă = do (Zz, at), we use a query network fo, mapping latent state-action tuples to h, and then minimize
LSPR [|sg(ee(sss)))
(z, + do a. ))1 + 3 1/12.112
We 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 environ- ments: i) acrobot-swingue due to its challenging and chaotic dynamics. üi) f finger-turn hard, which poses a challenging exploration problem that TD-MPC was found to struggle with. iī) quadruped-wer wwl.uudpdddlll,,,.. iv) quadruped run and v) humanoid d dwwakk due to the high-dimensional dynam-

---

A
DMC
acrobot-swingup
finger-turn hard
guadruped-walk
quadruped-run
humanoid-walk
1000
1000
Episode return
800
800
450
750
750
000
300
500
530
400
400
150
250
250
200
200
0.0
0.2
0.5 0.7
1.0
0.0
0.2
0.5 Step
0.7
1.0
3
2
5 1e5
0.0
0.3 0.6
0.9
12
Step
leb
le6
Step
Le5
Step
Step
leb
TD-MPC
TD-MPC
PLSM
Distracting Control Suite
B
walker-walk
cartpole-balance
hopper-stand
walker-run
finger-spin
1000
Episode return
800
630
400
450
800
600
300
400
200
200
150
200
100
Step
Step
Step
Step
Step
RePo
RePo
PLSM
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 dynamies, 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.
trained the latent dynamics and planning models the five tasks up to a million environment steps. Scores for TD-MPC are obtained from the original implementation provided by the authors' Again we used 3 = 0.1 for all tasks except for humanoid-wall where = 0.001 was more successful. Otherwise relied the standard hyperparameters from [6]. We see clear performance gains due in all tasks except quadruped-rum (Fig. 3A). Our results suggest that modeling the world with simple dynamics can be beneficial for RL and trajectory optimization.
3.2 Distracting visual control
Since our regularization compresses away aspects the environment whose dynamics are un- predictable given the agent's is actions, it could be beneficial in control tasks with distractors. We implemented PI.SM on top of RePo [12], relying on the authors' official implementation". RePo is 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 compressed query representation h, subject to L2 regularization instead of the full state representation Z,, resulting in recurrent dynamics that softy state-invariant. We then trained RePo with and without our regularization on the Distracting Control Suite [13], 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 '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 million environment steps across five seeds. We used regularization coefficient of 1e-7
'See https://github. com ni ck. htsaaaccckaaaaaaaeeeaaaaaa "See https://githu/ wwt.ub..co////oaaaaaa www/////miingrrepoooooooooo
5

---

Human Norm'd Score
1.6
c.8
UpNDown
Kangaroo
Krult
Freeway
Pong
Assault
Frostbite
RoadRunner
Seaquest
PrivateEye
Allan
-eero
Asterix
Obert
Anidar
DemonAllack
MsPacmmn
BattleZone
Gapher
ChopperCommaine
Jamesbord
Brealoutt
CrazyCllmbar llmber
RankHpiss
Baxing
Figure 4: Changing the dynamics model in SPR to PLSM increases score in several Atari games, with little implementation overhead. On average, human normalized scores are higher when using PLSM dynamics. Bars show difference in human normalized score between SPR with and without PLSM dynamics, averaged over five seeds.
for all tasks. This produced a considerable improvement over RePo more challenging tasks like hopper-stane walker- and finger-spin (see Fig. 3B). Our results suggest that encouraging the dynamics model represent effects actions more consistently can improve its generalization ability in environments with distractors.
3.3 Model-free RL
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 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 Atari [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, norm the loss function. Setting 3 5 and leaving all other hyperparameters as per the standard implementation, we train the PI.SM augmented SPR algorithm on 100k environment steps across seeds on all 26 Atari games. We use the SPR scores reported in [17] as our baseline.
Across several games we see substantial improvements human normalized score (see Fig. 4). Averaging over all games, using PLSM dynamics improves human normalized scores by 5.6 per- centage points (61.5 % for SPR 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 for the full score table.
4 Future state prediction
We found that our regularization improved model-free and model-based performance across several environments. Next we evaluated whether PLSSM also generally improves world models' long-horizon prediction accuracy 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 object in one of four cardinal directions. To make the learning tasks more challenging, we increased the number of movable objects from to 9. Additionally we created an environment based on the dSprite dataset [18], with four sprites
See https: ptttps://githee com wtt.a//ooo/maaaaii/aaa/eae

---

spaceinvaders
pong
cubes 9
shapes
dSprites
MOVi-e
100
Accuracy %
40
96
30
80
45
98
90
20
80
96
30
10
3
5
0
10
3
5
8
10
3 5
8
10
3 5
8
10
3
5
8
10
10
15
20
80
100
100
MRR %
99
96
90
70
45
98
86
75
60
80
30
50
3
3 Horizon
15 Horizon
Horizon
Horizon
Horizon
Horizon
CWM
PLSM
Figure 5: PLSM improves contrastive world models' accuracy in long- long-horizon latent prediction in five out of six environments. In the cubes and shapes dataset, the PLSM is close to perfect even when predicting as far as 10 timesteps the future. Lines show accuracy on entire test data averaged over five random seeds. The shaded region corresponds the standard error of the mean.
traversing latent generative factors on a random walk. The sprites had 6 generative factors: Spatial 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 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 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 Ze+1 and all other latent states in the training batch z up margin X.
Connrastivee || 2 max A || Ž
To apply our regularization on the contrastive dynamics model, we simply add the norm of the query representation the contrastive loss, similarly to equation (8). We fitted the regularization coefficient 8 with grid search and found value of 0.1 to work the best. As 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 in [3]: Given a state a sequence of N actions .c N and the resulting state St+N. we make the model predict its latent representation of st+N from the initial state and action sequence. The models were evaluated at several prediction horizons. Finally, we report the Hits at Rank accuracy for transitions in the test set, a common test metric for contrastive models [3, 20, 21].
4.1 PLSM improves long-horizon prediction accuracy
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, 9 objects can collide with each other and the grid boundaries depending on their position. Always considering all possible interactions makes challenging for model to generalize novel transitions. A more parsimonious solution is simply to represent whether not the object in question would collide or not if moved the direction specified by a, Our results suggest that our regularization can help leam such representations, affording better generalization to left-out transitions.
7

---

Novel cubes
Novel shapes
Noise: 0= 0.1
Noise: 0= 0.2
100
100
Accuracy %
40
80
80
40
20
20
60
60
40
40
PLSM
CWM
PLSM
CWM
1
3
5
8
10
1
3
5
8
10
Horizon
Horizon
CWM
PLSM
When exposed to scenes with fewer objects than trained on (cubes and shapes environment), or corrupted data (dSprite environment) from the test set, PLSM improves accuracy over the CWM. Lines represent the average of models trained across five seeds. Shaded regions and bars reflect the standard error of the mean.
Figure 6: PLSM improves generalization and robustness in contrastive models:
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 s paddle. This makes it challenging for the PLSM to capture all aspects of the environment state its dynamics. In environments with non-controle dynamics, we offer remedy by only enforcing half of the latent space to be governed by parsimonious dynamics, and allowing the other half to unconstrained. This hybrid model in turn shows the strongest performance the Atari environments. See Appendix G for details.
4.2
Generalization and robustness
Next, we evaluated the generalization and bustness 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 7. We also probed the mod- els' robustness to noise. We corrupted test data from the dSprite environment with noise sam- pled from Gaussian distributions. Models were tested both in low noise (a = 0.1) and a high noise (a 0.2) condition. See Supplementary Fig. for example data.
Generative factor decodability
0.7
0.6
0.5
R2
0.4
0.3
0.2
When tested on scenes with fewer objects than the training data, we see general decrease in accuracy because of the domain shift. Still, with our information-theoree the model generalizes significantly better to the -iitt--distribution transition data (Fig. 6). Since PI PM.SM 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 way that is generally invariant to the number of other blocks in the scene.
0.1
0.0
dynamics bottleneck
Latent state
Object matched query state
Query state
Figure 7: Latent states z, carry decodable infor- mation about the data generating factors, whereas query states h, do not. When conditioning on an action a,, query states carry more information about the object that the action changes.
Lastly, PLSM proved more robust to Gaussian noise than the unconstrained dynamics models: In the high noise condition, PLSSM dynamics still accurately predict almost 60% of the tran- sitions after 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, 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.
when we condition PLSM
In sum, parsimonious dynamies regularization improves both generalization and robustness properties in the three environments tested. The mutual information bottleneck on Å, therefore not only improves prediction accuracy for data in the training distribution, but may also allow the model to generalize better to ofitsiribbuuio data, and improve its robustness to noisy observations.
5 Related work
Our approach introduces a mutual information constraint between the latent states 2, and the latent dynamics Â 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 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-Leiblerr (KI.) divergence between latent states predicted by the model and latent states inferred from pixels. The KI. 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 A, themselves, but to the representations Z,. 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 '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.
variational
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 '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, 433. 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 Å itself as random variable, and minimize its conditional dependence on 2,. Furthermore, when A is fully disentangled from 2,, each action can be seen transformation that acts on the latent state space in the same way, invariantly of ot [44, 45J. We model the dynamics as softly state-invariant, allowing us to predict future latents both accurately and parsimoniously.
6 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 A, and the latent state representations 2, Optimizing this objective makes the effect of the actions on the agent's 's latent state more predictable. Combining our objective with different model classes contrastive 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 nigh-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 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 3 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 s history as input. Using our bottleneck here would correspond 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 A, transitions could be modelled by predicting transition matrices. Finally, recent recent advances in controllable gello-generation like Genie [47] and UniSim [48] both incorporate actions into their video models. Genie infers latent actions that explain transitions videos, and UniSim conditions on actions and language instructions generate controllable videos. Modelling the effect of actions in a parsimonious way could improve the accuracy and generation capabilities of such systems.
Acknowledgments
We thank the HCAI lab for valuable feedback and comments throughout the project. We especially grateful for feedback from Luca Schulze Buschoff and Can Demircan on an earlier version of the manuscript, and indebted to Julian Coda-Forno for helping name the model. This work was supported by the Institute for Human-Centered AI at the Helmholtz Center for Computational Health, the Volkswagen Foundation, the Max Planck Society, the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 011S18039A, and funded by the Deutsche Forschhungseemeinschatt (DFG, German Research Foundation) under Germany's 15/18.
Excellence
10

---

References
[1] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Ad- vances in neural information processing systems, 31, 2018.
[2] Danijar Hafner, Timothy Lillierap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. arXiv preprint arXiv:19/2.01603, 2019.
[3] Thomas Kipf, Elise Van der Pol, and Max Welling, Contrastive learning of structured world
models. arXiv preprint arXiv: 12247, 2019.
[4] Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference machine learning (ICML-11), pages 465 472.. Citeseer, 2011.
[5] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. arXiv preprint arXiv:2007.05922, 2020.
[6] Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203 04955, 2022.
[7] Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chandak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, et al. Understanding self- eeffprreictive learning for reinforcement learning. In International Conference Machine Learning, pages 33632 33656. PMLR, 2023.
[8] Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. arXiv preprint 1rXiv:190..12436 2019.
[9] Denis Yarats, Amy Zhang. Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency model-free reinforcement learning from images. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 10674-10681, 2021.
[10] Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 122) 1999.
[11] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint 180I 00690,, 2018.
[12] Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta. Repo: Resilient model- based reinforcement learning by regularizing posterior predictability. Advances in Neural Information Processing Systems, 36, 2024.
[13] Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting con- trol suite-a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722. 2021.
[14] Alex X Leee Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. Advances in Neural Information Processing Systems, 33:741-752, 2020.
[15] Deepak Pathak, Pulkit Agrawal, Alexci A Efros, and Trevor Darrell. -uriosity-driven exploration by self-supervised prediction. In International conference machine learning, pages 2778- 2787. PMLR, 2017.
[16] Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[17] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-29320 2021.

---

[18]
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentangle- ment testing sprites dataset, 2017.
[19]
Klaus Greff, Francois Belletti, Lucas Beyer, Carl Doersch, Yilun Du, Daniel Duckworth, David J Fleet, Dan Gnanapragasam, Florian Golemo, Charles Herrmann, Thomas Kipf, Abhijit Kundu, Dmitry Lagun, Issam Laradji, Hsueh-Ti (Derek) Liu, Henning Meyer, Yishu Miao, Derek Nowrouzezahrai, Cengiz Oztireli, Etienne Pot, Noha Radwan, Daniel Rebain, Sara Sabour, Mehdi S. M. Sajjadi, Matan Sela, Vincent Sitzmann, Austin Stone, Deqing Sun, Suhani Vora, Ziyu Wang. Tianhao Wu, Kwang Moo Yi, Fangcheng Zhong, and Andrea Tagliasacchi. Kubric: a scalable dataset generator. 2022.
[20] Zhen Wang, Jianwen Zhang, Jianlin Feng, and Zheng Chen. Knowledge graph embedding by translating on hyperplanes. In Proceedings of the AAAI conference on artificial intelligence, volume 28, 2014.
[21] Zhiqing Sun, Zhi-Hong Deng, Jun Nie, and Jian Tang. Rotate: Knowledge graph embed- ding by relational rotation in complex space. arXiv preprint arXiv: 10/97, 2019.
[22] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565. PMLR, 2019.
[23] Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, pages 1480-1490. PMLR, 2017.
[24] Arnb Kumar Mondal, Siba Smarak Panigrahi, Sai Rajeswar, Kaleem Siddiqi. and Siamak Ravanbakhsh. Efficient dynamics modeling in interactive environments with koopman theory. arXiv preprint arXiv:2306. / 1941, 2023.
[25] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.021933 2020.
[26] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv:1412.3555. 2014.
[27] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 13X.v.1322.6114, 2013.
[28] Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206. 15477, 2022.
[29] Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy, Deep variational information bottleneck. arXiv preprint arXiv:1612.000 2016.
[30] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pages 5171-5180. PMLR, 2019.
[31] Raef Bassily, Shay Moran, Ido Nachum, Jonathan Shafer, and Amir Yehudayoff. Learners that use little information. In Algorithmic Learning Theory, pages 25-55. PMLR, 2018.
[32] Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. The Journal of Machine Learning Research, (9(1) 2018.
[33] Alexander A Alemi. Variational predictive information bottleneck. In Symposium on Advances in Approximate Bayesian Inference, pages 1-6. PMLR, 2020.
[34] Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in dqn. arXiv preprint arXiv:1810.00123. 2018.
12

---

[35] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. In International conference on machine learning, pages 2170-2179. PMLR, 2019.
[36] Tankred Saanum, Noémi Éltetó, Peter Dayan, Marcel Binz, and Eric Schulz. Reinforcement learning with simple sequence priors. arXiv preprint arXiv:2305. /7109, 2023.
[37] Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Robust predictable control. Ad- vances in Neural Information Processing Systems, 34:27813-27825, 2021.
[38] Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew Botvinick, Hugo Larochelle, Yoshua Bengio, and Sergey Levine. Infobot: Transfer and exploration via the information bottleneck. arXiv preprint
2019.
[39] Chenjia Bai, Lingxiao Wang, Lei Han, Animesh Garg, Jianye Hao, Peng Liu, and Zhaoran Wang. Dynamic bottleneck for robust self-supervised exploration. Advances in Neural Information Processing Systems, 34:17007-17020, 2021.
[40] Felix Creutzig, Amir Globerson, and Naftali Tishby. Past-future information bottleneck in dynamical systems. Physical Review E, 79(4):044925, 2009.
[41] William Bialek. Ilya Nemenman, and Naftali Tishby. Predictability, complexity, and learning.
Neural computation, 13(11):2409-2485 2001.
[42] Naftali Tishby and Daniel Polani. Information theory of decisions and actions. In Perception- action cycle: Models, architectures, and hardware, pages 601-636. Springer, 2010.
[43] Nadav Amir, Stas Tiomkin, and Naftali Tishby. Past-future information bottleneck for linear feedback systems. In 2015 54th IEEE Conference on Decision and Control (CDC), pages 5737-5742. IEEE, 2015.
[44] Imre Risi Kondor. Group theoretical methods in machine learning. Columbia University, 2008.
[45] Robin Quessard, Thomas Barrett, and William Clements. Learning disentangled representations and group structure of dynamical environments. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 19727-19737. Curran Associates, Inc., 2020.
[46] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.00:: 000808088..0. 2022.
[47] Jake Bruce, Michael Dennis, Ashley Edwards, Jack Parker-Holder Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative interactive environments. arXiv preprint arXiv:2402.1533 15391, 2024.
[48] Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.06114, 2023.
[49] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data aata-auumented reinforcement learning. arXiv preprint arXiv:2107.00645, 2021.
[50] Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In lcml, 2010.
[51] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.690,, 2014.
[52] Jordi Pont-Tuset Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornun and Lue Van Gool. The 2017 davis challenge on video object segmentation. arXiv:1704.0067. 2017.
13

---

Appendix
A Mutual Information minimization
To simplify dynamics, our algorithm minimizes the mutual information between Z₁ and ÅS condi- tioned on ar (Eq. 8 & 3). With stochastic gradient descent, we search for parameters 0 that optimize the following objective:
min (Z; A, at)
(10)
To do so, we introduce a new variable h, that depends on 2, and a, (Bq. 6). We use this variable, together with a,, to produce Å,, instead of 24. For specific a,, the Markov chain behind A. can be written
Z, h - A.
(11)
Thus, to minimize the mutual information between and Á., instead minimize the mutual information between 2, and hy: ming (Ze; h be/a.). Following [29], we write the mutual information between 2, and h, follows:
/ da dh da p(z, h, a) log p(h|z, 22 a)
T(ze; helax) =
(12)
p(haa)
In case, p(h(z, a) is deterministic. The marginal distribution conditioned a, p(h(a), is challenging to obtain. We instead use a variational approximation: With our variational distribution (h a), we can upper bound the mutual information follows
= = f dz dh plz, a)p(hlz, a) log
p(h/z, 2 a)
I(ze; helas)
(13)
(hla)
First, this upper bound lets us see that the mutual information can be minimized by minimizing the number of bits of information z, contains about h, over and above a. This could potentially be done by introducing parameterized prior that depends a, go (he (a)) and minimize the KL divergence between these two probability distributions
K KLD(p(h,lzz a 1lgo (h, [a:)]
(14)
We opt for simpler approach, that allows us to do this without introducing new prior and parameterization. Following [8], we assume that gh,, aa)) is standard, d-dimensional, isotropic Gaussian distribution (O II). Assuming further that h, is another d-dimensional isotropic Gaussian, the KL divergence is a function of the mean and standard deviation of h, [8, 27]:
2KLD [|p((h))) d Σ o (h:); log (h;);
(15)
Here ((() and a(-) return the mean and standard deviation of the Gaussian, respectively. Importantly, we can minimize the KI. divergence by minimizing the first term the above equation, which is simply the norm of the mean of h. In our deterministic setup, we therefore simply minimize the norm of h, itself.
14

---

B PLSM vs L, and L2 norm regularization
CWM
1= 09
0=0.1
PLSM
-0.1
-14
-14
-19
-10
-35
2
B= 0.0005
0.001
= 0.005
L1
-07
-67
-03
-20
-21
-21
10
ß = 0.001
0=001
B = 0.1
L
-07
-67
-03
-22
-21
-23
37
To illustrate how our mutual information minimization impacts the latent space, we trained PLSM to learn the dynamics of dot moving in four directions gridworld with wall in the middle. We compared PLSM to contrastive models that regularize the L1 and L2 norms of the latent space, respectively. We observe that PLSM regularization leads to more regular representation of the states in the gridworld, whereas no regularization leads to a warped representation. Moreover, T. and La regularization simply shrinks the latent space. This type of shrinkage is not present in PLSM.
Figure 8:
C Ablations
We tested some alternative formulations and ablations of our PLSM objective subset of the datasets (see Fig. 9). One ablation removed the query representation h, and simply regularized the latent state 2,- This model achieved substantially worse performance in the three datasets tested it on. As an alternative to regularizing the L2 norm of h, we also tested a top-k bottleneck, which used the top 15 features from h, instead of the full representation. This model performed par with the original PLSM. Lastly, minimizing the L2 norm of h, could potentially be compensated for by the dynamics MLP by increasing the norm of the weight matrices. As control, added weight decay to the dynamics model to prevent this, and observe comparable performance.
D Is the L2 minimization effective?
One potential issue with minimizing the L2 norm of the query representation h, 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, of PLSM relative to the average norm of in the unregularized model. Comparing them we see indeed that h, 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, truly minimizes the amount information the model can extract from it to predict A.
15

---

cubes 9
shapes 9
sprites
100
100
Accuracy %
90
96
75
98
92
98
60
88
3
5
8
10
1
5
8
10
3
8
10
100
100
MRR %
98
100
90
95
75
92
98
60
8 10
5 Horizon
10
3
10
Horizon
Horizon
PLSM
Without ht
K-PLSM
PLSM (weight decay)
Figure PLSM and three alternative models, one dispensing of the query representation, one using top-k bottleneck, and one adding weight decay to the dynamics MLP. Removing the query representation decreases performance substantially. Performance intact in the two alternative formulations of PLSM.
spaceinvaders
pong
cubes
shapes 9
sprites
MOVI-A
122
1.0
10F
|linput|aa
10
20"11
20-3
-4
30-
30-
|[Wallz
10
0
Figure 10: While the norm of h, is orders of magnitude lower than the norm of zr, the norm of ensuing weights from the dynamics model are comparable and often not significantly different.
E
Atari
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 Atari. We see that in games where
regularization strenght of PLSS impacted performance important features of the environment are not controllable by the agent, weaker regularization is beneficial. For instance, in Boxing, the movement of the opponent hard to predict, and here having lower regularization strength is advantageous.
16

---

IQM
Mean
Optimality Gap
SPR PLSM SPR
DrQ
0.3
0.55
0.60
0.65
Figure 11: Median, IQM, Mean and Optimality Gap reported for the three model-free RL algorithms.
Human Norm'd d score
boxing
hero
0.3
3
0.2
2
0.1
1
0
0.0
5.0
0.1
1e-06
5.0
0.1
1e-06
B
Figure 12: In Boxing and Hero, lower regularization generally led better performance.
F
Datasets
See Fig. 13 for example frames from the training datasets and Fig. 14 for example frames from the generalization tests.
Figure 13: The six environments used to evaluate the accuracy of the latent dynamics. Top, from left right: dSprites 4, cubes 9m shapes 9, Bottom: pong, space invaders, and MOVi-e.
G Hybrid model
We propose hybrid model for the Atari environments. The hybrid model splits the latent state in two, one parsimonious state space z, and an unconstrained state space z? To predict the next state, the model uses two dynamics MLPs 2 zi + (z) 2?, a.) and 2 z! + dz (h,, az). The next latent prediction is then defined as 2,41 Concatenated ( 2 Results shown in Fig. 15.
H Slots with information regularization
We evaluated the C-SWM on our cubes task, both with and without our regularization. The performance of C-SWM dropper by small amount due to the increased difficulty of the task for
17

---

Figure 14: Datasets used to probe generalization and robustness. Left: Cubes and shapes data with varying numbers of objects. Right: dSprite images with mild and severe noise corruption.
spaceinvaders
pong
0.40
Accuracy
0.30
0.20
0.30
5
3
5
10
0.80
CWM
PLSM
MRR
0.70
0.40
0.50
0.30
Horizon
Horizon
Our hybrid model, which splits the latent space into parsimonious one and an uncon- strained one, outperforms the baseline on long-horizon latent prediction in the Atari environments.
Figure 15:
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 { 10, averaged over five seeds, plus and minus standard error of the mean.
Accuracy
PLSM
C-SWM
Cubes 99.05% #0.18
97.45%+1..
18

---

I Atari scores
Game
Random
Human
CURL
DrQ
SPR
SPR+PLSM
Alien
227.8
7127.7
711.0
865.2
841.9
832.2
Amidar
5.8
1719.5
113.7
137.8
179.7
152.
Assault
222.4
742.0
500.9
579.6
565.6
634.5
Asterix
210.0
8503.3
567.2
763.6
962.5
913.4
BankHeist
14.2
753.1
65.3
232.9
345.4
63.5
BattleZone
2360.0
37187.5
8997.8
10165.3
14834.1
13394.0
Boxing
0.1
12.1
0.9
9.0
35.7
28,4
Breakout
1.7
30.5
2.6
19.8
19.6
15.6
ChopperCommand
811.0
7387.8
783.5
844.6
946.3
450.6
CrazyClimbe
10780.5
35829.4
9154.4
21539.0
36700.5
30410.2
DemonAttack
152.1
1971.0
646.5
1321.5
517.6
477.2
Freeway
0.0
29.6
28.3
20.3
19.3
28.4
Frostbite
65.2
4334.7
1226.5
1014.2
1170.7
1371.9
Gopher
257.6
2412.5
400.9
621.6
660.6
557.0
Hero
1027.0
30826.4
4987.7
4167.9
5858.6
5763.2
Jamesbond
29.0
302.8
331.0
349.1
366.5
336.8
Kangaroo
52.0
3035.0
740.2
1088.4
3617.4
4886.6
Krull
1598.0
2665.5
3049.2
4402.1
3681.6
4043.4
KungFuMaster
258.5
22736.3
8155.6
11467.4
14783.2
15187.6
MsPacman
307.3
6951.6
1064.0
1218.1
1318.4
1156.7
Pong
-20.7
14.6
-18.5
-9.1
-5.4
1.1
PrivateFye
24.9
69571.3
81.9
3.5
86.0
85.8
Qbert
163.9
13455.0
727.0
1810.7
866.3
786.5
RoadRunner
11.5
7845.0
5006.1
11211.4
12213.1
12400.0
Seaquest
68.4
42054.7
315.2
352.3
558.1
561.3
UpNDown
533.4
11693.2
2646.4
4324.5
10859.2
29572.0
#Superhuman
0
N/A
2
3
6
6
Mean
0.0
1,000
0.261
0.465
0.616
0.671
19

---

Hyperparameters
J.1
Compute
We ran all experiments reported in the paper on compute nodes with 2 Nvidia A100 GPUs. On average, Atari runs lasted for 5 hours and DMC runs 4-9 hours depending the task.
J.2 Convolutional neural network architecture
To learn pixel representations we use a Convolutional Neural Network (CNN) architecture similar to the one used in [9] and [49]. For the Atari latent prediction experiments we stacked two frames to provide information about object movements. MOVi-E we stacked four frames.
import torch
from torch import mn encoder = nn Smquenttt (
mn nn. Conv2d (num_ channels, 32, 3, stride=2), ReLUeLU() (),
nn. Conv2d (32, 32, 3, stride-1), nn Re.U (),
Conv2d (32, 32, 3, stride=1), nn ReRU())))))))) (),
Conv2d (32, 32, 3, stride=1), ReLU
J.3
Contrastive models
Table 2: Contrastive model hyperparameters.
Hyperparameter
Value
Hidden units
512
Batch size
512
MLP hidden layers
2
Latent dimensions [ze
50
Query dimensions h
50
Regularization coefficient
0.1
Margin À Learning rate
0.001
Activation function Optimizer
ReLU [50]
Adam [51]
J.4 DeepMind Control Suite
We use the same hyperparameters [6] for the planning agents. The only addition is our PLSM regularization. We use regularization coefficient of 0.1 for all agents, and make the query net MLP fe the same size the dynamics network h, with [h, /zll-
Table 3: Action repeat values in the DeepMind Control Suite tasks.
Task
Action repeat
Acrobot Swingup
4
Finger Turn Hard
2
Quadruped Walk
4
Quadruped
Humanoid Walk
2
20

---

J.5 Distracting Control Suite
All models were trained with an action repeat value of in all environments. All model hyperparame- ters were the same is in [12]. The background videos were randomly sampled per episode from the DAVIS 2017 video dataset [52], projected in black and white.
21

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We include experiments and results for all claims made in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations of our model in the Discussion section.
22

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Mathematical details of the model are provided throughout the paper and in the appendix.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
All experiment details and hyperparameters included in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
23

---

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 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. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 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 NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please see the NeurIPS code and data submission guidelines (uttps://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. h/nips.cc/pppp wttp:cccpppppp.cccccceocoooppcpa httpa///ggiiiee///eebbbiiiiiiii Coodeebmeiiieeeeeeeeee sesiisssoosioioieoe
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
24

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
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 the core the paper to 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, supplemental material.
7. 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 accompanied by enror bars, confi- dence 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.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute details reported in the Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
25

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //pppr////uuips.... Spppblic//thissuiieeeee......e
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We found potential societal impacts worth discussing.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 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 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 models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. 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, scraped datasets)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
26

---

Released models that have high risk for misuse or dual-use 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset used and, if possible, include 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, paperswi thcode. www datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- 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 's creators.
13. 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 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.
14. 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 details about compensation (if any)?
Answer: [NA]
Justification:
27

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and 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 the institution conducting the review.
28


Reviews:
[{"review_id": "Xhh1oEP7xh", "reviewer": "Reviewer_PXJB", "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.", "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": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The methods and experimental setup were clear and well-presented. The authors performed extensive experiments both studying the properties of their learned latent space and on downstream task performance. Their method appears to perform extremely well when the state of the world is fully controllable by the agent. \n\nAlthough the method does not improve in all settings in which the authors evaluated it, they do not tune parameters for each environment, leaving room for potential improvement.", "weaknesses": "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. \n\nIn 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n>Have you run the TD-MPC experiments longer?\n\nWe 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.\n\nWe 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).\n\n>It would have been nice to see some explanation or hypothesis as to why some of the Atari games performance degraded\n\nWe 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:\n\n**\"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\"**\n\nReferences:\n\n[1] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021\n\n[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MWv2zty7iR", "reviewer": "Reviewer_gEcM", "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.", "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$\".\n- In Equation (5), should there be a comma between $z_t$ and $a_t$?\n- Line 111: Could the authors elaborate a bit on how (9) is paired with (8)?\n- 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$\".\n- Line 185: \"We evaluated the PLSM's effect\": remove \"the\" to be consistent with other places in the paper?", "limitations": "Limitations are commented above. No societal impact needs to be particularly addressed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The idea of learning state-invariant latent dynamics, and the way of introducing an information bottleneck to achieve it, are novel. This work contains an interesting exploration toward this idea.\n- Extensive experiments are done to support the proposed methods, including two settings in various environments.", "weaknesses": "- There seem to be important flaws in the proposed PLSM.\n  - 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?\n  - 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.\n- 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.\n- The writing has much room for improvement. \n  - 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.\n  - Some figures are dangling in the paper without being referenced, e.g., Figure 1.\n  - Equations are better referenced in parentheses like (1)(2); in LaTeX the command is `\\eqref`.\n  - Some other minor problems are detailed in the \"Questions\" section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n>First, the math derivation in Appendix A contains several loose ends\n\nWe 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).\n\n>I have concerns about the fundamental validity of the regularization on $||h_t||$\n\nThe 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).\n\n>The experiments do not appear to show that the proposed PLSM has a significant advantage\n\nWe 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).\n\n> It is unclear to me why Section 3 is a standalone section\n\nWe 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.\n\n>Some figures are dangling in the paper without being referenced\n\nThank you for pointing this out, we now reference the first figure in the following sentence in the introduction.\n\n**\"Here we explore the possibility of compressing states and dynamics jointly to learn systematic effects of actions (Fig 1.)\"**\n\n>Line 111: Could the authors elaborate a bit on how (9) is paired with (8)? \n\nTo clarify the workings of our model, we expounded on the connection between (8) and (9) in the paper.\n\n**\"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).\"**\n\nWe 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”.\n\nReferences:\n\n[1] From Variational to Deterministic Autoencoders, Ghosh et al., ICLR 2020\n\n[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023\n\n[3] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k7szflryAa", "reviewer": "Reviewer_itzz", "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.", "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.\n- 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": "Yes, the limitations are well-discussed by the authors.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem is well-motivated, and the solution is straightforward to implement.\n- The experiments cover a rather broad range of problems which showcase the general applicability of the method.", "weaknesses": "- While PLSM exhibits strong performance on specific prediction tasks, the improvements in RL tasks are relatively modest.\n- As is common with information bottleneck-based methods, PLSM's performance may be highly sensitive to the hyperparameter of the regularizer for the bottleneck.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n>While PLSM exhibits strong performance on specific prediction tasks, the improvements in RL tasks are relatively modest. \n\nWe 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).\n\n\n>Could you provide Interquartile Mean (IQM) results for the Atari experiments, as suggested by Agarwal et al. (2021)?\n\nWe 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.\n\n>Since the regularisor weights $\\beta$ is very important to the performance, I wonder how robust is PLSM to different $\\beta$?\n\nAfter 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$.   \n\nReferences:\n\n[1] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021\n\n[2] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QsGX1vms0Z", "reviewer": "Reviewer_hbwZ", "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.", "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": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The proposed method reasonably controls the information contained in the latent space, regularizing the behavior of the latent space dynamic model and allowing for strong generalization.", "weaknesses": "- The contribution is limited since the MI regularization for RL representation has been adopted in several previous works, including [1-3].\n\n- The empirical results presented in this paper do not show significant performance improvements. \n\n\n[1] Learning Robust Representations via Multi-View Information Bottleneck. ICLR 2020\n\n[2] Dynamic Bottleneck for Robust Self-Supervised Exploration. NeurIPS 2021\n\n[3] InfoBot: Structured Exploration in ReinforcementLearning Using Information Bottleneck. ICLR 2019", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n>The contribution is limited since the MI regularization for RL representation has been adopted in several previous works, including [1-3].\n\nWe have extended the related works section with the following paragraph:\n\n**\"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.\"**\n\n\n>Since the MI between representation and state is constrained, should the algorithm exhibit stronger robustness in various domains? \n\nWe 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. \n    \n>The empirical results presented in this paper do not show significant performance improvements. \n\nWe 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.\n\nReferences:\n\n[1] Generalization and Regularization in DQN, Farebrother et al., ArXiv 2018\n\n[2] DeepMDP: Learning Continuous Latent Space Models for Representation Learning\n\n[3] Learning Robust Representations via Multi-View Information Bottleneck, Federici et al.,  ICLR 2020\n\n[4] Robust Predictable Control, Eysenbach et al., NeurIPS 2022\n\n[5] Dynamic Bottleneck for Robust Self-Supervised Exploration, Bai et al., NeurIPS 2021\n\n[6] InfoBot: Structured Exploration in ReinforcementLearning Using Information Bottleneck, Goyal et al., ICLR 2019\n\n[7] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XMMqR9j2NK", "reviewer": "Reviewer_VZsj", "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.", "questions": "Please see the weaknesses section above.", "limitations": "The authors sufficiently address the limitations in Section 7.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Generally the paper is well-written, and I enjoyed reading the draft.\n* The related work section extensively discusses the difference between PLSM and previous approaches.\n* The proposed method is relatively easy to implement.\n* The authors experimentally show that PLSM improves long-time prediction accuracy.\n* The experiments are done across three very different settings (long-term prediction, model-based RL, and model-based RL).", "weaknesses": "* 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).\n* 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?\n* 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.\n* 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.).\n\nIn 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n>I'm not sure how simply minimizing the norm $h_t$ of leads to minimal representations\n\nWe 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).\n\n>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$? \n\nTo 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.\n\n>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\n\nWe 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:\n\n**\"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$.\"**\n\n>it's unclear how PLSM compares to other regularization methods (e.g., RePo, other regularization techniques, etc.)\n\nWe 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).\n\nReferences:\n\n[1] PILCO: A Model-Based and Data-Efficient Approach to Policy Search, Deisenroth et al., ICML 2011\n\n[2] Data efficient Reinforcement Learning With Self-Predictive Representations, Schwarzer et al., ICLR 2021\n\n[3] RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability, Zhu et al., NeurIPS 2023\n\n[4] The Distracting Control Suite – A Challenging Benchmark for Reinforcement Learning from Pixels, Stone et al., ArXiv, 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all of the reviewers for the time and effort put into providing thoughtful and constructive feedback on our paper. Reviewers generally found our method novel and interesting.

* Reviewer PXJB wrote that ‘The authors performed extensive experiments’ and that our method ‘appears to perform extremely well when the state of the world is fully controllable by the agent.’
* Reviewer gEcM wrote that ‘idea of learning state-invariant latent dynamics [...] are novel’ and that ‘Extensive experiments are done to support the proposed methods’

* Reviewer itzz said that our experiments ‘cover a rather broad range of problems which showcase the general applicability of the method.’

* Reviewer hbwZ wrote that our method is ‘regularizing the behavior of the latent space dynamic model and allowing for strong generalization.’

* Reviewer VZsj wrote that our paper ‘is well written’, including experimental evaluation ‘done across three very different settings’ 

At the same time, reviewers raised some important concerns. Reviewers hbwZ, gEcM and VZsj pointed out that PLSM's performance improvements are sometimes modest. Reviewers VZsj and gEcM also wondered if the sparsity of the query representation could be absorbed by having larger weights in the ensuing linear layer. Reviewers also suggested that we evaluate our regularization relative to other robust RL baselines, specifically RePo, ablate the query representations, run DMC experiments for more time steps, and analyze the model’s sensitivity to the regularization strength. We performed a series of experiments to address these points, and show the results in the attached supporting PDF document. 

* First, as requested by reviewers itzz, hbwZ, gEcM and VZsj, we show that PLSM can offer even stronger performance improvements in challenging visual control tasks with distracting visual cues (see Supporting Figure A).  We compare PLSM to RePO [1], a robust model-based RL algorithm, on tasks from the Distracting DeepMind Control Suite [2], where the background is replaced with a distracting, irrelevant video. We incorporated our regularization on top of RePO using the official implementation. Evaluating the models across five environments with five seeds each, we see that using our regularization offers considerable improvements in scores. This makes sense, since PLSM compresses away aspects of the environment that are hard to predict given the agent’s actions.

* We show empirically that PLSM does not compensate for the sparsity by adopting larger weights (requested by VZsj and gEcM, see Supporting Figure D). We also propose two alternative implementations of PLSM where artificially increasing the norm of the weights cannot compensate for our regularization. Both of these alternative implementations perform equally well, and sometimes better than the original PLSM (see Supporting Figure E).

* We extend the number of training steps in the DeepMind Control Suite experiments, showing faster and improved convergence (requested by PXJB, see Supporting Figure B).

* We provide more detailed results in Atari, along with Interquartile Mean estimates (requested by itzz, see Supporting Figure. C).

* We analyze how the regularization strength impacts performance in Atari, and offer an explanation as to why performance drops for some games (requested by itzz and PXJB, see Supporting Figure F).

* We show that removing the query net substantially decreases PLSM performance across several datasets (requested by VZsj, see supporting Figure E).

* We have cleaned up notation and clarified how we arrive at our regularization (requested by gEcM).



All of these analyses and results are outlined in detail in our responses to the individual reviews below, where we also address more specific questions reviewers had. We again want to thank the reviewers for their time and for actively engaging in the review process.

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

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

Meta-review: 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.
Decision: Accept (poster)