Title: STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: deep learning

reinforcement learning

model-based reinforcement learning

world model

learning in imagination

transformer

variational autoencoders

sequence modeling
Abstract: Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy  in visual input environments. These approaches begin by constructing a parameterized simulation world model of the real environment through self-supervised learning. By leveraging the imagination of the world model, the agent's policy is enhanced without the constraints of sampling from the real environment. The performance of these algorithms heavily relies on the sequence modeling and generation capabilities of the world model. However, constructing a perfectly accurate model of a complex unknown environment is nearly impossible. Discrepancies between the model and reality may cause the agent to pursue virtual goals, resulting in subpar performance in the real environment. Introducing random noise into model-based reinforcement learning has been proven beneficial.
In this work, we introduce Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture that combines the strong sequence modeling and generation capabilities of Transformers with the stochastic nature of variational autoencoders. STORM achieves a mean human performance of $126.7\%$ on the Atari $100$k benchmark, setting a new record among state-of-the-art methods that do not employ lookahead search techniques. Moreover, training an agent with $1.85$ hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only $4.3$ hours, showcasing improved efficiency compared to previous methodologies.

STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

Abstract

Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy in visual input environments. These approaches begin by constructing a parameterized simulation world model of the real environment through self-supervised learning. By leveraging the imagination of the world model, the agent's policy is enhanced without the constraints of sampling from the real environment. The performance of these algorithms heavily relies on the sequence modeling and generation capabilities of the world model. However, constructing a perfectly accurate model of a complex unknown environment is nearly impossible. Discrepancies between the model and reality may cause the agent to pursue virtual goals, resulting in subpar performance in the real environment. Introducing random noise into model-based reinforcement learning has been proven beneficial. In this work, we introduce Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture that combines the strong sequence modeling and generation capabilities of Transformers with the stochastic nature of variational autoencoders. STORM achieves a mean human performance of 126 . 7% on the Atari 100 k benchmark, setting a new record among state-of-the-art methods that do not employ lookahead search techniques. Moreover, training an agent with 1 . 85 hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only 4 . 3 hours, showcasing improved efficiency compared to previous methodologies.

We release our code at https://github.com/weipu-zhang/STORM.

1 Introduction

Deep reinforcement learning (DRL) has exhibited remarkable success across diverse domains. However, its widespread application in real-world environments is hindered by the substantial number of interactions with the environment required for achieving such success. This limitation becomes particularly challenging when dealing with broader real-world settings in e.g., unmanned and manufacturing systems [1, 2] that lack adjustable speed simulation tools. Consequently, improving the sample efficiency has emerged as a key challenge for DRL algorithms.

Popular DRL methods, including Rainbow [3] and PPO [4], suffer from low sample efficiency due to two primary reasons. Firstly, the estimation of the value function proves to be a challenging task. This

involves approximating the value function using a deep neural network (DNN) and updating it with n -step bootstrapped temporal difference, which naturally requires numerous iterations to converge [5]. Secondly, in scenarios where rewards are sparse, many samples exhibit similarity in terms of value functions, providing limited useful information for training and generalization of the DNN [6, 7]. This further exacerbates the challenge of improving the sample efficiency of DRL algorithms.

To address these challenges, model-based DRL algorithms have emerged as a promising approach that tackles both issues simultaneously while demonstrating significant performance gains in sampleefficient settings. These algorithms start by constructing a parameterized simulation world model of the real environment through self-supervised learning. Self-supervised learning can be implemented in various ways, such as reconstructing the original input state using a decoder [8-10], predicting actions between frames [7], or employing contrastive learning to capture the internal consistency of input states [6, 7]. These approaches provide more supervision information than conventional modelfree RL losses, enhancing the feature extraction capabilities of DNNs. Subsequently, the agent's policy is improved by leveraging the experiences generated using the world model, eliminating sampling constraints and enabling faster updates to the value function compared to model-free algorithms.

However, the process of imagining with a world model involves an autoregressive process that can accumulate prediction errors over time. In situations where discrepancies arise between the imagined trajectory and the real trajectory, the agent may inadvertently pursue virtual goals, resulting in subpar performance in the real environment. To mitigate this issue, introducing random noise into the world model has been proven beneficial [9-11, 14]. Variational autoencoders, capable of automatically learning low-dimensional latent representations of highdimensional data while incorporating reasonable random noise into the latent space, offer an ideal choice for image encoding.

Numerous endeavors have been undertaken to construct an efficient world model. For instance, SimPLe [11] leverages LSTM [15], while DreamerV3 [10] employs GRU [16] as the sequence model. LSTM and GRU, both variants of recurrent neural networks (RNNs), excel at sequence modeling tasks. However, the recurrent nature of RNNs impedes parallelized computing, resulting in slower training speeds [17]. In contrast, the Transformer architecture [17] has lately demonstrated superior performance over RNNs in various sequence modeling and generation tasks. It overcomes the challenge of forgetting long-term dependencies and is designed for efficient parallel computing. While several attempts have been made to incorporate Transformers into the world model [12, 13, 18], these works do not fully harness the capabilities of this architecture. Furthermore, these approaches require even longer training times and fail to surpass the performance of the GRUbased DreamerV3.

SimPLe

TWM

IRIS

STORM (ours)

0%

50%

100%

150%

33%

96%

104%

112%

127%

Mean human normalized score

SimPLe

TWM

IRIS

STORM (ours)

0%

20%

40%

60%

13%

51%

29%

49%

58%

