Simple, Good, Fast: Self-Supervised World Models Free of Baggage
Authors: Jan Robine
Marc Höftmann
Stefan Harmeling
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement learning
World models
Self-supervised learning
Atari 100k
Primary area: reinforcement learning
TL;DR: What are the essential components of world models? We present a simple, good, and fast world model, and evaluate it on the Atari 100k benchmark.

Abstract: What are the essential components of world models? How far do we get with world models that are not employing RNNs, transformers, discrete representations, and image reconstructions? This paper introduces SGF, a Simple, Good, and Fast world model that uses self-supervised representation learning, captures short-time dependencies through frame and action stacking, and enhances robustness against model errors through data augmentation. We extensively discuss SGF’s connections to established world models, evaluate the building blocks in ablation studies, and demonstrate good performance through quantitative comparisons on the Atari 100k benchmark. The code is available at https://github.com/jrobine/sgf.

Published as a conference paper at ICLR 2025
SIMPLE, GOOD, FAST: SELF-SUPERVISER WORLD MODELS FREE OF BAGGAGE
Jan Robine! Mare Höftmann" & Stefan Harmeling!
'TU Dortmund, LLamarr Institute for Machine Learning and Artificial Intelligence [jan. robine, ne, marc. hoentmann, stefan, hwreearrala........a. Somelinggl et de
ABSTRACT
What are the essential components of world models? How far do we get with world models that are not employing RNNs, transformers, discrete representations, and image reconstructions? This paper introduces SGF, a Simple, Good, and Fast world model that uses sulf--uue representation learning, captures short-time dependencies through frame and action stacking, and enhances robustness against model errors through data augmentation. We extensively discuss SGF's connections established world models, evaluate the building blocks in ablation studies, and demonstrate good performance through quantitative comparisons on the Atari 100k benchmark. The code available htttps: hats:///gthubbb/ com/ jrwbjjjobee/e.eee sgf
INTRODUCTION
Deep reinforcement learning has demonstrated remarkable success in solving challenging decision- making problems (Mnih et al. 2015 Schulman et al.) 2017; Mnih et al. 2016 Hessel et al, 2018
Badia et al. 2020 Schrittwieser al. 2020 Kapturowski et al. 2023 Hafner al. 2023). Despite these achievements, the primary challenge remains sample efficiency, i.c., the amount of data required to learn effective behaviors. Recent works have addressed this challenge by improving architectures and hyperparameters (van Hasselt et al.) [019 Schwarzer et al.] 2023). pretraining and fine- tuning (Schwarzer al. 2021b), applying data augmentation (Yarats et al., 2021 Laskin al. 2020a).
incorporating ideas from self-supervised representation learning (Laskin et al. | Schwarzer
al. 2021a b 2023), or learning a model of the environment (Kaiser et al. 22020] Ye al. 202
Robine et al. 2023] Micheli et al., 2023 Hafner et al. 2020 2023).
Several approaches have been proposed in the literature to leverage model of the environment. Improving the representations for the model- free agent can be achieved by learning (latent) transition model (Lee et al. 2020 Schwarzer et al. 2021a), reward model, or both (Gelada a (021). where the model serves auxiliary task. Aside from that, world model, i.e., a deep generative model of the environment, can be learned. World models find applications in learning in imagination, where model-free algorithm is applied to sequences generated by the world model (Sutton 1991
Information Maximization
Temporal Consistency
Information Muximization
Predictor ho
ct al. 2019 Zhang et al.)
Projector ge
Projector şin
Encoder fa
Encoder fo
ō
Ha & Schmidhuber 2018 Kaiser et al. 2020
Figure I: Our world model learns representations that are both temporally consistent and maximize the information about the observations.
Hatner et al. 2020 2021 2023) Micheli et al.
2023, Robine al.
2023), and in decision-time
planning, where the model is used for lookahead
search during action selection (Watter et al.) 20155 Banijamali et al. 2018 Chua et al., 2018 Hafner
et al. 2019) Another line of work performs decision-time planning without full world model, relying on value equivalence. In this paradigm, trajectories of the model are required to yield the same cumulative rewards those in the real environment, regardless of whether the produced hidden states correspond to any real environment states or not (Tamar et al. 2016 Silverr et al. 2017 Oh
al. 2017 Schrittwieser al. 2020 Ye al. 2021 Hansen et al. 2022).

---

Published as a conference paper at ICLR 2025
In computer vision, self-supervised learning of image representations has made significant progress in
recent years (Chen et al. 2020 He et al. 2020; Grill et al. 2020: Caron et al. 2020 Chen & He 2021
Zbontar al. 2021 C Caron et al. 2021 I Bardes et al. 2022). Many approaches are based Stamese
neural networks (Bromley et al. 1993) and can he categorized into contrastive and non-contrastive
methods. Contrastive methods (Chen et al. 2020 He et al. (220)) aim to learn representations that similar for different views (e.g. image augmentations) of the same image but dissimilar for different images to prevent the collapse of representations. Non-contrastive methods do not rely on negative samples. Instead, they prevent representation collapse by the design of the architecture (Grill et al.
2020 Chen & He (021) or by regularization of the representations (Zbontar et al. 2021 Bardes et al.
2022
Contributions: In this work, we explore simplifying world models while maintaining good per- formance. Many world models rely on RNNs or transformers to capture long-term dependencies, introducing computational complexity and instability. We focus problems where short-term dependencies might be sufficient, avoiding sequence models and instead leveraging self-supervised learning, data augmentation, and stacking. Our contributions
follows:
We present SGF, a world model based on a representation learning framework inspired by VICReg (Bardes al. 2022). While restricting ourselves simple world models (w/o RNNs and transformers) and choosing simple ingredients, we still have the essential properties of effective world models: maximum information and temporal consistency (Section 2.3)
We reduce the complexity of world models in terms of both methodology and implemen- tation: SGF does not require image reconstructions or discretization of representations. Besides avoiding sequence models, such as recurrent neural networks or transformers, we also do not use probabilistic predictions for deterministic environments. Instead, we employ simple techniques from model-free reinforcement learning, in particular data augmentation and stacking, which already have been successfully used before (Section 2.2).
We conduct several ablations and thoroughly discuss the similarities and differences between SGF and other world models that learn in imagination (Sections 4.2 and 5.2).
We demonstrate that our design choices lead to shorter training times compared other world models, while achieving good performance on the Atari 100k benchmark (Section 5.1).
SIMPLE, GOOD, AND FAST WORLD MODELS
While being increasingly powerful, model-based approaches have simultaneously grown in complex- ity and consist of more and more components that need to be adjusted to each other (Hafner et al. 2020; 2021 2023 Robine al.) 2023] Micheli al.) 2023). After introducing our notation, aim to distill minimal ingredients and objectives for world models that are easy to implement and computationally efficient.
2.1 PRELIMINARIES AND NOTATION
We formalize the environment in terms of a partially observable Markov decision process (POMDP) with discrete time steps, rewards E image observations o E xx and actions A, which are either discrete or continuous. Transitions within the environment are described by tuple (o, a, o', e), where o the current observation, o' is the next observation, and e {0, 1) indicates terminal states.
In its simplest form, world model generative model of the dynamics p(o', r, o, of POMDP. Given policy a = T()), iterative sampling from the world model generates trajectories without further real environment interactions. These trajectories can used for learning behaviors in imagination (Ha & Schmidhuber 2018 Hafner et al. 2020), c.g., via model-free RL.
To increase efficiency, world models should operate in a low-dimensional representation space (commonly known latent space) as opposed to the high-dimensional observation space (Ha & Schmidhuber 2018). For this we need two components: a representation model that maps image observations onto representations y. and a dynamics model that predicts the latent dynamics p(y', r, e y, a). Usually, the policy also operates the cow-dimensional space, i.c., a a enabling behavior learning with high computational efficiency (Ha & Schmidhuber 2018..
2

---

Published as a conference paper at ICLR 2025
A common assumption is the conditional independence of r and given y. a, and y', which is also employed by previous world models a..g...Mccher et al. 2023 Hafner et al. 2023). This leads to the following factorization of the latent dynamics
ply's r, y, a) p(y/ y, a) p(r y, a,
a, y') p(e y, a,y))
(I)
which consists of three conditional distributions: the transition distribution p(y' y.a)) which is only conditioned on y and a, the reward distribution p(r y: a. yy) and the terminal distribution p(e y, a, y'), both further conditioned on the next representation v'. This is natural choice for many POMDPs, where most of the complexity of the dynamics is captured by the transition distribution. This allows for learning three separate models rather than modeling the joint distribution.
2.2 INGREDIENTS LEADING TO SIMPLICITY
Stacking instead of memory. Model-free methods often assume that the observations of an POMDP approximately satisfy the Markov property. This means that the next observation o' (and consequently y') is independent of the preceding history of transitions given the current observation and action a. However, previous works on world models consider Markovinn observations, which might exhibit long-term dependencies. This is approached by adding a notion of memory to the dynamics model, which can be realized by introducing recurrent states (RNN-based, Ha Schmidhuber] [2088: Hafner et al. 2020) or by directly conditioning on the history of transitions (attention based Robine al. 2023] Micheli et al., 2023). However, the memory is typically a big computational burden and makes the model more complicated. We are asking whether can omit the memory to obtain a much faster world model.
To capture short-time dependencies with minimal computational overhead, we suggest to simply use frame and action stacking. Stacking the m most recent frames alleviates the problem of partial ob- servability. c.g., by capturing the velocity of objects in the scene. This is well-known preprocessing technique in model-free reinforcement learning (Mnih et al.) 2015), and has already been applied to world models by Robiie al. (2023). Additionally, stacking the most recent actions can be beneficial, considering potential delays in the effects of actions. In Section (..2, we show significant improvements of our world model through action stacking while being computationally cheap.
Augmentations instead of stochasticity. In deterministic POMDPs, executing a specific action in specific state consistently yields the same outcomes o', and However, prior world models stochastic even in deterministic POMDPs (see Section [..2). Ha & Schmidhuber (2018) argue that, due model errors, behaviors learned in imagination may perform poorly in the real environment. They propose that stochastic predictions can reduce the exploitability of an imperfect world model.
Similarly, introduce stochasticity through data augmentation, as demonstrated in previous model- free algorithms (Yarats et al. 2022]; Laskin et al., 2020a).. We investigate whether this helps to improve the robustness of our world model against model errors. In Section 4.2 we demonstrate that data augmentation significantly improves the performance of our world model.
2.3 ESSENTIAL PROPERTIES OF REPRESENTATIONS
Building meaningful representations of observations is crucial for dynamics modeling and behavior learning. In this work, argue that representations should possess two key properties: (1) the information about the observations should maximized, (2) they should be temporally consistent, ie., representations of two successive observations should be similar. We describe these properties in more detail below. Both properties are already present in other world models (see Section 5.2). however, in this work try to implement them in most simple and efficient manner.
Maximizing information. In latent world models, the representations are used for downstream behavior learning, any information not encoded in these representations is not accessible to the agent. Therefore, extracting maximum information from observations is necessary for learning optimal behaviors in latent space. This is often realized by reconstructing the input observations (i.e., autoencoder-style).. however, this can also be realized by self-supervised objectives, as we show in Section 3
3

