Title: Simplifying Latent Dynamics with Softly State-Invariant World Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: World model

latent dynamics

reinforcement learning

compression
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

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.

1 Introduction

In Reinforcement Learning (RL), the actions that an agent can use to solve tasks often have systematic and predictable effects on the state of the environment. When stepping on the gas pedal, the car tends to accelerate, when holding down the joystick in a given direction, the video game character tends to move in that direction, and so forth. These typical effects have exceptions that are predictable as well, e.g. stepping on the gas will not lead to acceleration if the car engine is turned off, and the video game character will not move if it is facing a wall. How can we learn world models that capture these systematic properties of actions? World models predict the agent's future states, given the current state and action [1, 2, 3]. Most approaches to world modelling represent high-dimensional observations (such as images) in compact, low-dimensional latent states z t , simplifying model-based prediction and control [4]. Here we explore the possibility of compressing states and dynamics jointly to learn systematic effects of actions (see Fig. 1). As is common practice in many dynamics model architectures [4, 5], we consider the case where the model predicts the next latent ˜ z t +1 state by predicting the difference ˜ ∆ a t t , or the change , between the current and future latent state, given an action a t .

Even if z t is low-dimensional, the effects of actions might not be represented parsimoniously within the world model: Performing the same action a t in two similar states z , z ′ might produce two very different deltas ∆ , ∆ ′ , due to how the observation encoder has constructed the latent state space (see

a1

a2

h1

1

h2

2

Distracting

Model-free RL

Control

(Atari)

Score

0.90

600

Return

Norm'd

0.75

400

200

Human

0.60

RePo+PLSM

0.45

Step

SPR

SPR+

PLSM

DrQ

Z3

Z1

DMC

750

Return

500

250

TDMPC

TDMPC+PLSM

Step

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 ˜ ∆ a t t depend on z t . This pushes the world model to represent states in a way that makes actions have coherent and predictable effects on the dynamics. We still allow the dynamics to vary depending on the state, but we penalize the extent of this dependence, resulting in dynamics that are softly state-invariant.

We combine PLSM with two classes of world models: Contrastive World Models (CWM) [3] for latent state prediction, and with Self Predictive Representations (SPR) for model-free and modelbased control (TD-MPC) [6, 5, 7]. Across control experiments and prediction experiments we see improvements in planning, representation learning for control, robustness to noise, world model accuracy, and generalization.

2 Latent dynamics

We assume that sequences of states, actions and rewards arise in a Markov Decision Process (MDP). AnMDPconsists of a state space S , an action space A , and transition dynamics s t +1 ∼ P ( s t +1 | s t , a t ) determining how the state evolves with the actions the agent performs. In RL, we additionally care about the reward function r ( s t , a t ) , which maps state-action pairs to a scalar reward term. Here, the goal is to learn the policy π θ ( a t | s t ) that maps states to the actions with the highest possible Q -values Q π θ ( s t , a t ) = E π θ [ ∑ T t =1 γ t r ( s t , a t ) ] , where γ is a discount factor. In this paper, we consider latent dynamics learning both in reward-free and RL settings.

Unconstrained

Parsimonious

0.35

0.18

21

0.00

0.00

-0.35

-0.18

-0.20

0.00

0.23

-0.14

0.00

0.13

22

22

To predict environment dynamics, the current state s t is first transformed using an encoder into a latent state z t that compactly represents the agent's sensory observation (2). The world model predicts the change in the latent state, ˜ ∆ a t t that is induced by the action a t (3) & (4).

We omit the action superscript from ˜ ∆ t for simplicity. Here, e θ is the encoder network mapping states to latent states, and d θ is the dynamics network mapping z t and a t to ˜ ∆ t .

2.1 Parsimonious latent dynamics

Consider the probability distribution of transitions P ( ˜ ∆ t | a t ) given the agent's actions, marginalizing across latent states. If actions have predictable effects on the state of the environment, the entropy of P ( ˜ ∆ t | a t ) will be low - knowing the latent state gives little information about the effect of a t . To make the world model handle actions more systematically, we propose to minimize the amount of information the world model needs from z t in order to predict correctly how an action changes the state of the world. This quantity is represented in the mutual information between the latent state z t and the dynamics ˜ ∆ t :

where H [ · ] denotes the Shannon entropy. If this quantity is 0, the latent dynamics ˜ ∆ t only depend on the agent's action a t and not z t . In this extreme, all ˜ ∆ t are predicted exclusively by the action. However, it is rarely the case that the action can capture an environment's full dynamics, and making dynamics contingent on states is often necessary to some degree.

To allow only the relevant information from z t to influence the dynamics, we introduce a query network f θ which maps latent state-action pairs to a latent code h t . We modify the next-step prediction components accordingly

We give the query network information about the action that the transition is conditioned on. This allows the network to attend to the relevant bits in z t to output an appropriate h t . Finally, to make

h t represent only the minimal amount of information needed to predict the next state, provided that a t is known, we penalize the norm of h t [8, 9]. This type of regularization has been used to constrain representations of deterministic Autoencoders in past work [8, 9], with [8] showing that it is equivalent to minimizing the KL divergence to a constant variance zero mean Gaussian. Other regularizers and stochastic formulations are also possible (see Appendix C). For simplicity we use the deterministic variant and leave stochastic versions for future work.

The strength of the penalization is controlled through a hyperparameter β . This regularization minimizes how much h t can vary with z t and hence their mutual information (see Appendix A). We then train our encoder e θ , query network f θ , and dynamics d θ jointly to minimize the following information regularized loss function.

Our loss function encourages that the mutual information term is kept as low as possible, while still allowing the model to predict the next latent state accurately. In contrast to information bottlenecks imposed on the latent states themselves, we apply an information bottleneck to the dynamics, making ˜ ∆ t easier to predict simply given a t . Regularizing h t differs from regularizing z t in important ways. In environments where the dynamics ∆ can be predicted perfectly from the actions and independently of the state, regularizing h t will not lead to a loss in information in the latent representation z t . This is because the bottleneck on h t only constrains the model in using information from z t to predict ∆ t , and not necessarily in predicting z t +1 . See Appendix B for a comparison of our method against L 1 and L 2 norm regularization on latent states. Notably, our method can also lead to state compression, in that e θ will be encouraged to omit features from s t +1 that cannot be predicted easily with little information from z t .

Unfortunately, the above loss function has a trivial solution: it can be minimized completely if d θ and e θ output a constant 0 vector for all states and state-action pairs [7]. This issue is referred to as representational collapse . Representational collapse can be remedied in various ways. To show the generality of our information bottleneck, we combine it with two different approaches for mitigating representational collapse, a self-supervised approach for model-based and model-free RL in Section 3, and a contrastive approach for future state prediction in Section 4.

3 Parsimonious dynamics for Reinforcement Learning

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-Oriented Latent Dynamics (TOLD) model. This dynamics model is trained to predict its own future state representations from an initial state and action sequence while making sure that the controller's policy π θ and Q -value function are decodable from the latent state (hence the name task-oriented latent dynamics). TOLD falls under the category of Self Predictive Representation (SPR) models, since it uses an exponentially moving target encoder e -θ with the stop-gradient operator to learn to predict its own representations. For planning TD-MPC uses the Cross-Entropy Method [10], searching for actions that maximize Q -values.

Only minimal adjustments to the TOLD model are necessary to attain parsimonious dynamics. Instead of predicting the next latent directly from the current latent and action ˜ z t +1 = d θ ( z t , a t ) , we use a query network f θ , mapping latent state-action tuples to h t and then minimize

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 environments: i ) acrobot-swingup , due to its challenging and chaotic dynamics. ii ) finger-turn hard , which poses a challenging exploration problem that TD-MPC was found to struggle with. iii ) quadruped-walk , iv ) quadruped-run and v ) humanoid-walk due to the high-dimensional dynam-