Median human normalized score

SimPLe

TWM

IRIS

STORM (ours)

0

5

10

0.5

5.6

0.7

9.3

11.9

Training FPS on V100

In this paper, we introduce the Stochastic Transformer-based wORld Model (STORM), a highly effective and efficient structure for model-based RL. STORM employs a categorical variational autoencoder (VAE) as the image encoder, enhancing the agent robustness and reducing accumulated autoregressive prediction errors. Subsequently, we incorporate the Transformer as the sequence

model, improving modeling and generation quality while accelerating training. STORM achieves a remarkable mean human normalized score of 126 . 7% on the challenging Atari 100 k benchmark, establishing a new record for methods without resorting to lookahead search. Furthermore, training an agent with 1 . 85 hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only 4 . 3 hours, demonstrating superior efficiency compared to previous methodologies. The comparison of our approach with the state-of-the-art methods is depicted in Figure 1.

2 Related work

Model-based DRL algorithms aim to construct a simulation model of the environment and utilize simulated experiences to improve the policy. While traditional model-based RL techniques like Dyna-Q have shown success in tabular cases [5], modeling complex environments such as video games and visual control tasks presents significant challenges. Recent advances in computing and DNNs have enabled model-based methods to learn the dynamics of the environments and start to outperform model-free methods on these tasks.

The foundation of VAE-LSTM-based world models was introduced by Ha and Schmidhuber [14] for image-based environments, demonstrating the feasibility of learning a good policy solely from generated data. SimPLe [11] applied this methodology to Atari games, resulting in substantial sample efficiency improvements compared to Rainbow [3], albeit with relatively lower performance under limited samples. The Dreamer series [8-10] also adopt this framework and showcase notable capabilities in Atari games, DeepMind Control, Minecraft, and other domains, using GRU [16] as the core sequential model. However, as discussed earlier, RNN structures suffer from slow training [17].

Recent approaches such as IRIS [13], TWM [12], and TransDreamer [18] incorporate the Transformer architecture into their world models. IRIS [13] employs VQ-VAE [19] as the encoder to map images into 4 × 4 latent tokens and uses a spatial-temporal Transformer [20] to capture information within and across images. However, the attention operations on a large number of tokens in the spatial-temporal structure can result in a significant training slowdown. TWM [12] adopts Transformer-XL [21] as its core architecture and organizes the sequence model in a structure similar to Decision Transformer [22], treating the observation, action, and reward as equivalent input tokens for the Transformer. Performing self-attention across different types of data may have a negative impact on the performance, and the increased number of tokens considerably slows down training. TransDreamer [18] directly replaces the GRU structure of Dreamer with Transformer. However, there is a lack of evidence demonstrating their performance in widely accepted environments or under limited sample conditions.

Other model-based RL methods such as MuZero [23], EfficientZero [7], and SpeedyZero [24] incorporate Monte Carlo tree search (MCTS) to enhance policy and achieve promising performance on the Atari 100 k benchmark. Lookahead search techniques like MCTS can be employed to enhance other model-based RL algorithms, but they come with high computational demands. Additionally, certain model-free methods [6, 25-27] incorporate a self-supervised loss as an auxiliary term alongside the standard RL loss, demonstrating their effectiveness in sample-efficient settings. Additionally, recent studies [28, 29] delve deeper into model-free RL, demonstrating strong performance and high data efficiency, rivaling that of model-based methods on several benchmarks. However, since the primary objective of this paper is to enhance the world model, we do not delve further into these methods.

We highlight the distinctions between STORM and recent approaches in the world model as follows:

Agent

Agent

C

a

1

o

1

ENC

Z

1

z

1

ˆ

Z

2

ˆ

o

1

DEC

h

1

ˆ

r

1

ˆ

c

1

a

2

o

2

Z

2

z

2

ˆ

Z

3

ˆ

o

2

h

2

ˆ

r

2

ˆ

c

2

ENC

DEC

C

Transformer Blocks

a

4

z

4

ˆ

Z

5

h

4

ˆ

r

4

ˆ

c

4

C

a

3

z

3

ˆ

Z

4

h

3

ˆ

r

3

ˆ

c

3

C

Legend:

Neural networks

Inner output

Environment data

Random sample

Policy

Concatenate

C

SimPLe [11] and Dreamer [10] rely on RNN-based models, whereas STORM employs a GPT-like Transformer [30] as the sequence model.

In contrast to IRIS [13] that employs multiple tokens, STORM utilizes a single stochastic latent variable to represent an image.

STORMfollows a vanilla Transformer [17] structure, while TWM [12] adopts a TransformerXL [21] structure.

In the sequence model of STORM, an observation and an action are fused into a single token, whereas TWM [12] treats observation, action, and reward as three separate tokens of equal importance.

Unlike Dreamer [10] and TransDreamer [18], which incorporate hidden states, STORM reconstructs the original image without utilizing this information.

3 Method

3.1 World model learning

Our approach adheres to the established framework of model-based RL algorithms, which focus on enhancing the agent's policy by imagination [5, 9-11, 13]. We iterate through the following steps until reaching the prescribed number of real environment interactions.

S1) Gather real environment data by executing the current policy for several steps and append them to the replay buffer.

S2) Update the world model using trajectories sampled from the replay buffer.

S3) Improve the policy using imagined experiences generated by the world model, with the starting points for the imagination process sampled from the replay buffer.

