STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning
Authors: Weipu Zhang
Gang Wang
Jian Sun
Yetian Yuan
Gao Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: deep learning
reinforcement learning
model-based reinforcement learning
world model
learning in imagination
transformer
variational autoencoders
sequence modeling
TL;DR: We introduce an efficient world model structrue that consists of a variational autoencoder and a GPT-like transformer for model-based 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.

STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning
Weipu Zhang, Gang Wang" Jian Sun, Yetian Yuan
National Key Lab of Autonomous Intelligent Unmanned Systems, Beijing Institute of Technology Beijing Institute of Technology Chongqing Innovation Center zhangup.. bit@gmail..... com, {gangwang, sunjian, ytyuun/abitttttmmmm yt. edu.
Gao Huang
Department of Automation, BNRist, Tsinghua University gaohuang@tsinghua. edu. en
Abstract
Recently, model-based reinforcement learning algorithms have demonstrated markable efficacy in visual input environments. These approaches begin by con- structing parameterized simulation world model of the real environment through self-supervised learning. By leveraging the imagination the world model, the agent's policy is enhanced without the constraints of sampling from the real en- vironment. The performance of these algorithms heavily relies on the sequence modeling and generation capabilities of the world model. However, constructing perfectly accurate model of complex unknown environment is nearly impossible. Discrepancies between the model and reality may cause the agent 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-bau efficient world model architecture that combines the strong sequence modeling and generation capabilities of Transformers with the stochastic nature of variational autoencoders. STORM achieves mean human performance of 126. on the Atari 100k benchmark, setting a new record among state-of-the-e art methods that not employ lookahead search techniques. Moreover, training an agent with 1.85 hours of real-time interaction experience single NVIDIA GeForce RTX 3090 graphics eard requires only 4.3 hours, showcasing improved efficiency compared to previous methodologies.
wORld Model (STORM), an
We release our code at htttgs///gghhbbbcoo////////////t
Introduction
Deep reinforcement learning (DRL) has exhibited remarkable success across diverse domains. How- ever, 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 manu- facturing systems [ 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 44, suffer from low sample efficiency due to two primary reasons. Firstly, the estimation of the value function proves be challenging task. This
Conresponding author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 converge [S. 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 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 sample- efficient settings. These algorithms start by constructing a parameterized simulation world model of the real environment through self- supervieed learning. Self-supervised learning can be implemented in various ways, such as reconstructing the original input state using a decoder [8 (8H1O), predicting actions between frames [7]. or employing contrastive learning to capture the internal consistency of
input states 66 These approaches provide more supervision information than conventional model- free RI. 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 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 en- vironment. To mitigate this issue, introducing random noise into the world model has been proven beneficial 1911 14. Variational au- toencoders, capable of automatically learning low-dimensional latent representations of high- dimensional data while incorporating reasonable random noise into the latent space, offer an ideal choice for image encoding.
Mean human normalized score
150%
100%
96%
104%
112%
50%
33%
0%
SimPLe
TWM
IRIS
DreamerV3
STORM (ours)
Median human normalized score
60%
51%
49%
58%
40%-
29%
20%
13%
0%
Numerous endeavors have been undertaken to construct efficient world model. For in-
SimPLe
TWM
IRIS
DreamerV3
STORM (ours)
stance, SimPLe [[] leverages LSTM [15). while DreamerV3 [TO] employs GRU [16] as sequence model. LSTM and GRU, both vari- ants of recurrent neural networks (RNNs), excel at sequence modeling tasks. However, the recur- rent nature of RNNs impedes parallelized com- puting, resulting in slower training speeds [17] In contrast, the Transformer architecture [17 has lately demonstrated superior performance RNNs in various sequence modeling and generation tasks. It overcomes the challenge of forgetting long-term dependencies and is de- signed for efficient parallel computing. While several attempts have been made incorporate Transformers into the world model [12 [3 [18), these works do not fully harness the capabil- ities of this architecture. Furthermore, these approaches require even longer training times and fail to surpass the performance of the GRU- based D3eamerV3.
Training FPS on V100
11.9
5.6
0.5
0.7
SimPLe
TWM
|R|S
DreamerV3
STORM (ours)
Comparison of methods Atari 100k. SimPLe [I] and DreamerV3 [IO] employ RNNs as their world models, whereas TWM [12). IRIS [3], and STORM use Transformers. The train- ing frames per second (FPS) results on single NVIDIA V100 GPU are extrapolated from other graphics cards for SimPLe, TWM, and IRIS, while DreamerV3 and STORM are directly evaluated.
Figure 1:
In this paper, we introduce the Stochastic Transformer-based wORld Model (STORM), a highly effective and efficient structure for model-based RL. STORM employs 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 remarkable mean human normalized score of 126.7% on the challenging Atari 100k benchmark, establishing 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-ar methods is depicted in Figure
Related work
Model-based DRL algorithms aim to construct 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 tabular cases [., 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-fire methods on these tasks.
The foundation of VAE-LSTM- world models was introduced by Ha and Schmidhuber [4 for image-based environments, demonstrating the feasibility of learning good policy solely from generated data. SimPLe [I]] 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 [O also adopt this framework and showcase notable capabilities in Atari games, Deep 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 113, TWM [2]. and TransDreamed [8] incorporate the Transformer architecture into their world models. IRIS [3 employs VQ-VAE [19 as the encoder to map images into 4 x 4 latent tokens and uses 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 significant training slowdown. TWM [2] 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. Trans 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.
[8]] directly replaces
24 incor- porate Monte Carlo tree search (MCTS) to enhance policy and achieve promising performance on the Atari 100k benchmark. Lookahead search techniques like MCTS can be employed enhance other model-based RL algorithms. but they come with high computational demands. Additionally, certain model-free methods 66 27 incorporate 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.
Other model-based RL methods such as MuZero
2 EfficientZero 77. and SpeedyZero
We highlight the distinctions between STORM and recent approaches in the world model as follows:
Table 1: Comparison between STORM and recent approaches. "Tokens" refers to the input tokens introduced to the sequence model during a single timestep. "Historical information" indicates whether the VAE reconstruction process incorporates historical data, such as the hidden states of an RNN.
Attributes
SimPl.e DI]
TWM 12
IRIS L
DreamerV3 10
STORM (ours)
Sequence model Tokens
LSTM [5]
Transformer XL.
21
Transformer [17]
GRU [T
Transformer
Latent
Latent, action, reward
Latent(4 x 4)
Latent
Latent
Latent representation
Binary-VAE
Categorical-Ve VAE No
VQ-VAE
Categorical-VAL
Categorical-Vas -AA
Histonical information
Yes
Yes
Yes
No
Agent state
Reconstricted image
Latent
Reconstructed image As DreamerV2 (
Latent, hidden
Latent, hidden
Agent training
PPO
As Dreamer V2 9
Dreamer
As Dreamer V33
3

---

f2
2
h2
C4
Transformer Blocks
C1
122
Agent
Agent
04
21
ENC
DEC
ENC
DEC
networks
legends
Inner ouaput
Random sample
01
O
Ö2
Policy
Conestieate
Figure Structure and imagination process of STORM. The symbols used in the figure are explained
in Sections 3.1 and 3.2 The Transformer blocks depict the sequence model fø in Equation (2). The Agent block, represented by a neural network, corresponds to mp (u(())))) Equation 6
SimPLe 11 and Dreamer 10 rely on RNN- bbased models, whereas STORM employs a
GPT-like Transformer
30
the sequence model.
contrast to IRIS 13 that employs multiple tokens, STORM utilizes single stochastic
latent variable to represent image.
STORM follows vanilla Transformer
17 structure, while TWM 12 adopts a Transformer-
XL [2]) structure.
In the sequence model of STORM, observation and action fused into single token, whereas TWM 12 treats observation, action, and reward as three separate tokens of equal importance.
Unlike Dreamer [10] and TransDreamer [8], which incorporate hidden states, STORM reconstructs the original image without utilizing this information.
Method
3.1 World model learning
Our approach adheres to the established framework of model-based RL algorithms, which focus enhancing the agent's policy by imagination (5] 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.
$2) 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 f, a data point comprises an observation or, an action de, a reward r:, and a continuation flag (a Boolean variable indicating whether the current episode is ongoing). The replay buffer maintains inst-in-first-out queue structure, enabling the sampling of consecutive trajectories from the buffer.
Section 3.1 provides a detailed description 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. 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 In our experiments, we focus on Atari games [3I], which generate image observations o, of the environment. Modeling the dynamics of the environment directly on raw images is computationally expensive and
prone to errors 7 23 address this, we leverage VAE [32] formulated in Equation (T) to convert Og into latent stochastic categorical distributions 2, Consistent with prior work [9] 10]
we set 2r as a stochastic distribution comprising 32 categories, each with 32 classes. The encoder (4a) and decoder (po) structures are implemented as convolutional neural networks (CNNs) [33]. Subsequently, we sample latent variable ot from 2, to represent the original observation or. Since sampling from distribution lacks gradients for backward propagation, we apply the traight-througee gradients trick [9 34 to preserve them.
Image encoder: 2t (o( 2₂
Image decoder: ō = [
Before entering the sequence model, we combine the latent sample 2 and the action d, into single token e using multi-layer perceptrons (MLPs) and concatenation. This operation, denoted as ma. prepares the imputs for the sequence model. The sequence model fo takes the sequence of e input and produces hidden states hę- We adopt a GPT-like Transformer structure [30] for the sequence model, where the self-attention blocks are masked with subsequent mask allowing e, to attend to the sequence é1 ,,,..., et. By utilizing MLPs se 9 and Pa C we rely on he to predict the current
reward fr. the continuation flag é, and the next distribution 2,41. The formulation of this part the world model is as follows
Action mixer:
mal a+)
Sequence model:
hEr fØ(Cn:)
Dynamics predictor:
2.41 (21t1l11)
Reward predictor:
ft (he)
Continuation predictor:
& = (ha).
Loss functions The world model is trained in self-supervised manner, optimizing end-to-end. The total loss function is calculated as in Equation (3) below, with fixed hyperparameters 33 = 0.5 and 82 0.1. In the equation, B denotes the batch size, and T denotes the batch length
1 BT
[
C.00 + (p) Cion (Ø) BI (2
(3)
The individual components of the loss function are defined follows: Cre(( represents the reconstruction loss of the original image, Lrew (Φ) represents the prediction loss of the reward, and Leou (6) represents the prediction loss of the continuation flag.
Lree (b)) 1,; 0r112
(4a)
2 (0) _aym (fr
(4b)
C.on (Ø) log è + (1 c.) log(1 -
(4c)
Additionally, Caym in Equation (4b) denotes the symlog two-hot loss, as described in [0] This loss
function transforms the regression problem into a classification problem, ensuring consistent loss scaling across different environments.
The losses cdy"(() and CT (Ø) are expressed as Kullback-Leible (KL) divergences but differ in their gradient backward and weighting. The dynamics loss L, zdy= (ω) guides the sequence model in predicting the next distribution, while the representation loss c.ep(o) 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.
(dym(6)= = max (1, KL ([LLLKL((((((((2+11o))))) 9 (2++1/h+))
(5a)
C.TP(() max KI (L. (g0 l0441) sg(o (2++1(h+)))
(5b)
where sg(:) denotes the operation of stop-gradients.
5

---

3.2 Agent learning
The agent's s learning is solely based on the imagination process facilitated by the world model, as illustrated in Figure To initiate the imagination process, a brief contextual trajectory is randomly selected from the replay buffer, and the initial posterior distribution 2, is computed. During infer- ence, rather than sampling directly from the posterior distribution 2,, we sample =, from the prior distribution Σ. To accelerate the inference, we employ the KV cache technique [35] within the Transformer structure.
The agent's state is formed by concatenating ž and hes as shown below:
State:
St = [Ze, [as]
8
Critie:
(8t) ETT.
k=0
Actor: a1 1o (ot((s)))
We adopt the actor learning settings from DreamerV3 10.
The complete loss of the actor-critic
algorithm is described by Equation 7 where f, corresponds the reward predicted the world model, and & represents the predicted continuation flag:
G (C-Vv((s)) 5))
) In HH (Toicaaa))n(((((asssss))))
((0)
sg
(7a)
BI.
n=
((1) BL
1
ΣΣ
(Vo(s))-s(() sg((:t) (Vo(s)) sg( Sg(VeRMa (8e))
(7b)
where H(-) denotes the entropy of the policy distribution, while constants y and L represent the coefficient for entropy loss and the imagination horizon, respectively. The X-return G [ 10 is recursively defined as follows
G re + ye [1 A)Vo()e41) + AG
(8a)
G VoLs))
(8b)
The normalization ratio S utilized actor loss 7a is defined in Equation ( which is computed as the range between the 95th and 5th percentiles of the X-return G. across the batch [10
$ percentile(C; 95) percentile( 5).
To regularize the value function, we maintain the exponential moving average (EMA) of th The EMA is defined in Equation (10), where dù represents the current critic parameters, is the decay rate, and v EMA denotes the updated critic parameters. This regularization technique aids stabilizing training and preventing overfitting
EMA
ovt EMA + (1 o))
(10)
Experiments
We evaluated the performance of STORM on the widely-used benchmark for sample-efficient RL, Atari 100k [3I]. 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 100k games are presented Section 4.2
4.1 Benchmark and baselines
Atari 100k consists of 26 different video games with discrete action dimensions of up to 18. The 100k sample constraint corresponds to 400k actual game frames, taking into account frame skip- ping (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 7 A-R
H-R
is calculated based on the score achieved by the agent, the score R obtained by random policy, and the average score H achieved by human player in specific environment. To determine the human player's performance H, 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 with model-based DRL algorithms that share similar training pipeline, discussed in Section However, similarly to 10 12 13 we do not directly compare our results with lookahead search methods like MuZero [23 and EfficientZero 77 as primary goal is 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 Atari 100k
Detailed results for each environment can be found in Table and the corresponding performance curve presented in Appendix due 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.
STORM demonstrates 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 be attributed 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 single small moving object, as observed 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.
Table 2: Game scores and overall human-normalized scores the 26 games in the Atari 100k benchmark. Following the conventions of [ scores that are the highest or within 5% of the highest score are highlighted in bold.
Game
Random
Human
SimPLe [I]
TWM 12
IRIS
[
DreamerV3
10
STORM (ours)
Alien
228
7128
617
675
420
959
984
Amidar
1720
74
122
143
139
205
Assault
222
742
527
683
1524
706
801
Asterix
210
8503
1128
1116
854
932
1028
Bank Heist
14
753
467
53
649
641
Battle Zone
2360
4031
13074
12250
13540
Boxing
78
70
78
80
Breakout
30
16
20
84
31
16
Chopper Command
811
7388
979
1697
1565
420
1888
Crazy Climber
10780
35829
62584
71820
59234
97190
66776
Demon Attack
152
1971
208
350
2034
303
165
Freeway
24
Freeway w/o
30
17
24
31
0
Frostbite
65
4335
237
1476
259
909
1316
Gopher
258
2413
597
1675
2236
3730
8240
Hero
30826
2657
7254
7037
11161
11044
James Bond
101
362
463
445
509
Kangaroo
52
3035
51
1240
838
4098
4208
Krull
1598
2666
2204
6349
6616
7782
8413
Kung Fu Master
256
22736
14862
24555
21760
21420
26182
Ms Pacman
307
6952
1480
1588
999
1327
2673
Pong
-21
19
15
Private Eye
69571
87
100
882
7781
Qbert
164
13455
1289
3331
746
3405
4522
Road Runner
12
7845
5641
9109
9615
15565
17564
Seaquest
68
42055
683
774
661
618
525
Up Down
533
11693
3350
15982
3546
7667
7985
Human Mean
0%
100%
33%
96%
105%
112%
126.7%
Human Median
0%
100%
13%
51%
29%
49%
58.4%

---

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.T] as well as on the agent's s design in Section [.2] Additionally, we propose 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 and Dreamer 9 10 can be formulated clearly using variational inference over time. However, the non-recursive Transformer-baset world model does not align with this practice and requires manual design. FagureBa shows alternative structures and their respective outcomes. In the "Decoder at rear" configuration, we employ 24 ~ 2, instead of 24 Z, 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, Equation 2 instead of h4 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 performance drop on tasks that require several contextual frames to predict the reward accurately (in c.g., Ms. Pacman).
utilize 2 as input for o (-) and og (.) in
By default, we configure our Transformer with 2 layers, which is significantly smaller than the 10 layers used IRIS [3] and TWM [12. Figure 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 100k to 400k, 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 (8), 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 complex model. Secondly, training a large model naturally requires a substantial amount of data, yet the Atari 100k 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 LeeP in Equation (56) directly influences the image encoder. The encoder may be overly influenced when tracking the output of large sequence model.
Ms Pacman
Pong
Ms Pacman
Pong
4000
4000
3000
10
10
2000
2000
1000
-00
1000
-10
-20
200k
400k
200k
400k
200k
400k
200k
400k
STORM
Decader at rear
Prodictor at front
STORM (2L)
(a) Rearranging the modules of STORM.
(b) Number of layers in the Transformer.
Figure 3: Ablation studies on the design and configuration of the STORM's world model.
5.2 Selection of the agent's state
The choice of the agent's 's state st offers several viable options: à 13
h4, 24 (as in TWM [2 ), or
the combination [he, 24] (as demonstrated by Dreamer, 10). In the case STORM, we employ 21]. as in Equation (6) Ablation studies investigating the selection of the agent's is state are presented in Figure 4 The results indicate that, environments where a good policy requires contextual information, such as in Ms. Pacman, the inclusion of h; leads to improved performance. However, in other environments like Pong and Kung Fu Master, this inclusion does not yield significant difference. When solely utilizing he in environments that evolve with the agent's policy,
8

---

like Pong, the agent may exhibit behaviors similar to catastrophic forgetting [39 due to the non- stationary and inaccurate nature of the world model. Consequently, the introduction of randomness through certain distributions like 2 proves to be beneficial.
Ms Pacman
Pong
Kung Fu Master
4000
20
30000
3000
10-
20000
2000
1000
-10-
10000
-20
0
200k
400k
0
200k
400k
200k
400k
STORM
Agent w/o Zt
Agent w/o
Figure 4: Ablation studies on the selection of the agent's 's state.
5.3 Impact of the demonstration trajectory
The inclusion of a demonstration trajectory is a straightforward implementation step when using world model, and 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 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 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 [3] 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-inter intensive games without prior knowledge is challenging. Typically, human players require instructions about the game's '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 demonstration trajectory train the world model and establish starting points for imagination. By leveraging sufficienlly robust world model, the utilization of limited offline information holds the potential to surpass specially designed curiosity- ddiven exploration strategies in the future.
Ms Pacman
Pong
Freeway
4000
3000
2000
1000
-10
-20
200k
200k
400k
200k
400k
w/o demonstration
w/ demonstration
Figure 5: Ablations studies on adding demonstration trajectory to the replay buffer.
As an integral part of our methodology, we integrate single trajectory from Freeway into our extensive results. Furthermore, to ensure fair comparisons in future research, we provide the results without incorporating Freeway's 's trajectory in Table and Figure 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%.

---

Conclusions and limitations
In this work, we introduce STORM, an efficient world model architecture for model-based RI,, surpassing previous methods in terms of both performance and training efficiency. STORM hamnesses 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 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 p(s) for these starting points is not explicitly considered, despite its significance the policy gradient formulation: (.) x E. m(s) La 9no (s, a) (a)(a [5].
Acknowledgments and Disclosure of Funding
We would like 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 20212X4100027.
References
[1] Jie Chen, Jian Sun, and Gang Wang. From unmanned systems autonomous intelligent systems. Engineering, 12:16-19, 2022.
[2] Boxi Weng, Jian Sun, Gao Huang. Fang Deng, Gang Wang, and Jie Chen. Competitive meta-learning. /EEE/CAA Journal of Automatica Sinica, 10(9):1902-1904, 2023.
[3] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improve- ments deep reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[4] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal 2017.
policy optimization algorithms.
[5] Richard S Sutton, Andrew G Barto, al. Introduction Reinforcement Learning, volume 135. MIT Press Cambridge, 1998.
[6] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: tations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
Contrastive unsupervised represen-
[7] 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.
[8] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Leurning Representa- tions, 2020. URL
[9] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In International Conference on Learning Representations, 2021. URL TtLhttts: htpp::::///peereeiiie http:p/wwweeeeeeeeeeeeeeeeeeaaaa
10

---

[10] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv:2301.04104, 2023.
[11] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, 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. URI. https
//openreview
[12] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In International Conference on Learning Representa- tions, 2023. URL
[13] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In International Conference on Learning Representations, 2023. URL https:// openreview
[14] David Ha and Jürgen Schmidhuber. World models. arXiv: /803. 2018.
[15] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9 (8):1735-1780, 1997.
[16] 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.
[17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[18] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv:2202.0943 2022.
[19] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017.
[20] Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: video vision transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6836-6846, 2021.
[21] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of Annual Meeting of the Association for Computational Linguistics, pages 2978 7788-2988,11 2019.
[22] 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, 14:15004-15099
2021.
[23] Julian Schritwieesrr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering Atari, Go, chess and Shogi by planning with a learned model. Nature, 588(7839):604-609, 2020.
[24] Yixuan Mei, Jiaxuan Gao, Weirui Ye, Shaohuai Liu, Yang Gao, and Yi Wu. SpeedyZero: Master- ing Atari with limited data and time. In International Conference on Learning Representations, 2023. URL http:ttttt///erreee http:////ppppiiiiee///////a//aaa
net
[25] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In Inter- national Conference on Learning Representations, 2021. URI. httpps://pperreviee.rt CorumPid=uCQgEPZwwaaii

---

[26] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. URI. htttps://openrevier: net/forumm2dd-06-666
[27] Hado P Van Hasselt, Matteo Hessel, and John Aslanides. When use parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 32, 2019.
[28] Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Mare G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In International Conference Learning Representations, 2023. URL htpps://///oooereerror
net/forum?id-OpC- DaBBBB
[29] Max Schwarzer, Johan Samir Obando Ceron, Aaron 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, pages 30365-30380, PMLR, 2023.
[30] Alec Radford, Karthik Narasimhan, tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[31] Marc 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.
[32] Diederik P Kingma and Max Welling, Auto-encoding variational Bayes. arXiv:1312.6114, 2013.
[33] 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, 144):541-551 1989.
[34] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. 13Xiv:1308.3432, 2013.
[35] Carol Chen. Transformer inference arithmetic. kipp.ly, 2022. URL htttps://kipp.1///00
[36] OpenAl. ChatGPT, https://openaat ccom/chat.............cc..cc techhtttttt//t 2021. Accessed: April 19, 2023.
[37]
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 scale. In International Conference on Learning Representations, 2021. URL htttss://operreviex www.:/eeeeeoommmmammoeeeeeraar
[38] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016
[39] 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.
[40] Sergey loffe 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.
[41] 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.
[42] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. 1rX1v:160.004450, 2016.
12

---

[43] 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.
[44] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization.
arXiv:1412.6980, 2014.
[45] Florin Gogianu, Tudor Berariu, Lucian Busoniu, and Elena Burceanu. Atari agents, 2022. URL
13

---

A Atari 100k curves
Alien
Amidar
Assault
Asterix
Bank Heist
1500
1000
200
1000
500
500
500
0
0
200k
400k
200k
400k
200k
400k
200k
400k
200k
400k
Battle Zone
Boxing
Breakout
Chopper Command
Crazy Climber
15000
10000
20
50000
200k
200k
Demon Attack
Freeway
Freeway w/o traj
Frostbite
Gopher
750
0.05
15000
500
10000
250
0.00
1000
5000
0.05
200k
400k
200k
400k
200k
400k
200k Krull
400k
200k
400k
Hero
James Bond
Kangaroo
Kung Fu Master
10000
750
30000
10000
500
5000
20000
5000
250
0
400k
200k
200k
400k
Pacman
Pong
Private Eye
Qbert
Road Runner
15000
2000
5000
2000
0
400k
Seaquest
Up N Down
Gamer Median
Gamer Mean
0.6
750
0.4
500
250
10000
0.2
0.5
0
0.0
200k
400k
200k
400k
200k
400k
200k
400k
STORM (Ours)
Dreamer V3
Figure 6: Performance comparison on the Atari 100k benchmark. Our method is represented blue, while DreamerV3 [10] is in orange. The solid line represents the average result over 5 seeds, and the filled area indicates the range between the maximum and minimum results across these 5 seeds.
14

---

B Details of model structure
Table 3: Structure of the image encoder. The size of the submodules is omitted and can be derived from the shape of the tensors. ReLU refers to the rectified linear units used for activation, while Linear represents a fully-connected layer. Flatten and Reshape operations are employed to alter the indexing method of the tensor while preserving the data and their original order. Conv denotes CNN layer [3], characterized by kernel 4, stride 2, and padding 1. BN denotes the batch
normalization layer 40
Submodule
Output tensor shape
Input image (ot)
3 x 64 64
Convl + BN1 ReLU
32 X 32
32
Conv2 + BN2 ReLU
64 16 16
Conv3 + BN3 + ReL.U
128 x 8
Conv4 BN4 ReLU
256 4 x 4
Flatten
4096
Linear
1024
Reshape (produce 2,)
32 x 32
Table Structure of the image decoder. DeConv denotes transpose CNN layer 41] characterized by kernel = 4, stride = 2, and padding = 1.
Submodule
Output tensor shape
Random sample (2+)
32 x 32
Flatten
1024
Linear + BNO ReLU
4096
Reshape
256 X 4
DeConyl + BÑI ReLU
128 x
x 8
DeConv2 + BN2 ReLU
64
16
16
DeConv3 + BN3 ReLU
32 x 32
x 32
DeConv4 (produce 6,)
3 64
64
Table 5: Action mixer = ma. at). Concatenate denotes combining the last dimension of two tensors and merging them into one new tensor. The variable represents the action dimension, which ranges from 3 to 18 across different games. D denotes the feature dimension of the Transformer. LN an abbreviation for layer normalization
Submodule
Output tensor shape
Random sample (ze), Action (a;) Reshape and concatenate
32 32, A
1024 + A
Linearl + LN1 ReLUU
D
Linear2 + LN2 (output ex)
D
15

---

Table 6: Positional encoding module. W1.7 is learnable parameter matrix with shape T D, and T refers to the sequence length.
Submodule
Output tensor shape
Input (e1:7) Add (e1:T + wi:T)
Tx D
LN
Table Transformer block. Dropout mechanism 44
can prevent overfitting.
Submodule
Module alias
Output tensor shape
Input features (label as x1)
T D
Multi-head self attention
Linearl + Dropout(p)
MHSA
Tx D
Residual (add xi)
LNI (label as x22)
Linear2 ReLU
T 2D
Linear3 + Dropout(p)
T D
FFN
Residual (add x2)
T D
LN2
T D
Table Transformer based sequence model h1.7 (e(:r)). Positional encoding is explained in Table and Transformer block is explained in Table
Submodule
Output tensor shape
Input (e1:T)
Positional encoding
Tx D
Transformer blocks x K Output (hirr)
Table 9: Pure MLP structures. A I-layer MLP corresponds to a fully-connected layer. 255 is the size of the bucket of symlog two-hot loss [10].
Module name
Symbol
MLP layers
Input/ MLP hidden/ Output dimension
Dynamics head
1
D/ -/ 1024
Reward predictor
3
D/ D/ 255
Continuation predictor
96
3
D/ D/
Policy network
To(a, (st)
3
D/ D/ A
Critic network
(sx)
3
D/ D/ 255
16

---

C Hyperparameters
Table 10: Hyerparameters. Note that the environment will provide "done" signal when losing life, but will continue running until the actual reset occurs. This life information configuration aligns with the setup used IRIS [3]. Regarding data sampling, each time we sample B₁ trajectories of length 7 for world model training, and sample Bg trajectories of length C' for starting the imagination process.
Hyperparameter
Symbol
Value
Transformer layers
K
2
Transformer feature dimension
D
512
Transformer heads
8
Dropout probability
0.1
World model training batch size World model training batch length
B1
16
64
Imagination batch size
B2
1024
Imagination context length
8
Imagination horizon
16
Update world model every env step
1
Update agent every step
1 16
Environment context length
Gamma
0.985
Lambda
0.95
Entropy coefficiency Critic EMA decay
7
3 10
0
0.98
Optimizer
Adam 44
World model learning rate
1.0 x
World model gradient clipping
1000
Actor-critic learning rate
3.0 x 10-5
Actor-critic gradient clipping
100
Gray scale input
False
Frame stacking
False
Frame skipping
4 (max over last frames)
Use of life information
True
17

---

D Demonstration trajectory information
Table 11: To account for frame skipping, the frame count multiplied by These trajectories were
gathered using pre-trained DQN agents 45
Game
Episode return
Frames
Ms Pacman
5860
1612 x
Pong
18
2079
Freeway
27
2048 x
18

---

E
Computational cost details and comparison
Table 12: Computational comparison. In the V100 column, an item marked with star indicates extrapolation based on other graphics cards, while items without star are tested using actual devices. The extrapolation method employed aligns with the setup used in DreamerV3 10, where it assumes the P100 is twice as slow and the A A100 is twice as fast.
Method
Original computing resource
V100 hours
SimPLe
NVIDIA P100, 20 days
240*
NVIDIA A A100, 10 hours
TWM 12
20
NVIDIA GeForce RTX 3090, 12.5 hours
IRIS 13
NVIDIA A100, days for two runs
168*
DreamerV3 10
NVIDIA V100, 12 hours
12
STORM
NVIDIA GeForce RTX 3090, 4.3 hours
9.3
19

---

Atari video predictions
Context Input
Imagination
Environment: Boxing
Real
Model
t
t=
t 15
= 23
31
39
t 47
55
63
Context Input
Imagination
Environment: ChopperCommand
Real
Model
t 23
31
t= 39
47
55
63
Context Input
Imagination
Environment: MsPacman
Real
Model
15
23
31
39
47
55
63
Context Input
Imagination
Environment: Pong
Real
Model
31
39
47
55
63
Context Input
Imagination
Environment: RoadRunner
Real
Model
63
Figure 7: Multi-step predictions on several environments in Atari games. The world model utilizes observations and actions contextual input, enabling the imagination of future events spanning 56 frames in an auto-regressive manner.
20


Reviews:
[{"review_id": "vcexTDfC2M", "reviewer": "Reviewer_Cg1v", "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.", "questions": "Please see the weaknesses.", "limitations": "yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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\n- It would also strengthen the paper if the authors could provide the converged results on the full Atari benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your effort spent reviewing our paper and providing many valuable suggestions. We will include the suggestions in the revised version. Below, we want to address the main concerns raised in the review.\n\n- Response to **Weakness 1**`criteria for task selection in the ablation studies`:\n  - Conducting ablation studies on the entire benchmark demands approximately **7 (ablation studies) $\\times$ 0.2 (days per game) $\\times$ 5 (seeds) $\\times$ 26 (games) = 182  (NVIDIA GeForce RTX 3090 days)**, which is resource-intensive and infeasible effort given our time and funding budget at present. Therefore, we have chosen representative environments that are sensitive to different configurations, based on our experience with STORM's development.\n  - Previous works like DreamerV3, IRIS, and TWM conduct ablation studies on their special contributions, such as policy training strategies, number of tokens, choice of policy Input, and other hyperparameters. In **Section 5.1**, we present the impact of model design and configuration, which are essential considerations when implementing a Transformer-based world model structure. In **Section 5.2**, we delve into the impact of policy input selection, an aspect also studied in TWM, providing researchers with valuable insights into STORM and model-based reinforcement learning. In **Section 5.3**, we analyze the influence of a single demonstration trajectory, highlighting that the challenging exploration issue in RL can potentially be addressed by integrating external knowledge and world models in future research. We believe this analysis could inspire further investigations in this field.\n  - Regarding concerns about performance (`The proposed method performs better in some tasks, while it doesn't in others`), we also address them in our response to **Reviewer rfCq Weakness 2**.\n  \n- Response to **Weakness 2** `It would also strengthen the paper if the authors could provide the converged results on the full Atari benchmark`: Please refer to our **Response 1** in the **Author Rebuttal** above."}, {"review_id": "mtL1f8aCzK", "reviewer": "Reviewer_4iYJ", "summary": "The 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.", "questions": "See **Weaknesses** section.", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.)", "weaknesses": "1. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your effort spent reviewing our paper and providing many valuable suggestions. We will include the suggestions in the revised version. Below, we want to address the main concerns raised in the review.\n\n- Response to **(1a, 1b)**: We sincerely appreciate your suggestion, and we will include a table of differences between these algorithms in the updated version to enhance the clarity and motivation of the modeling choices.\n\n- Response to **(1c)**: Transformer-XL was proposed to increase attention length and handle extremely long text in the NLP tasks. However, a vanilla Transformer structure like BERT was still proved to have a memory capacity of 512 tokens (as reported in their paper, Appendix A.2), which is much longer than the regular sequence modeling horizon in RL problems. STORM, Dreamer, IRIS, and TWM typically have a sequence modeling horizon ranging from 16 to 64. Given these considerations, the use of Transformer-XL is unnecessary and may lead to a drawback in performance and runtime.\n\n- Response to **(1d)**: We apologize for this confusion. As correctly pointed out,  TWM does provide different attention weights to different tokens. What we want to express here is that all observations, actions, and rewards are involved in the same self-attention process as input tokens, and performing self-attention across different types of data (observation, action, and reward of different physical meanings) may potentially negatively impact or limit the performance (**L112**). \n\n- Response to **(1e, 3)**, `Why ablations on \"decoder at rear\"`: Yes you are correct about this, \"decoder at rear\" means the the model is formulated as TSSM.\n\n  DreamerV3 is an RNN-based algorithm, that enables the fusion of observations and actions into the hidden state step by step, which naturally facilitates the observation reconstruction with historical information. In contrast, Transformer-based world models process a sequence of observations and actions simultaneously, leading to two options for the decoder position: **(1) decoder at rear** (like TSSM in TransDreamer) to reconstruct the original observation using the hidden state $h_t$ and prior $\\hat{z}_t$, and **(2) decoder at front** (STORM), which resembles the original structure of the variational autoencoder.\n\n  TransDreamer needs a great number of samples to converge in Atari games (see Figure 5 in their paper), which inspired us to further investigate it. As the ablation study presented in **Section 5.1**, we found that the structure of STORM is superior, and the \"decoder at rear\" approach struggles to converge with a 100k sample budget. This suggests that the variational reconstruction loss may not effectively drive the training of a Transformer-based model between the stochastic variable and the decoder, or the reconstruction loss might be excessively influenced by the KLDiv loss. In conclusion, this design is primarily for improving final performance, rather than lowering computational complexity.\n\n- Response to **(2a, 2b)**: Please refer to our **Response 2&3** in the **Author Rebuttal** above.\n\n- Response to **(4a)**: In our experiments, the samples in the demonstration buffer and the online buffer are uniformly sampled separately. For initiating the imagination process, a batch consists of 4 demonstration trajectories and 16 online trajectories.\n\n- Response to **(4b)**: Please refer to our **Response 4** in the **Author Rebuttal**.\n\n- Response to **(5)**: We appreciate your constructive suggestion, and we will explain the batches further in the header of Table 9 in the revised paper."}, {"review_id": "daYRwXRBvy", "reviewer": "Reviewer_rfCq", "summary": "This 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.", "questions": "### 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", "limitations": "The authors discussed some technical limitations in the conclusion.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your effort spent reviewing our paper and providing many valuable suggestions. We will include the suggestions in the revised version. Below, we want to address the main concerns raised in the review.\n\n- Response to **Weakness 1** and **Question 2** about `comparison with TWM`: For further experiments, please refer to our **Response 3** in the **Author Rebuttal** above. As you pointed out, STORM demonstrates a training speed twice as fast as that of TWM, while also achieving superior performance scores. We believe that this achievement represents a significant and meaningful innovation.\n\n- Response to **Weakness 2** `marginal improvements compared to DreamerV3 & not as good as BBF[2]`: Although the improvement in overall score was not significant (112% to 126%), STORM is equal or better than DreamerV3 on 21 out of 26 games.\n\n  While global human mean and median metrics offer an overall view of an RL algorithm's performance, it is common to observe variations in different environments and methods. To illustrate, let's consider BBF[2] for comparison:\n\n  | Algorithm | Human mean & human median | No. of envs that STORM $\\approx$ BBF | No. of envs that STORM > BBF |\n  | :-------: | :----------: | :-------------: | :-----------: |\n  |   STORM   |   1.267 & 0.584  |    4   |   5   |\n  |    BBF    |     2.247 & 0.917    |     -    |   -   |\n\n  Although BBF achieves nearly double the overall score compared to STORM, our method still outperforms it in five games, namely *Gopher*, *Hero*, *MsPacman*, *PrivateEye*, and *Qbert*. *Hero* and *PrivateEye* are two representative games in the Atari 100k benchmark that involve long-term exploration without rewards. In cases where BBF surpasses STORM, there are several environments where RL agents typically perform much better than humans. We believe that conducting research in such environments does not align with the expectations for future RL algorithms.\n\n- Response to **Weakness 3**: We mainly choose factors that may have a major impact on the performance of STORM for ablation study. Previous works like DreamerV3, IRIS, or TWM have also conducted ablation studies on their distinct contributions, such as policy training strategies, number of tokens, policy input selection, and other hyperparameters. Let us clarify your concerns:\n\n  - The ablation study on the number of layers is not meant for comparison with other methods, but to provide readers with insights into STORM's configuration. In the realm of deep learning, increasing the number of layers in neural networks may generally lead to improved performance, especially with residual connections. However, such a correlation is not evident in STORM. As discussed in **Sections 5.1 and 6**, we hope this analysis can inspire researchers interested in designing new model-based RL algorithms or exploring novel environments.\n  - Conducting ablation studies on the entire benchmark demands approximately **7 (ablation studies) $\\times$ 0.2 (days per game) $\\times$ 5 (seeds) $\\times$ 26 (games) = 182  (NVIDIA GeForce RTX 3090 days)** ,which is resource-intensive and infeasible effort given our time and funding budget at present. Therefore, we have chosen representative environments that are sensitive to different configurations, based on our experience with STORM's development. Of course, users should customize the configuration of STORM or the selection of the algorithms according to specific environments when employing as discussed above.\n  - For the demonstration trajectories, please refer to our **Response 1** in the **Author Rebuttal** above.\n\n- Response to **Question 1** about `the use of sticky actions`: \n\n  - It's correct that we use `v5` environments as described in `train.py` and `train.sh`. However, we will respectfully disagree with the result that you provided. We re-conducted experiments on your listed environments (`v5`) and obtained similar results to what we reported in the paper. We suspect that the discrepancy might arise from the single seed used for training in your case, whereas the reported results in the paper represent the average of 5 runs (retrained with different seeds). As plotted in **Appendix A**, all results in your table fall within a reasonable error range. If all 5 seeds produce such results, please kindly let us know, as it would be quite surprising.\n  - The use of sticky actions, recommended in [1], allows us to verify the algorithm's robustness. At the project's outset, we decided to adopt the latest environment configuration, using `v5` instead of `v4`.\n  - We have also tested our algorithm on your listed environments without sticky action and found that, while there was no significant difference in the other three environments, there is an improvement for the results on the game *Breakout*. This improvement can be attributed to the sensitivity of actions like catching, serving the ball, and moving the slider in this type of game.\n\n- Response to  **Question 3**, other concerns:\n  - Thanks for your helpful suggestion. Agreed. Indeed, performing additional experiments on new/complex environments would be helpful for sufficiently demonstrating the efficiency of STORM and understanding the performance limit of Transformer-based world models for RL. Nonetheless given our time and funding budget, we will leave it for future research. Please also refer to our **Response 1** in the **Author Rebuttal**.\n  - Thank you for your suggestion, and we will include these two papers in the related work section.\n\n\n---\n\n**References**\n\n[1] Machado, Marlos C., et al. \"Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents.\" *Journal of Artificial Intelligence Research* 61 (2018): 523-562.\n\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."}, {"review_id": "NgVwWVLZT4", "reviewer": "Reviewer_E5z3", "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.", "questions": "See weakness above", "limitations": "The generalizability of STORM could be further studied.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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.\n2. 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.\n3. It would be better if further analysis could be provided on the benefits of using a stochastic variational autoencoder.\n\nReference:\n\n[1]. Mei et.al. \"SpeedyZero: Mastering Atari with Limited Data and Time.\" In ICLR 2023 https://openreview.net/forum?id=Mg5CLXZgvLJ.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your effort spent reviewing our paper and providing many valuable suggestions. We will include the suggestions and the pointed concurrent work in the revised paper. Below, we want to address the main concerns raised in the review.\n\n- Response to **(1)**: We appreciate your observation regarding the missing paper, which we will include in the related work section **(Line 117)** rather than the experiments part for the reasons we elaborate next. As stated in **Section 2 (Lines 117-123)** and **Section 4.1 (Lines 221-226)**, the design and experimental validation of STORM have mainly followed the approach of previous works (DreamerV3, IRIS, TWM). Similarly to these works, we do not directly compare our results with **lookahead search based methods** such as MuZero, EfficientZero, and SpeedyZero, 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.\n- Response to **(2)**: Thanks for your suggestion on carrying out additional experiments using other benchmarks, which is interesting and indeed can further strengthen the contribution of the paper. Nonetheless, as explained in our **Response 1** in the **Author Rebuttal** above, the additional computing resources and time required for training and validation of STORM do not seem a feasible plan at present. We will explore it in the future. Please kindly refer to our **Response 1** in the **Author Rebuttal** for more details.\n- Response to **(3)**:  The use of stochastic representation is prevalent in reconstruction-based model-based RL algorithms, including STORM, Dreamer, TWM, and SimPLe. In our early experiments, we found that employing deterministic features to represent the observations (using e.g., a vanilla autoencoder) and replacing the KLDiv dynamics loss with L2 loss results in a significant inconsistency between the reconstructed environment and the original environment after a few steps. In contrast, a stochastic representation maintains stable reconstruction over a larger extended period (**Figure 7** in **Appendix F**). This phenomenon has also been observed by SimPLe [1] (in their Appendix A, ablations on models). Certainly, it should be remarked that this approach has limitations, for instance, contrastive learning methods, which have been proven effective for RL in EfficientZero/CURL, cannot be directly applied to STORM.\n\n---\n\n**References**\n\n[1] Kaiser, Lukasz, et al. \"Model-based reinforcement learning for atari.\" *arXiv preprint arXiv:1903.00374* (2019)."}]

Author rebuttal:
We express our gratitude for the valuable feedback and for the recognition of STORM's efficiency as emphasized by reviewers (4iYJ, rfCq, E5z3), along with the commendation for the paper's coherent structure and presentation, as indicated by reviewers (Cg1v, 4iYJ, rfCq). In the subsequent discussion, we address the main common points raised in the reviews.

1. Response to `STORM should be extended to other tasks like MuJoCo, Crafter, Minecraft, etc`: We appreciate the constructive suggestion and acknowledge the potential benefits of extending STORM to real-life applications like robot control or real-time game control learning tasks. However, due to the substantial computational resources required, evaluating our algorithm on a wide range of complex benchmarks, such as MuJoCo, Crafter, and Minecraft, is currently infeasible given our time and funding budget. For instance, to evaluate our algorithm on full Atari games with 200M sample steps, it demands approximately **5 (days per game) $\times$ 5 (seeds) $\times$ 55 (games) = 1375 (NVIDIA GeForce RTX 3090 days)**. While we aim to explore such extensions in the future, we think it is beyond the scope of this paper.

   Meanwhile, the selection of the 26 Atari 100k games in the experiments is routine and has been done in previous work such as IRIS, DreamerV3, TWM, SimPLe, the baseline methods used in our paper.  It is thus commonly acknowledged that such a selection and evaluation is sufficient to verify the effectiveness of these algorithms. We believe such an evaluation method is sufficient to verify the effectiveness of the algorithm.

2. Response to `Why IRIS, a Transformer based method, performs better than STORM on some tasks`:  IRIS is different from SimPLe, Dreamer, TWM, and STORM. It maps an image observation to 4$\times$4 or more tokens through a VQ-VAE structure. This design is capable of more precisely capturing the position and motion of small objects like in *Pong* and *Breakout*. However, the advantage comes with a trade-off in terms of slow training and inference speed.

3. Response to `Extra ablation studies on structures of STORM and TWM`: We made modifications to our code so that the world model is organized as TWM, and did **addtional experiments** on several environments. We trained the world model with a vanilla Transformer and trained the agent's policy with $s_t = [h_t, z_t]$. We did not use their "balanced dataset sampling" trick. The results obtained are as follows:
   | Environment | STORM (paper) | STORM modified as TWM | TWM (paper) |
   | :---------: | :-----------: | :-------------------: | :---------: |
   |    Alien    |      984      |          822          |     675     |
   |  BankHeist  |      641      |          394          |     467     |
   |  Breakout   |      16       |          12           |     20      |
   |    Pong     |      11       |           7           |     19      |
   |  MsPacman   |     2673      |         1829          |    1588     |
   |   UpNDown   |     7985      |         7092          |    15982    |
   
   Using the TWM modeling approach in several environments can harm the performance. Still, we acknowledge that comparing algorithms with different implementations can make fair ablation studies challenging.
   
   However, it is crucial to highlight that even though STORM and TWM achieve similar scores, the training cost differs significantly due to the $\mathcal{O}(n^2)$ complexity of self-attention operation, with a ratio of $3n_{\mathrm{STORM}}=n_{\mathrm{TWM}}$. Specifically, the training hours required on a GeForce RTX 3090 are as follows:
   
   |   Algorithm    | STORM | STORM modified as TWM | TWM  | Real sampling time |
   | :------------: | :---: | :-----------------: | :--: | :----------------: |
   | Training hours |  4.3  |        10.5         | 12.5 |        1.85        |
   
   It is noteworthy that current model-based RL algorithms exhibit slower convergence rates compared to real-time sampling, which means one will need several expensive computing devices to control a DayDreamer[1]-like robot. In contrast, using STORM is more economical and environment-friendly under such circumstances.
   
4. Response to `Demonstration trajectories`: We agree that the inclusion of demonstration trajectories may also benefit similar algorithms like TWM, IRIS, or Dreamer. While developing our approach, we found that certain games, such as *Freeway*, appear relatively easy for humans who understand or infer the rules quickly from past experience, but they are challenging for algorithms like STORM, DreamerV3, and IRIS. Interestingly, we experimentally observed that including even a single demonstration trajectory has a substantial impact on the performance of *Freeway*, confirming similar findings mentioned in the Appendix H of the IRIS paper.

   We believe that combining demonstration and model-based RL offers a promising solution to address the challenging exploration issue. However, we acknowledge that this method has some limitations, as discussed in **Section 5.3 (Lines 289-292)**, possibly due to the neglect of on-policy distribution, as mentioned in **Section 6 (Lines 321-325)**. As a result, we present the results on three typical games to demonstrate the potential benefits of demonstration trajectories while acknowledging these limitations. We hope to draw the attention of other researchers to further explore this avenue and evaluate the merits of integrating demonstration trajectories with model-based RL techniques.

---

**References**

[1] Wu, Philipp, et al. "Daydreamer: World models for physical robot learning." *Conference on Robot Learning*. PMLR, 2023.

Meta-review: The paper presents STORM, an architecture for model-based reinforcement learning that effectively combines Transformers and variational autoencoders. It demonstrates noteworthy performance improvements and time efficiency in Atari benchmarks, marking a valuable contribution to the field.

The manuscript is well-organized, technically sound, and offers meaningful advances in both performance and computational efficiency, as highlighted by the reviewers. However, there are suggestions for improvement, such as extending ablation studies, providing deeper motivations for modeling choices, and additional analysis on the use of stochastic variational autoencoders. Some of these concerns have been partially addressed in the rebuttal. Considering the overall quality of the work and the improvements made post-rebuttal, the paper is deemed worthy of an accept decision.
Decision: Accept (poster)