A

DMC

acrobot-swingup...

finger-turn hard

quadruped-walk

quadruped-run

humanoid-walk

1000

1000

Episode return

WW

800

800

450

750

750

600

600

300

500

500

400

400

150

250

250

200

200

0

0.00.2

0.5 0.7

1.0

0.00.2

0.5 0.7

1.0

0

0

123

5

0

123

5

0

0.00.30.60.91.2

Step

1e6

Step

1e6

Step

1e5

Step

1e5

Step

1e6

TD-MPC

TD-MPC+PLSM

Distracting Control Suite

B

walker-walk

cartpole-balance

hopper-stand

walker-run

finger-spin

1000

return

800

600

400

450

800

600

300

300

Episode

600

400

400

200

400

200

150

200

100

200

0

0

0

0

0.0

0.5

1.0

0

2

4

0.0

0.5

1.0

0.0

0.5

1.0

0.0

0.5

1.0

Step

1e6

Step

1e5

Step

1e6

Step

1e6

Step

1e6

RePo+PLSM

ics. These tasks have dynamics that appear complex in the original state-space but could potentially be simplified in an appropriate latent space by introducing a dynamics bottleneck.

We trained the latent dynamics and planning models in the five tasks up to a million environment steps. Scores for TD-MPC are obtained from the original implementation provided by the authors 1 . Again we used β = 0 . 1 for all tasks except for humanoid-walk , where β = 0 . 001 was more successful. Otherwise we relied on the standard hyperparameters from [6]. We see clear performance gains due in all tasks except quadruped-run (Fig. 3A). Our results suggest that modeling the world with simple dynamics can be beneficial for RL and trajectory optimization.