---

Published as a conference paper at ICLR 2025
Temporal consistency. Temporal consistency can be motivated by predictive coding, where the future or missing information is predicted. Predictive coding has been applied in information theory for data compression (Elias I955), and more recently in representation learning (Oord et al.) 2018 Hénaff et al. 2019). Also, the neuroscience literature suggests that the human brain learns internal representations of incoming sensory signals by minimizing prediction errors subject to particular constraints on the representations, in spatial and temporal domains (Rao & Ballard 1999]
Hosoya et al., 2005 Huang & Rao 2011). Furthermore, Hénaff al. (2019) proposed the temporal straightening hypothesis which suggests that inside the human brain visual inputs are transformed to follow straighter temporal trajectories, order make the stream visual inputs more predictable. These advantages of lemporal consistency for biological agents translate to advantages for agents in reinforcement learning that are based on latent world models:
Simpler dynamics prediction: successive observations are close in representation space, the dynamics model often only needs to predict minor changes and the danger of sudden jumps in representation space is reduced.
Improved behavior learning: both, the policy and value function benefit from temporally similar representations. This concept can be loosely connected to bisimulation metrics, where "behaviorally similar" states are grouped together (Zhang al. 202T). Further details are available in Appendix B.1
3 BUILDING BLOCKS FOR SGF
Having identified core ingredients, now describe the construction of simple, fast, and good world model. We will begin by outlining the representation model of SGF, followed by an description the dynamics model, implementational details, and our evaluation protocol. For representation learning we draw inspiration from VICReg (Bardes et al. 2022). Further connections to other self-supervised methods are discussed in Appendix B.2
3.1 LEARNING A WORLD MODEL
Representation learning. Given a POMDP transition (o, a, o', e). we apply random transforma- tions 1, ~ T from set of image augmentations to obtain augmented observations
((()) and o' ('o)... An encoder computes representations ÿ fō(ō) and ÿ fe(ō") with ÿ. ÿ Rd A projector network ge computes embeddings ă = go(g) and x = go (3 with ă, ø' E RD. An iition-conditioned predictor network ho predicts the next embedding z' ho a). An illustration can be seen in Figure
To achieve temporal consistency, we minimize the mean squared error between E and ž'. To maximize the information content, the embeddings are regularized using the variance and covariance regularization terms proposed by [Bardes al. (2022). The total representation loss is summarized by
Reppe (0) = E, B 11 (ž, a) - Z E
VC(2) + VC(Ž²)
(2)
Temporal Consistency
Information Maximization
where 7 is batch of transitions from replay buffer, 2 and Ž' are batches of embeddings, " controls the strength of the consistency loss, and VC (variance and covariance) is defined as
D = Σ
VC(Z)
p max (0, 1 Cov +8
+ w Σ kesf Cov
1=1
Variance Regularization
Covariance Regularization
where D is the dimensionality the embeddings, P, > 0 control the strength of variance and covariance regularization terms, respectively, and 1 x 10 prevents numerical instabilities. The goal the VC terms is to maximize information content and to prevent representation collapse. Variance regularization keeps the standard deviation of each embedding feature across the batch above 1 using hinge loss. Covariance regularization decorrelates the embedding features by attracting their covariances towards zero (Bardes et al. 2022).
4

---

Published as a conference paper at ICLR 2025
Dynamics learning. We build simple dynamics model and rely the capabilities of temporally consistent representations. Based on the dynamics factorization, we learn a transition distribution Po(y') y. a), reward distribution Po(r) y, a, y²), and terminal distribution po(e y, a, y'). In this work, we focus on deterministic prediction, i.e., we simply calculate the means for transitions and rewards, and the mode for terminals (more details on this in Appendix FD. Maximum likelihood estimation leads to the total dynamics loss
DDnn. (0) = E,,
log (y(s sg(y), a) (r y, y, log Pogpel (c) y, v,a,5))
Transition Distribation
Reward Distribution
Terminal Distribution
where sg(-) denotes the stop-gradient operator, meaning that the representations are not influenced by the loss of the transition distribution. This is because the transition distribution has moving targets, given that y originates from the representation model, which changes during training. The rewards and terminals provide stable signals from the POMDP. Note that we learn the transitions with con-augmented observations, i.c., y = fo(o) and y = fe(o').
3.2 LEARNING BEHAVIORS IN IMAGINATION
The representations y serve inputs to the policy Ta y). Through iterative dynamics prediction, batches of representations, actions, rewards, and terminals are generated and used to train the policy. The policy learns to maximize the expected return by performing approximate gradient ascent with the policy gradient (Sutton et al. 1999). To reduce the variance of the gradient estimates, we employ baseline, resulting in an advantage actor-critic approach (Mnih) et al.] (2166); the details are explained in Appendix F The pseudocode outlining our world model and policy training procedure is presented in Algorithm
learned value function uo (
3.3 EVALUATION PROTOCOL
We evaluate our world model on the Atari 100k benchmark, which was first proposed by Kaiser et al. (2020) and has been used to evaluate many sample-efficient reinforcement learning methods (Laskin et al. 2020b; Yarats et al. 2021, Schwarzer et al.) 202 2023 Micheli et al., 2023 Hafner et al. 2023). It includes subset of 26 Atari games from the Arcade Learning Environment (Bellemare et al. 2013) and is limited to 400k environment steps. which amounts to 100k steps after frame skipping or hours of human gameplay. Note that all games are deterministic (Machado et al.) 2018).
We perform 10 runs per game and for each run we compute the average score over 100 episodes the end of training. We follow Micheli et al. (2023) by selecting a random action with 1% probability inside the environment and using sampling temperature of 0.5 for the policy during evaluation. We also adapt their special handling of Freeway and use a sampling temperature of 0.01 for the policy.
4 EMPIRICAL STUDY
Behavior learning depends entirely on the quality of the world model. To assess our world model, we first analyze it qualitatively and then show how getting rid of individual design choices results in degradation of performance. Further analysis and ablations are presented in Appendices C CaandB
4.1 INSPECTING THE WORLD MODEL
To assess whether the learned representations contain relevant information and do not collapse, we train a separate decoder for analysis without affecting the world model's gradients. This allows us to visually interpret the imagined sequences of the world model. In Figure 2 we depict three exemplary sequences, demonstrating that the learned representations contain useful information without relying on image reconstructions. We can also see that the dynamics model can work for long sequences (30 time steps) without accumulating notable model errors, although being only feedforward model.
Furthermore, we study the temporal consistency of the learned representations. For that, we generate an episode in Pong following a policy trained with our approach. We then encode the observations and compute wwo-dimensional t-SNE embeddings (Van der Maaten & Hinton (008) of the represen- tations. The result can be seen in Figure 3 and it suggests that subsequent observations have similar
5

---

Published as a conference paper at ICLR 2025
Figure 2: Illustration of imagined sequences of length 30. Each frame in the frame stack is con- verted grayscale, and pixel changes are visual- ized in red, green, and blue. From top to bottom: Kung Fu Master, Ms Pacman, and Seaquest.
Figure 3: Two-ddmeesiooall t-SNE embeddings of the learned representations obtained by play- ing an episode in Pong. We highlight subse- quence of the episode, which starts with new ball and stops after point is scored.
representations and temporally consistency is successfully employed. Figure 7 we compare the learned embeddings when disabling temporal consistency.
4.2 ABLATING THE WORLD MODEL
We show empirically how effective the components presented in Section 2 are by performing five ablations. Each ablation assessed by the performance on five Atari games. The results are illustrated in Figure 4 Numerical results can be found in Table 3 We observe that data augmentation, action stacking, frame stacking, and temporal consistency are crucial:
1. No augmentations: omitting image augmentations leads to poor performance in all games.
2. No action stacking: stacking only the frames but not the actions decreases the overall performance for all five games.
3. No frame stacking: stacking only the actions but not the frames decreases the performance of all games, with complete failures Boxing and Breakout. However, Kung Fu Master suffers only a small degradation, possibly because all enemies face the direction they are heading, so their velocity identifiable from a single frame
4. No temporal consistency: setting the coefficient 7) zero leads significant decrease in performance for most games except Kung Fu Master.
5. Sample-contrastive: Garrido et al. (202)) show that VICReg can be also seen as dimension- contrastive method, as opposed cample-contrastive methods such SimCLR (Chen et al., 2020). They also show that VICReg can be converted to sample-contrastive method by transposing the embedding matrix Z in Equation B Making VICReg cumple-contrastive significantly worsens the performance in Breakout, but less so in the other games.
COMPARISONS
In this section we compare our method with previous world models regarding the results and the methodology. In Appendix B we discuss the relations to other methods.
5.1 RESULT COMPARISON
We compare our method with five baselines: the model-free algorithm SPR (Schwarzer et al. 2021a with updated scores from Agarwal et al. (2021), and the model-based methods EfficientZero (Ye ct al. 202T). IRIS (Micheli et 2023). and DreamerV3 (Hainer et al. 2023). The mean metric
across all games is calculated using human normalized scores (Mnih al. 2015)

---

Published as a conference paper at ICLR 2025
Boxing
Pong
Breakout
2.0
EfficientZ.cro
Normalized mean
1.5
Dreamer V3 TWM
Human SGF
IRIS
Ms Pacman
Kung Fu Master
SPR
0.5
SGF (default)
frame stacking
SimPLe
No augmentations
No temporal consistency
10a
10a
10
action stacking
Sample--ontzastive
Runtime in hours (log scale)
Figure 4: Ablations of SGF in five games. Hu- man normalized scores, normalized with the maximum value achieved per game.
Figure 5: Score and runtime comparison in the Atari 100k benchmark. SPR is model-free, Effi- cientZero performs lookahead.
In terms of performance, our simple world model achieves good scores with significantly faster training. Figure shows the scores relation runtime. Detailed scores can found Table 2 and numerical runtimes can be found in Table 4 Training SGF takes 1.5 hours a single NVIDIA A100 GPU. Obtaining precise training times for other methods challenging, as they depend on the GPU. Following Hafner al. (2023), we approximate runtimes for NVIDIA V100 GPU, assuming NVIDIA P100 GPUs are twice as slow and NVIDIA A100 GPUs are twice as fast. Notably, SGF's runtime is four times shorter than the runtime of DreamerV3, V3, despite both having the same number
of imagination steps (1.5 billion). In Table 5 we provide a breakdown of the runtime for certain components of method.
5.2 WORLD MODEL COMPARISON
We also compare the methodologies of SGF and state-of-the-arr world models used for learning in imagination: SimPLe (aiser et al. 2020). Dreamer (Hafner et al., 20200] [021 2023). IRIS (Micheli et al. 2023). TWM (Robine al. 2023 and the world model developed by Ha & Schmidhuber (2018). referred to as HS. In the Dreamer line work, our focus is on DreamerV2 and Dream- erV3, given their performance improvements over the initial version, notably achieved through the discretization of representations. A summarized comparison of the main differences is provided in Table (we highlight the components in the text).
Temporal consistency. We enforce successive representations be similar by the choice of our objective (Consistency).. IRIS and HS have no explicit concept of temporal consistency. Dreamer and TWM also seek temporal consistency and attract the representations slightly towards the outputs of the transition predictor from the previous time step. The transition predictor can be interpreted as time-dependent prior for variational auto-encoder. An illustration of this difference can be seen in Figure 6 Our approach is simpler for the following reasons: In Dreamer, the training of the representation model and the dynamics model is intertwined. Correctly balancing the representation loss and the dynamics loss is crucial for ensuring stable training. Our representation model learns in isolation, simplifying hyperparameter tuning while still achieving temporal consistency.
In Dreamer, consistency is imposed directly on the representations, whereas we maximize the similarity of the non-linear embeddings of the representations. In DreamerV3 (Hafner et al. 2023) the consistency loss is clipped when falls below a certain threshold, considering the similarity as sufficient (aka free bits). Our hypothesis is that maximizing the similarity of the embeddings offers a similar degree of freedom.
7

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of methodology with other world models used for imagination. DVx denotes
DreamerV2 and DreamerV3. Every component results in additional complexity.
Component
HS
SimPLe
IRIS
TWM
DVX
SGF
Augmentations
Information Maximization
X
Stacking
Consistency
Reconstructions
X
X
X
Discretization
X
X
X
Sequential Dynamics
X
X
X
Stochastic Transitions
X
X
X
Pixel Transitions
Pixel Dreams
X
Act with Memory
X
X
x
Information extraction. Previous world models depend pixel-wise image reconstruction for information extraction from observations (Reconstructions)) These auto-encoder architectures treat all pixels equally, including less important high-frequency details or noise. In contrast, adopt self-supervised objective and utilize data augmentation (Augmentations) to learn representations that maximize information (Information Maximization) and extract relevant features.
Discretization. Most previous methods learn discrete representations (Discretization). SimPLe discretizes representation values into bits. DreamerV2 (Hafner et al. 2021), DreamerV3 (Hafne al., 2023), and TWM (Robine et al. 2023) utilize softmax normalization to obtain a stack of independent categorical distributions. IRIS (Micheli et al. (023) converts each image observation into multiple discrete tokens. Discretization introduces additional complexity and requires techniques such straight- for the success of discretization in world models and explain how they are addressed in our approach:
gradient estimation (Bengio al. 2013). We propose two primary factors
Discretization significantly limits the information capacity of representations, potentially preventing collapse in auto-encoder architectures (LeCun 2022). Since our objective already prevents representation collapse, there is need for discretization that account.
Discretization potentially facilitates dynamics prediction by shrinking and stabilizing the support of p(y' y, a). However, we found that simple architectural choice, specifically adding layer normalization as the final layer of the encoder (as mentioned in Appendix (F). is sufficient to keep the mean and variance of the representations stable. Note that another common approach normalize the outputs lie in the interval [0, 1] (Schrittwieser et al. 2020 Schwarzer et al. 2021a).
Dynamics modeling. As we assume Markovian observations, our approach utilizes a feedforward dynamics model. Prior methods, with the exception of SimPLe, RNN-based (HS, DreamerV2. DreamerV3) or transformer-basse on image observations instead of operating tions). Additionally, our dynamics model is deterministic, whereas previous methods are stochastic (Stochastic Transitions), at least regarding transition prediction, while rewards and terminals typically remain deterministic.
(IRIS, TWM) (Sequential Dynamics). SimPLe operates directly low-dimensional representation space (Pixel Transi-
Behavior learning. Efficient world models, such ours, usually train policy based on the low- dimensional representations. However, IRIS decodes the representations back to pixels, slowing down training significantly (Pixel Dreams). Since SimPLe predicts the pixels directly, their policy has to operate on pixels as well. Moreover, previous methods that use a sequence model usually equip the policy with memory, since the representations encapsulate the history of transitions either via attention or through a compressed recurrent state (Act with Memory). Due to our feedforward dynamics model, policy is memoryless.
8

---

Published as a conference paper at ICLR 2025
2
ō'
y
y
H
y
a
ö'
a
(a) SGF (ours)
(b) Dreamer
Figure 6: High-level illustration how SGF and Dreamer ensure temporal consistency. The dashed lines denote that two variables are attracted towards each other by loss function. White nodes indicate inputs, blue nodes indicate representations, and green nodes indicate variables used for representation learning. The red node indicates that Dreamer depends on the dynamics model for representation learning. We omit Dreamer's recurrent states for simplicity.
Stacking. Previous world models considered in this section only employ the usual preprocessing. e.g., conversion grayscale or downscaling the observations. Our method applies frame and action stacking (Stacking); TWM applies only frame stacking.
LIMITATIONS
The insights presented in our paper provide a solid foundation for new world models that are simple, good, and fast. For this paper, we limited our method to deterministic MDPs. To also include non- deterministic POMPDs, the transition distribution needs to be stochastic, allowing for the prediction of multiple possible outcomes. The predictor must also be stochastic to account for non-deterministic information between and o'. Both networks would need to make stochastic predictions, e.g., by modeling the mean and variance of independent normal distributions or using Gaussian mixtures.
Having avoided sequence models such as RNN and transformers, we did limit ourselves to envi- ronments with mainly short-term dependencies, which might be the reason that did not reach state-of-the-art performance. Future work could replace the MLPs used for the transition, reward, and terminal distributions with a sequence model, requiring some implementation effort and resulting in increased computation time. Since the transition distribution is independent of representation learning. this change would not affect the encoder. We evaluated a preliminary version of this in Appendix Another more sophisticated approach would involve using a sequence model for the predictor, which would also affect the features extracted by the encoder. However, this would significantly increase the complexity of the model, which we aimed to avoid.
Another current limitation of our approach that VICReg requires the observations to images. In principle, VICReg could be applied to other modalities if reasonable augmentations available. We could also combine SGF with other self-supervised learning methods.
7 CONCLUSION
The starting point of our work are the questions: What are the essential components of world models? How far do we get with world models that are not employing RNNs, transformers, discrete representations, and image reconstructions? We demonstrate that representations learned in self- supervised fashion using VICReg combined with aecion-conditionel predictor network and applied to stacked observations can learn latent representations without resorting to resource-intensive sequence models. Self-supervised learning. coupled with augmentations and frame and action stacking, proves effective in building good world model. Applying SGF to the Atari 100k benchmark, attained good results with significantly reduced training times. We advocate for future research in model-bar based reinforcement learning to focus on sparingly adding new components and analyzing their necessity under varying circumstances.

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research has been funded/supported by the Federal Ministry of Education and research of Germany and the state of North Rhine- Seine-Westpphliaa as part of the Lamarr Institute for Machine Learning and Artificial Intelligence.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Mare G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NewrlPS 2021, December 6-14, 2021, virtual, pp. 29304-29320, 2021.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. (bss/600.004450 2016.
Layer normalization. CoRR,
Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020. 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 507-517. PMLR, 2020.
Ershad Banijamali, Rui Shu, Hung Bui, Ali Ghodsi, et al. Robust locally-linear controllable embed- ding. International Conference on Artificial Intelligence and Statistics, pp. 1751-1759, PMLR, 2018.
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variane-iivarrnee--vvaaacee regularization for self supervuper learning. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Mare G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: An evaluation platform for general agents. J. Artif. Intell. Res., 47:253-279, 2013. doi: 10.1613/jair.912.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville, Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint 308. 2013.
Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah. Signature verification using siamese time delay neural network. In Advances in Neural Information Processing Systems 6, [7th NIPS Conference, Denver, Colorado, USA, 1993], pp. 737-744. Morgan Kaufmann, 1993.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Syslems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10- 17, 2021, pp. 9630-9640 IEEF, 2021. doi: 10. /00////CCC49922222....000
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations, Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In IEEE Confer- ence on Computer Vision and Pattern Recognition, CVPR 2021, virtual. June 19-25, 2021, pp. 15750-15758. Computer Vision Foundation IEEF, 2021. doi: 10.1 N..119/CVVR44433..000 2021 .01549.
10

---

Published as a conference paper at ICLR 2025
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Peter Elias. Predictive coding-i. IRE transactions on information theory, 1(1):16-44, 1955.
Norman Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In UAI, pp- 210-219, 2014.
Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann Lecun. On the duality between contrastive and non-contrastive self-supervised learning, arXiv preprint arXiv:2206.022. 02574,, 2022.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2170-2179. PMLR, 2019.
Jean-Bastien Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent A new approach self-supervised learning. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 2455-2467, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2555-2565. PMLR, 2019.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. CoRR, Abs/2301.04400 2023. doi: 10.48350/axxi..220..44444 144144.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. International Conference on Machine Learning, ICMI. 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 8387-8406. PMLR, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He. Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEFF/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp. 9726-9735. Computer Vision Foundation IEEE, 2020. doi: 1.L109CCVPR226002000000000
Olivier Hénaff, Robbe LT Goris, and Eero P Simoncelli. Perceptual straightening of natural videos.
Nature neuroscience, 22(6):984-99) 2019.
11

---

Published as a conference paper at ICLR 2025
Olivier J. Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, S. M. Ali Eslami, and Aäron van den Oord. Data-efficient image recognition with contrastive predictive coding. CoRR, abs/1905.09272, 2019. URL
org. abs abs/1905....... 09011199550002222222 09272
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/1606.08415. 2016.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining improve- ments in deep reinforcement learning. In Proceedings of the Thirty-Secone AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI- 18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 3215-3222. AAAI Press, 2018.
Toshihiko Hosoya, Stephen A Baccus, and Markus Meister. Dynamic predictive coding by the retina. Nature, 436(7047):71-77, 2005.
Yanping Huang and Rajesh PN Rao. Predictive coding. Wiley Interdisciplinary Reviews: Cognitive Science, 2(5):580-593, 2011.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 448 456. JMLR.org. 2015.
Lukasz Kaiser, Mohammad Babacizadeh. Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In 8th International Conference Leurning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Steven Kapturowski, Victor Campos, Ray Jiang, Nemanja Rakicevic, Hado van Hasselt, Charles Blundell, and Adrià Puigdomènech Badia. Human-level atari 200x faster. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020a.
Michael Laskin. Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representations for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650 PMLR, 2020b.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2022-06-27. Open Review, 62(1), 2022.
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.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. Open net, 2023.
12

---

Published as a conference paper at ICLR 2025
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nat., 518(7540):529- 533, 2015. doi: 10. 10.1038/nature1443 14236.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the 33nd International Conference on Machine Learning, ICML. 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928-1937. JMLR.org, 2016.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems, 30, 2017.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint
2018.
Rajesh PN Rao and Dana H Ballard. Predictive coding in the visual cortex: functional interpretation of some extra-classical receptive-field effects. Nature neuroscience, 2(1):79-87, 1999.
Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-basee world models are happy with 100k interactions. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap. and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nat., 588 (7839):604-600) 2020. doi: 10. 10.1038/4415860000-0000 1586002000355---
John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Picter Abbeel. High- dimensional continuous control using generalized advantage estimation. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. CoRR, Abs//777.00637, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self- predictee representations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R. De- von Hjelm, Philip Bachman, and Aaron C. Courville. Pretraining representations for data-efficient reinforcement learning. CoRR, abs/2 100.0499, 2021b.
Max Schwarzer, Johan Samir Obando-Ceron, Aaron C. Courville, Mare G. Bellemare, Rishabh Agarwal, and Pablo Samuel Castro. Bigger, better, faster: Human-level atari with human-level efficiency, In International Conference on Machine Learning, ICML 2023, 23-29 July 2023. Hon- olulu, Hawaii USA, volume 202 of Proceedings of Machine Learning Research, pp. 30365-30380. PMLR, 2023.
David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Amold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-end learning and planning. In International Conference on Machine Learning, pp. 3191-3199, PMLR, 2017.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull., 160) 1991. doi: 10. 10.1145//22344412237
Richard S. Sutton and Andrew G. Barto. Reinforcement Learnigg An Introduction. A Bradford Book, Cambridge, MA, USA, 2018. ISBN 0262039249.
13

---

Published as a conference paper at ICLR 2025
Richard S. Sutton, David A. McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in Neural Information Processing Systems 12, [NIPS Conference, Denver, Colorado, USA, November 29 December 1999], pp. 1057-1063. The MIT Press, 1999.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Hado van Hasselt, Matteo Hessel, and John Aslanides. When parametric models in reinforce- ment learning? In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 14322-14333, 2019.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances neural information processing systems, 28, 2015.
Ronald J. Williams and Jing Peng. Function optimization using connectionist reinforvement learning algorithms. Connection Science, 3(3):241-266) 1991. doi: 10.1089/09540099 108946887.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OeenReview.net, 2021.
Weirui Ye. Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp- 25476-25488, 2021.
Jure Zbontar, Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 Proceedings of Machine Learning Research, pp. 12310-12320. PMLR, 2021.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference on Learning Representations, ICLR 1021, Virtual Event, Austria, May 3-7, 2021. OeenReview.net, 2021.
14

---

Published as a conference paper at ICLR 2025
A ETHICS STATEMENT
The advancements in developing Simple, Good, and Fast (SGF) world models for reinforcement learning can significantly enhance various fields by making advanced techniques more accessible and reducing computational demands. While this democratization can drive innovation in areas like robotics and autonomous systems, also raises ethical concerns, such as potential misuse in surveillance or autonomous weaponry. Therefore, it is crucial for the research community to address these risks and develop guidelines to ensure responsible use.
B RELATIONS TO OTHER METHODS
B.1 RECONSTRUCTION- MODELS
Previous world models learning in imagination rely on image reconstructions. However, there are other approaches that learn a model without relying on image reconstructions, which we will discuss in this section.
Value equivalence. In the value equivalence paradigm, trajectories of the model must achieve the same cumulative rewards as those in the real environment, regardless of whether the produced hidden states correspond any real environment states not. These models are used for decision-time planning. For instance, MuZero (Schrittwisser et al. (2020) trains a model with hidden states by predicting the policy, the value function, and the reward. EfficientZero (Ye al., 202T) extends this objective by introducing self-super consistency loss with projector and a predictor network, which shares similarities with our temporal consistency loss. However, they predict the next representation before feeding it into the projector, rather than employing an action-conditioned predictor. Additionally, their approach bears more resemblance to SimSiam (Chen & He 2021), as they utilize the noo-ggraieent operation and lack explicit information maximization.
Auxiliary tasks and bisimulation metrics. Learning model as an auxiliary task can improve the representations of the model-free agent. Although these approaches only share loose connection to generative world models, they typically are reconstruction-Fit freee. For instance, SPR (Schwarzer al, 2021a) learns a transition model predict the latent states of future time steps using a projector and predictor network, incorporating data augmentation. Their architecture shares similarities with ours, however, their transition model is convolutional, and they predict the next representation before feeding it into the projector (similar to Efficien(Zero). Moreover, their methodology is influenced by BYOL (Grill et al., 2020), utilizing a momentum encoder and lacking explicit information maximization. SPI (Schwarzer al. 2021b) combines SPR with goal-conditioned RL and inverse dynamics modelling, i... predicting the action ag from states St and St+1. They pretrain an encoder on unlabelled data, which is later finetuned on task-specific data.
special type of auxiliary tasks are connected to bisimulation metrics, where "behaviorally similar" states grouped together (Ferns & Precup 2014). Similar to approach, this also amounts to learning a reward model and a (distributional) latent transition model (by minimizing the Wasser-
stein distance). Prominent works include DeepMDP (Gelada et al., 2019), which still requires reconstructions for good results on Atari, and DBC (Zhang el al. 2021
B.2 SELF-SUPERVISED REPRESENTATION LEARNING
Our self-supervised representation learning framework similar to existing visual representation learning methods. We describe the differences the most related works. Note that a common difference is our architecture, as we employ layer normalization instead of batch normalization (Toffe & Szegedy 2015), SiLU instead of ReL.U nonlinearities, and no ResNet-based encoder (He al. 2016).
Relation to VICReg (Bardes et al., 2022). Our work is greatly inspired by VICReg, which is used to learn representations of (stationary) images. Originally, the same image is augmented and fed into both branches of the Siamese neural network. We augment two successive image observations, so the two branches get different inputs, which nonetheless related. Furthermore, VICReg maximizes
15

---

Published as a conference paper at ICLR 2025
the similarity between the embeddings and z' directly, whereas integrate action-conditen conditioned predictor network, since the observations o and might have a more complicated connection.
The basis of our self- supervised learning setup is a combination of ideas from VICReg and BYOL. Specifically, we adopt the predictor network concept from BYOL but utilize the regularization terms from VICReg, omitting BYOL's momentum encoder. Additionally, our predictor network is action-come conditiionddd. Similarly, our method is related to SimSiam (Chen He 2021), which also employs a predictor network, but we do not need the stop-gradient operation.
Relation BYOL (Grill al. 2020 and SimSiam (Chen He 2021
From a practical perspective, these methods could likely achieve comparable performance with appro- priate hyperparameter tuning. However, our decision to use VICReg was motivated by its conceptual advantages, which we believe make it particularly suitable for our framework: Specifically, VICReg offers two key advantages over BYOL and SimSiam. First, it avoids the need for additional target networks updated through moving averages (as in BOL). Second, VICReg has a more established theoretical foundation for its loss functions, leveraging variance and covariance regularization to prevent representation collapse. In contrast, YOL and SimSiam rely on mechanisms like target networks top-gradient operations, which are more heuristic in nature.
B.3 MODEL-FREE DATA AUGMENTATION
Our method is model-based, but there are model-free methods that also use data augmentation. Laskin al. (2020a) also use data augmentation to increase the sample efficiency, but by augmenting the image observations passed to the model-free agent. Yarats et al. (2021) additionally regularize the value function such that it is invariant to the augmentations.
16

---

Published as a conference paper at ICLR 2025
C ADDITIONAL ANALYSIS
(a) With temporal consistency.
(b) Without temporal consistency.
Figure 7: Comparison of two-dimensional t-SNE embeddings of the learned representations with and without temporal consistency. When only maximizing information, the representations are arranged in Gaussian blobs, which are harder predict. We show the representations of one episode in Pong.
Kung Fu Master
Ms Pacman
Seaquest
0.8
urcuracy (%)
0.4
Top-1
0.2
Top-5
Top-10
0.0
5K
20K
40K
60K
BOK
100K
5K
20K
40K
60K
BOK
100K
5K
20K
40K
60K
80K
100K
Environment Steps
(a) Top-k accuracies the reconstructions for different values of k. The accuracy is determined by first encoding and reconstructing an observation, then computing the MSE between the reconstruction and all ground truth observations in the replay buffer, and finally testing whether the input observation is among the k-nearest neighbors. We caleulate the mean over batch of 512 observations. Note that the observations the replay buffer can be very similar or even identical, so top-1 accuracy is not as expressive as top-5 and top-10.
Kung Fu Master
Ms Pacman
Seaquest
8o.4
Mean Squared Error
6o.4
4o-4
2e:4
20K
40K
60K
8OK
100K
20K
40K
60K
80K
100K
5K
20K
40K
60K
80K
100K
Environment Steps
(b) Reconstruction loss of the decoder.
Figure 8: Additional analysis of the decoder from Section 4.1 environment steps.
We start training after collecting 5000
17

---

Published as a conference paper at ICLR 2025
D DETAILED RESULTS
Table 2: Comparison with other methods on the Atari100k benchmark. Averaged over 10 seeds.
Model-froc
Lookahead
Learning in imagination
Game
Random
Human
SPR
Eff. Zero
IRIS
DreamerV3
SGF
(std. dev.)
Alien
227.8
7127.7
841.9
808.5
420.0
959
518.8
(116.7)
Amidar
5.8
1719.5
179.7
148.6
143.0
139
62.7
(19.4)
Assault
222.4
742.0
565.6
1263.1
1524.4
706
850.1
(250.3)
Asterix
210.0
8503.3
962.5
25557.8
853.6
932
802.5
(317.5)
Bank Heist
14.2
753.1
345.4
351.0
53.1
649
58.7
(38.9)
Battle Zone
2360.0
37187.5
14834.1
13871.2
13074.0
12250
3747.0
(1240.3)
Boxing
0.1
12.1
35.7
52.7
70.1
78
83.4
(10.7)
Breakout
1.7
30.5
19.6
414.1
83.7
31
50.7
(37.6)
Chopper Cmd.
811.0
7387.8
946.3
1117.3
1565.0
420
1775.4
(593.9)
Crazy Climber
10780.5
35829.4
36700.5
83940.2
59324.2
97190
15751.3
(5488.9)
Demon Attack
152.1
1971.0
$17.6
13003.9
2034.4
303
2809.5
(749.2)
Freeway
0.0
29.6
19.3
21.8
31.1
11.9
(4.7)
Frostbite
65.2
4334.7
1170.7
296.3
259.1
909
265.6
(8.6)
Gopher
257.6
2412.5
660.6
3260.3
2236.1
3730
416.4
(133.7)
Hero
1027.0
30826.4
5858.6
9315.9
7037.4
11161
1522.9
(1513.1)
James Bond
29.0
302.8
366.5
517.0
462.7
445
280.9
(60.9)
Kangaroo
52.0
3035.0
3617.4
724.1
838.2
4098
271.2
(298.0)
Krull
1598.0
2665.5
3681.6
5663.3
6616.4
7782
7813.7
(1598.0)
Kung Fu Master
258.5
22736.3
14783.2
30944.8
21759.8
21420
20169.8
(8206.7)
Ms Pacman
307.3
6951.6
1318.4
1281.2
999.1
1327
1356.8
(775.8)
Pong
-20.7
14.6
-5.4
20.1
14.6
18
12.6
(10.0)
Private Eye
24.9
69571.3
86.0
96.7
100.0
882
405.5
(1144.3)
Qbert
163.9
13455.0
866.3
13781.9
745.7
3405
685.0
(68.3)
Road Runner
11.5
7845.0
12213.1
17751.3
9614.6
15565
8164.2
(4066.8)
Seaquest
68.4
42054.7
558.1
1100.2
661.3
618
476.8
(88.7)
UP n' Down
533.4
11693.2
10859.2
17264.2
3546.2
N/A
7745.0
(8515.7)
Normalized mean
0.000
1.000
0.616
1,943
1.046
1.12
0.884
Normalized median
0.000
1.000
0.396
1.090
0.289
0.49
0.152
Table 3: Mean scores for the ablation studies.
Ablation
Boxing
Breakout
Kung Fu Master
Ms Paeman
Pong
SGF (default)
83.4
50.7
20169.8
1356.8
12.6
No augmentations
9.1
6031.4
847.6
-20.6
No action stacking
63.3
22.9
11981.2
651.4
-3.1
No frame stacking
12.1
8.2
17833.2
708.4
3.2
No temporal consistency
0.0
6.4
17008.8
754.3
-20.6
Sample-contrastive
77.4
25.1
19237.2
1162.6
Table 4: Total training times of various methods the Atari 100k benchmark. They approxi- mated for an NVIDIA V100 GPU.
Table 5: Detailed time breakdown. Percentages in the lower half are relative the default setting, obtained by enabling or disabling components.
Method
Runtime (hours)
Component
Percentage
SPR
2.3
Total training
100 %
SGF (ours)
3
World model training
63 %
DreamerV3
12
Policy training
37 %
TWM
20
No augmentations
16 %
EfficientZero
29
No action stacking
-0.1 %
IRIS
168
frame stacking
-0.1 %
SimPLE
240
+ With decoder
+19 %
18

---

Published as a conference paper at ICLR 2025
E ADDITIONAL ABLATIONS
Table 6: Mean scores for additional ablation studies.
Ablation
Boxing
Breakout
Kung Fu Master
Ms Pacman
Pong
SGF Projector x 4
83.9
42.2
22626.2
1134.0
14.8
79.1
23.7
22481.4
1176.1
6.7
Transition x 4
82.9
27.1
23779.4
1208.7
13.5
Actor/critic 4
77.9
29.7
17627.8
787.5
1.7
Horizon = 15
57.9
30,9
21579.4
978.1
-0.4
Training 4
64.1
29.0
14035.8
572.4
13.3
Stack size 8
68.1
21.8
19171.2
1154.7
11.5
Stack size 12
35.3
14.4
16423.6
1042.8
12.6
Recurrent transition
71.7
11.8
18639.4
1105.5
13
Recurrent predictors
87.8
9.4
16086.2
1032.1
10.5
In this section we provide additional ablation studies to analyze the effect of increasing the model size and training time. The results are shown in Table 6 We evaluated these additional ablation studies on 5 instead of 10 random seeds. We observe that the default configuration of SGF performs best most cases. The ablations are as follows:
1. Projector x 4: We increase the hidden dimension of the projector network from 2048 8192. This increases the number of parameters this network from 9.5M to 88M.
2. Transition X We increase the hidden dimension of the transition network from 1024 to 4096. This increases the number of parameters of this network from 5.2M 71.3M.
3. Actor/critic x 4: We increase the hidden dimension of the actor and critic networks from 512 to 2048. This increases the total number of parameters of the agent from 0.5M 8.4M.
4. Horizon 15: We increase the imagination horizon H from 10 15 steps, and reduce the imagination batch size to 2048 to keep the effective batch size constant.
5. Training x 4: We train the world model and the agent with two batches per environment step instead of one batch every second step. This effectively multiplies the training time by four and is similar to the training time of Dreamer V3.
6. Stack size = We stack 8 frames and actions instead of 4.
7. Stack size 12: We stack 12 frames and actions instead of 4.
8. Recurrent transition: We incorporate recurrent layer (LSTM) into the transition network, placing it after the five hidden linear layers and before the final output layer. This requires several changes in the implementation, since the recurrent states must be maintained and passed between steps.
9. Recurrent predictors: We make the transition, reward, and terminal distributions recurrent by introducing shared three-layer MLP followed by LSTM layer. The output of the LSTM is then fed into a two-layer transition head, as well as the reward and terminal networks.
19

---

Published as a conference paper at ICLR 2025
F IMPLEMENTATION DETAILS
Implementing a world model involves numerous design choices, many of which may seem arbitrary the first glance or are obscured in the source code. In the following, we explain all of our implementation details.
Stacking and preprocessing As detailed Section 2.2 we stack the m most recent observations and actions, with m = 4. Frame stacking also plays role in our representation learning approach, with observations o and o' sharing information from three subsequent frames. For data augmentation. we apply the transformations proposed for Atari by Yarats et al. (2021), i.c., random shifts and imagewise intensity jittering.
Distributions. We model the transition distribution using independent normal distributions with unit variance, i.c., Pe(y') y, a) = NWy' n a), Ia), where μe is a neural network computing the mean vector. We chose this distribution since the loss function reduces to minimizing the mean squared error 2 || a) y I12. Also, the mean that we use for prediction is available without further computation. We model the reward distribution po(r y, a, y') using discrete regression with two-hot encoded targets and symlog predictions, as recently proposed by Hafner et al. (2023). Al- though not yet being a widely used approach, it makes reward prediction stable across different scales without the need for domain-specific reward normalization or hyperparameter tuning. We model the terminal distribution using a Bernoulli distribution, i.c., (e y. a, yy) = Bernouliie | (3). a, y')), where dg is a neural network computing the terminal probability. We chose this distribution since is a common choice for distributions with binary support: the loss function reduces to the binary cross-entropy, and the mode can be computed by [oo a. y') 2 0.5) with the squared brackets being Iverson brackets.
Architecture. All networks use SiLU nonlincarities (Hendrycks & Gimpel 2016) to prevent dead ReLUs, especially given that the data is coming form an ever-changing replay buffer. Furthermore, we employ layer normalization (Ba et al. (016) in all networks. The encoder fo consists of four convolutional layers with kernel size of 4. stride of 2, and padding size of 1, followed by a linear layer that computes representations of dimension normalized using layer normalization; refer to Section 5.2 for our motivation. The projector network is an MLP with two hidden layers of dimension 2048, computing embeddings of dimension D = 2048; these dimension have been proven to strike good balance between qualitative performance and architecture size for VICReg (Garrido et al. 2022). The predictor network uses the same architecture the projector network. The network of the transition distribution is an MLP with five hidden layers of dimension 1024, and residual connection from the input to the output. The networks of the reward distribution, terminal distribution, policy, and value function MLPs with two hidden layers of dimension 1024. We use the AdamW optimizer (Loshchilov & Hutter 2019) for all networks and loss functions.
512. To stabilize training, the representations are also
Actor-critic. We estimate advantages using generalized advantage estimation (Schulman et al. 2016) and calculate multi-step truncated X-returns (Sutton & Barto 2018) as the target for the value function. To improve exploration and prevent early convergence suboptimal policies, add the entropy of the policy to the objective (Williams & Pengl 1991 Mnih et al.) (2166). Additionally, adopt the following strategies from DreamerV3 (Hafner al. (002), which have demonstrated success across various environments and reward scales without domain-specific fine-tuning. For advantage computation, the returns are normalized by mapping the 5th and the 95th percentile and 1, respectively. The value function utilizes the same discrete regression approach as the reward predictor, i.c., two-hot encoded targets and symlog predictions. target network, which is the exponential moving average of the online value network, computes additional targets for the value function. This allows for estimating returns using the online network instead of the target network.
20

---

Published as a conference paper at ICLR 2025
Algorithm 1: SGF's main training procedure.
environment environment steps M, imagination horizon initialize replay buffer D initialize networks the world model and the policy for i fl....................... MM) do execute action in & according to policy T store observed transition in D
Input:
# all following computations are batch-wise if world model update then
sample batch of transitions 7 D estimate Repe. (0) and Dyn (f) according
update 0 to minimize the losses if policy update then
# sample from arbitrary time steps sample batch of observations ~ encode observations Yı fo(O)
for E f H do
select actions A, TT Yi)
predict Y++1 ~ (Y(+1 Y,, Ad)
predict T141 ~ (P Yr, Ar, Y411 predict 6r+1 ~ Po Y,, Ar, Y++1) update 6 actor-critic style using trajectories
Table 7: Summary of all hyperparameters. Note that we use the original coefficients for VICReg.
Hyperparameter
Symbol
Value
Dimensionality of y Dimensionality of z Consistency coefficient Covariance coefficient Variance coefficient Frame resolution
d
512
D
2048
n
12.5
P
1.0
25.0
64 64
Grayscale frames Terminal on loss
No
Yes
Frame and action stacking
m
Random shifts Discount factor
0-3 pixels 0.997
A-return parameter
X
Entropy cocfficient
x 10-3
Target network decay
0.98
World model training interval
Every 2nd environment step
Policy training interval
Every 2nd environment step
Environment steps Initial random steps World model batch size World model learning rate World model warmup steps World model weight decay World model gradient clipping Imagination batch size Imagination horizon Actor-critic learning rate Actor-critic gradient clipping
M
100 000
5000
1024
x 10
5000
10
10.0
3072
H
10
2.4 10-4
100.0
Policy temperature for evaluation
0.5 (0.01 for Freeway)
Random actions during collection
1%
21


Reviews:
[{"review_id": "Q8ENn7anYw", "reviewer": "Reviewer_hwrw", "summary": "This paper discusses the design of a world model - a parametric model that predicts the transitions probabilities, reward functions and terminal state distributions to serve as a simulation environment for reinforcement learning. The paper proposes to simplify some existing elements in recent world-models, such as sequence models (i.e. RNNs, transformers), to focus on simple ingredients (e.g. frame stacking). It instead proposes to keep the maximum information and temporal consistency formulation as the most essential properties of effective world models. The resultant model is dubbed \"Simple, Good and Fast World Models\" (SGF) which demonstrates somewhat competitive performance against other existing world models in the Atari 100k benchmark.", "questions": "I think it will be useful if the paper can present accuracy numbers in similar training time, or training time numbers at roughly the same accuracy numbers. That will be very useful in positioning this work against existing methods.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "## Presentation\nThis paper is well-written with thorough discussion on the related works, the design philosophy and the precise formulations of the proposed modeling. The discussions are usually precise and insightful. The important elements in building the proposed world-models, such as the POMDP formulation, the representation learning (including sufficient details, such as image augmentations, temporal consistency, covariance regularization), the dynamics learning (the conditional independence assumption and the resultant factorization) are all presented clearly, leaving the readers with no doubts on the technical components and the underlying reasoning.\n\nIn many cases, the discussions have involved clear contrast with prior works, for example in Table 1 where the design choices in a good selection of prior works are presented, and compared with the design choices in the proposed SGF.\n\n## Empirical Studies\nThe paper is tested on the standard Atari 100k tasks that is standard for this kind of work. The empirical study seems sounds and demonstrates a few interesting properties of this work that may be useful for researchers in this area. In particular\n- It has presented through ablation studies in Section 4.2. showing the importance of various design components, such as state (image) augmentations, action/frame stacking, temporal consistency and sample contrastive formulation. Additional details are also included in the appendix, such as in Table 3 in Appendix D.\n- The paper has presented detailed comparisons with other methods, as presented in Appendix D Table 2.", "weaknesses": "There are a few limitations that seem to limit the contributions of this work.\n\n- As discussed thoroughly in the related work section in the paper, world model in training reinforcement learning agents is not a new idea. In such cases, it is useful to establish that this work is addressing a significant weakness in prior works, without sacrificing other important metrics. In this case, the main motivation of SGF seems to be presenting a simple, fast, yet accurate method to train a good world model. While SGF certainly fit the bill for the first two, it seems to fall short significantly in the third. The result in Table 2 seems to suggest that SGF is much weaker than prior works such as IRIS and DreamerV3, which are all based on learning on imagination and hence are arguably fairly related.\n\n-  The choice to drop sequence models appears to be a fairly significant limiting factor. This was acknowledged clearly in Line 462-464, which lists the decision to drop sequence models as a potential reasons why this work does not reach SOTA performance on Atari 100k. This seems like a fairly trivial observation - certain games in Atari 100k requires long term reasoning, and it is indeed one of the most challenging aspects of world models for RL agents. It should be expected that removing a components specifically designed to address this important challenge will result in inferior performance. It seems unlikely that researchers in this field will learn much more about how to design better world models if they are simply presented with results comparing methods with and without a sequence model - they probably already know that it is going to be much worse for certain tasks that require long term reasoning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "mrPr24NXMs", "reviewer": "Reviewer_LRfe", "summary": "This paper proposes a new approach to world modeling for the training of RL-based agents. The current state of literature employs different mechanisms to achieve a correct modeling of long sequences, including reconstruction, recurrent modules, memory, etc. The authors propose to simplify the world modeling and just use self-supervised techniques, inspired by VicReg, and a simple learning strategy for the system dynamics. There is a comparison of existing methods and the proposed one on Atari100k, and appropriate ablations of the components.", "questions": "1. How do the method perform with long-term dependencies?\n2. How do feature extraction perform if we change the SSL training objective?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. I believe these kinds of works are important. It is easy to just incrementally propose new components to improve the performance of systems while considerably increasing the engineering complexity. This does not give a clear view of the actual importance of the components included in the SOTA of world modeling. Going in a completely different direction is in my opinion a needed move sometimes, and it will help to shape new design choices for world modeling. Hence, the motivation is strong, and the reasoning behind that is coherent.\n2. The results are compelling. I believe that on Atari100K the simple method proposed performs quantitatively well, even compared to baselines that are far more complex, and with considerably lower runtime.\n3. The comparison with existing methods is nontrivial and requires a proper analysis of the literature. Table 1 is also interesting and will be useful for future work.\n4. The paper is well-written and well-motivated, all introduced explanations are useful and the writing is compact enough. The proposed experiments are interesting.", "weaknesses": "1. I believe that while the proposed method focuses on short-term dependencies, as correctly stated in the limitation, how much performance degrades with an increasingly long-term dependency on actions would be important to quantify. This will allow us to assess the limitations of the proposed method in a more robust manner, for people to build upon.\n2. It is not clear to me why only VICReg is chosen for representation extraction. There are relationships with BYOL and SimSiam as reported in the appendix, but it is unclear how performance would change if these approaches were instead used for representation extraction rather than VICReg.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "zgDVeANqLJ", "reviewer": "Reviewer_DZ8x", "summary": "Authors introduce model-based method: SGF (a Simple, Good, and Fast) world model.\nThe simplicity of the model is that: it stacks 4 previous frames and actions to capture short-time dependencies, and it uses data augmentation to enhance robustness. It is trained to provide maximum information and temporal consistency.\nFor simplicity, it does not use: image reconstructions, discretization of representations, sequence models (RNN, transformers, ...),  probabilistic predictions for deterministic environments.\n\nThey use several loss functions:\n- MSE-loss for temporal consistency (between next embedding and action-conditioned next embedding)\n- Variance and Covariance regularization loss for Information Maximization (use the current and next batches of embeddings as input)\n- Loss for 3 distributions for Dynamics learning: \n    - transition - probability of transition to y'-representation if a-action is applied to y-representation (gradient is only propagated to a-action)\n    - reward - probability of r-reward if a-action is applied to y-representation\n    - terminal - probability of e-terminal-action if a-action is applied to y-representation\n\nThey train several models using AdamW optimizer: \n- Encoder (f computes representations from current and next input image observations): 4 convolutional layers (kernel size 4, stride 2, pad 1), linear layer dim = 512, norm-layer, SiLU activation\n- Projector (g computes current and next embeddings from current and next representations): MLP with 2 hidden layers of dim = 2048\n- Predictor (h predicts action-conditioned next embedding): MLP with 2 hidden layers of dim = 2048\n- Transition network: MLP with 5 hidden layers of dim = 1024, and a residual connection from the input to the output (suggested to use the sequence model for future work)\n- The networks of the reward distribution, terminal distribution, policy, and value function: MLPs with 2 hidden layers of dim = 1024\n\nThis simple approach achieves shorter training times compared to other world models and good performance on the Atari 100k benchmark.", "questions": "Although this is partially explained in the Limitations, could you go into more detail about: for more complex tasks than Atari-games, where longer term knowledge needs to be remembered (where orders of magnitude more memory of early events is required), what parts of the models should be scaled up or significantly modified / replaced (e.g. using  sequence model: RNN, Transformers, ...) to be optimal in terms of accuracy and training time?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "In this work, authors try to find the most necessary components that are the most optimal in terms of accuracy and training time.\n\n1. The presented SGF approach lies on the Pareto optimality curve on the chart of Accuracy (normalized mean score) and Training time (hours) - Figure 5, where the other points on the Pareto optimality curve are: SPR, DreamerV3, EfficientZero (performs lookahead)\n\n2. The optimal combination of improvements (frame stacking, action stacking, temporal consistency, augmentations, sample-contrastive) has been found to achieve the highest accuracy in five games - Figure 4\n\n3. Optimal sizes of models and training times have been found in Table 6 to achieve the highest possible mean scores\n\n4. The presented experimental results show the necessity of temporal consistency in Figure 7", "weaknesses": "1. While optimal sizes of models and training times have been found in Table 6 to achieve the highest possible mean scores, this may mean that either the scalability of the approach is limited, or the approach must scale in many directions simultaneously to achieve even higher mean scores. Although the approach lies on the Pareto optimality curve on the Accuracy vs Training time chart, i.e. it is one of many optimal options, it is not shown how this approach can be scaled or improved to achieve the highest accuracy with increasing Training time. \nOr it requires more serious architectural changes, f.e. as the authors suggest for future work - to use the sequence model for transition network.\n\n2. The proposed method is optimal for tasks that are probably simple enough and do not require remembering very old events, so it is sufficient to have a stack of 4 previous frames and actions, and it is not shown how this approach can be transferred to more complex tasks where orders of magnitude more memory of early events is required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "HwCzVWqLmX", "reviewer": "Reviewer_iALA", "summary": "This paper introduces a world model that uses self-supervised representation learning, captures short-time dependencies through frame and action stacking, and enhances robustness against model errors through data augmentation. This paper is based on a partially observable Markov decision process. It stacks the recent observations and actions to capture short-time dependencies. It introduces stochasticity through data augmentation. To build meaningful representations of observations, this paper enforeces information maximization and temporal consistency of the features.", "questions": "Considering its simplicity, it would be nice and help the paper strongly if scalability can be demonstrated. Can it be scaled?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper contains thorough experiments.\n2. The method is simple and fast.", "weaknesses": "1. The strategy of stacking observations and actions seem overfitting to the chosen dataset, Atari 100k benchmark, rather than fitting to the real-world, where long-term dependencies are common.\n2. Related to the first weakness, while the games in the dataset are deterministic, real-world can be very stochastic. It is questionable whether the model can be applied in real-world cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)