At each time t , a data point comprises an observation o t , an action a t , a reward r t , and a continuation flag c t (a Boolean variable indicating whether the current episode is ongoing). The replay buffer maintains a first-in-first-out queue structure, enabling the sampling of consecutive trajectories from the buffer.

Section 3.1 provides a detailed description of the architecture and training losses employed by STORM. On the other hand, Section 3.2 elaborates on the imagination process and the training methodology employed by the agent. It provides an thorough explanation of how the agent leverages the world model to simulate experiences and improve its policy.

Model structure The complete structure of our world model is illustrated in Figure 2. In our experiments, we focus on Atari games [31], which generate image observations o t of the environment. Modeling the dynamics of the environment directly on raw images is computationally expensive and prone to errors [7-11, 13, 23]. To address this, we leverage a V AE [32] formulated in Equation (1) to convert o t into latent stochastic categorical distributions Z t . Consistent with prior work [9, 10, 12], we set Z t as a stochastic distribution comprising 32 categories, each with 32 classes. The encoder ( q ϕ ) and decoder ( p ϕ ) structures are implemented as convolutional neural networks (CNNs) [33]. Subsequently, we sample a latent variable z t from Z t to represent the original observation o t . Since sampling from a distribution lacks gradients for backward propagation, we apply the straight-through gradients trick [9, 34] to preserve them.

Before entering the sequence model, we combine the latent sample z t and the action a t into a single token e t using multi-layer perceptrons (MLPs) and concatenation. This operation, denoted as m ϕ , prepares the inputs for the sequence model. The sequence model f ϕ takes the sequence of e t as input and produces hidden states h t . We adopt a GPT-like Transformer structure [30] for the sequence model, where the self-attention blocks are masked with a subsequent mask allowing e t to attend to the sequence e 1 , e 2 , . . . , e t . By utilizing MLPs g D ϕ , g R ϕ , and g C ϕ , we rely on h t to predict the current reward ˆ r t , the continuation flag ˆ c t , and the next distribution ˆ Z t +1 . The formulation of this part of the world model is as follows

Loss functions The world model is trained in a self-supervised manner, optimizing it end-to-end. The total loss function is calculated as in Equation (3) below, with fixed hyperparameters β 1 = 0 . 5 and β 2 = 0 . 1 . In the equation, B denotes the batch size, and T denotes the batch length

The individual components of the loss function are defined as follows: L rec t ( ϕ ) represents the reconstruction loss of the original image, L rew t ( ϕ ) represents the prediction loss of the reward, and L con t ( ϕ ) represents the prediction loss of the continuation flag.

Additionally, L sym in Equation (4b) denotes the symlog two-hot loss, as described in [10]. This loss function transforms the regression problem into a classification problem, ensuring consistent loss scaling across different environments.

The losses L dyn t ( ϕ ) and L rep t ( ϕ ) are expressed as Kullback-Leibler (KL) divergences but differ in their gradient backward and weighting. The dynamics loss L dyn t ( ϕ ) guides the sequence model in predicting the next distribution, while the representation loss L rep t ( ϕ ) allows the output of the encoder to be weakly influenced by the sequence model's prediction. This ensures that the learning of distributional dynamics is not excessively challenging.

where sg( · ) denotes the operation of stop-gradients.

3.2 Agent learning

The agent's learning is solely based on the imagination process facilitated by the world model, as illustrated in Figure 2. To initiate the imagination process, a brief contextual trajectory is randomly selected from the replay buffer, and the initial posterior distribution Z t is computed. During inference, rather than sampling directly from the posterior distribution Z t , we sample z t from the prior distribution ˆ Z t . To accelerate the inference, we employ the KV cache technique [35] within the Transformer structure.

The agent's state is formed by concatenating z t and h t , as shown below:

We adopt the actor learning settings from DreamerV3 [10]. The complete loss of the actor-critic algorithm is described by Equation (7), where ˆ r t corresponds to the reward predicted by the world model, and ˆ c t represents the predicted continuation flag:

where H ( · ) denotes the entropy of the policy distribution, while constants η and L represent the coefficient for entropy loss and the imagination horizon, respectively. The λ -return G λ t [5, 10] is recursively defined as follows

The normalization ratio S utilized in the actor loss (7a) is defined in Equation (9), which is computed as the range between the 95 th and 5 th percentiles of the λ -return G λ t across the batch [10]

To regularize the value function, we maintain the exponential moving average (EMA) of ψ . The EMA is defined in Equation (10), where ψ t represents the current critic parameters, σ is the decay rate, and ψ EMA t +1 denotes the updated critic parameters. This regularization technique aids in stabilizing training and preventing overfitting

4 Experiments

We evaluated the performance of STORM on the widely-used benchmark for sample-efficient RL, Atari 100 k [31]. For detailed information about the benchmark, evaluation methodology, and the baselines used for comparison, please refer to Section 4.1. The comprehensive results for the Atari 100 k games are presented in Section 4.2.

4.1 Benchmark and baselines

Atari 100 k consists of 26 different video games with discrete action dimensions of up to 18 . The 100 k sample constraint corresponds to 400 k actual game frames, taking into account frame skipping ( 4 frames skipped) and repeated actions within those frames. This constraint corresponds to

approximately 1 . 85 hours of real-time gameplay. The agent's human normalized score τ = A -R H -R is calculated based on the score A achieved by the agent, the score R obtained by a random policy, and the average score H achieved by a human player in a specific environment. To determine the human player's performance H , a player is allowed to become familiar with the game under the same sample constraint.