3.2 Distracting visual control

Since our regularization compresses away aspects of the environment whose dynamics are unpredictable given the agent's actions, it could be beneficial in control tasks with distractors. We implemented PLSM on top of RePo [12], relying on the authors' official implementation 2 . RePo is a model-based RL algorithm based on the Dreamer [2] architecture. Here the environment dynamics are represented through a GRU network which is updated recurrently with latent state and action variables. To make the dynamics parsimonious, we update the GRU using a compressed query representation h t subject to L 2 regularization instead of the full state representation z t , resulting in recurrent dynamics that are softy state-invariant. We then trained RePo with and without our regularization on the Distracting Control Suite [13], a challenging visual control benchmark based on DMC, where the background is replaced with a random, distracting video from the 2017 Davis video dataset. These videos are independent of the agent's actions, irrelevant for rewards, and change from episode to episode.

We trained RePo with and without the PLSM objective in five Distracting Control Suite tasks for one million environment steps across five seeds. We used a regularization coefficient of β = 1 e7

Human Norm'd Score

1.6

0.8

0.0

-0.8

UpNDown

Kangaroo

Krull

Freeway

Pong

Assault

Frostbite

RoadRunner

KungFuMaster-

Seaquest

PrivateEye

Alien

Hero

Asterix

Amidar

DemonAttack

MsPacman

BattleZone

Gopher

ChopperCommand

Jamesbond

Breakout

CrazyClimber

BankHeist

Boxing

for all tasks. This produced a considerable improvement over RePo in more challenging tasks like hopper-stand , walker-run and finger-spin (see Fig. 3B). Our results suggest that encouraging the dynamics model to represent the effects of actions more consistently can improve its generalization ability in environments with distractors.

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 in way that makes state transitions more predictable, it may discover useful state features and ignore aspects of the environments that would make the dynamics unpredictable otherwise. We build upon the SPR implementation for Atari 3 [16], which uses a latent dynamics model for next latent state prediction. We alter the architecture of this latent dynamics model in the same way we did for TD-MPC, and add the h t norm to the loss function. Setting β = 5 and leaving all other hyperparameters as per the standard implementation, we train the PLSM augmented SPR algorithm on 100k environment steps across 5 seeds on all 26 Atari games. We use the SPR scores reported in [17] as our baseline.

Across several games we see substantial improvements to human normalized score (see Fig. 4). Averaging over all games, using PLSM dynamics improves human normalized scores by 5.6 percentage points (61.5 % for SPR vs 67.1% with PLSM). While we see improvements in games such as UpNDown and Kangaroo, there are other games where the regularization impacts performance negatively. Performance could potentially improve by fine-tuning the regularization strength for these domains. See Appendix I for the full score table.

4 Future state prediction

We found that our regularization improved model-free and model-based performance across several environments. Next we evaluated whether PLSM also generally improves world models' long-horizon prediction accuracy in latent space. Using the evaluation framework and environments from [3] (see Fig. 13 for example observations), we generated datasets of image, action, next-image triplets from two Atari games (Pong and Space Invaders), and grid worlds with moving 3D cubes and 2D shapes, where each action corresponds to moving an object in one of four cardinal directions. To make the learning tasks more challenging, we increased the number of movable objects from 5 to 9. Additionally we created an environment based on the dSprite dataset [18], with four sprites

spaceinvaders

pong

cubes 9

shapes 9

dSprites 4

MOVi-e

100

%

40

60

96

66

96

30

80

45

90

