SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fei Deng
fei.deng@rutgers.edu
Junyeong Park
jyp10987@kaist.ac.kr
Sungjin Ahn*
sungjin.ahn@kaist.ac.kr
https://fdeng18.github.io/s4wm
World models are a fundamental component in model-based reinforcement learning (MBRL). To perform temporally extended and consistent simulations of the future in partially observable environments, world models need to possess long-term memory. However, state-of-the-art MBRL agents, such as Dreamer, predominantly employ recurrent neural networks (RNNs) as their world model backbone, which have limited memory capacity. In this paper, we seek to explore alternative world model backbones for improving long-term memory. In particular, we investigate the effectiveness of Transformers and Structured State Space Sequence (S4) models, motivated by their remarkable ability to capture long-range dependencies in low-dimensional sequences and their complementary strengths. We propose S4WM, the first world model compatible with parallelizable SSMs including S4 and its variants. By incorporating latent variable modeling, S4WM can efficiently generate high-dimensional image sequences through latent imagination. Furthermore, we extensively compare RNN-, Transformer-, and S4-based world models across four sets of environments, which we have tailored to assess crucial memory capabilities of world models, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. Our findings demonstrate that S4WM outperforms Transformer-based world models in terms of long-term memory, while exhibiting greater efficiency during training and imagination. These results pave the way for the development of stronger MBRL agents.
The human brain is frequently compared to a machine whose primary function is to construct models of the world, enabling us to predict, plan, and react to our environment effectively [51, 39]. These mental representations, referred to as world models, are integral to essential cognitive functions like decision-making and problem-solving. Similarly, one of the pivotal tasks in artificial intelligence (AI) systems that aim for human-like cognition is the development of analogous world models.
Model-Based Reinforcement Learning (MBRL) [42] has emerged as a promising approach that builds world models through interaction with the environment. As a fundamental component of MBRL, these world models empower artificial agents to anticipate the consequences of their actions and plan accordingly, leading to various advantages. Notably, MBRL offers superior sample efficiency, mitigating the high data requirements commonly associated with model-free methods. Moreover, MBRL exhibits enhanced exploration, transferability, safety, and explainability [42], making it well-suited for complex and dynamic environments where model-free methods tend to struggle.
*Correspondence to sungjin.ahn@kaist.ac.kr.
The effectiveness and characteristics of world models crucially depend on their backbone neural network architecture. In particular, the backbone architecture dictates the model's capabilities of capturing long-term dependencies and handling stochasticity in the environment. Additionally, it affects the compactness of memory footprint and the speed of future prediction rollouts. Furthermore, in visual MBRL that finds extensive practical applications, the backbone architecture holds even greater significance than it does in state-based MBRL. This is due to the need to deal with high-dimensional, unstructured, and temporal observations.
Nevertheless, choosing the appropriate backbone architecture for visual MBRL has become a considerable challenge due to the rapidly evolving landscape of deep architectures for temporal sequence modeling. This includes the recent advancements in major backbone architecture classes, notably Transformers [61, 5] and the Structured State Space Sequence models such as S4 [21] and S5 [57].
Traditionally, Recurrent Neural Networks (RNNs) [7, 33] have been the go-to backbone architecture, thanks to their efficient use of computational resources in processing sequential data. However, RNNs tend to suffer from vanishing gradient issues [49], limiting their long-term memory capacity. Recently, Transformers [61] have demonstrated superior sequence modeling capabilities in multiple domains, including natural language processing and computer vision [12]. Their self-attention mechanism grants direct access to all previous time steps, thereby enhancing long-term memory. Moreover, Transformers offer parallel trainability and exhibit faster training speeds than RNNs. However, their quadratic complexity and slow generation speed pose challenges when dealing with very long sequences. To address these limitations, the S4 model has been proposed, offering both parallel training and recurrent generation with sub-quadratic complexity. In the Long Range Arena [59] benchmark consisting of low-dimensional sequence modeling tasks, the S4 model outperforms many Transformer variants in both task performance and computational efficiency.
In this paper, we present two primary contributions. Firstly, we introduce S4WM, the first and general world model framework that is compatible with any Parallelizable SSMs (PSSMs) including S4, S5, and other S4 variants. This is a significant development since it was unclear whether the S4 model would be effective as a high-dimensional visual world model, and if so, how this could be achieved. To this end, we instantiate the S4WM with the S4 architecture to manage high-dimensional image sequences, and propose its probabilistic latent variable modeling framework based on variational inference. Secondly, we conduct the first empirical comparative study on the three major backbone architectures for visual world modeling—RNNs, Transformers, and S4. Our results show that S4WM outperforms RNNs and Transformers across multiple memory-demanding tasks, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. In terms of speed, S4WM trains the fastest, while RNNs exhibit significantly higher imagination throughput. We believe that by shedding light on the strengths and weaknesses of these backbones, our study contributes to a deeper understanding that can guide researchers and practitioners in selecting suitable architectures, and potentially inspire the development of novel approaches in this field.
Structured State Space Sequence (S4) Model. Originally introduced in [21], S4 is a sequence modeling framework that solves all tasks in the Long Range Arena [59] for the first time. At its core is a structured parameterization of State Space Models (SSMs) that allows efficient computation and exhibits superior performance in capturing long-range dependencies both theoretically and empirically. However, the mathematical background of S4 is quite involved. To address this, a few recent works seek to simplify, understand, and improve S4 [23, 20, 40, 22, 57, 24, 47]. It has been discovered that S4 and Transformers have complementary strengths [70, 40, 15, 34, 24]. For example, Transformers can be better at capturing local (short-range) information and performing context-dependent operations. Therefore, hybrid architectures have been proposed to achieve the best of both worlds. Furthermore, S4 and its variants have found applications in various domains, such as image and video classification [43, 36, 34, 62], audio generation [17], time-series generation [69], language modeling [70, 15, 40], and model-free reinforcement learning [10, 38]. Our study introduces the first world model compatible with S4 and its variants (more generally, parallelizable SSMs) for improving long-term memory in MBRL. We also investigate the strengths and weaknesses of S4 and Transformers in the context of world model learning.
World Models. World models [25] are typically implemented as dynamics models of the environment that enable the agent to plan into the future and learn policies from imagined trajectories. RNNs have been the predominant backbone architecture of world models. A notable example is RSSM [28], which has been widely used in both reconstruction-based [29-31, 65, 56, 16, 35, 63, 64, 68] and reconstruction-free [44-46, 11, 26] MBRL agents. With the advent of Transformers [61], recent works have also explored using Transformers as the world model backbone [5, 41, 54]. While Transformers are less prone to vanishing gradients [49] than RNNs, their quadratic complexity limits their applicability to long sequences. For example, recent works [41, 54] use a short imagination horizon of  $\sim 20$  steps. In contrast, S4WM can successfully imagine hundreds of steps into the future with sub-quadratic complexity. We also develop an improved Transformer-based world model that can deal with long sequences by employing Transformer-XL [9].
Agent Memory Benchmarks. While many RL benchmarks feature partially observable environments, they tend to evaluate multiple agent capabilities simultaneously [2, 8, 27] (e.g., exploration and modular skill learning), and may be solvable with a moderate memory capacity [14, 48]. Additionally, some benchmarks are designed for model-free agents [66, 37, 52], and may contain stochastic dynamics that are not controlled by the agents, making it hard to separately assess the memory capacity of world models. The recently proposed Memory Maze [50] focuses on measuring long-term memory and provides benchmark results for model-based agents. We build upon Memory Maze and introduce additional environments and tasks to probe a wider range of memory capabilities. Another recent work, TECO [67], also introduces datasets and a Transformer-based model for evaluating and improving long-term video prediction. Our work has a different focus than TECO in that we stress test models on extremely long sequences (up to 2000 steps), while TECO considers more visually complex environments, with sequence lengths capped at 300. Our experiment setup allows using relatively lightweight models to tackle significant challenges involving long-term memory. We include a comparison with TECO in Appendix F.
S4 [21] and its variants [20, 40, 57, 15] are specialized parameterizations of linear state space models. We first present relevant background on linear state space models, and then introduce the S4 model.
Linear State Space Models (SSMs) are a widely used sequence model that defines a mapping from a 1-D input signal  $u(t)$  to a 1-D output signal  $y(t)$ . They can be discretized into a sequence-to-sequence mapping by a step size  $\Delta$ . The continuous-time and discrete-time SSMs can be described as:
$$
\text {(c o n t i n u o u s - t i m e)} \quad \begin{array}{l} \boldsymbol {s} ^ {\prime} (t) = \boldsymbol {A} \boldsymbol {s} (t) + \boldsymbol {B} \boldsymbol {u} (t) \\ y (t) = \boldsymbol {C} \boldsymbol {s} (t) + \boldsymbol {D} \boldsymbol {u} (t) \end{array} , \quad \text {(d i s c r e t e - t i m e)} \quad \begin{array}{l} \boldsymbol {s} _ {k} = \overline {{\boldsymbol {A}}} \boldsymbol {s} _ {k - 1} + \overline {{\boldsymbol {B}}} \boldsymbol {u} _ {k} \\ y _ {k} = \overline {{\boldsymbol {C}}} \boldsymbol {s} _ {k} + \overline {{\boldsymbol {D}}} \boldsymbol {u} _ {k} \end{array} . \tag {1}
$$
Here, the vectors  $s(t)$  and  $s_k$  are the internal hidden state of the SSM, and the discrete-time matrices  $\overline{A}, \overline{B}, \overline{C}, \overline{D}$  can be computed from their continuous-time counterparts  $A, B, C, D$  and the step size  $\Delta$ . We will primarily deal with the discrete-time SSMs, which allow efficient autoregressive generation like RNNs due to the recurrence in  $s_k$ .
Unlike RNNs, however, linear SSMs can offer parallelizable computation like Transformers. That is, given the input sequence  $u_{1:T}$ , the output sequence  $y_{1:T}$  can be computed in parallel across time steps by a discrete convolution [21] or a parallel associative scan [57, 4]. In this work, we define the class of Parallelizable SSMs (PSSMs) to be the SSMs that provide the following interface for both parallel and single-step computation:
$$
\text {(p a r a l l e l)} \quad y _ {1: T}, s _ {T} = \operatorname {P S S M} \left(u _ {1: T}, s _ {0}\right), \quad \text {(s i n g l e s t e p)} \quad y _ {k}, s _ {k} = \operatorname {P S S M} \left(u _ {k}, s _ {k - 1}\right), \tag {2}
$$
where the inputs  $\pmb{u}_k$  and outputs  $\pmb{y}_k$  can be vectors.
The S4 model aims to use SSMs for deep sequence modeling, where the matrices  $A, B, C, D$  and the step size  $\Delta$  are learnable parameters to be optimized by gradient descent. Because SSMs involve computing powers of  $\overline{A}$ , which is in general expensive and can lead to the exploding/vanishing gradients problem [49], SSMs with a randomly initialized  $A$  perform very poorly in practice [19].
To address these problems, S4 parameterizes  $A$  as a Diagonal Plus Low-Rank (DPLR) matrix [21, 17]:  $A = \Lambda - PP^*$ , where  $\Lambda$  is a diagonal matrix,  $P$  is typically a column vector (with rank 1), and  $P^*$  is the conjugate transpose of  $P$ . This parameterization allows efficient computation of powers of  $\overline{A}$ ,
Figure 1: We propose S4WM, the first S4-based world model for improving long-term memory. S4WM efficiently models the long-range dependencies of environment dynamics in a compact latent space, using a stack of S4 blocks. This crucially allows fully parallelized training and fast recurrent latent imagination. S4WM is a general framework that is compatible with any parallelizable SSM including S5 and other S4 variants.
Algorithm 1 S4WM Training
Input:  $(\mathbf{x}_0,\mathbf{a}_1,\mathbf{x}_1,\dots ,\mathbf{a}_T,\mathbf{x}_T)$  ▷Obtain posterior latents
1: for time step  $t = 0, \dots, T$  parallel do
2:  $\mathbf{z}_t\sim q(\mathbf{z}_t|\mathbf{x}_t) = \mathrm{Encoder}(\mathbf{x}_t)$
3: end for  
 $\triangleright$  Prepare inputs to S4 blocks
