SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
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 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.
Table 2: Game scores and overall human-normalized scores on the 26 games in the Atari 100k benchmark. Following the conventions of [9], scores that are the highest or within  $5\%$  of the highest score are highlighted in bold.
Game	Random	Human	SimPLe [11]	TWM [12]	IRIS [13]	DreamerV3 [10]	STORM (ours)
Alien	228	7128	617	675	420	959	984
Amidar	6	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	34	467	53	649	641
Battle Zone	2360	37188	4031	5068	13074	12250	13540
Boxing	0	12	8	78	70	78	80
Breakout	2	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	0	30	17	24	31	0	34
Freeway w/o traj	0	30	17	24	31	0	0
Frostbite	65	4335	237	1476	259	909	1316
Gopher	258	2413	597	1675	2236	3730	8240
Hero	1027	30826	2657	7254	7037	11161	11044
James Bond	29	303	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	15	13	19	15	18	11
Private Eye	25	69571	35	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 N 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.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} \sim \hat{\mathcal{Z}}_{t}$  instead of  $z_{t} \sim \mathcal{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_{\phi}^{R}(\cdot)$  and  $g_{\phi}^{C}(\cdot)$  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 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 [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 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  $\mathcal{L}^{\mathrm{rep}}$  in Equation (5b) directly influences the image encoder. The encoder may be overly influenced when tracking the output of a large sequence model.
(a) Rearranging the modules of STORM.
Figure 3: Ablation studies on the design and configuration of the STORM's world model.
(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:  $\hat{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  $\mathcal{Z}_t$  proves to be beneficial.
Figure 4: Ablation studies on the selection of the agent's state.
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.
Figure 5: Ablations studies on adding a demonstration trajectory to the replay buffer.
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\%$ .
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  $\mu(s)$  for these starting points is not explicitly considered, despite its significance in the policy gradient formulation:  $\nabla J(\theta) \propto \sum_{s} \mu(s) \sum_{a} q_{\pi_{\theta}}(s, a) \nabla \pi_{\theta}(a|s)$  [5].
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.
[1] Jie Chen, Jian Sun, and Gang Wang. From unmanned systems to 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. IEEE/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 improvements in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[4] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv:1707.06347, 2017.
[5] Richard S Sutton, Andrew G Barto, et al. Introduction to Reinforcement Learning, volume 135. MIT Press Cambridge, 1998.
[6] 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.
[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 Learning Representations, 2020. URL https://openreview.net/forum?id=S110TC4tDS.
[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 https://openreview.net/forum?id=OoabwyZb0u.
[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 Buşoniu, and Elena Burceanu. Atari agents, 2022. URL https://github.com/floringogianu/atari-agents.
A Atari 100k curves
Figure 6: Performance comparison on the Atari 100k benchmark. Our method is represented in 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.
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 a CNN layer [33], characterized by kernel  $= 4$ , stride  $= 2$ , and padding  $= 1$ . BN denotes the batch normalization layer[40].
Submodule	Output tensor shape
Input image (ot)	3 × 64 × 64
Conv1 + BN1 + ReLU	32 × 32 × 32
Conv2 + BN2 + ReLU	64 × 16 × 16
Conv3 + BN3 + ReLU	128 × 8 × 8
Conv4 + BN4 + ReLU	256 × 4 × 4
Flatten	4096
Linear	1024
Reshape (produce Zt)	32 × 32
Table 4: Structure of the image decoder. DeConv denotes a transpose CNN layer [41], characterized by kernel  $= 4$ , stride  $= 2$ , and padding  $= 1$ .
Submodule	Output tensor shape
Random sample (zt)	32 × 32
Flatten	1024
Linear + BN0 + ReLU	4096
Reshape	256 × 4 × 4
DeConv1 + BN1 + ReLU	128 × 8 × 8
DeConv2 + BN2 + ReLU	64 × 16 × 16
DeConv3 + BN3 + ReLU	32 × 32 × 32
DeConv4 (produce \(\hat{\sigma}_{t}\))	3 × 64 × 64
Table 5: Action mixer  $e_t = m_\phi(z_t, a_t)$ . Concatenate denotes combining the last dimension of two tensors and merging them into one new tensor. The variable  $A$  represents the action dimension, which ranges from 3 to 18 across different games.  $D$  denotes the feature dimension of the Transformer. LN is an abbreviation for layer normalization [42].
Submodule	Output tensor shape
Random sample (zt), Action (at)	32 × 32, A
Reshape and concatenate	1024 + A
Linear1 + LN1 + ReLU	D
Linear2 + LN2 (output et)	D
Table 6: Positional encoding module.  $w_{1:T}$  is a learnable parameter matrix with shape  $T \times D$ , and  $T$  refers to the sequence length.
Submodule	Output tensor shape
Input (e1:T)
Add (e1:T + w1:T)	T × D
LN
Table 7: Transformer block. Dropout mechanism [43] can prevent overfitting.
Submodule	Module alias	Output tensor shape
Input features (label as x1)	T × D
Multi-head self attention
Linear1 + Dropout(p)	MHSA	T × D
Residual (add x1)
LN1 (label as x2)
Linear2 + ReLU	FFN	T × 2D
Linear3 + Dropout(p)	T × D
Residual (add x2)	T × D
LN2	T × D
Table 8: Transformer based sequence model  $h_{1:T} = f_{\phi}(e_{1:T})$ . Positional encoding is explained in Table 6 and Transformer block is explained in Table 7.
Submodule	Output tensor shape
Input (e1:T)
Positional encoding	T × D
Transformer blocks × K
Output (h1:T)
Table 9: Pure MLP structures. A 1-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	gDφ	1	D/ -1 1024
Reward predictor	gRφ	3	D/ D/ 255
Continuation predictor	gCφ	3	D/ D/ 1
Policy network	πθ(at|st)	3	D/ D/ A
Critic network	Vψ(st)	3	D/ D/ 255
Table 10: Hyerparameters. Note that the environment will provide a "done" signal when losing a life, but will continue running until the actual reset occurs. This life information configuration aligns with the setup used in IRIS [13]. Regarding data sampling, each time we sample  $B_{1}$  trajectories of length  $T$  for world model training, and sample  $B_{2}$  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	p	0.1
World model training batch size	B1	16
World model training batch length	T	64
Imagination batch size	B2	1024
Imagination context length	C	8
Imagination horizon	L	16
Update world model every env step	-	1
Update agent every env step	-	1
Environment context length	-	16
Gamma	γ	0.985
Lambda	λ	0.95
Entropy coefficiency	η	3 × 10-4
Critic EMA decay	σ	0.98
Optimizer	-	Adam [44]
World model learning rate	-	1.0 × 10-4
World model gradient clipping	-	1000
Actor-critic learning rate	-	3.0 × 10-5
Actor-critic gradient clipping	-	100
Gray scale input	-	False
Frame stacking	-	False
Frame skipping	-	4 (max over last 2 frames)
Use of life information	-	True
D Demonstration trajectory information
Table 11: To account for frame skipping, the frame count is multiplied by 4. These trajectories were gathered using pre-trained DQN agents [45].
Game	Episode return	Frames
Ms Pacman	5860	1612 × 4
Pong	18	2079 × 4
Freeway	27	2048 × 4
E Computational cost details and comparison
Table 12: Computational comparison. In the V100 column, an item marked with a star indicates extrapolation based on other graphics cards, while items without a 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 A100 is twice as fast.
Method	Original computing resource	V100 hours
SimPLe [11]	NVIDIA P100, 20 days	240*
TWM [12]	NVIDIA A100, 10 hours	20*
NVIDIA GeForce RTX 3090, 12.5 hours
IRIS [13]	NVIDIA A100, 7 days for two runs	168*
DreamerV3 [10]	NVIDIA V100, 12 hours	12
STORM	NVIDIA GeForce RTX 3090, 4.3 hours	9.3
F Atari video predictions
Figure 7: Multi-step predictions on several environments in Atari games. The world model utilizes 8 observations and actions as contextual input, enabling the imagination of future events spanning 56 frames in an auto-regressive manner.