98

90

20

60

84

96

84

30

10

1

35

810

135

810

135

810

1

3

5

810

135

810

15101520

80

60

100

100

96

%

96

99

90

70

MRR

45

92

98

88

75

60

80

30

88

98

60

50

1

35

810

135

810

1

35

810

1

35

810

1

35

810

15101520

Horizon

Horizon

Horizon

Horizon

Horizon

Horizon

CWM

PLSM

traversing latent generative factors on a random walk. The sprites had 6 generative factors: Spatial x, y coordinates, scale, rotation, color, and shape. The sprites could vary in coordinates, scale, and rotation within episodes, and additionally vary in color across episodes. Lastly, we evaluate PLSM on a dynamic object interaction dataset with realistic textures and physics without actions, MOVi-E [19], to see if our method can be beneficial in action-free settings.

Following [3], we pair PLSM with a contrastive loss function to mitigate representational collapse: The contrastive loss encourages that different states are distinguishable in the latent space. Given a latent state and action, the model should minimize the distance between the predicted and true future latent state, while maximizing the distance between the predicted future latent state z t +1 and all other latent states in the training batch z -, up to a margin λ .

To apply our regularization on the contrastive dynamics model, we simply add the norm of the query representation to the contrastive loss, similarly to equation (8). We fitted the regularization coefficient β with a grid search and found a value of 0 . 1 to work the best. As a baseline we used the unregularized contrastive model from [3], referred to as CWM (for Contrastive World Model), and its dynamics are defined through Equation (2), (4) and (9). We also combine PLSM with the slot-based version of this model, called C-SWM (see Appendix H for results).

The models were scored based on their ability to correctly predict future states (e.g. latent prediction accuracy). The models were trained and evaluated using the same parameters and metrics as in [3]: Given a state s t , a sequence of N actions a t , ..., a t + N -1 , and the resulting state s t + N , we make the model predict its latent representation of s t + N from the initial state and action sequence. The models were evaluated at several prediction horizons. Finally, we report the Hits at Rank 1 accuracy for transitions in the test set, a common test metric for contrastive models [3, 20, 21].

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, all 9 objects can collide with each other and the grid boundaries depending on their position. Always considering all possible interactions makes it challenging for a model to generalize to novel transitions. A more parsimonious solution is simply to represent whether or not the object in question would collide or not if moved in the direction specified by a t . Our results suggest that our regularization can help learn such representations, affording better generalization to left-out transitions.

Novel cubes

Novel shapes

Noise: 0= 0.1

Noise: 0= 0.2

1001

1001

60

%

40

80

80

40

20

20

60

60

0

0

40

40

PLSM

CWM

PLSM

CWM

1

3

5

8

10

1

3

5

8

10

Horizon

Horizon

CWM

PLSM

In one environment, Pong, we do not see an advantage in encouraging parsimonious dynamics. Here, various components are outside of the agent's sphere of influence, for instance, the movement of the opponent's paddle. This makes it challenging for the PLSM to capture all aspects of the environment state in its dynamics. In environments with non-controllable dynamics, we offer a remedy by only enforcing half of the latent space to be governed by parsimonious dynamics, and allowing the other half to be unconstrained. This hybrid model in turn shows the strongest performance in the Atari environments. See Appendix G for details.

4.2 Generalization and robustness

Generative factor decodability

0.7

0.6

0.5

2

0.4

R

0.3

0.2

0.1

0.0

Latent

Objectmatched

Query

state

querystate

state

Next, we evaluated the generalization and robustness properties of PLSM. For the cubes and shapes environments, we generated novel datasets where the number of moving objects was lower than in the original training data. In these datasets, the number of moving objects varied from 1 to 7. We also probed the models' robustness to noise. We corrupted test data from the dSprite environment with noise sampled from Gaussian distributions. Models were tested both in a low noise ( σ = 0 . 1 ) and a high noise ( σ = 0 . 2 ) condition. See Supplementary Fig. 14 for example data.

When tested on scenes with fewer objects than the training data, we see a general decrease in accuracy because of the domain shift. Still, with our information-theoretic dynamics bottleneck the model generalizes significantly better to the out-of-distribution transition data (Fig. 6). Since PLSM seeks to predict the next state using as little information from its latent representations as possible, it is more likely to learn that the blocks move in a way that is generally invariant to the number of other blocks in the scene.