4: for time step  $t = 1, \dots, T$  parallel do
5:  $\pmb{g}_t = \mathrm{MLP}(\mathrm{concat}[\mathbf{z}_{t-1}, \pmb{a}_t])$
6: end for  
▷ Encode history by S4 blocks
7:  $\pmb{h}_{1:T}, \pmb{s}_T = \mathbf{S4Blocks}(\pmb{g}_{1:T}, \pmb{s}_0)$ $\triangleright$  Compute prior and decode posterior latency.
8: for time step  $t = 1, \dots, T$  parallel do
9:  $p(\mathbf{z}_t|\mathbf{z}_{< t},\mathbf{a}_{\leq t}) = \mathrm{MLP}(\mathbf{h}_t)$
10:  $\hat{\mathbf{x}}_t = \mathrm{Decoder}(\mathrm{concat}[h_t, \mathbf{z}_t])$
11: end for
12: Compute objective by Equation (9)
Algorithm 2 S4WM Imagination
Input: context  $(\mathbf{x}_{0:C},\mathbf{a}_{1:C})$  , query  $a_{C + 1:T}$  -Encodecontext
1: for context step  $t = 0, \dots, C$  parallel do
2:  $\mathbf{z}_t\sim q(\mathbf{z}_t|\mathbf{x}_t) = \mathrm{Encoder}(\mathbf{x}_t)$
3:  $\pmb{g}_{t + 1} = \mathrm{MLP}(\mathrm{concat}[\mathbf{z}_t,\pmb{a}_{t + 1}])$
4: end for
5:  $\pmb{h}_{1:C+1}, \pmb{s}_{C+1} = \mathbf{S4Blocks}(\pmb{g}_{1:C+1}, \pmb{s}_0)$   
> Imagine in latent space
6:  $\mathbf{z}_{C + 1}\sim p(\mathbf{z}_{C + 1}|\mathbf{z}_{0:C},\pmb{a}_{1:C + 1}) = \mathrm{MLP}(\pmb{h}_{C + 1})$
7: for query step  $t = C + 2,\dots ,T$  do
8:  $\pmb{g}_t = \mathrm{MLP}(\mathrm{concat}[\mathbf{z}_{t-1}, \pmb{a}_t])$
9:  $\pmb{h}_t, \pmb{s}_t = \mathrm{S4Blocks}(\pmb{g}_t, \pmb{s}_{t-1})$
0:  $\mathbf{z}_t\sim p(\mathbf{z}_t|\mathbf{z}_{< t},\pmb{a}_{\leq t}) = \mathrm{MLP}(\pmb {h}_t)$
11: end for  
▷ Decode imagined latents
12: for query step  $t = C + 1, \dots, T$  parallel do
13:  $\hat{\mathbf{x}}_t = \mathrm{Decoder}(\mathrm{concat}[h_t, \mathbf{z}_t])$
14: end for
while also including the HiPPO matrices [18], which are theoretically derived based on continuous-time memorization and empirically shown to better capture long-range dependencies. In practice, S4 initializes  $\mathbf{A}$  to the HiPPO matrix. To cope with vector-valued inputs and outputs (i.e.,  $\mathbf{u}_k, \mathbf{y}_k \in \mathbb{R}^H$ ), S4 makes  $H$  copies of the SSM, each operating on one dimension, and mixes the outputs by a position-wise linear layer. Follow-up works further simplify the parameterization of  $\mathbf{A}$  to a diagonal matrix [20], and use a multi-input, multi-output SSM for vector-valued sequences [57].
4 S4WM: A General World Model for Parallelizable SSMs
We consider an agent interacting with a partially observable environment. At each time step  $t$ , the agent receives an image observation  $\mathbf{x}_t$ . It then chooses an action  $\mathbf{a}_{t+1}$  based on its policy, and receives the next observation  $\mathbf{x}_{t+1}$ . For simplicity, we omit the reward here.
We aim to model  $p(\mathbf{x}_{1:T} \mid \mathbf{x}_0, \mathbf{a}_{1:T})$ , the distribution of future observations given the action sequence. We note that it is not required to model  $p(\mathbf{x}_0)$ , as world model imagination is typically conditioned
on at least one observation. While S4 and its variants have shown remarkable abilities to model long-range dependencies, they operate directly in the observation space. For example, S4 models images as sequences of pixels, and directly learns the dependencies among individual pixels. This is hard to scale to high-dimensional sequences, such as the sequences of images that we aim to model.
Inspired by RSSM [28], we propose S4WM, the first PSSM-based world model that learns the environment dynamics in a compact latent space. This not only allows fast parallelizable training, but also enables efficient modeling of long-range dependencies in the latent space. Importantly, S4WM is a general framework that can incorporate not only the specific S4 model [21] but also any PSSM defined by Equation (2), including S5 [57] and other variants [20, 40, 15]. It models the observations and state transitions through a probabilistic generative process:
$$
p \left(\mathbf {x} _ {1: T} \mid \mathbf {x} _ {0}, \boldsymbol {a} _ {1: T}\right) = \int p \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right) \prod_ {t = 1} ^ {T} p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) d \mathbf {z} _ {0: T}, \tag {3}
$$
where  $\mathbf{z}_{0:T}$  are the stochastic latent states. We note that computing the likelihood  $p(\mathbf{x}_t\mid \mathbf{z}_{\leq t},\mathbf{a}_{\leq t})$  and the prior  $p(\mathbf{z}_t\mid \mathbf{z}_{< t},\mathbf{a}_{\leq t})$  requires extracting relevant information from the history  $(\mathbf{z}_{< t},\mathbf{a}_{\leq t})$ . Therefore, it is crucial to maintain a long-term memory of the history. To this end, we use a stack of PSSM blocks to encode the history  $(\mathbf{z}_{< t},\mathbf{a}_{\leq t})$  into an embedding vector  $\boldsymbol{h}_t$  for each  $t$ . This can be done in parallel during training and sequentially during imagination:
$$
\text {(p a r a l l e l)} \quad h _ {1: T}, s _ {T} = \operatorname {P S S M} _ {-} \operatorname {B l o c k s} \left(g _ {1: T}, s _ {0}\right), \tag {4}
$$
$$
\text {(s i n g l e s t e p)} \quad h _ {t}, s _ {t} = \operatorname {P S S M} _ {-} \operatorname {B l o c k s} \left(g _ {t}, s _ {t - 1}\right). \tag {5}
$$
Here,  $g_{t} = \mathrm{MLP}(\mathrm{concat}[\mathbf{z}_{t - 1},\mathbf{a}_{t}])$  is the input to the PSSM blocks,  $s_t$  is the collection of all internal hidden states, and  $h_t$  is the final output. In our experiments, we use S4 blocks (shown in Figure 1 Right), a particular instantiation of the PSSM blocks using the S4 model. We find that adding the final MLP in each S4 block can improve generation quality (see Appendix B.1 for an ablation and Figure 21 for a detailed illustration of S4 block architecture).
After obtaining  $h_t$ , we use it to compute the sufficient statistics of the prior and likelihood:
$$
p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) = \operatorname {M L P} \left(\boldsymbol {h} _ {t}\right), \tag {6}
$$
$$
p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \mathbf {a} _ {\leq t}\right) = \mathcal {N} \left(\hat {\mathbf {x}} _ {t}, \mathbf {1}\right), \quad \hat {\mathbf {x}} _ {t} = \operatorname {D e c o d e r} \left(\operatorname {c o n c a t} \left[ \mathbf {h} _ {t}, \mathbf {z} _ {t} \right]\right). \tag {7}
$$
For training, we use variational inference. The approximate posterior is defined as:
$$
q \left(\mathbf {z} _ {0: T} \mid \mathbf {x} _ {0: T}, \boldsymbol {a} _ {1: T}\right) = \prod_ {t = 0} ^ {T} q \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {t}\right), \quad \text {w h e r e} q \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right) = p \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right). \tag {8}
$$
We use a CNN encoder to compute the sufficient statistics of the posterior from image observations. This allows all posterior samples  $\mathbf{z}_{0:T}$  to be obtained in parallel, thereby fully leveraging the parallel computation ability offered by PSSMs during training. We also provide an alternative design of the posterior in Appendix B.1. It is conditioned on the full history, and can obtain better generation quality at the cost of more computation.
The training objective is to maximize the evidence lower bound (ELBO):
$$
\log p \left(\mathbf {x} _ {1: T} \mid \mathbf {x} _ {0}, \boldsymbol {a} _ {1: T}\right) \geq \mathbb {E} _ {q} \left[ \sum_ {t = 1} ^ {T} \log p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) - \mathcal {L} _ {\mathrm {K L}} \left(q \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {t}\right), p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right)\right) \right]. \tag {9}
$$
In our experiments, we instantiate S4WM using the S4 model [21]. Figure 1 provides an illustration of training and imagination procedures, and Algorithms 1 and 2 provide detailed descriptions. Hyperparameters and further implementation details can be found in Appendix J.
Unlike previous works [14, 48, 37, 52] that primarily evaluate the final performance of model-free agents on memory-demanding tasks, we seek to understand the memory capabilities of world models
Figure 2: Partially observable 3D (Top) and 2D (Bottom) environments for evaluating memory capabilities of world models, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning.
Figure 3: Comparison of speed and memory usage during training and imagination. S4WM is the fastest to train, while RSSM-TBTT is the most memory-efficient during training and has the highest throughput during imagination.
in model-based agents in terms of long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. We believe that our investigation provides more insights than the final performance alone, and paves the way for model-based agents with improved memory.
To this end, we develop a diverse set of environments shown in Figure 2, each targeting a specific memory capability. The environments are based on the 3D Memory Maze [50] and the 2D Mini-Grid [6], both with partial observations. The world models are learned from an offline dataset collected by a scripted policy for each environment. This allows the world models to be evaluated independently of the policy learning algorithms.
Specifically, for each episode, the environment is regenerated. To simplify the evaluation of world model imagination, we design the data collecting policy to consist of a context phase and a query phase. In the context phase, the policy fully traverses the environment, while in the query phase, the policy revisits some parts of the environment. For evaluation, we use unseen episodes collected by the same policy as training. The world model observes the context phase, and is then evaluated on its imagination given the action sequence in the query phase. Because the environments are deterministic and have moderate visual complexity, and the context phase fully reveals the information of the environment, it suffices to use the mean squared error (MSE) as our main evaluation metric.
In the following, we first motivate our choice of the baselines through a comparison of speed and memory consumption, and then introduce and present the results for each environment in detail.