Facing Off World Model Backbones: RNNs, Transformers, and S4
Authors: Fei Deng
Junyeong Park
Sungjin Ahn
Venue: NeurIPS 2023 poster (2023)
Keywords: world models
structured state space sequence models
S4
long-term memory
model-based reinforcement learning
TL;DR: We propose the first world model compatible with parallelizable SSMs including S4 and its variants, and extensively compare it with RNN- and Transformer-based world models across four sets of environments tailored to assess memory capabilities.

Abstract: 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.

Facing Off World Model Backbones:
RNNs, Transformers, and S4
Fei Deng Rutgers University fei dei.deng@rutgerr... edu
Junyeong Park
Sungjin Ahn" KAIST
KAIST jyp10987@kaist.ac.k5
sungjin aungjin.ahn@kaiss wwn@aaintttaaeeemm
https://fdeng18 github 1o/s4wm
Abstract
World models are 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 possess long-term memory. However, state-of 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) mod- els, 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-dimenn dimensional image sequences through latent imagination. Further- more, 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.
Introduction
The human brain is frequently compared to a machine whose primary function is to construct models of the world, enabling to predict, plan, and react mental representations, referred to as world models, are integral to essential cognitive functions like decision-making and problem-solving. Similarly, systems that aim for human-like cognition is the development of analogous world models.
environment effectively [51, 39]. These
of the pivotal tasks in artificial intelligence (AI)
Model-Based Reinforcement Learning (MBRL) [42] has emerged promising approach that builds world models through interaction with the environment. As a fundamental component 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 well-suited for complex and dynamic environments where model-free methods tend struggle.
"Conresoondence sungj in, ahn@kaist.cm ac. kr.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 consid- erable 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 a significant development since it was unclear whether the S4 model would be effective as a high-dimennn fith-ddmmnionall visual world model, and if so, how this could be achieved. To this end, we instantiate the S4WM with the S4 architecture manage eigh-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 $4. Our results show that S4WM outperforms RNNs and Transformers across multiple memory-demand imagination, context-depender 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.
tasks, including long-term
2 Related Work
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-deeenden 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 $4 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 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-base
[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 ~0 steps. In contrast, S4WM 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-X XL [9].
Agent Memory Benchmarks. While many RL. benchmarks feature partially observable environ- ments, 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]. Addition- ally, some benchmarks are designed for model-free agents [66, 37, 52], and may contain stochastic dynamics that are not controlled by the agents, making 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- -baee agents. We build upon Memory Maze and introduce additional environments and tasks to probe wider range of memory capabilities. Another recent work, TECO [67]. also introduces datasets and a Transformer baasdd model for evaluating and improving long-ter term video prediction. Our work has different focus than TECO in that 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 comparison with TECO in Appendix F.
3 Background
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 1-D output signal y(t). They can be discretized into sequence-to-sequeeee mapping by a step size A. The continuous-ting time and discrete-time SSMs can be described as:
8' (t) = As(t) + Bu(() y() Cs(() + Du(t)'
(discrete-time) 8% = As + Bus
(continuous-time)
(1)
Cs k + Duk
Here, the vectors st) and 8% are the internal hidden state of the SSM, and the discrete-ting matrices A, B, C, D can be computed from their continuous-timt counterparts A, B.C. C, D and the step size A. We will primarily deal with the discrete-time SSMs, which allow efficient autoregressive generation like RNNs due to the recurrence in Sy.
Unlike RNNs, however, linear SSMs can offer parallelizable computation like Transformers. That is, given the input sequence u1.7, the output sequence f1.7 can be computed in parallel across time steps by 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:
(parallel) y1.7 8T = PSSM(u 80)
(single step) yv, 8%- = uu 8k-1)
(2)
where the inputs un and outputs yk can be vectors.
The $4 model aims to use SSMs for deep sequence modeling, where the matrices A, B, C, and the step size A are learnable parameters to be optimized by gradient descent. Because SSMs involve computing powers of A, which is in general expensive and can lead to the exploding/vamishing gradients problem [49], SSMs with randomly initialized A perform very poorly in practice [19].
To address these problems, S4 parameterizes as Diagonal Plus Low-Rank (DPLR) matrix [21, 17]: A A - PP*, where A is diagonal matrix, P is typically column vector (with rank 1), and P is the conjugate transpose of P. This parameterization allows efficient computation of powers of A.
3

---

Training
Imagination
54 Block
Next Step Prediction
Generation
MU
Lavendarm
Ax
(pa allif
Gingle stop!
Gangle step!
Context
Figure 1: propose S4WM, the first S4-hased 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 general framework that is compatible with any parallelizable SSM including S5 and other S4 variants.
Algorithm S4WM Training
Algorithm S4WM Imagination
Input: (Xo, x1; Xi....y QT, XT)
Input: context (Xo.c, a1:c), query aC+1.
Obtain posterior latents
Encode context
1: for time step t = T parallel do
1: for context step t O.....
, C parallel do 2, ~ (z, X₂ Encoder(x,))
2: Z: ~ (zz x,) Encoder(x;)
3: end for
3: gi+1 MLP concat [z, a(+1))
Prepare inputs to blocks
end for
4: for time step 1 1, 1..., T parallel do
5: 11c+1, 8C+1 S4Blocks(g..+4),
5: 9: MI=MMMP(cnnctt/ 12-1, at))
Imagine in latent space
6: end for
6: zc+1 p(zc+1 Zo:c; a1:c+1) = MLP(hc+1)
Encode history by $4 blocks
7: for query step = C. + , do
7: HLT, 8T S4Blocks(gI.T,
8: g: MMLL(ccncat/zz-11
b Compute prior and
9: he, $4
decode posterior latents
10:
2, p(z スc agt) = MLP(h.)
8: for time step t = 1, . parallel do
11: end for
9:
p(z, Z agr)) MLP(h,,
Decode imagined latents
10: X, = Decoder( condderroonctt/a... 2:))
12: for query step 1 = C + 1,..., T parallel do Decoder( (h,, 2:))
11: end for
13:
12: Compute objective by Equation (9)
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 A to the HiPPO matrix. To cope with vector-valued inputs and outputs (i.e., ug yk E IR Hf) S4 makes H copies of the SSM, each operating position-wwi wise linear layer. Follow-up works further simplify the parameterization of A to a diagonal matrix [20], and use multi-input, multi-output SSM for vector-valued sequences [57].
one dimension, and mixes the outputs by
4 S4WM: A General World Model for Parallelizable SSMs
We consider agent interacting with partially observable environment. At each time step the agent receives an image observation x,, It then chooses an action a,+1 based on its policy, and receives the next observation X,+1- For simplicity, we omit the reward here.
We aim to model p(xET Xo, a1:T), the distribution of future observations given the action sequence. We note that it is not required to model p(xo), 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 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 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 probabilistic generative process:
T p(XIT Xo. 01:T) = p(zo xo) II p(x agts agi) Zet,, age) dzo:T
(3)
where Zor are the stochastic latent states. We note that computing the likelihood p(x; Egee ass)) and the prior p(z, Z<t; a-ge) requires extracting relevant information from the history (Z<z, agt). Therefore, is crucial to maintain long-term memory of the history. To this end, we use a stack of PSSM blocks to encode the history (z<t, acz) into embedding vector h, for each t. This can be done in parallel during training and sequentially during imagination:
(parallel)
h1Tr: ST PSSM Blockss(a.2,.8))
(single step)
h PSSM_Blockee BLMM__locss(g.... $t-1)
Here, 9: = MLP(concat/z 21-1, arl is the input to the PSSM blocks, 8, is the collection of all internal hidden states, and h, is the final output. In our experiments, Figure Right), a particular instantiation of the PSSM blocks using the S4 model. We find that adding the final MIP in each S4 block can improve generation quality (see Appendix B.1 for ablation and Figure 21 for detailed illustration of S4 block architecture).
S4 blocks (shown in
After obtaining he,
it to compute the sufficient statistics of the prior and likelihood:
p(z: Zet: age) = MLP(he), p(x;
(6)
Zgt; age) Nxx 1). &, Decoder(connat((h1, z,])
(7)
For training, use variational inference. The approximate posterior is defined as:
((zo:T
Xo:r, a.:))
II g(z xx)), where q(z0) xo) p(zo Xn) t=0
(8)
We use CNN encoder compute the sufficient statistics of the posterior from image observations. This allows posterior samples Zo.T to be obtained 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. quality at the cost of more computation.
conditioned on the full history, and can obtain better generation
The training objective to maximize the evidence lower bound (ELBO):
T
log p(x; lige ast) - LKI. (g(z)) Xx), p(zr [zzc, as:)) (9)
logp(xx:r No, I.)) N E,
Σ
In our experiments, we instantiate S4WM using the S4 model [21]. Figure provides an illustration of training and imagination procedures, and Algorithms and 2 provide detailed descriptions. Hyperparameters and further implementation details can be found in Appendix J.
Experiments
5.1 Environments
Unlike previous works [14, 48, 37, 52] that primarily evaluate the final performance of model-free agents on memory- memory-demanding tasks, seek to understand the memory capabilities of world models
5