Lastly, PLSM proved more robust to Gaussian noise than the unconstrained dynamics models: In the high noise condition, PLSM dynamics still accurately predict almost 60% of the tran- sitions after 10 steps, whereas the unconstrained dynamics models, more susceptible to use noisy information in the state, predict less than 50% accurately.

We also investigated the representations learned by the query network (see Fig. 7). To verify that PLSM learns objects' attributes (such as position and orientation), we attempted to decode ground truth object positions, scales and orientations in the dSprite dataset from PLSM latent states. We find that one can decode these ground truth factors with high accuracy from the latents. However, trying to decode these variables from the query states yielded substantially lower accuracy, indicating that they contain less information about the generative factors of the environment. Interestingly, when we condition PLSM on an action that affects only one object, we can decode attributes of that object more accurately from the query state than average. We observe this because the query state is designed to only encode information that is relevant to predict the effect of individual actions.

In sum, parsimonious dynamics regularization improves both generalization and robustness properties in the three environments tested. The mutual information bottleneck on ˜ ∆ t therefore not only improves prediction accuracy for data in the training distribution, but may also allow the model to generalize better to out-of-distribution data, and improve its robustness to noisy observations.

5 Related work

Our approach introduces a mutual information constraint between the latent states z t and the latent dynamics ˜ ∆ t inferred by the model. Several methods focus on state compression for dynamics modeling [1, 22, 5, 6, 23, 24]: The Recurrent State Space Model (RSSM) [22, 2, 25], uses a variational Auto-Encoder in combination with a recurrent model (e.g. a GRU [26]) to infer compact latent states in partially observable environments. Latent consistency is enforced by minimizing the Kullback-Leibler (KL) divergence between latent states predicted by the model and latent states inferred from pixels. The KL term regularizes the latent state not to contain more information than can be predicted by the dynamics model [2, 27]. Unlike our approach, this information bottleneck is not applied to the dynamics ˜ ∆ t themselves, but to the representations z t . Expanding on this line of work, RePo [12] discards image reconstruction from the RSSM, and simply enforces that the model can reconstruct the environment's reward function, leading to stronger compression and improved performance in tasks with unpredictable elements. Similar approaches like Denoised MDP [28] only model controllable and reward-relevant aspects of the environment. While simplified latents can make transition dynamics more tractable to model, they do not necessarily give rise to the systematic action representations that we are interested in. Lastly, Self Predictive Representation (SPR) models [7, 6, 5] learn dynamics by predicting the future representations of a target encoder. SPR models have been used both for model-free and model-based control.

Mutual information minimization is used in many deep learning frameworks more generally: [29] and [30] use variational methods for minimizing the mutual information between the network's inputs X and latent representations Z while maximizing the mutual information between representations Z and outputs Y . Mutual information minimization also has links to generalization ability [31, 32, 33, 34, 35], robustness in RL [36, 37], and exploration [38, 39]. Our information bottleneck differs from previous approaches in that it directly constrains the effect the latent state can have on the residual term in the latent dynamics over and above the agent's actions.

Closest to our regularization method is the past-future information bottleneck [40, 41]. Here the mutual information between sequences of past states and future states is minimized [42, 43]. While this method simplifies dynamics, our approach differs in important ways: Rather than representing the environment's dynamics, say, using a low number of its principal components, we treat the dynamics operator ˜ ∆ t itself as a random variable, and minimize its conditional dependence on z t . Furthermore, when ˜ ∆ t is fully disentangled from z t , each action can be seen as a transformation that acts on the latent state space in the same way, invariantly of z t [44, 45]. We model the dynamics as softly state-invariant, allowing us to predict future latents both accurately and parsimoniously.

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 ˜ ∆ t and the latent state representations z t . Optimizing this objective makes the effect of the actions on the agent's latent state more predictable. Combining our objective with different model 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 in high-dimensional and pixel-based environments with complex dynamics. Overall, our results suggest that systematic action representations can offer important improvements to the generalization ability and data-efficiency of world models.

Limitations: Our model, in its current formulation, assumes that actions have predictable and deterministic effects on the environment. Aspects of the environment that do not behave predictably conditioning on actions are susceptible to be ignored by the model, even if they are relevant for the downstream task. While performance in Atari was improved on average, this aspect led to reduced performance in some tasks. Similarly, the degree of regularization β needs to be tuned to achieve the right level of dynamics compression for different environments.