To demonstrate the efficiency of our proposed world model structure, we compare it with model-based DRL algorithms that share a similar training pipeline, as discussed in Section 2. However, similarly to [10, 12, 13], we do not directly compare our results with lookahead search methods like MuZero [23] and EfficientZero [7], as our primary goal is to refine the world model itself. Nonetheless, lookahead search techniques can be combined with our method in the future to further enhance the agent's performance.

4.2 Results on Atari 100k

Detailed results for each environment can be found in Table 2, and the corresponding performance curve is presented in Appendix A due to space limitations. In our experiments, we trained STORM using 5 different seeds and saved checkpoints every 2 , 500 sample steps. We assessed the agent's performance by conducting 20 evaluation episodes for each checkpoint and computed the average score. The result reported in Table 2 is the average of the scores attained using the final checkpoints.

STORMdemonstrates superior performance compared to previous methods in environments where the key objects related to rewards are large or multiple, such as Amidar , MsPacman , Chopper Command , and Gopher . This advantage can be attributed to the attention mechanism, which explicitly preserves the history of these moving objects, allowing for an easy inference of their speed and direction information, unlike RNN-based methods. However, STORM faces challenges when handling a single small moving object, as observed in Pong and Breakout , due to the nature of autoencoders. Moreover, performing attention operations under such circumstances can potentially harm performance, as the randomness introduced by sampling may excessively influence the attention weights.

5 Ablation studies

In our experiments, we have observed that the design and configuration choices of the world model and the agent can have significant impacts on the final results. To further investigate this, we conduct ablation studies on the design and configuration of the world model in Section 5.1, as well as on the agent's design in Section 5.2. Additionally, we propose a novel approach to enhancing the exploration efficiency through the imagination capability of the world model using a single demonstration trajectory, which is explained in Section 5.3.

5.1 World model design and configuration

The RNN-based world models utilized in SimPLe [11] and Dreamer [9, 10] can be formulated clearly using variational inference over time. However, the non-recursive Transformer-based world model does not align with this practice and requires manual design. Figure 3a shows alternative structures and their respective outcomes. In the 'Decoder at rear' configuration, we employ z t ∼ ˆ Z t instead of z t ∼ Z t for reconstructing the original observation and calculating the loss. The results indicate that the reconstruction loss should be applied directly to the output of the encoder rather than relying on the sequence model. In the 'Predictor at front' setup, we utilize z t as input for g R ϕ ( · ) and g C ϕ ( · ) in Equation (2), instead of h t . These findings indicate that, while this operation has minimal impact on the final performance for tasks where the reward can be accurately predicted from a single frame (in e.g., Pong ), it leads to a performance drop on tasks that require several contextual frames to predict the reward accurately (in e.g., Ms. Pacman ).

By default, we configure our Transformer with 2 layers, which is significantly smaller than the 10 layers used in IRIS [13] and TWM [12]. Figure 3b presents the varied outcomes obtained by increasing the number of Transformer layers. The results reveal that increasing the layer count does not have a positive impact on the final performance. However, in the case of the game Pong , even when the sample limit is increased from 100 k to 400 k, the agent still achieves the maximum reward in this environment regardless of whether a 4 -layer or 6 -layer Transformer is employed. This scaling discrepancy, which differs from the success observed in other fields [36-38], may be attributed to three reasons. Firstly, due to the minor difference between adjacent frames and the presence of residual connections in the Transformer structure [17], predicting the next frame may not require a complex model. Secondly, training a large model naturally requires a substantial amount of data, yet the Atari 100 k games neither provide images from diverse domains nor offer sufficient samples for training a larger model. Thirdly, the world model is trained end-to-end, and the representation loss L rep in Equation (5b) directly influences the image encoder. The encoder may be overly influenced when tracking the output of a large sequence model.

0

200k

400k

0

1000

2000

3000

4000

0

200k

400k

20

10

0

10

20

Pong

STORM

Decoder at rear

Predictor at front

0

200k

400k

0

1000

2000

3000

4000

0

200k

400k

20

10

0

10

20

Pong

STORM (2L)

4L

6L

(a) Rearranging the modules of STORM.

(b) Number of layers in the Transformer.

5.2 Selection of the agent's state

The choice of the agent's state s t offers several viable options: ˆ o t [13], h t , z t (as in TWM [12]), or the combination [ h t , z t ] (as demonstrated by Dreamer, [9, 10]). In the case of STORM, we employ s t = [ h t , z t ] , as in Equation (6). Ablation studies investigating the selection of the agent's state are presented in Figure 4. The results indicate that, in environments where a good policy requires contextual information, such as in Ms. Pacman , the inclusion of h t leads to improved performance. However, in other environments like Pong and Kung Fu Master , this inclusion does not yield a significant difference. When solely utilizing h t in environments that evolve with the agent's policy,

like Pong , the agent may exhibit behaviors similar to catastrophic forgetting [39] due to the nonstationary and inaccurate nature of the world model. Consequently, the introduction of randomness through certain distributions like Z t proves to be beneficial.

0

200k

400k

0

1000

2000

3000

4000

0

200k

400k

20

10

0

10

20

Pong

0

200k

400k

0

10000

20000

30000

STORM

Agent w/o

zt

Agent w/o

ht

5.3 Impact of the demonstration trajectory