---

Two Rooms
Four Rooms
Ten Rooms
Distracting Memory
Multi Doors Keys
Figure 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-bant reasoning.
100k
20K
10
SOUM
1S:M
TESM XL
TSSH XL
RSSMIIBITY
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.
model-based agents terms 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 regenerated. To simplify the evaluation of world model imagination, we design the data collecting policy to consist of a context phase and query phase. In the context phase, the policy fully traverses the environment, while 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 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.
5.2 Baselines
RSSM-TBTT. RSSM [28] is an RNN-based world model backbone used in state-of-the-art MBRL agents [29-31]. Recently, [50] show that training RSSM with truncated backpropagation through time (TBTT) can lead to better long-term memory ability. We follow their implementation and denote the model RSSM-TBTT.

---

Table 1: Evaluation of long-term imagination. Each environment is labeled with (context steps I query steps). All models obtain good reconstruction, while S4WM much better long-term generation up 500 steps. models struggle in the Ten Rooms environment.
Two Rooms (301
200)
Four Rooms (501
(500)
Ten Rooms (1101 900))
Recon. MSE (U)
Gen. MSE (L)
Recon. MSE (L)
Gen. MSE (4)
Recon. MSE (J)
Gen. MSE ())
RSSM-TBIT
1.7
62.2
1.5
219.4
1.5
323.1
TSSM-XL.
2.5
62.9
224.4
2.6
360.4
S4WM
1.7
44.0
1.8
224.4
t- 50
100
150
200
250
300
350
400
450
500
Truth
RSSM
SAWM T9SM
Figure 4: Long-term imagination in the Four Rooms environment. While RSSM-TB and TSSM-XL make many mistakes wall colors, object colors, and object positions, S4WM is able to generate much more accurately, with only minor errors in object positions.
TSSM-XL. TSSM [5] is the first Transformer-basee world model for improving long-term memory. It was originally evaluated on sequences of length ~000. In this paper, we seek to evaluate on much longer sequences (up to 2000 steps), and is impractical to feed the entire sequence the vanilla Transformer [61]. Therefore, we use Transformer-X. XL [9] as the backbone and denote the model as TLSM-XL. It divides the full sequence into chunks, and maintains cache of the intermediate hidden states from processed chunks. This cache serves an extended context that allows modeling longer-term dependencies.
Speed and Memory Usage. We note that the cache length m is a crucial byperparameter of TSSM- XL. A larger can potentially improve the memory capacity, at the cost slower training and higher memory consumption. To ensure a fair comparison in our experiments in the next few sections, we first investigate the speed and memory usage of S4WM, ESSM-TBTT and TSSM-XL XL with several m values. Details can be found in Appendix I. As shown in Figure 3, S4WM and TSSM-XL trains much faster than RSSM-TB due to their parallel computation during training. while RSSM-TBTT is much more memory-efficient. For imagination, RSSM-TBTT achieves ~10 throughput compared to S4WM and TSSM-XL. While S4WM also uses recurrence during imagination, its multi-layered recurrence structure with MI.ss in between slows down its performance. As for memory usage, the decoder takes up most of the memory decoding steps parallel, leading to similar memory consumption of all models.
Based on our investigation, TSSM-XL with cache length m 128 is the closest to S4WM in terms of speed and memory usage. Therefore, we use TSSM-XL with m 128 for all subsequent experiments. We provide more thorough investigation with larger cache lengths in Appendix B.2.
5.3 Long-Term Imagination
The ability of world models to perform long-term imagination is crucial to long-horizon planning. While many RL benchmarks can be tackled with short-term imagination of ~15 steps [3 [31], here we seek to understand the long-term imagination capability of world models and explore their limits by letting the world models imagine hundreds of steps into the future.
7

---

Two Rooms (301 200)
Four Rooms (501 500)
Ten Rooms (1101 900)
Generation MSE
300
600
50
100
200
50
100
150
2.00
125
250
375
500
0
225
450
675
900
Imagination Steps
Imagination Steps
Imagination Steps
54WM
TSSM-XL
RSSM-TBIT
Figure 5: Generation MSE per imagination step. Each environment is labeled with (context steps query steps). S4WM maintains relatively good generation quality for up 500 steps, while RSSM-TBTT and TSSM-XL make large generation errors within 50 steps.
Two Rooms (151 100)
Four Rooms (301 200)
Rooms 400)
550
Generation MSE
450
350
250
15
20
Contexts After Teleport
# Contexts After Teleport
# Contexts After Teleport
S4WM
TSSM-XL
RSSM-TBIT
Figure Evaluation context-dependent recall in teleport environments. Each environment labeled with (context steps query steps). We provide up to 20 observations after the teleport as additional contexts. TSSM- XL performs the best in the Two Rooms environment where the context phase is short, and is able to recall successfully without additional observations. When the context phase is longer, S4WM performs the best.
To this end, we develop three environments with increasing difficulty, namely Two Rooms, Four Rooms, and Ten Rooms, based on the 3D Memory Maze [50]. The top-down views are shown in Figure 2. In the context phase, the data collecting policy starts from random room, sequentially traverses all rooms, and returns to the starting room. the query phase, the policy revisits each room in the same order as the context phase.
As shown in Table I, all models obtain good reconstruction, while S4WM is much better in the Two Rooms and Four Rooms environment for long-tern term generation up to 500 steps. We demonstrate the superior generation quality of S4WM in Figure All models are able to capture the high-level maze layout. However, RSSM-TBTT and TSSM-XL XL make many mistakes in details such as wall colors, object colors, and object positions, while S4WM is able generate much more accurately, with only minor errors in object positions. We further show the per step generation MSE in Figure 5. S4WM able maintain a relatively good generation quality for up 500 steps, while RSSM-TBTT and TSSM-XL make large generation errors even within 50 steps. We notice periodic drop in the generation MSE. This when the agent moves from one room to another through a narrow corridor where the action sequence is less diverse.
We also find that all models struggle the Ten Rooms environment where the context length 1101 and the query length is 900. This likely reaches the sequence modeling limits of the S4 model, and we leave the investigation of more sophisticated model architectures to future work.
5.4 Context- teetext-Deeendent Recall
Humans are able recall past events in great detail. This has been compared to "mental time travel" [60, 58, 37]. Motivated by this, we develop a "teleport" version of the Two Rooms, Four Rooms, and Ten Rooms environments. After the initial context phase, the agent is teleported to random point in history, and is asked to recall what happened from that point onwards, given the exact same action sequence that the agent took.
To succeed in this task, the agent needs figure out where it is teleported by comparing the new observations received after the teleport to its own memory of the past. In other words, the content

---