Future directions: A promising future direction is to combine our mutual information regularization with recurrent models. In partially observable, non-Markovian environments, next-state prediction is often done with a recurrent model that uses the agent's history as input. Using our bottleneck here would correspond to the assumption that the effect of the agent's actions are softly history -invariant. Another promising avenue of research is to combine PLSM with discrete dynamics. Many successful world modelling approaches assume that the latent state space is categorical [25, 46]. Instead of modelling transitions using affine transformations ∆ , transitions could be modelled by predicting transition matrices. Finally, recent recent advances in controllable video-generation like Genie [47] and UniSim [48] both incorporate actions into their video models. Genie infers latent actions that explain transitions in videos, and UniSim conditions on actions and language instructions to generate controllable videos. Modelling the effect of actions in a parsimonious way could improve the accuracy and generation capabilities of such systems.

Acknowledgments

We thank the HCAI lab for valuable feedback and comments throughout the project. We are 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 us 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: 01IS18039A, and funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's Excellence Strategy-EXC2064/1-390727645.15/18.

References

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems , 31, 2018.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 , 2019.

Thomas Kipf, Elise Van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint arXiv:1911.12247 , 2019.

Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11) , pages 465-472. Citeseer, 2011.

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.05929 , 2020.

Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.04955 , 2022.

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-predictive learning for reinforcement learning. In International Conference on Machine Learning , pages 33632-33656. PMLR, 2023.

Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. arXiv preprint arXiv:1903.12436 , 2019.

Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10674-10681, 2021.

Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability , 1(2):127-190, 1999.

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 arXiv:1801.00690 , 2018.

Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta. Repo: Resilient modelbased reinforcement learning by regularizing posterior predictability. Advances in Neural Information Processing Systems , 36, 2024.

Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite-a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722 , 2021.

Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances in Neural Information Processing Systems , 33:741-752, 2020.

Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning , pages 27782787. PMLR, 2017.

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

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset, 2017.

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.

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.

Zhiqing Sun, Zhi-Hong Deng, Jian-Yun Nie, and Jian Tang. Rotate: Knowledge graph embedding by relational rotation in complex space. arXiv preprint arXiv:1902.10197 , 2019.

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.

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.

Arnab 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.11941 , 2023.

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193 , 2020.

Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 , 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

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.

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410 , 2016.

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.

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.

Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep representations. The Journal of Machine Learning Research , 19(1):1947-1980, 2018.

Alexander A Alemi. Variational predictive information bottleneck. In Symposium on Advances in Approximate Bayesian Inference , pages 1-6. PMLR, 2020.

Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in dqn. arXiv preprint arXiv:1810.00123 , 2018.

Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International conference on machine learning , pages 2170-2179. PMLR, 2019.

Tankred Saanum, Noémi Éltet˝ o, Peter Dayan, Marcel Binz, and Eric Schulz. Reinforcement learning with simple sequence priors. arXiv preprint arXiv:2305.17109 , 2023.

Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Robust predictable control. Advances in Neural Information Processing Systems , 34:27813-27825, 2021.

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 arXiv:1901.10902 , 2019.

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.

Felix Creutzig, Amir Globerson, and Naftali Tishby. Past-future information bottleneck in dynamical systems. Physical Review E , 79(4):041925, 2009.

William Bialek, Ilya Nemenman, and Naftali Tishby. Predictability, complexity, and learning. Neural computation , 13(11):2409-2463, 2001.

Naftali Tishby and Daniel Polani. Information theory of decisions and actions. In Perceptionaction cycle: Models, architectures, and hardware , pages 601-636. Springer, 2010.

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.

Imre Risi Kondor. Group theoretical methods in machine learning . Columbia University, 2008.

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.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.00588 , 2022.

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.15391 , 2024.

Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.06114 , 2023.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021.

Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Icml , 2010.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornung, and Luc Van Gool. The 2017 davis challenge on video object segmentation. arXiv:1704.00675 , 2017.

Appendix

A Mutual Information minimization

To simplify dynamics, our algorithm minimizes the mutual information between z t and ˜ ∆ t conditioned on a t (Eq. 8 & 3). With stochastic gradient descent, we search for parameters θ that optimize the following objective:

To do so, we introduce a new variable h t that depends on z t and a t (Eq. 6). We use this variable, together with a t , to produce ˜ ∆ t , instead of z t . For a specific a t , the Markov chain behind ˜ ∆ a t t can be written as

