SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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
zhangwp.bit@gmail.com, {gangwang,sunjian,tyyuan}@bit.edu.cn
Gao Huang
Department of Automation, BNRist, Tsinghua University
gaohuang@tsinghua.edu.cn
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 100k benchmark, setting a new record among state-of-the-art methods that do not employ lookahead search techniques. Moreover, training an agent with 1.85 hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only 4.3 hours, showcasing improved efficiency compared to previous methodologies.
We release our code at https://github.com/weipu-zhang/STORM.
Deep reinforcement learning (DRL) has exhibited remarkable success across diverse domains. However, its widespread application in real-world environments is hindered by the substantial number of interactions with the environment required for achieving such success. This limitation becomes particularly challenging when dealing with broader real-world settings in e.g., unmanned and manufacturing systems [1, 2] that lack adjustable speed simulation tools. Consequently, improving the sample efficiency has emerged as a key challenge for DRL algorithms.
Popular DRL methods, including Rainbow [3] and PPO [4], suffer from low sample efficiency due to two primary reasons. Firstly, the estimation of the value function proves to be a challenging task. This
involves approximating the value function using a deep neural network (DNN) and updating it with  $n$ -step bootstrapped temporal difference, which naturally requires numerous iterations to converge [5]. Secondly, in scenarios where rewards are sparse, many samples exhibit similarity in terms of value functions, providing limited useful information for training and generalization of the DNN [6, 7]. This further exacerbates the challenge of improving the sample efficiency of DRL algorithms.
To address these challenges, model-based DRL algorithms have emerged as a promising approach that tackles both issues simultaneously while demonstrating significant performance gains in sample-efficient settings. These algorithms start by constructing a parameterized simulation world model of the real environment through self-supervised learning. Self-supervised learning can be implemented in various ways, such as reconstructing the original input state using a decoder [8-10], predicting actions between frames [7], or employing contrastive learning to capture the internal consistency of input states [6, 7]. These approaches provide more supervision information than conventional model-free RL losses, enhancing the feature extraction capabilities of DNNs. Subsequently, the agent's policy is improved by leveraging the experiences generated using the world model, eliminating sampling constraints and enabling faster updates to the value function compared to model-free algorithms.
However, the process of imagining with a world model involves an autoregressive process that can accumulate prediction errors over time. In situations where discrepancies arise between the imagined trajectory and the real trajectory, the agent may inadvertently pursue virtual goals, resulting in subpar performance in the real environment. To mitigate this issue, introducing random noise into the world model has been proven beneficial [9-11, 14]. Variational autoencoders, capable of automatically learning low-dimensional latent representations of high-dimensional data while incorporating reasonable random noise into the latent space, offer an ideal choice for image encoding.
Numerous endeavors have been undertaken to construct an efficient world model. For instance, SimPLe [11] leverages LSTM [15], while DreamerV3 [10] employs GRU [16] as the sequence model. LSTM and GRU, both variants of recurrent neural networks (RNNs), excel at sequence modeling tasks. However, the recurrent nature of RNNs impedes parallelized computing, resulting in slower training speeds [17]. In contrast, the Transformer architecture [17] has lately demonstrated superior performance over RNNs in various sequence modeling and generation tasks. It overcomes the challenge of forgetting long-term dependencies and is designed for efficient parallel computing. While several attempts have been made to incorporate Transformers into the world model [12, 13, 18], these works do not fully harness the capabilities of this architecture. Furthermore, these approaches require even longer training times and fail to surpass the performance of the GRU-based DreamerV3.
Figure 1: Comparison of methods on Atari 100k. SimPLe [11] and DreamerV3 [10] employ RNNs as their world models, whereas TWM [12], IRIS [13], and STORM use Transformers. The training frames per second (FPS) results on a single NVIDIA V100 GPU are extrapolated from other graphics cards for SimPLe, TWM, and IRIS, while DreamerV3 and STORM are directly evaluated.
In this paper, we introduce the Stochastic Transformer-based wORld Model (STORM), a highly effective and efficient structure for model-based RL. STORM employs a categorical variational autoencoder (VAE) as the image encoder, enhancing the agent robustness and reducing accumulated autoregressive prediction errors. Subsequently, we incorporate the Transformer as the sequence
model, improving modeling and generation quality while accelerating training. STORM achieves a remarkable mean human normalized score of  $126.7\%$  on the challenging Atari 100k benchmark, establishing a new record for methods without resorting to lookahead search. Furthermore, training an agent with 1.85 hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only 4.3 hours, demonstrating superior efficiency compared to previous methodologies. The comparison of our approach with the state-of-the-art methods is depicted in Figure 1.
Model-based DRL algorithms aim to construct a simulation model of the environment and utilize simulated experiences to improve the policy. While traditional model-based RL techniques like Dyna-Q have shown success in tabular cases [5], modeling complex environments such as video games and visual control tasks presents significant challenges. Recent advances in computing and DNNs have enabled model-based methods to learn the dynamics of the environments and start to outperform model-free methods on these tasks.
The foundation of VAE-LSTM-based world models was introduced by Ha and Schmidhuber [14] for image-based environments, demonstrating the feasibility of learning a good policy solely from generated data. SimPLe [11] applied this methodology to Atari games, resulting in substantial sample efficiency improvements compared to Rainbow [3], albeit with relatively lower performance under limited samples. The Dreamer series [8-10] also adopt this framework and showcase notable capabilities in Atari games, DeepMind Control, Minecraft, and other domains, using GRU [16] as the core sequential model. However, as discussed earlier, RNN structures suffer from slow training [17].
Recent approaches such as IRIS [13], TWM [12], and TransDreamer [18] incorporate the Transformer architecture into their world models. IRIS [13] employs VQ-VAE [19] as the encoder to map images into  $4 \times 4$  latent tokens and uses a spatial-temporal Transformer [20] to capture information within and across images. However, the attention operations on a large number of tokens in the spatial-temporal structure can result in a significant training slowdown. TWM [12] adopts Transformer-XL [21] as its core architecture and organizes the sequence model in a structure similar to Decision Transformer [22], treating the observation, action, and reward as equivalent input tokens for the Transformer. Performing self-attention across different types of data may have a negative impact on the performance, and the increased number of tokens considerably slows down training. TransDreamer [18] directly replaces the GRU structure of Dreamer with Transformer. However, there is a lack of evidence demonstrating their performance in widely accepted environments or under limited sample conditions.
Other model-based RL methods such as MuZero [23], EfficientZero [7], and SpeedyZero [24] incorporate Monte Carlo tree search (MCTS) to enhance policy and achieve promising performance on the Atari 100k benchmark. Lookahead search techniques like MCTS can be employed to enhance other model-based RL algorithms, but they come with high computational demands. Additionally, certain model-free methods [6, 25-27] incorporate a self-supervised loss as an auxiliary term alongside the standard RL loss, demonstrating their effectiveness in sample-efficient settings. Additionally, recent studies [28, 29] delve deeper into model-free RL, demonstrating strong performance and high data efficiency, rivaling that of model-based methods on several benchmarks. However, since the primary objective of this paper is to enhance the world model, we do not delve further into these methods.
We highlight the distinctions between STORM and recent approaches in the world model as follows:
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	SimPLe [11]	TWM [12]	IRIS [13]	DreamerV3 [10]	STORM (ours)
Sequence model	LSTM [15]	Transformer-XL [21]	Transformer [17]	GRU [16]	Transformer
Tokens	Latent	Latent, action, reward	Latent(4 × 4)	Latent	Latent
Latent representation	Binary-VAE	Categorical-VAE	VQ-VAE	Categorical-VAE	Categorical-VAE
Historical information	Yes	No	Yes	Yes	No
Agent state	Reconstructed image	Latent	Reconstructed image	Latent, hidden	Latent, hidden
Agent training	PPO [4]	As DreamerV2 [9]	As DreamerV2 [9]	DreamerV3	As DreamerV3
Figure 2: 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_{\phi}$  in Equation (2). The Agent block, represented by a neural network, corresponds to  $\pi_{\theta}(a_t|s_t)$  in Equation (6).
- SimPLe [11] and Dreamer [10] rely on RNN-based models, whereas STORM employs a GPT-like Transformer [30] as the sequence model.
- In contrast to IRIS [13] that employs multiple tokens, STORM utilizes a single stochastic latent variable to represent an image.
- STORM follows a vanilla Transformer [17] structure, while TWM [12] adopts a Transformer-XL [21] structure.
- In the sequence model of STORM, an observation and an action are fused into a single token, whereas TWM [12] treats observation, action, and reward as three separate tokens of equal importance.
- Unlike Dreamer [10] and TransDreamer [18], which incorporate hidden states, STORM reconstructs the original image without utilizing this information.
3.1 World model learning
Our approach adheres to the established framework of model-based RL algorithms, which focus on enhancing the agent's policy by imagination [5, 9-11, 13]. We iterate through the following steps until reaching the prescribed number of real environment interactions.
S1) Gather real environment data by executing the current policy for several steps and append them to the replay buffer.
S2) Update the world model using trajectories sampled from the replay buffer.
S3) Improve the policy using imagined experiences generated by the world model, with the starting points for the imagination process sampled from the replay buffer.
At each time  $t$ , a data point comprises an observation  $o_t$ , an action  $a_t$ , a reward  $r_t$ , and a continuation flag  $c_t$  (a Boolean variable indicating whether the current episode is ongoing). The replay buffer maintains a first-in-first-out queue structure, enabling the sampling of consecutive trajectories from the buffer.
Section 3.1 provides a detailed description of the architecture and training losses employed by STORM. On the other hand, Section 3.2 elaborates on the imagination process and the training methodology employed by the agent. It provides an thorough explanation of how the agent leverages the world model to simulate experiences and improve its policy.
Model structure The complete structure of our world model is illustrated in Figure 2. In our experiments, we focus on Atari games [31], which generate image observations  $o_{t}$  of the environment. Modeling the dynamics of the environment directly on raw images is computationally expensive and prone to errors [7-11, 13, 23]. To address this, we leverage a VAE [32] formulated in Equation (1) to convert  $o_{t}$  into latent stochastic categorical distributions  $\mathcal{Z}_{t}$ . Consistent with prior work [9, 10, 12], we set  $\mathcal{Z}_{t}$  as a stochastic distribution comprising 32 categories, each with 32 classes. The encoder  $(q_{\phi})$  and decoder  $(p_{\phi})$  structures are implemented as convolutional neural networks (CNNs) [33]. Subsequently, we sample a latent variable  $z_{t}$  from  $\mathcal{Z}_{t}$  to represent the original observation  $o_{t}$ . Since sampling from a distribution lacks gradients for backward propagation, we apply the straight-through gradients trick [9, 34] to preserve them.
$$
\text {I m a g e e n c o d e r :} \quad z _ {t} \sim q _ {\phi} \left(z _ {t} \mid o _ {t}\right) = \mathcal {Z} _ {t} \tag {1}
$$
Image decoder:  $\hat{o}_t = p_\phi (z_t)$
Before entering the sequence model, we combine the latent sample  $z_{t}$  and the action  $a_{t}$  into a single token  $e_{t}$  using multi-layer perceptrons (MLPs) and concatenation. This operation, denoted as  $m_{\phi}$ , prepares the inputs for the sequence model. The sequence model  $f_{\phi}$  takes the sequence of  $e_{t}$  as input and produces hidden states  $h_{t}$ . We adopt a GPT-like Transformer structure [30] for the sequence model, where the self-attention blocks are masked with a subsequent mask allowing  $e_{t}$  to attend to the sequence  $e_{1}, e_{2}, \ldots, e_{t}$ . By utilizing MLPs  $g_{\phi}^{D}$ ,  $g_{\phi}^{R}$ , and  $g_{\phi}^{C}$ , we rely on  $h_{t}$  to predict the current reward  $\hat{r}_{t}$ , the continuation flag  $\hat{c}_{t}$ , and the next distribution  $\hat{\mathcal{Z}}_{t+1}$ . The formulation of this part of the world model is as follows
$$
\text {A c t i o n m i x e r :} \quad e _ {t} = m _ {\phi} \left(z _ {t}, a _ {t}\right)
$$
$$
\text {S e q u e n c e m o d e l :} \quad h _ {1: T} = f _ {\phi} \left(e _ {1: T}\right)
$$
$$
\text {D y n a m i c s p r e d i c t o r :} \quad \hat {\mathcal {Z}} _ {t + 1} = g _ {\phi} ^ {D} \left(\hat {z} _ {t + 1} \mid h _ {t}\right) \tag {2}
$$
$$
R e w a r d \quad \hat {r} _ {t} = g _ {\phi} ^ {R} (h _ {t})
$$
$$
\text {C o n t i n u a t i o n p r e d i c t o r :} \quad \hat {c} _ {t} = g _ {\phi} ^ {C} (h _ {t}).
$$
Loss functions The world model is trained in a self-supervised manner, optimizing it end-to-end. The total loss function is calculated as in Equation (3) below, with fixed hyperparameters  $\beta_{1} = 0.5$  and  $\beta_{2} = 0.1$ . In the equation,  $B$  denotes the batch size, and  $T$  denotes the batch length
$$
\mathcal {L} (\phi) = \frac {1}{B T} \sum_ {n = 1} ^ {B} \sum_ {t = 1} ^ {T} \left[ \mathcal {L} _ {t} ^ {\operatorname {r e c}} (\phi) + \mathcal {L} _ {t} ^ {\operatorname {r e w}} (\phi) + \mathcal {L} _ {t} ^ {\operatorname {c o n}} (\phi) + \beta_ {1} \mathcal {L} _ {t} ^ {\mathrm {d y n}} (\phi) + \beta_ {2} \mathcal {L} _ {t} ^ {\operatorname {r e p}} (\phi) \right]. \tag {3}
$$
The individual components of the loss function are defined as follows:  $\mathcal{L}_t^{\mathrm{rec}}(\phi)$  represents the reconstruction loss of the original image,  $\mathcal{L}_t^{\mathrm{rew}}(\phi)$  represents the prediction loss of the reward, and  $\mathcal{L}_t^{\mathrm{con}}(\phi)$  represents the prediction loss of the continuation flag.
$$
\mathcal {L} _ {t} ^ {\text {r e c}} (\phi) = \left| \left| \hat {o} _ {t} - o _ {t} \right| \right| _ {2} \tag {4a}
$$
$$
\mathcal {L} _ {t} ^ {\text {r e w}} (\phi) = \mathcal {L} ^ {\text {s y m}} \left(\hat {r} _ {t}, r _ {t}\right) \tag {4b}
$$
$$
\mathcal {L} _ {t} ^ {\mathrm {c o n}} (\phi) = c _ {t} \log \hat {c} _ {t} + (1 - c _ {t}) \log (1 - \hat {c} _ {t}). \tag {4c}
$$
Additionally,  $\mathcal{L}^{\mathrm{sym}}$  in Equation (4b) denotes the symlog two-hot loss, as described in [10]. This loss function transforms the regression problem into a classification problem, ensuring consistent loss scaling across different environments.
The losses  $\mathcal{L}_t^{\mathrm{dyn}}(\phi)$  and  $\mathcal{L}_t^{\mathrm{rep}}(\phi)$  are expressed as Kullback-Leibler (KL) divergences but differ in their gradient backward and weighting. The dynamics loss  $\mathcal{L}_t^{\mathrm{dyn}}(\phi)$  guides the sequence model in predicting the next distribution, while the representation loss  $\mathcal{L}_t^{\mathrm{rep}}(\phi)$  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.
$$
\mathcal {L} _ {t} ^ {\mathrm {d y n}} (\phi) = \max  \left(1, \mathrm {K L} \left[ \operatorname {s g} \left(q _ {\phi} \left(z _ {t + 1} \mid o _ {t + 1}\right)\right) \mid \mid g _ {\phi} ^ {D} \left(\hat {z} _ {t + 1} \mid h _ {t}\right) \right]\right) \tag {5a}
$$
$$
\mathcal {L} _ {t} ^ {\operatorname {r e p}} (\phi) = \max  \left(1, \mathrm {K L} \left[ q _ {\phi} \left(z _ {t + 1} \mid o _ {t + 1}\right) \mid \mid \operatorname {s g} \left(g _ {\phi} ^ {D} \left(\hat {z} _ {t + 1} \mid h _ {t}\right)\right) \right]\right) \tag {5b}
$$
where  $\mathrm{sg}(\cdot)$  denotes the operation of stop-gradients.
3.2 Agent learning
The agent's learning is solely based on the imagination process facilitated by the world model, as illustrated in Figure 2. To initiate the imagination process, a brief contextual trajectory is randomly selected from the replay buffer, and the initial posterior distribution  $\mathcal{Z}_t$  is computed. During inference, rather than sampling directly from the posterior distribution  $\mathcal{Z}_t$ , we sample  $z_t$  from the prior distribution  $\hat{\mathcal{Z}}_t$ . To accelerate the inference, we employ the KV cache technique [35] within the Transformer structure.
The agent's state is formed by concatenating  $z_{t}$  and  $h_{t}$ , as shown below:
State:  $s_t = [z_t, h_t]$
$\begin{array}{rl}{V_{\psi}(s_t)\approx \mathbb{E}_{\pi_\theta ,p_\phi}\Big[\sum_{k = 0}^{\infty}\gamma^k r_{t + k}\Big]} & {} \end{array}$  (6)
Actor:  $a_{t}\sim \pi_{\theta}(a_{t}|s_{t}).$
We adopt the actor learning settings from DreamerV3 [10]. The complete loss of the actor-critic algorithm is described by Equation (7), where  $\hat{r}_t$  corresponds to the reward predicted by the world model, and  $\hat{c}_t$  represents the predicted continuation flag:
$$
\mathcal {L} (\theta) = \frac {1}{B L} \sum_ {n = 1} ^ {B} \sum_ {t = 1} ^ {L} \left[ - \operatorname {s g} \left(\frac {G _ {t} ^ {\lambda} - V _ {\psi} (s _ {t})}{\max  (1 , S)}\right) \ln \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) - \eta H \left(\pi_ {\theta} \left(a _ {t} \mid s _ {t}\right)\right) \right] \tag {7a}
$$
$$
\mathcal {L} (\psi) = \frac {1}{B L} \sum_ {n = 1} ^ {B} \sum_ {t = 1} ^ {L} \left[ \left(V _ {\psi} (s _ {t}) - \operatorname {s g} \left(G _ {t} ^ {\lambda}\right)\right) ^ {2} + \left(V _ {\psi} (s _ {t}) - \operatorname {s g} \left(V _ {\psi^ {\mathrm {E M A}}} (s _ {t})\right)\right) ^ {2} \right] \tag {7b}
$$
where  $H(\cdot)$  denotes the entropy of the policy distribution, while constants  $\eta$  and  $L$  represent the coefficient for entropy loss and the imagination horizon, respectively. The  $\lambda$ -return  $G_t^\lambda$  [5, 10] is recursively defined as follows
$$
G _ {t} ^ {\lambda} \doteq r _ {t} + \gamma c _ {t} \left[ (1 - \lambda) V _ {\psi} \left(s _ {t + 1}\right) + \lambda G _ {t + 1} ^ {\lambda} \right] \tag {8a}
$$
$$
G _ {L} ^ {\lambda} \doteq V _ {\psi} \left(s _ {L}\right). \tag {8b}
$$
The normalization ratio  $S$  utilized in the actor loss (7a) is defined in Equation (9), which is computed as the range between the 95th and 5th percentiles of the  $\lambda$ -return  $G_{t}^{\lambda}$  across the batch [10]
$$
S = \operatorname {p e r c e n t i l e} \left(G _ {t} ^ {\lambda}, 9 5\right) - \operatorname {p e r c e n t i l e} \left(G _ {t} ^ {\lambda}, 5\right). \tag {9}
$$
To regularize the value function, we maintain the exponential moving average (EMA) of  $\psi$ . The EMA is defined in Equation (10), where  $\psi_t$  represents the current critic parameters,  $\sigma$  is the decay rate, and  $\psi_{t+1}^{\mathrm{EMA}}$  denotes the updated critic parameters. This regularization technique aids in stabilizing training and preventing overfitting
$$
\psi_ {t + 1} ^ {\mathrm {E M A}} = \sigma \psi_ {t} ^ {\mathrm {E M A}} + (1 - \sigma) \psi_ {t}. \tag {10}
$$
We evaluated the performance of STORM on the widely-used benchmark for sample-efficient RL, Atari 100k [31]. For detailed information about the benchmark, evaluation methodology, and the baselines used for comparison, please refer to Section 4.1. The comprehensive results for the Atari 100k games are presented in 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 skipping (4 frames skipped) and repeated actions within those frames. This constraint corresponds to
approximately 1.85 hours of real-time gameplay. The agent's human normalized score  $\tau = \frac{A - R}{H - R}$  is calculated based on the score  $A$  achieved by the agent, the score  $R$  obtained by a random policy, and the average score  $H$  achieved by a human player in a specific environment. To determine the human player's performance  $H$ , a player is allowed to become familiar with the game under the same sample constraint.
To demonstrate the efficiency of our proposed world model structure, we compare it with model-based DRL algorithms that share a similar training pipeline, as discussed in Section 2. However, similarly to [10, 12, 13], we do not directly compare our results with lookahead search methods like MuZero [23] and EfficientZero [7], as our primary goal is to refine the world model itself. Nonetheless, lookahead search techniques can be combined with our method in the future to further enhance the agent's performance.
4.2 Results on Atari 100k