Table 2: Reward prediction accuracy in the Distracting Memory environments. Each environment is labeled with (context steps query steps). S4WM succeeds in environments. TSSM-XL has limited success when observing the full sequence, but fails to predict rewards within imagination. RSSM-TBTT completely fails.
Width = 100 (199 | 51)
Width = 200 (399 101)
Width = 400 (799 201)
Inference
Imagination Accuracy (1)
Inference Accuracy (1)
Imagination Accuracy (1)
Inference Accuracy (t)
Imagination Accuracy (1)
Accuracy (t)
RSSM-TBTT
47.9% 100.0%
49.6%
48.7%
48.4%
50.4%
52.2%
TSSM-XL
51.2%
99.9%
51.3%
50.4%
51.0%
S4WM
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
32
18
24
27
33
36
39
42
48
Truth
RSSM
54VIM TS5M
Figure 7: Imagination in the Distracting Memory environment of width 100. RSSM-TBTT and TSSM-XL fail keep track of the agent's position, leading inaccurate reward prediction within imagination.
to recall depends on the new observations. Transformers have been shown to be better than S4 at performing such context-dependent operations in low-dimensional sequence manipulation tasks [24] and synthetic language modeling tasks [15]. We investigate this in the context of world model learning, with high-dimensional image inputs.
To help the model retrieve the correct events in history, we provide up to 20 observations after the teleport as additional contexts. The generation MSE of the recall reported in Figure TSSM-XL performs the best in the Two Rooms environment where the context phase is short, and is able to recall successfully without additional observations. When the context phase longer as in Four Rooms and Ten Rooms, S4WM performs the best. We visually show the recall quality with 20 observations after the teleport in Figure 18. In the Two Rooms environment, both TSSM-XL and S4WM are able to recall accurately. However, only S4WM is able to maintain such recall quality in the more challenging Four Rooms environment.
5.5 Reward Prediction
To facilitate policy learning within imagination, world models need to accurately predict the rewards. In this section, we evaluate the reward prediction accuracy over long time horizons. To decouple the challenges posed by 3D environments from long-term reward prediction, use the visually simpler 2D MiniGrid [6] environment.
Specifically, we develop the Distracting Memory environment, which is more challenging than the original MiniGird Memory environment, due distractors of random colors being placed in the hallway. A top-down view is shown in Figure 2. Each episode terminates when the agent reaches one of the squares on the right. A reward of given if the square reached of the same color as the square in the room on the left. Otherwise, no reward is given. In the context phase, the data collecting policy starts in the middle of the hallway, then traverses the hallway and returns to the starting position. In the query phase, the policy goes to of the two squares on the right uniformly at random. To accurately predict the reward, the world model must learn to ignore the distractors while keeping track of the agent's 's position.
We report two types of reward prediction accuracy in Table The inference accuracy is measured when the model takes the full sequence of observations as input (including both the context and the query phases). This evaluates the model's ability to capture long-range dependencies independently of the imagination quality. In contrast, the imagination accuracy is evaluated within the model's 's imagination, conditioned on the observations in the context phase and additionally the action sequence in the query phase.

---

Table 3: Memory-based reasoning in the Multi Doors Keys environments. Each environment is labeled with (context steps query steps). S4WM performs well on all environments, while others struggle.
Three Keys (76
114)
Five Keys (170 380)
Seven Keys (296
65-4)
Recon. MSF (1)
Gen. MSE (1)
Recon. MSF (1)
Gen. MSE (4)
Recon. MSE ())
Gen. MSF ())
RSSM-TBTT
0.05
5.16
0.04
6.36
0.03
6.28
TSSM-XL
0.05
1.27
0.03
6.05
0.02
9.21
S4WM
0.01
0.04
0.02
0.27
0.02
0.10
Our results show that only S4WM able to accurately predict rewards within imagination. TSSM- XL has limited success when observing the full sequence, but fails to imagine future rewards accurately. RSSM-IBTT completely fails, and its reward prediction is close to random guessing. Our visualization of model imagination in Figure 7 reveals that the failure of TSSM-XL and RSSM-TBTT mainly due to their inability to keep track of the agent's ' position.
5.6 Memory-Based Reasoning
In the previous experiments, the model's memory of the environment can largely be kept fixed after the context phase. In this section, we explore the setting where the memory needs to be frequently updated in order to reason about the future.
We develop the Multi Doors Keys environment, where the agent collects keys to unlock doors. A top-down view is shown in Figure 2. Each time a door is unlocked, the corresponding key will be consumed, so cannot be used to unlock other doors of the same color. The agent allowed to possess multiple keys. In the context phase, the agent visits all keys and doors without picking up any keys. In the query phase, the agent attempts to unlock two random doors after picking up each key. After all keys are picked up, the agent will try to unlock each door once again. To successfully predict the outcome when the agent attempts to unlock door, the world model must constantly update its memory when key is picked up or consumed.
Since the environment is visually simple, we find the generation MSE to be good indicator of how well the model predicts the future door states. As reported in Table 3 and visually shown in Figure 19, S4WM performs well all environments, demonstrating its ability to keep updating the memory, while both RSSM-TBTT and TSSM-XL struggle.
6 Conclusion
In this paper, we introduced S4WM, the first PSSM-based visual world model that effectively expands the long-range sequence modeling ability of $4 and its variants from low-dimensional inputs to high-dimensional images. Furthermore, we presented the first comparative investigation of major world model backbones in a diverse set of environments specifically designed to evaluate critical memory capabilities. Our findings demonstrate the superior performance of S4WM over RNNs and Transformers across multiple tasks, including long-term imagination, context-dependent recall, reward prediction, and memory- reasoning.
Limitations and Future Work. We primarily focused on visually simple and deterministic environ- ments to limit the computation cost and simplify the evaluation process. Future work could explore more sophisticated model architectures and proper evaluation metrics for complex and stochastic environments. In addition, we mainly evaluated imagination quality and did not test world models in conjunction with policy learning. Future work could develop and thoroughly test MBRL agents based on S4WM. To demonstrate the potential of S4WM for policy learning, we provide offline probing results in Appendix D and conduct skill-level MPC experiment in Appendix E. We find that S4WM outperforms RSSM in offline probing when it is instantiated with S5, and leads to higher task success rates when used for planning. In Appendix G. we additionally demonstrate that the long-term imagination quality can be further improved by instantiating S4WM with S5, showing the potential of our general S4WM framework for incorporating more advanced parallelizable SSMs.
10

---