Thus, to minimize the mutual information between z t and ˜ ∆ t , we can instead minimize the mutual information between z t and h t : min θ I ( z t ; h t | a t ) . Following [29], we write the mutual information between z t and h t as follows:

In our case, p ( h | z , a ) is deterministic. The marginal distribution conditioned on a , p ( h | a ) , is challenging to obtain. We instead use a variational approximation: With our variational distribution q ( h | a ) , we can upper bound the mutual information as follows

First, this upper bound lets us see that the mutual information can be minimized by minimizing the number of bits of information z t contains about h t over and above a t . This could potentially be done by introducing a parameterized prior that depends on a t , q θ ( h t | a t ) and minimize the KL divergence between these two probability distributions

We opt for a simpler approach, that allows us to do this without introducing a new prior and parameterization. Following [8], we assume that q ( h t | a t ) is a standard, d -dimensional, isotropic Gaussian distribution N ( 0 , I ) . Assuming further that h t is another d -dimensional isotropic Gaussian, the KL divergence is a function of the mean and standard deviation of h t [8, 27]:

Here µ ( · ) and σ ( · ) return the mean and standard deviation of the Gaussian, respectively. Importantly, we can minimize the KL divergence by minimizing the first term in the above equation, which is simply the norm of the mean of h t . In our deterministic setup, we therefore simply minimize the norm of h t itself.

B PLSM vs L 1 and L 2 norm regularization

CWM

β=1e-09

β=0.1

1.4

3.7

3.6

PLSM

0.3

1.2

1.1

2

2

-0.8

1.4

1.4

-1.9

22

0.1

0.1

2.2

4.0

2.0

0.8

0.9

3.9

-2.4

.0.8

0.8

Z1

Z1

Z1

β=0.0005

β=0.001

β=0.005

2.11

2.11

2.11

L1

0.7

0.7

0.7

2

0.7

0.7

0.7

2.1

2.

0.1

0.1

2.1

-2.1

0.1

0.1

21

2.1

2.4

0.1

0.1

Z1

Z1

Z1

β=0.001

β=0.01

β=0.1

2.11

2.11

2.11

L2

0.7

0.7

0.7

2

2

0.7

0.7

0.7

2.1

2.1

0.1

0.1

2.

-2.1

0.1

0.1

2.

2.

0.1

0.1

Z1

Z1

Z1

C Ablations

We tested some alternative formulations and ablations of our PLSM objective on a subset of the datasets (see Fig. 9). One ablation removed the query representation h t and simply regularized the latent state z t . This model achieved substantially worse performance in the three datasets we tested it on. As an alternative to regularizing the L 2 norm of h t , we also tested a top -k bottleneck, which used the top 15 features from h t instead of the full representation. This model performed on par with the original PLSM. Lastly, minimizing the L 2 norm of h t could potentially be compensated for by the dynamics MLP by increasing the norm of the weight matrices. As a control, we added weight decay to the dynamics model to prevent this, and observe comparable performance.

D Is the L 2 minimization effective?

One potential issue with minimizing the L 2 norm of the query representation h t is that the model can compensate for this by increasing the magnitude of the weights in the ensuing dynamics module. While this can be prevented by adding weight decay to the dynamics model's weights, we show that PLSM does not suffer from this empirically in the datasets we evaluated it on. We calculated the average norm of h t of PLSM relative to the average norm of z t in the unregularized model. Comparing them we see indeed that h t has several orders of magnitude lower norm. However, comparing the norm of the ensuing linear layers, we see that they most often do not differ significantly (see Fig. 10). This suggests that shrinking h t truly minimizes the amount of information the model can extract from it to predict ∆ .

cubes 9

shapes 9

sprites

100

100

%

90

2

96

99

pi

75

98

92

C

AC

98

60

88

i

3

35

810

i

3

5

810

1

3

5

810

100

100

%

98

100

90

R

MRI

95

99

75

92

98

60

35

810

35

810

135

810

Horizon

Horizon

Horizon

PLSM

Without ht

K-PLSM

PLSM (weight decay)

spaceinvaders

pong

cubes9

shapes9

sprites

MOVi-A

101

100

100

100

100

100

100

Ilinputll2,

10-1

10-1

10-1

10-1

10-1

10-1

10-2

10-2

10-2