The inclusion of a demonstration trajectory is a straightforward implementation step when using a world model, and it is often feasible in real-world settings. Figure 5 showcases the impact of incorporating a single demonstration trajectory in the replay buffer. Details about the provided trajectory can be found in Appendix D. In environments with sparse rewards, adding a trajectory can improve the robustness, as observed in Pong , or the performance, as seen in Freeway . However, in environments with dense rewards like Ms Pacman , including a trajectory may hinder the policy improvement of the agent.

Freeway serves as a prototypical environment characterized by challenging exploration but a simple policy. To receive a reward, the agent must take the 'up' action approximately 70 times in a row, but it quickly improves its policy once the first reward is obtained. Achieving the first reward is extremely challenging if the policy is initially set as a uniform distribution of actions. In the case of TWM [12], the entropy normalization technique is employed across all environments, while IRIS [13] specifically reduces the temperature of the Boltzmann exploration strategy for Freeway . These tricks are critical in obtaining the first reward in this environment. It is worth noting that even for most humans, playing exploration-intensive games without prior knowledge is challenging. Typically, human players require instructions about the game's objectives or watch demonstrations from teaching-level or expert players to gain initial exploration directions. Inspired by this observation, we aim to directly incorporate a demonstration trajectory to train the world model and establish starting points for imagination. By leveraging a sufficiently robust world model, the utilization of limited offline information holds the potential to surpass specially designed curiosity-driven exploration strategies in the future.

0

200k

400k

0

1000

2000

3000

4000

0

200k

400k

20

10

0

10

20

Pong

0

200k

400k

0

10

20

30

Freeway

w/o demonstration

w/ demonstration

As an integral part of our methodology, we integrate a single trajectory from Freeway into our extensive results. Furthermore, to ensure fair comparisons in future research, we provide the results without incorporating Freeway 's trajectory in Table 2 and Figure 6. It is important to highlight that even without this trajectory, our approach consistently outperforms previous methods, attaining a mean human normalized score of 122 . 3% .

6 Conclusions and limitations

In this work, we introduce STORM, an efficient world model architecture for model-based RL, surpassing previous methods in terms of both performance and training efficiency. STORM harnesses the powerful sequence modeling and generation capabilities of the Transformer structure while fully exploiting its parallelizable training advantages. The improved efficiency of STORM broadens its applicability across a wider range of tasks while reducing computational costs.

Nevertheless, it is important to acknowledge certain limitations. Firstly, both the world model of STORM and the compared baselines are trained in an end-to-end fashion, where the image encoder and sequence model undergo joint optimization. As a result, the world model must predict its own internal output, introducing additional non-stationarity into the optimization process and potentially impeding the scalability of the world model. Secondly, the starting points for imagination are uniformly sampled from the replay buffer, while the agent is optimized using an on-policy actor-critic algorithm. Although acting in the world model is performed on-policy, the corresponding on-policy distribution µ ( s ) for these starting points is not explicitly considered, despite its significance in the policy gradient formulation: ∇ J ( θ ) ∝ ∑ s µ ( s ) ∑ a q π θ ( s, a ) ∇ π θ ( a | s ) [5].

Acknowledgments and Disclosure of Funding

We would like to thank anonymous reviewers for their constructive comments. The work was supported partially by the National Key R&D Program of China under Grant 2021YFB1714800, partially by the National Natural Science Foundation of China under Grants 62173034, U23B2059, 61925303, 62088101, and partially by the Chongqing Natural Science Foundation under Grant 2021ZX4100027.

References

Jie Chen, Jian Sun, and Gang Wang. From unmanned systems to autonomous intelligent systems. Engineering , 12:16-19, 2022.

Boxi Weng, Jian Sun, Gao Huang, Fang Deng, Gang Wang, and Jie Chen. Competitive meta-learning. IEEE/CAA Journal of Automatica Sinica , 10(9):1902-1904, 2023.

Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 32, 2018.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv:1707.06347 , 2017.

Richard S Sutton, Andrew G Barto, et al. Introduction to Reinforcement Learning , volume 135. MIT Press Cambridge, 1998.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering Atari games with limited data. Advances in Neural Information Processing Systems , 34:25476-25488, 2021.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=S1lOTC4tDS .

Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=0oabwyZbOu .

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv:2301.04104 , 2023.

Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Bła˙ zej Osi´ nski, 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 International Conference on Learning Representations , 2020. URL https: //openreview.net/forum?id=S1xCPJHtDB .

Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=TdBaDGCpjly .

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In International Conference on Learning Representations , 2023. URL https:// openreview.net/forum?id=vhFu1Acb0xb .

David Ha and Jürgen Schmidhuber. World models. arXiv:1803.10122 , 2018.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation , 9 (8):1735-1780, 1997.

Kyunghyun Cho, Bart Merrienboer, Caglar Gulcehre, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In EMNLP , 2014.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv:2202.09481 , 2022.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in Neural Information Processing Systems , 30, 2017.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 6836-6846, 2021.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of Annual Meeting of the Association for Computational Linguistics , pages 2978-2988, 2019.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems , 34:15084-15097, 2021.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering Atari, Go, chess and Shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

Yixuan Mei, Jiaxuan Gao, Weirui Ye, Shaohuai Liu, Yang Gao, and Yi Wu. SpeedyZero: Mastering Atari with limited data and time. In International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=Mg5CLXZgvLJ .

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations , 2021. URL https://openreview.net/ forum?id=uCQfPZwRaUu .

Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=GY6-6sTvGaf .

Hado P Van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? Advances in Neural Information Processing Systems , 32, 2019.

Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Marc G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In International Conference on Learning Representations , 2023. URL https://openreview. net/forum?id=OpC-9aBBVJe .