Acknowledgments and Disclosure of Funding
This work is supported by Brain Pool Plus Program (No. 2021HID342A0310345 through the National Research Foundation of Korea (NRF) funded by the Ministry of Science and ICT. We thank Jurgis Pašukonis, Danijar Hafner, Chang Chen, Jaesik Yoon, and Caglar Gulcehre for insightful discussions.
References
[1] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[2] Charles Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al. DeepMind Lab. arXiv preprint arXiv: 16/2.03801, 2016.
[3] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 1rXiv:1308.344 2013.
[4] Guy Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Carnegie Mellon University, 1990.
[5] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with Transformer world models. In Deep RL Workshop NeurIPS 2021, 2021.
[6] Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for Gymnasium, 2018. URL https://gittub: com/ htte//arraaaaaaae Fouaaaaaoooounattit//iiiiee
[7] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: /412.3555 2014.
[8] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International Conference on Machine Learning. 2020.
[9] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov: Transformer-XL: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[10] Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: Effi- cient sequence-based RL via state spaces layers. In International Conference on Learning Representations, 2023.
[11] Fei Deng, Ingook Jang, and Sungjin Ahn. DreamerPro: Reconstruction-frer model-based reinforcement learning with prototypical representations. In International Conference on Machine Learning, 2022.
[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[13] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming Transformers for high-resolution image synthesis. In CVPR, 2021.
[14] Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomênech Badia, Gavin Buttimore, Charles Deck, Joel Z Leibo, and Charles Blundell. Generalization of rein- forcement learners with working and episodic memory. Advances Neural Information Processing Systems, 2019.
[15] Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. International Conference Learning Representations, 2023.

---

[16] Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, 2021.
[17] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! Audio generation with state-space models. In International Conference on Machine Learning, 2022.
[18] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems, 2020.
[19] Albert Gu, Isys Johnson, Karan Goel, Khaled Kamal Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous- models with linear state space layers. In Advances in Neural Information Processing Systems, 2021.
[20] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, 2022.
[21] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022.
[22] Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your HiPPO: State space models with generalized orthogonal basis projections. In International Conference on Learning Representations, 2023.
[23] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Advances in Neural Information Processing Systems, 2022.
[24] Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear RNNs. arXiv preprint arXiv:2212.00766 2022.
[25] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, 2018
[26] Jeongsoo Ha, Kyungsoo Kim, and Yusung Kim. Dream to generalize: Zero-shot model-based reinforcement learning for unseen visual distractions. AAAI, 2023.
[27] Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations, 2022.
[28] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, 2019.
[29] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- tions, 2020.
[30] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In International Conference on Learning Representations, 2021.
[31] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint PrXivv2301.04100 2023.
[32] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Advances in Neural Information Processing Systems, 2022.
[33] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
[34] Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. In ECCV, 2022.
12

---

[35]
Arna Kumar Jain, Shivakanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. In Advances in Neural Information Processing Systems, 2022.
[36] David M Knigge, David W Romero, Albert Gu, Efstratios Gavves, Erik Bekkers, Jakub Mikolaj Tomezak, Mark Hoogendoom, and Jan Jan-jakob Sonke. Modelling long range dependencies in ND: From saesispecifiic to general purpose CNN. In International Conference on Learning Representations, 2023.
[37] Andrew Kyle Lampinen, Stephanie C.Y. Chan, Andrea Banino, and Felix Hill. Towards mental time travel: a hierarchical memory for reinforcement learning agents. In Advances in Newral Information Processing Systems, 2021.
[38] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning, arXiv preprint arXiv:2303.00 039882 2023.
[39] Marcelo G Mattar and Máté Lengyel. Planning in the brain. Newron, 10(6):914-934, 2022.
[40] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In International Conference on Learning Representations, 2023.
[41] Vincent Micheli. Eloi Alonso, and François Fleuret. Transformers sample-cfficient world models. In International Conference on Learning Representations, 2023.
[42] Thomas Moerland, Joost Broekens, Aske Plaat, and Catholijn M Jonker. Model-based reinforcement learning: A survey. Foundations and Trends@ in Machine Learning, 16(1):1-118, 2023.
[43] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4ND: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, 2022.
[44] Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning, 2021.
[45] Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. In ICRA, 2021.
[46] Masashi Okada and Tadahiro Taniguchi. DreamingV2 V2: Reinforcement learning with discrete world models without reconstruction. In IROS, 2022.
[47] Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcchre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, 2023.
[48] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing Transformers for reinforcement learning. In International Conference on Machine Learning, 2020.
[49] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference Machine Learning, 2013.
[50] Jurgis Pašukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-ten memory in 3D mazes. In International Conference on Learning Representations, 2023.
[51] Joel Pearson. The human imagination: the cognitive neuroscience of visual mental imagery. Nature Reviews Neuroscience, 20 20(10):624-634 2019.
[52] Marco Pleines, Matthias Pallasch, Frank Zimmer, and Mike Preuss. Memory Gym: Partially observable challenges to memory-based agents. In International Conference on Learning Representations, 2023.
13

---

[53] Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv: 1710.05941, 2017.
[54] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformerrbased world models are happy with 100k interactions. In International Conference on Learning Representa- tions, 2023.
[55] Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. In Advances in Neural Information Processing Systems, 2021.
[56] Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Picter Abbeel. Masked world models for visual control. In CoRL, 2023.
[57] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. International Conference on Learning Representations, 2023.
[58] Thomas Suddendorf, Donna Rose Addis, and Michael C Corballis. Mental time travel and the shaping of the human mind. Philosophical Transactions of the Royal Society B: Biological Sciences, 364(1521)11317-1124, 2009.
[59] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A benchmark for efficient Transformers. In International Conference on Learning Representations, 2021.
[60] Endel Tulving. Memory and consciousness. Canadian Psychology/ssycooooic canadienne, 26 (1):1, 1985.
[61] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, t. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
[62] Jue Wang. Wentao Zhu, Pichao Wang. Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In CVPR, 2023.
[63] Junjie Wang, Yao Mu, Dong Li, Qichao Zhang, Dongbin Zhao, Yuzheng Zhuang. Ping Luo, Bin Wang, and Jianye Hao. Prototypical context-aware dynamics generalization for high- dimensional model-based reinforcement learning. arXiv preprint arXiv:2211.12774 2022.
[64] Jialong Wu, Haoyu Ma, Chaoyi Deng. and Mingsheng Long. Pre-training contextualized world models with mn-the-wild videos for reinforcement learning. arXiv preprint arXiv:2305.18449, 2023.
[65] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Day- Dreamer: World models for physical robot learning. In CoRL, 2023.
[66] Marek Wydmuch, Michat Kempka, and Wojeiech Jaškowski. ViZDoom competitions: Playing Doom from pixels. IEEE Transactions on Games, 11(3):248-25 2598,22992200 2018.
[67] Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel. Temporally consistent Trans- formers for video generation. In International Conference on Machine Learning. 2023.
[68] Chengyang Ying, Zhongkai Hao, Xinning Zhou, Hang Su, Songming Liu, Jialian Li, Dong Yan, and Jun Zhu Reward informed dreamer for task generalization in reinforcement learning. arXiv preprint arXiv:2303..m 05092, 2023.
[69] Linqi Zhou, Michael Poli, Winnie Xu, Stefano Massaroli, and Stefano Ermon. Deep latent state space models for time-series generation. In International Conference on Machine Learning, 2023.
[70] Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented Transformer. arXiv preprint arXiv:2212.081: 36,,, 2022.
14

---

A Environment and Dataset Details
For each 3D environment (i.e., Two Rooms, Four Rooms, and Ten Rooms), we generate 30K trajectories using scripted policy, of which 28K are used for training. 1K for validation, and 1K for testing. For each 2D environment (i.e., Distracting Memory and Multi Doors Keys). we generate 10K trajectories using a scripted policy, of which 8K are used for training, 1K for validation, and 1K for testing. All reported results are obtained from the test trajectories, using the model checkpoints that achieve the best validation loss. The image observations are of size 64 x64x3 for 3D environments, and 40x3 for 2D environments.
B Ablation Study
In this section, investigate alternative architectural choices for S4WM and different cache lengths m for TSSM-XI.. We conduct these ablation studies on the Four Rooms and Ten Rooms environments.
B.1 Alternative Architectures of S4WM
S4WM-Full-Posterior, In our main experiments, we have chosen to use the factorized posterior
((Zrrr Xo:T, (1:T)
II ==
[[z xx)
(10)
for simplicity and parallel training ability. However, note that it is possible to condition on the full history while maintaining the parallel training ability:
T
(Zo.T Xa.T, (.T) II a() X age)) x=0
(11)
We illustrate this architecture Figure 8. Here, we first use a CNN encoder to obtain a deterministic embedding e, for each image observation Xt, and then use stack of S4 blocks to encode the history and compute the sufficient statistics the posterior for all in parallel:
qZ: Xge asi) MLP(h;), hu.T. ST S4Blocks(go.T: 8-1). gt MLP(concatle,, a:l)
(12)
We have defined the dummy action ag and the initial $4 hidden state 8 which are both implemented vectors of all zeros. We note that the S4 blocks in the posterior are not shared with those in the prior.
S4WM-NoMMLP. In our implementation, each S4 block consists of two S4 layers and one MLP. We that the MLP not used in the original S4 [21] for the Long Range Arena tasks [59], but is commonly used in language modeling and audio generation [17]. Hence, we consider a model variant without the MLP in the S4 blocks investigate the importance of this MLP in world model learning.
Results. We report results on the non-teleport Four Rooms and Ten Rooms environments in Table and Figures 9 and 10. Results teleport environments are reported in Figures 11 and 12. We also show a comparison of speed and memory usage in Figure 13. Our results suggest that S4W M-Ful- Posterior performs similarly S4WM on the Four Rooms environment, and becomes better in the more challenging Ten Rooms environment where the episode length is longer. However, it is more computationally demanding than S4WM during training. contrast, while 54WM-No-MLP the most computationally efficient, its performance is much worse than S4WM, indicating the importance of the MLP in S4 blocks to both long-term imagination and context-dependent recall.
15

---

Training
Imagination
Next Step Prediction
Generation
7/2
fg
frz
hu
Prior
S4 Blocks (parallel)
54 Blacks (parallel)
54 Blocks Laingle step)
54 Blocks (single stea)
R2
x2
X,
(paraliei)
Posterior
EO
X
Xa
X
X
Context
Figure 8: Architecture of S4WM-FFull Full-Posterior. It maintains the parallel training ability while computing the posterior from the full history. denotes dummy actions.
Table Comparison of alternative S4WM architectures on long-term imagination. labeled with (context steps query steps). S4WM-FFll-Posterinn is comparable to S4WM on the Four Rooms environment, but is better the more challenging Ten Rooms environment. In contrast, S4WM-No-MI.P performs much worse, suggesting the importance of the MLP in the S4 blocks to long-ter fengtterm imagination.
environment
Four Rooms (501 500)
Ten Rooms (1101 900)
Recon. MSE (L)
Gen. MSE (L)
Recon. MSE (1)
Gen.
MSE (L)
S4WM
1.7
44.0
1.8
224.4
HAWM-Full-Posterree
1.7
38.0
2.0
171.1
S4WM-No-MLP
2.3
68.6
2.5
277.0
16

---

Four Rooms (501 500)
Ten Rooms (1101 900)
450
Generation MSE
100
Ywel
300
50
150
0
125
250
375
500
225
450
675
900
Imagination Steps
Imagination Steps
S4WM
S4WM-Full-Posterior
54WM-No-MLP
Figure 9: Generation MSE per imagination step of alternative S4WM architectures. Each environment labeled with (context steps query steps). S4WM-FulllPPsteriin Posterir performs the best as imagination horizon increases.
100
300
350
450
Truth
54WM
No-MLP Full Post.
S4WM. S4WM-
Figure 10: Long-term imagination from alternative S4WM architectures in the Four Rooms environment. S4WM and S4WM-Full Full- Posttrootr have similar imagination quality in this environment, while S4WM-No- MLP makes more mistakes in wall colors and object positions.
17

---

Four Rooms (301 200)
Ten Rooms (601 400)
Generation MSE
120
400
320
240
10
15
20
10 15
20
# Contexts After Teleport
Contexts After Teleport
S4WM
S4WM-Full Posterior
S4WM-No-MLL
Figure 11: Comparison of alternative S4WM architectures on context-dependent recall in teleport environments. Each environment is labeled with (context steps query steps). We provide up to 20 observations after the teleport additional contexts. S4WM-Full-Posiede short, becomes betler importance the MLP in the S4 blocks to context-dependent recall.
performs similarly S4WM when the context phase the context phase longer. S4WM-No-MLP performs the worst, suggesting
90
105
120
135
150
Truth
S4WM
No-MLP Full Post.
S4WM- S4WM-
Figure 12: Context-dependent recall from alternative S4WM architectures in the Four Rooms teleport environ- ment. observations after the teleport are provided as additional contexts. S4WM and S4WM -uulFull-- Full-Posterio perform similarly in this environment, while S4WM-No -o-MoL makes more mistakes in wall colors and object positions.
Training Episodes per Second
Imagination Frames per Second
Training Memory Usage (G8)
Imagination Memory Usage (GB)
15K
20
5K
10:
S4WM WWW.-Ril-Poaee Full-Pottrrior
54WM-NO-MLP
Figure 13: Comparison of speed and memory usage for alternative S4WM architectures. S4WM-Full-Posterie more computationally demanding than S4WM during training, while P4WM-No-MI. is the most efficient.
18

---

Table 5: Comparison of TSSM-XL with different cache lengths mn long-term imagination. Each environment labeled with (context steps query steps). Larger cache lengths show better generation quality, the cost more computation.
Four Rooms (501
500)
Ten Rooms (1101
900)
Recon. MSE (4)
Gen.
Recon. MSE (L)
Gen. MSE (L)
MSE (1)
TSSM-XL (m 128)
2.4
224.4
2.6
360.4
TSSM-XL (m 256)
2.4
124.4
2.5
295.1
TSSM-XL (m 512)
2.4
38.4
2.5
277.6
Four Rooms (501 500)
Ten Rooms (1101) 9000
600
Generation MSE
300
200
400
100
200
125
250
375
500
225
450
675
900
Imagination Steps
Imagination Steps
TSSM-XL (m = 128)
TSSM-XL (m = 256)
TSSM-XL (m 512)
Figure 14: Generation MSE per imagination step of TSSM-XL with different cache lengths m. Each envi- ronment is labeled with (context steps query steps). Increasing cache length helps improving long-term imagination.
200
250
300
350
400
450
500
Truth
(m= (mm2mmmmm66)) (m= (m= 128)
TSSM TSSM TSSM
Figure 15: Long-term imagination from TSSM-XL with different cache lengths m in the Four Rooms environ-
ment. TSSM-XL 512) works well this environment and is comparable S4WM, while TSSM-XL. (m 128) and TSSM-XI. (m 256) make many mistakes.
19

---

Four Rooms (301 200)
Ten Rooms (601 400)
210
Generation MSE
140
360
240
70
120
# Contexts After Teleport
Contexts After Teleport
TSSM-XL (m 128)
TSSM-LL XL (m 256)
TSSM-XL (m 512)
Figure 16: Comparison of TSSM-XL XL with different cache lengths context-dependen recall in teleport environments. Each environment is labeled with (context steps query steps). We provide up observations after teleport additional contexts, Increasing cache length improves context-dependent recall. sufficiently large cache length, TSSM-XL without observing additional contexts can outperform S4WM.
75
90
105
120
135
150
Truth
(m=512) (m=256) (m= 128)
TSSM
TSSM TSSM
Figure 17: Context-dependent recall from TSSM-XL XL with different cache lengths in the Four Rooms teleport environment. 20 observations after the teleport are provided as additional contexts. TSSM-XL 256) and TSSM-XL (m = 512) perform similarly this environment, and are both better than TSSM-XL (m 128).
B.2 Cache Length of TSSM-XL
our main experiments, we have used the cache length m TSSM-XL, because it is close to S4WM in tenms of computational cost. Here provide more thorough investigation with larger cache lengths.
We report results on the non-teleport Four Rooms and Ten Rooms environments in Table and Fig- ures 14 and 15. Results on teleport environments are reported in Figures 16 and 17. We find that increasing the cache length generally improves generation quality, at the cost of slower training and imagination speed. Notably, TSSM-XL with m 512 shows better context-depente deeeedttd recall than S4WM on the Ten Rooms teleport environment, consistent with the findings in previous work [24, 15] that Transformers are better than $4 at performing context-dependent operations.
20

---

C
Additional Experiment Figures
t=
15
30
45
60
75
t=
30
60
90
120
150
Truth
RSSM
TSSM
S4WM
Two Rooms
Four Rooms
Figure 18: depeexe-deeendent recall in teleport environments. observations teleport are provided as additional contexts. TSSM-XI M- and S4WM perform similarly in the Two Rooms environment, while only S4WM able to maintain good recall quality in the Four Rooms environment.
105
525
S4WM TSSM RSSM Truth
Figure 19: Imagination in the Multi Doors Keys environment with seven keys.
D
Offline Probing on Memory Maze
Pašukonis et al. [50] recently proposed the Memory Maze offline probing benchmark for evaluating the representation learning ability of world models. For completeness, we report the benchmark results in Table 6. Our implementation is based on the newer DreamerV3 [31], and the results of RSSM-TBTT are slightly better than reported in the original Memory Maze paper.
For RSSM-TBTT, TSSM-XI, and S4WM, we use concat [h,, 2,] as the input to the probing network, where h2 is the output of the final Transformer/S4 block for TSSM-XI. and S4WM. The probing network is an MI. with 4 hidden layers, each consisting of 1024 hidden units and followed by layer normalization [1] and SiLU [53] nonlinearity.
Table 6: Memory Maze offline probing benchmark results.
Memory 9x9
Memory 15x15
Memory 9x9 Objects (MSE
Memory 15x
Walls (Aec. 1)
Walls (Acc. 1)
Objects (MSE })
Constant Baseline
80.8%
78.3%
23.9
64.8
RSSM-TBTT
95.0%
81.7%
5.4
32.6
RS
86.4%
79.8%
10.2
35.1
S4WM
88.4%
80.2%
8.9
34.8
S5WM
98.3%
81.8%
1.8
25.3
21

---

Both TSSM-XL and S4WM underperform RSSM-TBTT. We conjecture that this is because in RSSM-TBTT. h, captures more global information, while in TSSM-XL and S4WM, h, captures more local information for each I.
Comparing RSSM-TBTT and S4WM, find that the S4 hidden state 8, plays similar role the h; in RSSM-TBTT by carrying over information from previous time steps. Therefore, we conjecture that 8, will capture more global information and concat [&,, Zt can be more suitable embedding for probing. However, the S4 hidden state 8, has a much higher (typically 64x) dimension than ht, making the above solution impractical.
Fortunately, S4WM is a general framework that works with many S4 variants. In particular. S5 [57] uses hidden state 8; that is of similar dimension h2. Hence, replace the S4 layers in S4WM with S5 layers, and call this model S5WM. We extract the S5 hidden state of the last S5 layer each S5 block, and concatenate them with 2, form the input the probing network. Our results show that this model outperforms RSSM-TBTT by a large margin on the 9x9 maze. Nevertheless, the larger 15x 15 maze remains challenging for all models.
E Skill-Level Model-Predictive Control
To show the potential of S4WM for planning, we consider a skill-level Model-Predictive Control (MPC) agent, with predefined skills and an dffline-trained world model. Specifically, we consider the Multi Doors Keys environments. The predefined skills are (1) picking up a key at specific position, and (2) going to a door at a specific position and attempting to unlock The agent will return to its starting position at the end of each skill. The task is to unlock door specified by a goal image, which shows the door in unlocked state (see Figure 20 for an illustration). The agent is allowed to execute two skills: pick up one key and then attempt to unlock one door. Hence, planning has two stages. First, the agent enumerates all allowed skill sequences (of length 2), and uses the world model to predict which skill sequence is most likely to unlock the correct door. Then the agent executes the first skill the environment and replans. We show the success rate of RSSM-TBTT TSSM-XL. and S4WM in Table 7. The MPC agent equipped with S4WM is able unlock every door in all three environments, while RSSM-TBTT and TSSM-XL succeed roughly half of the time.
Start
Goal
(First Person View)
(Top-down View)
Figure 20: Illustration of the task solved by the skill-level MPC agent. The task is unlock the door specified by the goal image the left. The agent is allowed to execute skills: pick up one key and then attempt unlock one door. We use MPC to plan the optimal skill sequence with the offline-trained world model.
Table 7: Success rate of skill-level MPC agents in the Multi Doors Keys environments.
Three Keys
Five Keys
Seven Keys
RSSM-TBTT
0%
60% 40%
42.86%
TSSM-XL
100%
57.14%
S4WM
100%
100%
100%
22

---

Table 8: Comparison of generation quality DMLab following TECO evaluation protocol.
DMLab
PSNR (1)
SSIM (t)
LPIPS (J)
#Params
CW-VAE [55]
12.6
0.372
0.465
111M
Latent FDM [32]
17.8
0.588
0.222
31M
S4WM
20.6
0.667
0.196
41M
TECO [67]
21.9
0.703
0.157
169M
Table 9: Evaluation of long-term imagination. Each environment is labeled with (context steps query steps).
Two Rooms (301
200)
Four Rooms (501
500)
Ten Rooms (1101
900)
Recon. MSE (4)
Gen. MSE (L)
Recon. MSE (J)
Gen. MSE (J)
Recon. MSE (1)
Gen. MSE (1)
S4WM
1.8
27.3
1.7
44.0
1.8
224.4
S5WM
1.5
10.6
1.3
19.2
1.3
159.2
F
Comparison with TECO
We provide quantitative comparison with TECO [67] and relevant baselines in Table 8. This comparison is done on DMLab following the TECO evaluation protocol. TECO and other baselines use pretrained VQGAN [13] encoder/decoder, while S4WM uses jointly trained simple CNN encoder/decoder (see Appendix J and Table 10 for architecture details and hyperparameters), S4WM significantly outperforms Clockwork VAE [55] and Latent FDM [32] (a diffusion model), and close to TECO despite having much fewer parameters.
G Additional Results of Instantiating S4WM with S5
We show a comparison of long-term imagination quality between S4WM and S5WM in Table 9, where S5WM is our instantiation of S4WM with the S5 model. S5WM demonstrates even better imagination quality than S4WM, showing the potential of our general framework for incorporating more advanced parallelizable SSMs.
H Extended Background
In this section, we briefly introduce RSSM and TSSM for completeness. We denote the sequence of observations and actions as (Xo, a1; X1, 022 X ar, XT). Namely, the agent takes action ar+1 after observing X,, and receives the next observation X We omit the reward for simplicity.
H.1 RSSM
RSSM [28] models the observations and state transitions through the following generative process:
/ II p(x: Zgg; agx) p(z, Zets agt) dzo.T
p aI:T)
(13)
where oo.r the stochastic latent states. The approximate posterior is defined
g(zo.T XorT a1.T)
II g(z Zet agt, Xt)
(14)
The conditioning on previous states Z and actions a- appears multiple times. RSSM uses a shared GRU [7] to compress z< and ac: into deterministic encoding h,:
h, GRU(h MLP(concate 2,-1, a.)))
(15)
23