10-2

40

40

40

40

30

40

30

30

30

30

20

30

IWdell2

20

20

20

20

20

10

10

10

10

10

10

0

0

n

0

CWM

PLSM

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 regularization strenght of PLSM impacted performance in Atari. We see that in games where important features of the environment are not controllable by the agent, weaker regularization is beneficial. For instance, in Boxing, the movement of the opponent is hard to predict, and here having a lower regularization strength is advantageous.

Median

IQM

Mean

OptimalityGap

SPR+PLSM

SPR

DrQ

0.6

0.8

0.3

0.4

0.6

0.8

0.55

0.60

0.65

score

boxing

hero

4

0.3

Norm'd

3

0.2

2

0.1

Human

1

0

0.0

5.0

0.1

1e-06

5.0

0.1

1e-06

β

β

F Datasets

See Fig. 13 for example frames from the training datasets and Fig. 14 for example frames from the generalization tests.

G Hybrid model

We propose a hybrid model for the Atari environments. The hybrid model splits the latent state in two, one parsimonious state space z 1 t and an unconstrained state space z 2 t . To predict the next state, the model uses two dynamics MLPs ˜ z 2 t +1 = z 2 t + d 1 θ ( z 1 t , z 2 t , a t ) and ˜ z 1 t +1 = z 1 t + d 2 θ ( h t , a t ) . The next latent prediction is then defined as ˜ z t +1 = Concatenate (˜ z 1 t +1 , ˜ z 2 t +1 ) . Results are shown in Fig. 15.

H Slots with information regularization

We evaluated the C-SWM on our cubes 9 task, both with and without our regularization. The performance of C-SWM dropper by a small amount due to the increased difficulty of the task for

T.

spaceinvaders

pong

0.40

2

0.60

ccurac

0.30

0.45

0.20

0.30

1

3

5

8

10

1

3

5

8

10

0.80

0.60

CWM

PLSM*

0.70

0.50

R

MRF

0.60

0.40

0.30

0.50

1

35

8

10

1

35

8

10

Horizon

Horizon

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.

I Atari scores

J 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 on 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. In MOVi-E we stacked four frames.

J.3 Contrastive models

J.4 DeepMind Control Suite

We use the same hyperparameters as [6] for the planning agents. The only addition is our PLSM regularization. We use a regularization coefficient of 0 . 1 for all agents, and make the query net MLP f θ the same size as the dynamics network h t , with | h t | = | z t | .

J.5 Distracting Control Suite

All models were trained with an action repeat value of 2 in all environments. All model hyperparameters 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.

NeurIPS Paper Checklist

The checklist is designed to 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 is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a 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 it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a 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 "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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.

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 in 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 a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text 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 fairness.

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 important 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 a 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 not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

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 if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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 are included in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will release code with model implementation upon publication.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Details are found in Experiment section and Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

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 are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Compute details reported in the Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 is a special consideration 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 no potential societal impacts worth discussing.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All datasets, models and assets used in the paper are cited properly.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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 was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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 as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 to 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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).\n\nWEAKNESSES:\n* 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).\n\nQUESTIONS:\nPlease see the weaknesses section above.\n\nLIMITATIONS:\nThe authors sufficiently address the limitations in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-21"}

comments: 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.

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

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


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

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

[3] InfoBot: Structured Exploration in ReinforcementLearning Using Information Bottleneck. ICLR 2019

QUESTIONS:
- 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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-09

comments: 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.

STRENGTHS:
- The problem is well-motivated, and the solution is straightforward to implement.
- 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.
- As is common with information bottleneck-based methods, PLSM's performance may be highly sensitive to the hyperparameter of the regularizer for the bottleneck.

QUESTIONS:
- Could you provide Interquartile Mean (IQM) results for the Atari experiments, as suggested by Agarwal et al. (2021) in "Deep Reinforcement Learning at the Edge of the Statistical Precipice"? This would offer a more robust measure of aggregate performance across multiple tasks.
- Since the regularisor weights $\beta$ is very important to the performance, I wonder how robust is PLSM to different $\beta$? Could you also share some insights about tuning the parameter? This can be a valuable information to the community.

LIMITATIONS:
Yes, the limitations are well-discussed by the authors.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations are commented above. No societal impact needs to be particularly addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nHave you run the TD-MPC experiments longer? This is a super important experiment to run and should be included in the paper\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}