Max Schwarzer, Johan Samir Obando Ceron, Aaron Courville, Marc G Bellemare, Rishabh Agarwal, and Pablo Samuel Castro. Bigger, better, faster: Human-level Atari with human-level efficiency. In International Conference on Machine Learning , pages 30365-30380. PMLR, 2023.

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.

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.

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

Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation , 1(4):541-551, 1989.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv:1308.3432 , 2013.

Carol Chen. Transformer inference arithmetic. kipp.ly , 2022. URL https://kipp.ly/blog/ transformer-inference-arithmetic/ .

OpenAI. ChatGPT. https://openai.com/chat-gpt/ , 2021. Accessed: April 19, 2023.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=YicbFdNTTy .

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 , pages 770-778, 2016.

Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation , volume 24, pages 109-165. Elsevier, 1989.

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning , pages 448-456. PMLR, 2015.

Matthew D Zeiler, Dilip Krishnan, Graham W Taylor, and Rob Fergus. Deconvolutional networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 2528-2535. IEEE, 2010.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv:1607.06450 , 2016.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research , 15(1):1929-1958, 2014.

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

Florin Gogianu, Tudor Berariu, Lucian Bus , oniu, and Elena Burceanu. Atari agents, 2022. URL https://github.com/floringogianu/atari-agents .

A Atari 100k curves

0

200k

400k

0

500

1000

1500

Alien

0

200k

400k

0

100

200

Amidar

0

200k

400k

500

1000

Assault

0

200k

400k

500

1000

Asterix

0

200k

400k

0

500

1000

Bank Heist

0

200k

400k

5000

10000

15000

0

200k

400k

0

50

Boxing

0

200k

400k

0

20

40

60

Breakout

0

200k

400k

1000

2000

Chopper Command

0

200k

400k

0

50000

100000

0

200k

400k

0

250

500

750

0

200k

400k

0

10

20

30

Freeway

0

200k

400k

0.05

0.00

0.05

Freeway w/o traj

0

200k

400k

0

1000

2000

Frostbite

0

200k

400k

0

5000

10000

15000

Gopher

0

200k

400k

0

5000

10000

Hero

0

200k

400k

0

250

500

750

0

200k

400k

0

5000

Kangaroo

0

200k

400k

0

5000

10000

Krull

0

200k

400k

0

10000

20000

30000

0

200k

400k

0

2000

4000

0

200k

400k

20

0

20

Pong

0

200k

400k

0

5000

10000

15000

Private Eye

0

200k

400k

0

2000

4000

0

200k

400k

0

10000

20000

0

200k

400k

250

500

750

Seaquest

0

200k

400k

0

10000

20000

Up N Down

0

200k

400k

0.0

0.2

0.4

0.6

0

200k

400k

0.0

0.5

1.0

Gamer Mean

STORM (Ours)

B Details of model structure

×

C Hyperparameters

D Demonstration trajectory information

×

E Computational cost details and comparison

F Atari video predictions

Real

Model

t = 0

t = 7

Imagination

t = 15

t = 23

t = 31

t = 39

t = 47

t = 55

Environment: Boxing

t = 63

Real

Model

t = 0

t = 7

Imagination

t = 15

t = 23

t = 31

t = 39

t = 47

t = 55

Environment: ChopperCommand

t = 63

Real

Model

t = 0

t = 7

Imagination

t = 15

t = 23

t = 31

t = 39

t = 47

t = 55

t = 63

Real

Model

t = 0

t = 7

Imagination

t = 15

t = 23

t = 31

t = 39

t = 47

t = 55

Environment: Pong

t = 63

Real

Model

t = 0

t = 7

Imagination

t = 15

t = 23

t = 31

t = 39

t = 47

t = 55

Environment: RoadRunner

t = 63

号

5-+

E3

940

+1

Peer reviews:
comments: SUMMARY:
The paper introduces the Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture. STORM proposes to encode image inputs using a stochastic variational autoencoder, and predicts latent state using a GPT-like sequential model. It then trains dynamics and policy based on the outputs of these stochastic variational autoencoder and sequential model. The authors conducted experimental comparisons with some classic baseline methods in the Atari100k environment. The results indicate that STORM has a faster speed and achieves better performance.

STRENGTHS:
1. Although the individual components proposed by STORM have been introduced before, the authors designed a solid structure to combine these parts into an effective method.
2. The experimental results demonstrate the effectiveness of STORM, and more importantly, they improve the time efficiency of Model-Based Reinforcement Learning methods.

WEAKNESSES:
1. In the experimental section, please add a comparison with Speedy Zero [1]. Speedy Zero is also a model-based RL method  which is proposed recently and has also achieved good time efficiency and performance on Atari.
2. Beyond Atari, can STORM be applied to other tasks, such as MuJoCo, DMC, MetaWorld, etc.? Dreamerv3 can achieve very good results in a wide range of different environments. If STORM cannot, the practical significance of this paper will be greatly reduced.
3. It would be better if further analysis could be provided on the benefits of using a stochastic variational autoencoder.

Reference:

[1]. Mei et.al. "SpeedyZero: Mastering Atari with Limited Data and Time." In ICLR 2023 https://openreview.net/forum?id=Mg5CLXZgvLJ.

QUESTIONS:
See weakness above