---

This is then used to compute the sufficient statistics of the prior, likelihood, and posterior:
p(z: Zet, agi) MLP(h;),
(16)
p(x: Rigt, age) (N; 1). X, Decoder(concat (he; ze))
(17)
g(2z Zec, agt, Xt) MLP(concat[hy, erl) er = Encoder(x:) (x+).
(18)
The training objective to maximize the evidence lower bound (ELBO):
T
loga(Xoor (1.1) 2 Bg
Σ
log p(x (og, agi) - LKL (1(2s R21s age, xi), p(ze Be aze)
1=
(19)
H.2 TSSM
Our implementation of TSSM [5] uses the same generative process, approximate posterior, and training objective as S4WM. For convenience, we repeat them below. The generative process
T
p(X1TT Xo: a1:r)= = p(zo xo)
II p(x: Zgt, agi) p(z: Zet, age) dzo.r.
(20)
where Zo.T are the stochastic latent states. The approximate posterior is defined as:
T
g(Zo.T Xo.T, a..) = a( x where g(zo Xo) p((o xo)
(21)
t=
The training objective to maximize the evidence lower bound (ELBO):
1 Xg. a1r)) 2 Eg Σ p(x agg, age) - LKI (oc(992zz (9(2e)2x xx), p(zz Z 2,,, azi))
(22)
The main difference from S4WM is that in TSSM the prior p(z: Zzci, agi) is computed by stack of Transformer [61] blocks. Specifically, the Transformer blocks output an embedding vector he through self-attention over the history:
h TransforneBBlooksggna)) gr ML.P(concat(/x-11..
(23)
The hr is then used for predicting the next latent state Z, and decoding the latent state into image &y:
p(zz Zict, age) = MLP(h,,, & Decoder(concat|/k, 2:))
(24)
I Speed and Memory Usage Evaluation Details
All results in Figure are obtained on single NVIDIA RTX A6000 GPU. We make sure that the models have comparable number of parameters. For training, we use batch size of 8, sequence length 1000, and image size 64x64, We report the number of episodes per second processed by each model, averaged over 100 batches, and also the peak memory usage. For imagination, use a batch size of 64, context length of 500, generation length of 500, and image size of 64x64. We report the number of frames per second, averaged over batches.
J Implementation Details and Hyperparameters
We base our implementation on the publicly available code of S4 [21] and DreamerV3 [31]. We provide the hyperparameters used for 3D and 2D environments in Tables 10 and 11 respectively. For 3D environments, we largely follow the hyperparameters used in Memory Maze [50]. For 2D environments, we follow the architecture DreamerV3-S. We use linear warmup (1000 gradient steps) and cosine anneal learning rate schedule for S4WM. Figure 21 shows the detailed architecture of an S4 block. For TSSM-XI. and RSSM-TBTT, we use constant learning rates following the original papers. To ensure fair comparison, all models the same CNN encoder and decoder,

---

with convolutional and transposed convolutional layers of kernel size 4 and stride 2, layer normaliza- tion [1], and SiLU [53] nonlinearity. The latent states Zo:T are vectors of categorical variables [30] optimized by straight-through gradients [3]. To facilitate stable training, we parameterize the cate- gorical distributions as mixtures of 1% uniform and 99% neural network output [31]. We use KI.
balancing [30] to scale the gradient of the KL. loss LKL (g. p) with respect to the posterior and prior
p:
(q, o- KL(stop_grrd(() p] (1 - a)- || stop-grad(p)]
(25)
Here, stop-grad denotes the stop-gradient operator, and we set a 0.8 to put more emphasis on learning the prior toward the posterior than the other way around.
K Broader Impact
The proposed S4WM and other world models investigated in this paper fundamentally deep gener- ative models. Therefore, they inherit the potential negative social impacts that deep generative models may have, such generating fake images and videos that can contribute to digital misinformation and deception. Caution must be exercised in the application these models, adhering to ethical guidelines and regulations to mitigate the risks.
25

---

Dropout
Linear
Dropout
GELU
Linear
LayerNorm
x2
Dropout
GLU
Linear
Dropout
GELU
S4
LayerNorm
Figure 21: Detailed architecture of an S4 block.
26

---

Table 10: Hyperparameters for environments.
S4WM
TSSM-XL
RSSM-TBTT
Training
Optimizer
AdamW
AdamW
AdamW
Batch size
8
8 57
32
Epochs
57
Learning rate
1 10-3
1 10
3 x 10
Weight decay
1 X 10-2
1 X 10
1 x 10-2
Gradient clipping
1000
1000
200
TBTT steps
48
Model
Stochastic discrete latent size
32 32
32 x 32
32 32
History encoding blocks
6
12
History encoding dmodel
512
512
2048
History encoding du
2048
512
History encoding cache length m
128
256 512
CNN layers
4
4
CNN multiplier
48
48
MLP hidden units
1000
1000
1000
Total parameters
41.1
43.8
54.1
Objective
KL balancing a
0.8
0.8
0.8
Table 11: Hyperparameters for environments.
S4WM
TSSM-XI.
RSSM-TBTT
Training
Optimizer
Adam W
AdamW 8 100
AdamW W
Batch size
8 100
32
Epochs
100
Learning rate
1 1 10 3
1 x 10 -4
3 x 10
Weight decay
1 x 10-2
1 x 10
1 x 10-2
Gradient clipping
1000
1000
200
TBTT steps
50
Model
Stochastic discrete latent size
32 32
32 x 32
32 32
History encoding blocks
6
10
History encoding dmodal
512
512
2048
History encoding da
2048
512
History encoding cache length m
128
CNN layers
3
3
3
CNN multiplier
32
32
MLP hidden units
512
512
512
Total parameters
28.1
27.1 M
31.2 M
Objective
KL balancing a
0.8
0.8
0.8
27


Reviews:
[{"review_id": "jkvyVnCAOz", "reviewer": "Reviewer_hTvu", "summary": "This work aims to better model based RL agents by focusing on developing robust world models. They propose S4WM which uses the S4 model as backbone and incorporates actions as context in the hidden state. They evaluate on prediction tasks for image sequences through latent imagination and show experimental results to compare with other common sequence methods such as RNNs and Transformers.", "questions": "* why did the author choose the games that were used for evaluation? Are there potential nuances in stochastic vs non-stochastic dynamics that would also affect how easily the method can generalize?\n* in terms of long range dependency modeling, why would action-dependent modeling affect the ability to evaluate the memorization capabilities of a model? \n* Is there a set of synthetic tasks that may be a sufficient indicator for the MBRL problem itself?\n* In Figure 4, what makes it harder for other models to reconstruct the correct color information? Is that a latent variable that is difficult to store efficiently in memory? \n* Is it possible to set query steps beyond the context steps to check for some notion of generalization or circularity? Could be interesting to investigate the inductive biases of the recurrence versus other non-recurrent modules.\n* Does good recall in the two rooms correlate to good recall in the four and proportionally in the ten as well?", "limitations": "Not applicable.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* easy application of vanilla S4 SSM without much modification, results appear well in terms of matching colors and scenes in the Four Rooms environment\n* it recognizes and seeks to address the issue of conditional generation of sequences, which is a current limitation of S4 formulated SSMs and an interesting direction to consider\n* results shown make sense, showing higher throughout from recurrent model and memory efficiency for the SSM / Transformer based models\n* paper well-written with few typos and simple summary of related works", "weaknesses": "* Would be good to add to background section for those less familiar with world model learning the common techniques and current state of progress in terms of long term memory\n* I'm also curious why we only do action conditioning -- would it be effective to also evaluate the effects of conditioning on rewards / hallucinated rewards given the access to the emulator in this task setting?\n* This work seems to be a direct application of the S4 framework in the MBRL domain with few comparisons to other SOTA MBRL methods. I also wonder if there needs to be different approaches for conditioning, perhaps a hybrid attention based mechanism instead of direct concatenation and mixing, which has been common in some of the recent conditional diffusion architectures, in order to improve the outcomes of latent space conditioning?\n* It seems as well that as the context length is increase, the trend of MSE across the three models appear more consistent with each other. Is there a threshold where one would catch up and/or overtake the other? If so, can this be another limitation of the method themselves at scale?\n* Are there additional baselines for showing the two / four / ten rooms and how they scale w.r.t. the same combinations of query / context steps? Could be important in ensuring that the results observed are robust and validated in a gradual manner.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "VcNoP9QWzp", "reviewer": "Reviewer_ALhP", "summary": "This paper introduces S4WM, the first world model based on Structured State Space Sequence (S4) models for MBRL agents. Specifically, the proposed world model leverages a stack of S4 blocks for temporal prediction in latent space. CNN encoder and decoder are used to map the image to the latent space. Compared to RNNs and Transformers, S4WM demonstrates superior long-term memory.", "questions": "- As mentioned in the weakness section, the impact of the improved world model on MBRL algorithms.\n- It seems that during training, the model only performs one-step prediction. I wonder whether multi-step prediction during training can further benefit the generation process.", "limitations": "The limitation stated is the lack of evaluation in more complex and stochastic environments. ", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The method extends the application of S4 in world models, and demonstrates better performance than RNN and transformer-based world models, especially in capturing long-term memory.\n- The authors make full comparisons among RNNs, Transformers, and S4WMs in different environments.", "weaknesses": "- The contribution of the proposed backbone is incremental as it heavily relies on S4.\n- The efficiency of the proposed method in the imagination process is worse than that of the widely used RSSM, which will bring additional overhead for MBRL agents.\n- As the proposed method aims to provide a better world model for MBRL agents, a natural question is how the  ability to capture long-term dependency can benefit MBRL agents. The authors may evaluate  advanced MBRL algorithms with your improved world model.\n- The experiments are all conducted in visually simple environments. I am curious whether the proposed method can handle more complex dynamics, such as robot manipulation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "7VYyQlV5Jw", "reviewer": "Reviewer_TTTx", "summary": "This paper studies the backbones of world models. An S4-based world model S4WM is proposed by replacing RNNs in RSSM with S4 Blocks, which is known to be good at long-term memory. An empirical study has been conducted to compare three major backbones, RNN, Transformers, and S4, on speed, memory usage, and performance of multiple memory-demanding tasks.", "questions": "Only one minor question:\n\n1. The large memory usage of RSSM for imagination (Fig 3d) is attributed to the image decoder. Why does not the decoder also make a high memory usage in RSSM for training (Fig 3c)?", "limitations": "This work has discussed its limitations and future work in the conclusion. I recommend discussing weaknesses 1 and 2 mentioned above in the part on limitations.\n\nThere does not seem to be any negative social impact of this paper that should be discussed. ", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Introduction of S4 into the field of world models and MBRL is appreciated\n2. Carefully designed memory-demanding tasks and extensive experiments\n3. This paper is well-written and easy to follow", "weaknesses": "1. Eq. (6), in order to obtain parallel trainability, approximates posterior only based on the current observation, regardless of previous latent states. This approximation, as pointed out in the paper of TSSM, is myopic and may hurt memory ability, despite satisfactory empirical performance in evaluated environments. \n2. This work does not evaluate MBRL equipped with S4WM on memory-demanding tasks, such as DMLab.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "EgpNPvQB0H", "reviewer": "Reviewer_vwzx", "summary": "Model-Based RL algorithms rely on learning efficient world models for learning optimal behaviors. Existing agents like Dreamer rely on an RNN for learning world models which have limited memory capacity. S4 has shown promising results when compared with RNNs and Transformers in terms of both complexity and efficiency on various tasks. This work proposes a method called S4WM that uses S4 as the backbone of world models. The method is proposed for high-dimensional inputs (images) under a probabilistic learning framework and uses variational inference for learning. The architecture seems similar to Dreamer, where the deterministic component (RNN) in Dreamer is replaced with a stack of S4 layers. The paper also proposes a few benchmarks for evaluation in an offline learning setting and shows that S4WM outperforms RNN and Transformers-based baselines.", "questions": "1. Line 171 in the paper mentions that- “ We make sure that the models have comparable number of parameters.”. What were the hyperparameters used for RSSM/TSSM-XL/S4WM to have a similar number of parameters? Was the same number of hyperparameters only used to run the validation for memory usage and speed (in Fig 3) or for other experiments too? \n2. There have been more recent works on improving RSSM to have better memory capacity- Clockwork VAEs [1] and VSG [2]. How well does S4WM compare with such methods? Clockwork VAEs did not do the action-conditioned generation of frames but should be mentioned in related works, however, they are used in [4] for evaluation. \n3. Although the paper proposed benchmarks for evaluation, the methods should be evaluated on offline datasets used in [3]. Section 5 discusses some offline experiments and uses probing for evaluation. Rather than using MSE as the evaluation metric, it would be interesting to observe the results of probing to test the quality of learned embeddings. Is there a way to do probing for the benchmarks used in the paper?\n4. TECO [4] is another baseline that uses Transformers for action-conditioned video generation and reported results on benchmarks like Habitat and Minecraft, and should be added for comparison.\n5. How do these methods compare when only a few initial frames (maybe 20) are provided and open-loop prediction is done conditioned on them?\n6. Is there a way to do some kind of probing to evaluate the embeddings of the world models rather than using MSE?\n7. The details of the S4 block are not mentioned in the paper. If there are multiple layers, was an experiment conducted to study the effect of the number of layers in the S4 block?\n \\\n### References\n[1] Saxena, Vaibhav, Jimmy Ba, and Danijar Hafner. \"Clockwork variational autoencoders.\" Advances in Neural Information Processing Systems 34 (2021): 29246-29257.\\\n[2] Jain, Arnav Kumar, et al. \"Learning robust dynamics through variational sparse gating.\" Advances in Neural Information Processing Systems 35 (2022): 1612-1626.\\\n[3] Pasukonis, Jurgis, Timothy Lillicrap, and Danijar Hafner. \"Evaluating Long-Term Memory in 3D Mazes.\" arXiv preprint arXiv:2210.13383 (2022).\\\n[4] Yan, Wilson, et al. \"Temporally consistent video transformer for long-term video prediction.\" arXiv preprint arXiv:2210.02396 (2022).", "limitations": "The only limitation mentioned in the paper is to extend it to more complex environments. More limitations of the current algorithm should be discussed. For example, the paper needs to discuss how to extend this to online learning and if there will be any bottlenecks when moving to harder tasks or stochastic environments.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The problem statement of improving the memory capacity of world models while being computationally efficient by deploying S4 layers is interesting.\n2. The experiments thoroughly study different aspects to test the ability of proposed methods on tasks like reward prediction, memory, and reasoning.", "weaknesses": "1. Although the paper introduces new benchmarks, it is important to evaluate methods on the existing benchmarks (Questions 3, 4).\n2. Some of the baselines and related works are missing in the paper (Question 2) which focused on a similar problem of improving the long-term memory of world models.\n3. Section 4 is not well explained. The paper should highlight the differences with the Dreamer architecture. From my understanding, S4 layers are used to replace the GRU in the deterministic component.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "WiqnKvANIv", "reviewer": "Reviewer_kwZk", "summary": "This paper presents a novel architecture for world-models in deep RL, which is based on S4, a generative sequence model. Furthermore, the paper presents an evaluation to two other world-model architectures: RSSM's and transformer-based world-models.", "questions": "- I am willing to reconsider my evaluation score if there is some discussion of the included reference (Yan et al., 2022) and how this study sets itself apart.\n- Equation 5: This equation states that $h_{t-1}$ is fed into the S4 block. Figure 1. however does not show this. ", "limitations": "The limitation of the benchmark environment simplicity was acknowledged.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Integrating more powerful models into existing world-model algorithms can lead to more performant and easier to train algorithms. Hence, the general idea is valid.\n- The benchmark contribution sheds some light on which models are useful as a backbone for world-model RL and is certainly beneficial for the RL community.\n- The experimental evaluation is insightful and appears to be of high-quality. Especially useful are the throughput comparisons during training and inference. ", "weaknesses": "- The idea to use novel sequence model architectures as backbones for world-models in RL is somewhat incremental. \n- There is already another (more extensive) benchmark in video prediction models [(Yan et al., 2022)](https://arxiv.org/pdf/2210.02396.pdf). While this study does not explicitly target world-models, it contains models that have been applied as such. Furthermore, those models are also action-conditioned, which means they should be able to serve as world-models as well. This paper misses this, quite important, reference.\n- The benchmark lacks some depth and detail. While the general setup, environment selection and metrics are suitable, it would be great to have some more direct insight. For instance, show also the path that was taken during trajectory reconstruction. This would make it easier to understand how the agent's path influences the prediction. Also, it would give some insight on how the reconstruction behaves on loop-closures (revisiting old locations) and identical looking positions. In general, different types of runs (loops, straight lines, etc) would benefit the evaluation.\n- The approach is only tested only on offline data. While the rationale behind it is clear (comparing w.r.t model quality only), i think it would still make sense to use it actually as a world model for RL. I acknowledge that training world-model RL algorithms is not easy and can be brittle. However, it not necessary to present SOTA results for some world in this regard, just to have some evidence and insight on how well it actually does in the task for this model was developed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}]

Author rebuttal:
We thank all reviewers for their insightful and constructive feedback. We are encouraged that they appreciate our contribution of introducing S4 into the field of world models to improve long-term memory capacity (@TTTx, @vwzx). Many reviewers praised our thorough and high-quality experiments that evaluate several aspects of world models across different backbones and environments (@ALhP, @TTTx, @vwzx, @kwZk), which can be beneficial for the RL community. We address the main concerns from reviewers below.

> Evaluate MBRL equipped with S4WM (@ALhP, @ TTTx, @ kwZk)

We appreciate the suggestion. With S4WM, we aim to make a step towards tackling environments with long-range dependencies. Such environments entangle several challenges such as exploration, long-horizon planning and skill learning in addition to the world model learning. An agent with any of these pieces missing may fail in such environments. Therefore, successfully training an MBRL agent from scratch in such environments is extremely difficult, even if the agent is equipped with better world models such as S4WM.

To address the reviewer concern, we instead consider a skill-level Model-Predictive Control (MPC) agent, with predefined skills and an offline-trained world model. This will show whether the world model can benefit RL tasks through planning. Specifically, we consider the Multi Doors Keys environments. The predefined skills are (1) picking up a key at a specific position, and (2) going to a door at a specific position and attempting to unlock it. The agent will return to its starting position at the end of each skill. The task is to unlock a door specified by a goal image, which shows the door in unlocked state. The agent is allowed to execute two skills: pick up one key and then attempt to unlock one door. Hence, planning has two stages. First, the agent enumerates all allowed skill sequences (of length 2), uses the world model to predict which skill sequence is most likely to unlock the correct door. Then the agent executes the first skill in the environment and replan. We show the success rate of RSSM, TSSM, and S4WM in Table 1 of upload pdf. The MPC agent equipped with S4WM is able to unlock every door in every environment, while RSSM and TSSM succeed roughly half of the time.

> Compare with additional baselines such as TECO (@hTvu, @vwzx, @kwZk)

We thank reviewers for pointing to this relevant paper, which is published at ICML 2023. We note that TECO considers more visually complex environments, with sequence length capped at 300. Therefore, TECO uses the quite heavy VQGAN encoder/decoder, and a large model of 169M parameters. In contrast, we consider extremely long sequences (up to 2000 steps), where using heavy models becomes too costly. For this reason, we mainly considered visually simpler environments that can be tackled with lightweight models (S4WM has 41M parameters).

Following the reviewer suggestions, we provide a quantitative comparison with additional baselines in Table 2 of uploaded pdf. This comparison is done on DMLab following the evaluation protocol of the TECO paper. All baselines use a pretrained VQGAN encoder/decoder, while S4WM uses a jointly trained simple CNN encoder/decoder (we use the same hyperparameter as our 3D experiments). S4WM significantly outperforms Clockwork VAE and Latent FDM (a diffusion-based model), and is close to TECO.

We further showcase the training efficiency of S4WM in Table 3 of uploaded pdf. Here we train both S4WM and TECO for only 80K steps in the Two Rooms environment (in Table 2, the models are trained for 1M steps). We find that S4WM performs well while TECO fails to learn as quickly as S4WM. This suggests that S4WM may better suit online MBRL agents, where the world model needs to constantly adapt to the data collected by the agent.

> The contribution of combining S4 with world models is somewhat incremental (@ALhP, @kwZk)

We acknowledge that the proposed model S4WM is not completely new, but a novel combination of existing components. In addition to proposing S4WM as the first S4-based world model, we also contribute an extensive empirical study of the three major world model backbones in terms of several important memory capabilities evaluated in carefully designed environments. We believe that this thorough investigation presents new knowledge and can be valuable to the community.

Meta-review: This paper proposes to use S4 as a backbone for world models. Authors compare RNN, Transformers, and S4 on various prediction tasks based on the learnt world model and show that S4 is superior to both RNNs and Transformers.

The authors have addressed most of the reviewers' concerns during the rebuttal. There was a question of applying the proposed world model to MBRL. The authors added some MPC results in the rebuttal. However, I don't think it is very convincing application of the proposed world model. Hence, I suggest the authors make it explicit in the abstract that this is not a paper about using S4 for MBRL. It is only about using S4 for supervised learning of the world model.
Decision: Accept (poster)