LIMITATIONS:
The generalizability of STORM could be further studied.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a world model architecture (STORM) to train RL agents in imagination. The world model is composed of an autoencoder with categorical latents and a Transformer. These modules are trained jointly with a reconstruction loss, a next latent state prediction loss, as well as reward, episode termination, and representation losses. Experiments in the Atari 100k benchmark indicate that the approach is effective.\n\nSTRENGTHS:\n- The method is technically sound with empirical results to back its effectiveness (outperforms other methods based on learning in imagination).\n- It is faster to train than similar methods (e.g. twice as fast as TWM).\n- The paper is well-written and easy to follow.\n\nWEAKNESSES:\n- The method lacks novelty and the incremental improvements over previous work are not properly explored.\n    - STORM is a variant of TWM with two minor modifications: (1) a vanilla Transformer is used instead of a TransformerXL, (2) latent state and actions tokens are fused instead of having separate tokens.\n    - Currently, it is not clear why STORM achieves better results than TWM. Is it due to implementation details or the proposed modifications actually matter? And if they matter, why?\n    \n- The results have limited significance. Although STORM outperforms TWM, it only yields marginal improvements over DreamerV3, which was not specifically optimized for Atari100k. Moreover, recent model-free methods [1, 2] achieve similar or better results in the benchmark.\n\n- Ablations and additional experiments fail to provide additional insights:\n    - lines 237-238: IRIS also relies on an autoencoder trained with a reconstruction loss and it still obtains good results on Breakout and Pong.\n    - l261: comparing the number of layers for STORM and IRIS/TWM is not straightforward since the hidden dimension is not the same (256 vs 512).\n    - 5.2: it is hard to draw any conclusions as only a few environments were considered and we do not know how they were picked. Also it is not clear if including z_t yields statistically signficant improvements.\n    - 5.3: the premise of this section is interesting but again too few environments are considered and results have limited statistical significance. The improvement on Freeway over previous methods relies on the addition of an extra demonstration while the other methods use exploration strategies that do not involve expert data. It would be interesting to know whether STORM leverages demonstrations better than other world models.\n\nQUESTIONS:\n### Do you use sticky actions?\n\n---\n\nI ran your code, on the 4 following games, with and without sticky actions:\n\n  \n\n| Environment | Reported | Without sticky actions | With sticky actions |\n| --- | --- | --- | --- |\n| BankHeist | 641 | 1044.0 | 218.5 |\n| Breakout | 16 | 29.3 | 11.2 |\n| MsPacman | 2673 | 2942.0 | 1921 |\n| PrivateEye | 7781 | 4458.4 | 100.0 |\n\nCan you clarify whether you used the Atari environments with (`ALE/<envname>-v5`) or without (`<envname>NoFrameskip-v4`) sticky actions? It is not mentioned in your paper but the code seems to use the v5 environments by default (TWM, IRIS and DreamerV3 do *not* use sticky actions). \n\nIt seems that my `v5` results are significantly below the reported results, and that the `v4` are on-par/slightly better. \n\n\n### Ablations to investigate the differences with TWM\n\n---\n\nThe paper would greatly benefit from a thorough investigation of the differences with TWM. It should be made clear why STORM performs better than TWM as it seems like a variant with two minor modifications. \n\n- Is the substantial performance improvement explained by implementation differences (in particular, there is [an open issue about reproducibility](https://github.com/jrobine/twm/issues/3) on TWM’s repo), or by the two modifications?\n- If these modifications matter, can you run ablations to demonstrate their effectiveness?\n- Maybe it is due to the incorporation of tricks for policy learning from DreamerV3?\n\nI am keen to significantly increase my overall rating if this concern is properly addressed during the rebuttals.\n\n### Other concerns\n\n---\n\n- In my opinion, there would be a substantial gain to experiment in more complex environments, e.g. Crafter [3], Minecraft [4], or Memory Maze [5]. Such results would expand our knowledge of what Transformer-based world models can achieve.\n\n- Can you include the recent work on the benchmark [1, 2] in the related work section?\n\n---\n\n### References\n\n- [1] *Sample-Efficient Reinforcement Learning by Breaking the Replay Ratio Barrier*. D'Oro, Pierluca and Schwarzer, Max and Nikishin, Evgenii and Bacon, Pierre-Luc and Bellemare, Marc G and Courville, Aaron. The Eleventh International Conference on Learning Representations, 2022.\n- [2] *Bigger, Better, Faster: Human-level Atari with human-level efficiency*. Schwarzer, Max and Obando-Ceron, Johan and Courville, Aaron and Bellemare, Marc and Agarwal, Rishabh and Castro, Pablo Samuel. arXiv preprint arXiv:2305.19452, 2023.\n- [3] *Benchmarking the Spectrum of Agent Capabilities*. Hafner, Danijar. International Conference on Learning Representations, 2021.\n- [4] *Minerl diamond 2021 competition: Overview, results, and lessons learned*. Kanervisto, Anssi and Milani, Stephanie and Ramanauskas, Karolis and others. NeurIPS 2021 Competitions and Demonstrations Track, 2022.\n- [5] *Evaluating Long-Term Memory in 3D Mazes*. Pasukonis, Jurgis and Lillicrap, Timothy P and Hafner, Danijar. The Eleventh International Conference on Learning Representations, 2022.\n\nLIMITATIONS:\nThe authors discussed some technical limitations in the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose several modifications to the recently proposed Transformer based world model for Model-based reinforcement learning. Specifically, they come up with a single latent stochastic state and treat action as an explicit input to the state as opposed to a token (as in previous works) and show significant performance and speed in limited regime of Atari 100k tasks.\n\nSTRENGTHS:\n1. Improved performance as well as significant time reduction in STORM's agent training which is crucial in MBRL agents alongside sample efficiency.\n2. Well written paper overall and was easy for me to read in a single go (NOTE: I'm very familiar with this MBRL space so that is another reason that has contributed to this but overall the flow of the paper was very intuitive.)\n\nWEAKNESSES:\n1. **Highlighting differences between TWM, Dreamerv3, Transdreamer, IRIS, SimPLe**: \n\n\t(a) It would be much better to have a table indicating different aspects such as (a) RNN (b) multiple tokens (c) Transformer variant etc. This would be much easier for the reader to read. Additionally, I did not find the motivation behind some of these modeling choices (see below for detailed descriptions).\n\n\t(b) On L127 the authors write \"In contrast to IRIS [12] that employs multiple tokens, STORM utilizes a single stochastic latent variable to represent an image.\" This is correct but TWM encoding style and STORM's encoding style are the same -- so I find it a bit misleading to omit TWM as all three (TWM, IRIS and STORM) are transformer models\n\n\t(c) On L129, re: \"STORM follows a vanilla Transformer [15] structure, while TWM [11] adopts a Transformer- XL [19] structure.\" -- it would be helpful why one would like to use a vanilla Transformer as opposed to Transformer-XL? What are the benefits?\n\n\t(d) On L131, \"TWM [11] treats observation, action, and reward as three separate tokens of equal importance.\" -- I am not sure if this statement is true. The input to the transformer in the case of TWM is (obs, action, rew) but that doesn't imply that they are equal -- the attention weights of the transformer would be the ones deciding whether to consider these tokens or not. See Figure 6 of TWM paper showing the attention map over $(s, a, r)$ and it is clear that not all are weighted equally.\n\n\t(e) On L134, \"Unlike Dreamer [8] and TransDreamer [16], which incorporate hidden states, STORM reconstructs the original image without utilizing this information.\" -- What is the additional benefit of not using the hidden state? I don't think there is a significant reduction in time for the reconstruction.\n\n2. **Performance**: \n\t(a) What specific component according to the authors attributes to the superior performace of STORM. L235-237 mention the self-attention mechanism -- which is a valid argument against the RNN based methods. However, it is unclear to me what is helpful in STORM when compared to IRIS or TWM.\n\n\t(b) L237-238: what do authors mean by the \"nature of autoencoders?\" For example, the encoding style in TWM and STORM is very similar. It is unclear to me why specifically to single object games does STORM perform poorly. I'd like the authors to elaborate on this.\n\n\n3. **\"Decoder at rear\" experiment**: I am not sure what exactly the purpose of this experiment was. From what I understand, the model can be formulated as TSSM (as mentioned in TransDreamer paper), and hence the the reconstruction would be using the posterior $z_t$ and not the prior $\\hat{z}_t$. Is there something that I'm missing here?\n\n4. **Impact of trajectory**: \n    (a) How exactly was the trajectory used in the replay buffer? Was there any prioritized replay buffer or was the sampling of trajectories for the training of world model uniform? \n    (b) Inclusion of trajectory is *not* specific to STORM -- so I'm not sure how this experiment validates the usefulness of STORM specifically and not apply to other world models (TWM, IRIS etc).\n\n\n**[Very minor comment -- not considered for rating the paper]**\n\n5. In appendix, Table 9, it would be more clearer to denote the \"Imagination batch size\" by another symbol or $B \\times T$ as that is effectively what happens during imagination.\n\n\n----\n**Rationale for my rating**\nI think the model details mentioned in the paper are important to share with the broader community, however I do think that motivation for several choices made in the modeling of STORM that has *not* been explicitly provided.\n\nAdditional it is not clear to me that experiments such as \"Decoder at rear\" or the inclusion of trajectory are fundamentally because of STORM and seem more of a study on Dreamer-*like* paradigm.\n\nI would like the authors to address my concerns above and for the time being, I have leaned towards Borderline reject. I will update my rating post-rebuttal and discussion with the authors.\n\n-----\n**Post-rebuttal rating**\n\nBased on the rebuttal provided by the authors, I've decided to bump up my rating to *Weak Accept* as they have addressed a majority of my concerns. The only section that I'm currently unconvinced is the inclusion of demonstrations which is not the major contribution of this work and would require its own analysis instead (a single section won't suffice in my opinion). I once again thank the authors for their comprehensive rebuttal.\n\nQUESTIONS:\nSee **Weaknesses** section.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper presented a Transformer-based model-based RL framework. As with earlier approaches, online data is gathered into the replay buffer with the learned reactive policy, the Transformer-based world model is trained with segments sampled from the replay buffer, then the policy is optimized with imaged data generated from the world model. The improved results were reported on the Atari 100K benchmark.  The authors also conduct a thorough ablation study to justify heir architectural design choices.

STRENGTHS:
This paper is well-organized and easy to follow.

WEAKNESSES:
- What criteria were employed for task selection in the ablation studies?  They are not consistent over all the subsections, nor do they have the same trend in Table 1 compared to baseline models. The proposed method performs better in some tasks, while it doesn't in others. To provide a more robust evaluation of their proposed method, it would be better to consider extending the ablation studies to a wider range of tasks
- It would also strengthen the paper if the authors could provide the converged results on the full Atari benchmark.

QUESTIONS:
Please see the weaknesses.

LIMITATIONS:
yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07