Title: Facing Off World Model Backbones: RNNs, Transformers, and S4
Venue: NeurIPS 2023
Accepted: yes
Keywords: world models

structured state space sequence models

S4

long-term memory

model-based reinforcement learning
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:

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.

1 Introduction

The human brain is frequently compared to a machine whose primary function is to construct models of the world, enabling us to predict, plan, and react to our environment effectively [51, 39]. These mental representations, referred to as world models, are integral to essential cognitive functions like decision-making and problem-solving. Similarly, one of the pivotal tasks in artificial intelligence (AI) systems that aim for human-like cognition is the development of analogous world models.

Model-Based Reinforcement Learning (MBRL) [42] has emerged as a promising approach that builds world models through interaction with the environment. As a fundamental component of MBRL, these world models empower artificial agents to anticipate the consequences of their actions and plan accordingly, leading to various advantages. Notably, MBRL offers superior sample efficiency, mitigating the high data requirements commonly associated with model-free methods. Moreover, MBRL exhibits enhanced exploration, transferability, safety, and explainability [42], making it well-suited for complex and dynamic environments where model-free methods tend to struggle.

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 highdimensional, unstructured, and temporal observations.

Nevertheless, choosing the appropriate backbone architecture for visual MBRL has become a considerable challenge due to the rapidly evolving landscape of deep architectures for temporal sequence modeling. This includes the recent advancements in major backbone architecture classes, notably Transformers [61, 5] and the Structured State Space Sequence models such as S4 [21] and S5 [57].

Traditionally, Recurrent Neural Networks (RNNs) [7, 33] have been the go-to backbone architecture, thanks to their efficient use of computational resources in processing sequential data. However, RNNs tend to suffer from vanishing gradient issues [49], limiting their long-term memory capacity. Recently, Transformers [61] have demonstrated superior sequence modeling capabilities in multiple domains, including natural language processing and computer vision [12]. Their self-attention mechanism grants direct access to all previous time steps, thereby enhancing long-term memory. Moreover, Transformers offer parallel trainability and exhibit faster training speeds than RNNs. However, their quadratic complexity and slow generation speed pose challenges when dealing with very long sequences. To address these limitations, the S4 model has been proposed, offering both parallel training and recurrent generation with sub-quadratic complexity. In the Long Range Arena [59] benchmark consisting of low-dimensional sequence modeling tasks, the S4 model outperforms many Transformer variants in both task performance and computational efficiency.

In this paper, we present two primary contributions. Firstly, we introduce S4WM, the first and general world model framework that is compatible with any Parallelizable SSMs (PSSMs) including S4, S5, and other S4 variants. This is a significant development since it was unclear whether the S4 model would be effective as a high-dimensional visual world model, and if so, how this could be achieved. To this end, we instantiate the S4WM with the S4 architecture to manage high-dimensional image sequences, and propose its probabilistic latent variable modeling framework based on variational inference. Secondly, we conduct the first empirical comparative study on the three major backbone architectures for visual world modeling-RNNs, Transformers, and S4. Our results show that S4WM outperforms RNNs and Transformers across multiple memory-demanding tasks, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. In terms of speed, S4WM trains the fastest, while RNNs exhibit significantly higher imagination throughput. We believe that by shedding light on the strengths and weaknesses of these backbones, our study contributes to a deeper understanding that can guide researchers and practitioners in selecting suitable architectures, and potentially inspire the development of novel approaches in this field.

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-dependent operations. Therefore, hybrid architectures have been proposed to achieve the best of both worlds. Furthermore, S4 and its variants have found applications in various domains, such as image and video classification [43, 36, 34, 62], audio generation [17], time-series generation [69], language modeling [70, 15, 40], and model-free reinforcement learning [10, 38]. Our study introduces the first world model compatible with S4 and its variants (more generally, parallelizable SSMs) for improving long-term memory in MBRL. We also investigate the strengths and weaknesses of S4 and Transformers in the context of world model learning.

World Models. World models [25] are typically implemented as dynamics models of the environment that enable the agent to plan into the future and learn policies from imagined trajectories. RNNs have been the predominant backbone architecture of world models. A notable example is RSSM [28], which has been widely used in both reconstruction-based [29-31, 65, 56, 16, 35, 63, 64, 68] and reconstruction-free [44-46, 11, 26] MBRL agents. With the advent of Transformers [61], recent works have also explored using Transformers as the world model backbone [5, 41, 54]. While Transformers are less prone to vanishing gradients [49] than RNNs, their quadratic complexity limits their applicability to long sequences. For example, recent works [41, 54] use a short imagination horizon of ∼ 20 steps. In contrast, S4WM can successfully imagine hundreds of steps into the future with sub-quadratic complexity. We also develop an improved Transformer-based world model that can deal with long sequences by employing Transformer-XL [9].

Agent Memory Benchmarks. While many RL benchmarks feature partially observable environments, they tend to evaluate multiple agent capabilities simultaneously [2, 8, 27] ( e.g. , exploration and modular skill learning), and may be solvable with a moderate memory capacity [14, 48]. Additionally, some benchmarks are designed for model-free agents [66, 37, 52], and may contain stochastic dynamics that are not controlled by the agents, making it hard to separately assess the memory capacity of world models. The recently proposed Memory Maze [50] focuses on measuring long-term memory and provides benchmark results for model-based agents. We build upon Memory Maze and introduce additional environments and tasks to probe a wider range of memory capabilities. Another recent work, TECO [67], also introduces datasets and a Transformer-based model for evaluating and improving long-term video prediction. Our work has a different focus than TECO in that we stress test models on extremely long sequences (up to 2000 steps), while TECO considers more visually complex environments, with sequence lengths capped at 300 . Our experiment setup allows using relatively lightweight models to tackle significant challenges involving long-term memory. We include a comparison with TECO in Appendix F.

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 a 1 -D output signal y ( t ) . They can be discretized into a sequence-to-sequence mapping by a step size ∆ . The continuous-time and discrete-time SSMs can be described as:

Here, the vectors s ( t ) and s k are the internal hidden state of the SSM, and the discrete-time matrices A , B , C , D can be computed from their continuous-time counterparts A , B , C , D and the step size ∆ . We will primarily deal with the discrete-time SSMs, which allow efficient autoregressive generation like RNNs due to the recurrence in s k .

Unlike RNNs, however, linear SSMs can offer parallelizable computation like Transformers. That is, given the input sequence u 1: T , the output sequence y 1: T can be computed in parallel across time steps by a discrete convolution [21] or a parallel associative scan [57, 4]. In this work, we define the class of Parallelizable SSMs (PSSMs) to be the SSMs that provide the following interface for both parallel and single-step computation:

where the inputs u k and outputs y k can be vectors.

The S4 model aims to use SSMs for deep sequence modeling, where the matrices A , B , C , D and the step size ∆ 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/vanishing gradients problem [49], SSMs with a randomly initialized A perform very poorly in practice [19].

To address these problems, S4 parameterizes A as a Diagonal Plus Low-Rank (DPLR) matrix [21, 17]: A = Λ -PP ∗ , where Λ is a diagonal matrix, P is typically a column vector (with rank 1 ), and P ∗ is the conjugate transpose of P . This parameterization allows efficient computation of powers of A ,

<latexi sh

1_b

64="7

OMEQ

I3

5

U0

u

ZNF

/grL

>A

c

VD

S

J

2vW

d

k

w

R

y

KCf

T

m

Y

H

p

X

P

n

9

j

G

8

q

o

z

<latexi sh

1_b

64="Q

d/u7

WPL3U+yc

9T

IXE

Cw

>A

n

VD

S

N

F

2pr

f

Z

Bovgq

0

m

M

R

O

K

Y

H

5

J

z

8

j

G

k

<latexi sh

1_b

64="F

wmMOgz

CB

NK+XW

>A

/n

c

V

9S

EJ2L

R

o

7

G

U

8

H

fY

3

jd0

I

Zr

k

5

P

u

y

p

D

Q

<latexi sh

1_b

64="KEQ

Hw0

I5XjTM/

kLR+

27A

>

C

n

c

VD

S

N

F

pr

f

ZduBovgq

U

m

O

Y

9

W

J

z

3

8

P

G

y

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="K

pM

I9

kRGF

H

J

w

>A

C

n

c

VDL

N

3

rPU

YB

W

d

q

Z

E

m

0

u

O

7j

5y

z

g

/

2

v

o

X

Q

8

<latexi sh

1_b

64="f3gk

JW

rvK2

m7c

Dj

GA

>

C

n

V

LS

N

F

p

ZduBo

q

T

w0

U

+

M

R

O

/

Y

H

5

9

Q

z

8

E

X

P

y

I

<latexi sh

1_b

64="T

G72gf

V

uSj

/DUdv5

>

B

n

L

N

EOy

r

Y9

oP

K

X8

m

k

zC

+

pH

J

W

Q

3

F

0

M

I

q

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="

0

gk qGo

S

7K

/C+zuQ5

>

B

n

c

VDL

N

EOy

r

fU

9

P

w

X8

jm

d

Z

v

pH

RJ

TW

3

F

2

M

I

<latexi sh

1_b

64="zdkSR

7g

XfJ

yM

>A

C

n

c

VD

F

3

W

ur

8

Q

H9

E

p

O2

UvI

m

w

Y

K0

Z

/

5

j

B

P

+

T

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="Q

d/u7

WPL3U+yc

9T

IXE

Cw

>A

n

VD

S

N

F

2pr

f

Z

Bovgq

0

m

M

R

O

K

Y

H

5

J

z

8

j

G

k

<latexi sh

1_b

64="KEQ

Hw0

I5XjTM/

kLR+

27A

>

C

n

c

VD

S

N

F

pr

f

ZduBovgq

U

m

O

Y

9

W

J

z

3

8

P

G

y

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="f3gk

JW

rvK2

m7c

Dj

GA

>

C

n

V

LS

N

F

p

ZduBo

q

T

w0

U

+

M

R

O

/

Y

H

5

9

Q

z

8

E

X

P

y

I

<latexi sh

1_b

64="K

pM

I9

kRGF

H

J

w

>A

C

n

c

VDL

N

3

rPU

YB

W

d

q

Z

E

m

0

u

O

7j

5y

z

g

/

2

v

o

X

Q

8

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

<latexi sh

1_b

64="T

G72gf

V

uSj

/DUdv5

>

B

n

L

N

EOy

r

Y9

oP

K

X8

m

k

zC

+

pH

J

W

Q

3

F

0

M

I

q

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="

0

gk qGo

S

7K

/C+zuQ5

>

B

n

c

VDL

N

EOy

r

fU

9

P

w

X8

jm

d

Z

v

pH

RJ

TW

3

F

2

M

I

<latexi sh

1_b

64="Cg/w

MY

GvfOE7I

yS

A

>

n

c

VDL

N

FJ3

rPU

9

B

T

u

m

2

0ok+

H

z

Q

5

d

q

p

K

W

R

j

X

Z

8

<latexi sh

1_b

64="z87FV2dUc

g

BR0q

Y

X

o

>A

3

DLS

M

j

9Z

W

nfy

K

v

Gk mE5

I

+

r

N

Hu

J

/O

C

w

P

T

Q

p

<latexi sh

1_b

64="PfoX

O

N

Kg+j

E

>A

B/n

c

VDLS

y

r

UY9

85

C

Z

k

M7v

p

W

m3

J0

q

T

F

d

Q

R

H

G

I

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="A8C

gr

MR

P/y

V+

WzoBY

>

n

c

DLS

N

EO

fU

9

K

5

j

Zwu

k

7v

p

m3

J0

q

T

X

F

d

Q

H

G

I

2

<latexi sh

1_b

64="

/DuQ

T

r

H9IvnFwE

W

>A

C

c

V

LS

N

J3U

2

d

Bo gq

0

+

m

M

Y

Z

f

X7

P

G

K

j

kR

8

5

z

O

y

p

<latexi sh

1_b

64="n

Y/2

+5

IpR

P

Qm

8KySf

3

>A

B

c

VDL

gM

F

r

Zdu

k

w

W

E

7

9

G

o

T

J

q

ONC0X

j

z

U

v

H

<latexi sh

1_b

64="z87FV2dUc

g

BR0q

Y

X

o

>A

3

DLS

M

j

9Z

W

nfy

K

v

Gk mE5

I

+

r

N

Hu

J

/O

C

w

P

T

Q

p

<latexi sh

1_b

64="I

m

H

U

fyjRKWC/

T

M+

N

>A

B

n

c

VDLSg

EO

r

Y9

oP

w

X

Q

GF2

p

Z

d

z

7

8

5

k

3

0

u

J

q

v

<latexi sh

1_b

64="rqUKgmMVy/v nTc2+k3Z

JQ

>A

B

DLS

N

EO

f

Y9

oP

w

X

GF

p

d

Cz

7

8

R

5

W

I

0

u

j

H

<latexi sh

1_b

64="k

Bz

Z

d

fIG

M

ycK

rXQ

>A

/H

VDLSgN

EO

UY9

oP

2

R

J

j

m

v

3

0

qp

u

T

+

F7

nw

5

C

8

W

<latexi sh

1_b

64="Q

uUFk

Y

Xjz

7q

T

H+

Lc

>A

C

VD

SgM

3

9Z

W

BE

o

yIr2

p

/

J

G

f

N

n

P

5

Od

w

0

R

m

K

v

8

<latexi sh

1_b

64="Q

uUFk

Y

Xjz

7q

T

H+

Lc

>A

C

VD

SgM

3

9Z

W

BE

o

yIr2

p

/

J

G

f

N

n

P

5

Od

w

0

R

m

K

v

8

<latexi sh

1_b

64="Q

uUFk

Y

Xjz

7q

T

H+

Lc

>A

C

VD

SgM

3

9Z

W

BE

o

yIr2

p

/

J

G

f

N

n

P

5

Od

w

0

R

m

K

v

8

Training

Context

Imagination

(parallel)

ˆ

x

1

h

1

x

0

Enc

z

0

a

1

ˆ

x

2

h

2

x

1

Enc

z

1

a

2

x

2

Enc

z

2

Prior

Posterior

(parallel)

x

0

Enc

z

0

a

1

x

1

Enc

z

1

a

2

ˆ

x

2

h

2

(single step)

a

3

ˆ

x

3

h

3

(single step)

a

4

ˆ

x

4

h

4

Generation

· · ·

Next Step Prediction

Dec

Dec

Dec

Dec

Dec

z

1

z

2

z

2

z

2

z

3

z

3

z

4

+

LayerNorm

Linear

LayerNorm

Linear

+

LayerNorm

MLP

+

Algorithm 1 S4WM Training

Input:

(

x

0

,

a

1

,

x

1

, . . . ,

a

T

,

x

T

)

▷

Obtain

posterior

latents

1: for time step t = 0 , . . . , T parallel do

2: z t ∼ q ( z t | x t ) = Encoder ( x t )

3: end for ▷ Prepare inputs to S4 blocks

4: for time step t = 1 , . . . , T parallel do

5: g t = MLP ( concat [ z t - 1 , a t ])

6: end for

▷

Encode

history by

blocks

7: h 1: T , s T = S4Blocks ( g 1: T , s 0 ) ▷ Compute prior and decode posterior latents

8: for time step t = 1 , . . . , T parallel do

9: p ( z t | z <t , a ≤ t ) = MLP ( h t )

10: ˆ x t = Decoder ( concat [ h t , z t ])

11: end for

12: Compute objective by Equation (9)

Algorithm 2 S4WM Imagination

Input: context ( x 0: C , a 1: C ) , query a C +1: T

▷

Encode

context

1: for context step t = 0 , . . . , C parallel do

2: z t ∼ q ( z t | x t ) = Encoder ( x t )

3: g t +1 = MLP ( concat [ z t , a t +1 ])

4: end for

5: h 1: C +1 , s C +1 = S4Blocks ( g 1: C +1 , s 0 ) ▷ Imagine in latent space

6: z C +1 ∼ p ( z C +1 | z 0: C , a 1: C +1 ) = MLP ( h C +1 )

7: for query step t = C +2 , . . . , T do

8: g t = MLP ( concat [ z t - 1 , a t ])

9: h t , s t = S4Blocks ( g t , s t - 1 )

10: z t ∼ p ( z t | z <t , a ≤ t ) = MLP ( h t )

11:

end for

▷

Decode

imagined latents

12: for query step t = C +1 , . . . , T parallel do

13: ˆ x t = Decoder ( concat [ h t , z t ])

14: end for

while also including the HiPPO matrices [18], which are theoretically derived based on continuoustime 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. , u k , y k ∈ ❘ H ), S4 makes H copies of the SSM, each operating on one dimension, and mixes the outputs by a position-wise linear layer. Follow-up works further simplify the parameterization of A to a diagonal matrix [20], and use a multi-input, multi-output SSM for vector-valued sequences [57].

4 S4WM: A General World Model for Parallelizable SSMs

We consider an agent interacting with a partially observable environment. At each time step t , the agent receives an image observation x t . It then chooses an action a t +1 based on its policy, and receives the next observation x t +1 . For simplicity, we omit the reward here.

We aim to model p ( x 1: T | x 0 , a 1: T ) , the distribution of future observations given the action sequence. We note that it is not required to model p ( x 0 ) , as world model imagination is typically conditioned

on at least one observation. While S4 and its variants have shown remarkable abilities to model long-range dependencies, they operate directly in the observation space. For example, S4 models images as sequences of pixels, and directly learns the dependencies among individual pixels. This is hard to scale to high-dimensional sequences, such as the sequences of images that we aim to model.

Inspired by RSSM [28], we propose S4WM, the first PSSM-based world model that learns the environment dynamics in a compact latent space. This not only allows fast parallelizable training, but also enables efficient modeling of long-range dependencies in the latent space. Importantly, S4WM is a general framework that can incorporate not only the specific S4 model [21] but also any PSSM defined by Equation (2), including S5 [57] and other variants [20, 40, 15]. It models the observations and state transitions through a probabilistic generative process:

where z 0: T are the stochastic latent states. We note that computing the likelihood p ( x t | z ≤ t , a ≤ t ) and the prior p ( z t | z <t , a ≤ t ) requires extracting relevant information from the history ( z <t , a ≤ t ) . Therefore, it is crucial to maintain a long-term memory of the history. To this end, we use a stack of PSSM blocks to encode the history ( z <t , a ≤ t ) into an embedding vector h t for each t . This can be done in parallel during training and sequentially during imagination:

Here, g t = MLP ( concat [ z t -1 , a t ]) is the input to the PSSM blocks, s t is the collection of all internal hidden states, and h t is the final output. In our experiments, we use S4 blocks (shown in Figure 1 Right), a particular instantiation of the PSSM blocks using the S4 model. We find that adding the final MLP in each S4 block can improve generation quality (see Appendix B.1 for an ablation and Figure 21 for a detailed illustration of S4 block architecture).

After obtaining h t , we use it to compute the sufficient statistics of the prior and likelihood:

For training, we use variational inference. The approximate posterior is defined as:

We use a CNN encoder to compute the sufficient statistics of the posterior from image observations. This allows all posterior samples z 0: T to be obtained in parallel, thereby fully leveraging the parallel computation ability offered by PSSMs during training. We also provide an alternative design of the posterior in Appendix B.1. It is conditioned on the full history, and can obtain better generation quality at the cost of more computation.

The training objective is to maximize the evidence lower bound (ELBO):

In our experiments, we instantiate S4WM using the S4 model [21]. Figure 1 provides an illustration of training and imagination procedures, and Algorithms 1 and 2 provide detailed descriptions. Hyperparameters and further implementation details can be found in Appendix J.

5 Experiments

5.1 Environments

Unlike previous works [14, 48, 37, 52] that primarily evaluate the final performance of model-free agents on memory-demanding tasks, we seek to understand the memory capabilities of world models

TwoRooms

FourRooms

TenRooms

Distracting Memory

Multi Doors Keys

5

10

15

Training Episodes per Second

10K

100K

Imagination Frames per Second

5

15

25

Training Memory Usage (GB)

10

20

30

Imagination Memory Usage (GB)

S4WM

TSSM-XL

(m=128)

TSSM-XL

(m=256)

TSSM-XL

(m=512)

RSSM-TBTT

in model-based agents in terms of long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. We believe that our investigation provides more insights than the final performance alone, and paves the way for model-based agents with improved memory.

To this end, we develop a diverse set of environments shown in Figure 2, each targeting a specific memory capability. The environments are based on the 3D Memory Maze [50] and the 2D MiniGrid [6], both with partial observations. The world models are learned from an offline dataset collected by a scripted policy for each environment. This allows the world models to be evaluated independently of the policy learning algorithms.

Specifically, for each episode, the environment is regenerated. To simplify the evaluation of world model imagination, we design the data collecting policy to consist of a context phase and a query phase. In the context phase, the policy fully traverses the environment, while in the query phase, the policy revisits some parts of the environment. For evaluation, we use unseen episodes collected by the same policy as training. The world model observes the context phase, and is then evaluated on its imagination given the action sequence in the query phase. Because the environments are deterministic and have moderate visual complexity, and the context phase fully reveals the information of the environment, it suffices to use the mean squared error (MSE) as our main evaluation metric.

In the following, we first motivate our choice of the baselines through a comparison of speed and memory consumption, and then introduce and present the results for each environment in detail.

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 as RSSM-TBTT.

t=50

100

150

200

250

300

350

400

450

500

Truth

RSSM

TSSM

S4WM

TSSM-XL. TSSM [5] is the first Transformer-based world model for improving long-term memory. It was originally evaluated on sequences of length ∼ 100 . In this paper, we seek to evaluate on much longer sequences (up to 2000 steps), and it is impractical to feed the entire sequence to the vanilla Transformer [61]. Therefore, we use Transformer-XL [9] as the backbone and denote the model as TSSM-XL. It divides the full sequence into chunks, and maintains a cache of the intermediate hidden states from processed chunks. This cache serves as an extended context that allows modeling longer-term dependencies.

Speed and Memory Usage. We note that the cache length m is a crucial hyperparameter of TSSMXL. A larger m can potentially improve the memory capacity, at the cost of 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, RSSM-TBTT, and TSSM-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-TBTT 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 MLPs in between slows down its performance. As for memory usage, the decoder takes up most of the memory decoding all steps in parallel, leading to similar memory consumption of all models.

Based on our investigation, TSSM-XL with a 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 a 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 [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.

0

50

100

150

200

Imagination Steps

0

50

100

150

Generation MSE

Two Rooms (301 | 200)

0

125

250

375

500

Imagination Steps

0

100

200

300

Four Rooms (501 | 500)

0

225

450

675

900

Imagination Steps

0

200

400

600

Ten Rooms (1101 | 900)

S4WM

TSSM-XL

RSSM-TBTT

0

5

10

15

20

# Contexts After Teleport

0

100

200

300

Generation MSE

Two Rooms (151 | 100)

0

5

10

15

20

# Contexts After Teleport

0

100

200

300

Four Rooms (301 | 200)

0

5

10

15

20

# Contexts After Teleport

250

350

450

550

Ten Rooms (601 | 400)

S4WM

TSSM-XL

RSSM-TBTT

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 a random room, sequentially traverses all rooms, and returns to the starting room. In the query phase, the policy revisits each room in the same order as the context phase.

As shown in Table 1, all models obtain good reconstruction, while S4WM is much better in the Two Rooms and Four Rooms environment for long-term generation up to 500 steps. We demonstrate the superior generation quality of S4WM in Figure 4. All models are able to capture the high-level maze layout. However, RSSM-TBTT and TSSM-XL make many mistakes in details such as wall colors, object colors, and object positions, while S4WM is able to generate much more accurately, with only minor errors in object positions. We further show the per step generation MSE in Figure 5. S4WM is able to maintain a relatively good generation quality for up to 500 steps, while RSSM-TBTT and TSSM-XL make large generation errors even within 50 steps. We notice a periodic drop in the generation MSE. This is 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 in the Ten Rooms environment where the context length is 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-Dependent Recall

Humans are able to 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 a 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 to 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

= 100

(

199

|

51

)

= 200

(

399

|

101

)

= 400

(

799

|

201

)

Inference

↑

)

Imagination

↑

)

Inference

↑

)

Imagination

↑

)

Inference

↑

)

Imagination

↑

)

RSSM-TBTT

47

.

9%

49

.

6%

48

.

7%

48

.

4%

50

.

4%

52

.

2%

TSSM-XL

100

.

0

%

51

.

2%

99

.

9%

51

.

3%

50

.

4%

51

.

0%

S4WM

100

.

0

%

100

.

0

%

100

.

0

%

100

.

0

%

100

.

0

%

100

.

0

%

t=3

6

9

12

15

18

21

24

27

30

33

36

39

42

45

48

51

RSSM

S4WM TSSMF

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 is reported in Figure 6. 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 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, we 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 to 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 1 is given if the square reached is 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 one 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 position.

We report two types of reward prediction accuracy in Table 2. 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 imagination, conditioned on the observations in the context phase and additionally the action sequence in the query phase.

Our results show that only S4WM is able to accurately predict rewards within imagination. TSSMXL has limited success when observing the full sequence, but fails to imagine future rewards accurately. RSSM-TBTT 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 is 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 it cannot be used to unlock other doors of the same color. The agent is 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 a door, the world model must constantly update its memory when a key is picked up or consumed.

Since the environment is visually simple, we find the generation MSE to be a 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 on 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 S4 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-based reasoning.

Limitations and Future Work. We primarily focused on visually simple and deterministic environments 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 a 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.

Acknowledgments and Disclosure of Funding

This work is supported by Brain Pool Plus Program (No. 2021H1D3A2A03103645) 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

Jimmy Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al. DeepMind Lab. arXiv preprint arXiv:1612.03801 , 2016.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Guy Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Carnegie Mellon University, 1990.

Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with Transformer world models. In Deep RL Workshop NeurIPS 2021 , 2021.

Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for Gymnasium, 2018. URL https://github.com/Farama-Foundation/Minigrid .

Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 , 2014.

Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International Conference on Machine Learning , 2020.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , 2019.

Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: Efficient sequence-based RL via state spaces layers. In International Conference on Learning Representations , 2023.

Fei Deng, Ingook Jang, and Sungjin Ahn. DreamerPro: Reconstruction-free model-based reinforcement learning with prototypical representations. In International Conference on Machine Learning , 2022.

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.

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming Transformers for high-resolution image synthesis. In CVPR , 2021.

Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charles Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. In Advances in Neural Information Processing Systems , 2019.

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. In International Conference on Learning Representations , 2023.

Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning , 2021.

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.

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.

Albert Gu, Isys Johnson, Karan Goel, Khaled Kamal Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. In Advances in Neural Information Processing Systems , 2021.

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.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations , 2022.

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.

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.

Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear RNNs. arXiv preprint arXiv:2212.00768 , 2022.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems , 2018.

Jeongsoo Ha, Kyungsoo Kim, and Yusung Kim. Dream to generalize: Zero-shot model-based reinforcement learning for unseen visual distractions. In AAAI , 2023.

Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations , 2022.

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.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2020.

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In International Conference on Learning Representations , 2021.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

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.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8): 1735-1780, 1997.

Md Mohaiminul Islam and Gedas Bertasius. Long movie clip classification with state-space video models. In ECCV , 2022.

Arnav 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.

David M Knigge, David W Romero, Albert Gu, Efstratios Gavves, Erik J Bekkers, Jakub Mikolaj Tomczak, Mark Hoogendoorn, and Jan-jakob Sonke. Modelling long range dependencies in N D: From task-specific to a general purpose CNN. In International Conference on Learning Representations , 2023.

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 Neural Information Processing Systems , 2021.

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.03982 , 2023.

Marcelo G Mattar and Máté Lengyel. Planning in the brain. Neuron , 110(6):914-934, 2022.

Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In International Conference on Learning Representations , 2023.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In International Conference on Learning Representations , 2023.

Thomas M 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.

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.

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.

Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. In ICRA , 2021.

Masashi Okada and Tadahiro Taniguchi. DreamingV2: Reinforcement learning with discrete world models without reconstruction. In IROS , 2022.

Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning , 2023.

Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, 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.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning , 2013.

Jurgis Pašukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-term memory in 3D mazes. In International Conference on Learning Representations , 2023.

Joel Pearson. The human imagination: the cognitive neuroscience of visual mental imagery. Nature Reviews Neuroscience , 20(10):624-634, 2019.

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.

Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941 , 2017.

Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In International Conference on Learning Representations , 2023.

Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. In Advances in Neural Information Processing Systems , 2021.

Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In CoRL , 2023.

Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations , 2023.

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):1317-1324, 2009.

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.

Endel Tulving. Memory and consciousness. Canadian Psychology/Psychologie canadienne , 26 (1):1, 1985.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , 2017.

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.

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 highdimensional model-based reinforcement learning. arXiv preprint arXiv:2211.12774 , 2022.

Jialong Wu, Haoyu Ma, Chaoyi Deng, and Mingsheng Long. Pre-training contextualized world models with in-the-wild videos for reinforcement learning. arXiv preprint arXiv:2305.18499 , 2023.

Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. DayDreamer: World models for physical robot learning. In CoRL , 2023.

Marek Wydmuch, Michał Kempka, and Wojciech Ja´ skowski. ViZDoom competitions: Playing Doom from pixels. IEEE Transactions on Games , 11(3):248-259, 2018.

Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel. Temporally consistent Transformers for video generation. In International Conference on Machine Learning , 2023.

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.05092 , 2023.

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.

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.08136 , 2022.

A Environment and Dataset Details

For each 3D environment ( i.e. , Two Rooms, Four Rooms, and Ten Rooms), we generate 30 K trajectories using a scripted policy, of which 28 K are used for training, 1 K for validation, and 1 K for testing. For each 2D environment ( i.e. , Distracting Memory and Multi Doors Keys), we generate 10 K trajectories using a scripted policy, of which 8 K are used for training, 1 K for validation, and 1 K 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 × 64 × 3 for 3D environments, and 40 × 40 × 3 for 2D environments.

B Ablation Study

In this section, we investigate alternative architectural choices for S4WM and different cache lengths m for TSSM-XL. 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

for simplicity and parallel training ability. However, we note that it is possible to condition on the full history while maintaining the parallel training ability:

We illustrate this architecture in Figure 8. Here, we first use a CNN encoder to obtain a deterministic embedding e t for each image observation x t , and then use a stack of S4 blocks to encode the history and compute the sufficient statistics of the posterior for all t in parallel:

We have defined the dummy action a 0 = ∅ and the initial S4 hidden state s -1 , which are both implemented as vectors of all zeros. We note that the S4 blocks in the posterior are not shared with those in the prior.

S4WM-No-MLP. In our implementation, each S4 block consists of two S4 layers and one MLP. We note that the MLP is 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 to 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 4 and Figures 9 and 10. Results on 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 S4WM-FullPosterior performs similarly as 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. In contrast, while S4WM-No-MLP is 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.

<latexi sh

1_b

64="7

OMEQ

I3

5

U0

u

ZNF

/grL

>A

c

VD

S

J

2vW

d

k

w

R

y

KCf

T

m

Y

H

p

X

P

n

9

j

G

8

q

o

z

<latexi sh

1_b

64="Q

d/u7

WPL3U+yc

9T

IXE

Cw

>A

n

VD

S

N

F

2pr

f

Z

Bovgq

0

m

M

R

O

K

Y

H

5

J

z

8

j

G

k

<latexi sh

1_b

64="I

3Y

pg

w

9jqz+u

CKJ8Ly

E

>A

B/n

c

V

S

N

2

X

F

R

Do

7

r

WU

P

f

m

d0

H

M

T

vG

5

Z

Q

O

k

<latexi sh

1_b

64="F

wmMOgz

CB

NK+XW

>A

/n

c

V

9S

EJ2L

R

o

7

G

U

8

H

fY

3

jd0

I

Zr

k

5

P

u

y

p

D

Q

<latexi sh

1_b

64="KEQ

Hw0

I5XjTM/

kLR+

27A

>

C

n

c

VD

S

N

F

pr

f

ZduBovgq

U

m

O

Y

9

W

J

z

3

8

P

G

y

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

Training

Imagination

(parallel)

ˆ

x

1

h

1

z

0

a

1

ˆ

x

2

h

2

z

1

a

2

z

2

Prior

Posterior

Next Step Prediction

Dec

Dec

z

1

z

2

x

0

Enc

x

1

Enc

x

2

Enc

(parallel)

e

0

e

1

e

2

∅

a

1

a

2

(parallel)

z

0

a

1

z

1

a

2

ˆ

x

2

h

2

(single step)

a

3

ˆ

x

3

h

3

(single step)

a

4

ˆ

x

4

h

4

Generation

· · ·

Dec

Dec

Dec

z

2

z

2

z

3

z

3

z

4

a

1

x

0

Enc

x

1

Enc

(parallel)

∅

e

0

e

1

Context

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="wSoXJOkvE

Cu2

AW5j

U

IQ

K

>

H

c

VDL

N

F

3

q

g0

R

y

G

+

m

M

d

Z

r

/

n

T

z

79

f

8

B

p

Y

P

<latexi sh

1_b

64="K

pM

I9

kRGF

H

J

w

>A

C

n

c

VDL

N

3

rPU

YB

W

d

q

Z

E

m

0

u

O

7j

5y

z

g

/

2

v

o

X

Q

8

<latexi sh

1_b

64="f3gk

JW

rvK2

m7c

Dj

GA

>

C

n

V

LS

N

F

p

ZduBo

q

T

w0

U

+

M

R

O

/

Y

H

5

9

Q

z

8

E

X

P

y

I

<latexi sh

1_b

64="T

G72gf

V

uSj

/DUdv5

>

B

n

L

N

EOy

r

Y9

oP

K

X8

m

k

zC

+

pH

J

W

Q

3

F

0

M

I

q

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="zwP

fq

Kp

+M

5

I

Vnu

R

Tg

>A

c

9S

N

EJ2LX

F

Do

7

r

WU

Y

8m

3

jd0

y

H

vG

C

Z

Q

O

k

<latexi sh

1_b

64="

0

gk qGo

S

7K

/C+zuQ5

>

B

n

c

VDL

N

EOy

r

fU

9

P

w

X8

jm

d

Z

v

pH

RJ

TW

3

F

2

M

I

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="zdkSR

7g

XfJ

yM

>A

C

n

c

VD

F

3

W

ur

8

Q

H9

E

p

O2

UvI

m

w

Y

K0

Z

/

5

j

B

P

+

T

<latexi sh

1_b

64="LQ

Jj85S

dDMn/

Yc vrT

Wg

>A

B

V

N

9

oP

K

wGX

z

I

7m

2

u

C

q

Z

R

+

0

3

p

F

k

H

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="

0

gk qGo

S

7K

/C+zuQ5

>

B

n

c

VDL

N

EOy

r

fU

9

P

w

X8

jm

d

Z

v

pH

RJ

TW

3

F

2

M

I

<latexi sh

1_b

64="Q

d/u7

WPL3U+yc

9T

IXE

Cw

>A

n

VD

S

N

F

2pr

f

Z

Bovgq

0

m

M

R

O

K

Y

H

5

J

z

8

j

G

k

<latexi sh

1_b

64="I

3Y

pg

w

9jqz+u

CKJ8Ly

E

>A

B/n

c

V

S

N

2

X

F

R

Do

7

r

WU

P

f

m

d0

H

M

T

vG

5

Z

Q

O

k

<latexi sh

1_b

64="KEQ

Hw0

I5XjTM/

kLR+

27A

>

C

n

c

VD

S

N

F

pr

f

ZduBovgq

U

m

O

Y

9

W

J

z

3

8

P

G

y

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="wSoXJOkvE

Cu2

AW5j

U

IQ

K

>

H

c

VDL

N

F

3

q

g0

R

y

G

+

m

M

d

Z

r

/

n

T

z

79

f

8

B

p

Y

P

<latexi sh

1_b

64="K

pM

I9

kRGF

H

J

w

>A

C

n

c

VDL

N

3

rPU

YB

W

d

q

Z

E

m

0

u

O

7j

5y

z

g

/

2

v

o

X

Q

8

<latexi sh

1_b

64="f3gk

JW

rvK2

m7c

Dj

GA

>

C

n

V

LS

N

F

p

ZduBo

q

T

w0

U

+

M

R

O

/

Y

H

5

9

Q

z

8

E

X

P

y

I

<latexi sh

1_b

64="w+YgWRE

CX

zq9d

0kH/V

fM

Q

>A

n

c

DLS

N

F

2pr

Z

uBov

T

U

m

K

O7

5

J

3

8

j

P

G

y

I

<latexi sh

1_b

64="T

G72gf

V

uSj

/DUdv5

>

B

n

L

N

EOy

r

Y9

oP

K

X8

m

k

zC

+

pH

J

W

Q

3

F

0

M

I

q

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="zwP

fq

Kp

+M

5

I

Vnu

R

Tg

>A

c

9S

N

EJ2LX

F

Do

7

r

WU

Y

8m

3

jd0

y

H

vG

C

Z

Q

O

k

<latexi sh

1_b

64="

0

gk qGo

S

7K

/C+zuQ5

>

B

n

c

VDL

N

EOy

r

fU

9

P

w

X8

jm

d

Z

v

pH

RJ

TW

3

F

2

M

I

<latexi sh

1_b

64="

SpV

0RqH/oZ9TyG+

M

FA

kY

>

B

n

c

wN

EJ2LXz

g

7

mU

8

f

3

jd

I

rO

5

P

u

W

v

C

D

Q

K

<latexi sh

1_b

64="Cg/w

MY

GvfOE7I

yS

A

>

n

c

VDL

N

FJ3

rPU

9

B

T

u

m

2

0ok+

H

z

Q

5

d

q

p

K

W

R

j

X

Z

8

<latexi sh

1_b

64="z87FV2dUc

g

BR0q

Y

X

o

>A

3

DLS

M

j

9Z

W

nfy

K

v

Gk mE5

I

+

r

N

Hu

J

/O

C

w

P

T

Q

p

<latexi sh

1_b

64="PfoX

O

N

Kg+j

E

>A

B/n

c

VDLS

y

r

UY9

85

C

Z

k

M7v

p

W

m3

J0

q

T

F

d

Q

R

H

G

I

<latexi sh

1_b

64="H

J

G

+Ap3N

q

y

P8uQCoU

w

>

n

c

VDLS

F

W

r

g

9

EM

O2

zk

I

m

Yf

K0

Zd7

X

/

5

j

B

R

v

T

<latexi sh

1_b

64="A8C

gr

MR

P/y

V+

WzoBY

>

n

c

DLS

N

EO

fU

9

K

5

j

Zwu

k

7v

p

m3

J0

q

T

X

F

d

Q

H

G

I

2

<latexi sh

1_b

64="

/DuQ

T

r

H9IvnFwE

W

>A

C

c

V

LS

N

J3U

2

d

Bo gq

0

+

m

M

Y

Z

f

X7

P

G

K

j

kR

8

5

z

O

y

p

<latexi sh

1_b

64="n

Y/2

+5

IpR

P

Qm

8KySf

3

>A

B

c

VDL

gM

F

r

Zdu

k

w

W

E

7

9

G

o

T

J

q

ONC0X

j

z

U

v

H

<latexi sh

1_b

64="z87FV2dUc

g

BR0q

Y

X

o

>A

3

DLS

M

j

9Z

W

nfy

K

v

Gk mE5

I

+

r

N

Hu

J

/O

C

w

P

T

Q

p

<latexi sh

1_b

64="I

m

H

U

fyjRKWC/

T

M+

N

>A

B

n

c

VDLSg

EO

r

Y9

oP

w

X

Q

GF2

p

Z

d

z

7

8

5

k

3

0

u

J

q

v

<latexi sh

1_b

64="rqUKgmMVy/v nTc2+k3Z

JQ

>A

B

DLS

N

EO

f

Y9

oP

w

X

GF

p

d

Cz

7

8

R

5

W

I

0

u

j

H

<latexi sh

1_b

64="k

Bz

Z

d

fIG

M

ycK

rXQ

>A

/H

VDLSgN

EO

UY9

oP

2

R

J

j

m

v

3

0

qp

u

T

+

F7

nw

5

C

8

W

0

125

250

375

500

Imagination Steps

0

50

100

150

Generation MSE

Four Rooms (501 | 500)

0

225

450

675

900

Imagination Steps

0

150

300

450

Ten Rooms (1101 | 900)

S4WM

S4WM-Full-Posterior

S4WM-No-MLP

t=50

100

150

200

250

300

350

400

450

500

Truth

S4WM

Full-Post.

S4WM-

S4WM-

No-MLP

0

5

10

15

20

# Contexts After Teleport

0

60

120

180

Generation MSE

Four Rooms (301 | 200)

0

5

10

15

20

# Contexts After Teleport

240

320

400

480

Ten Rooms (601 | 400)

S4WM

S4WM-Full-Posterior

S4WM-No-MLP

t=15

30

45

60

75

90

105

120

135

150

Truth

S4WM

Full-Post.

S4WM-

S4WM-

No-MLP

5

10

15

Training Episodes per Second

5K

10K

15K

Imagination Frames per Second

5

15

25

Training Memory Usage (GB)

10

20

30

Imagination Memory Usage (GB)

S4WM

S4WM-Full-Posterior

S4WM-No-MLP

0

125

250

375

500

Imagination Steps

0

100

200

300

Generation MSE

Four Rooms (501 | 500)

0

225

450

675

900

Imagination Steps

0

200

400

600

Ten Rooms (1101 | 900)

TSSM-XL

(

m

=128)

TSSM-XL

(

m

=256)

TSSM-XL

(

m

=512)

t=50

100

150

200

250

300

350

400

450

500

Truth

(m=512)(m=256)(m=128)

TSSM

TSSM

TSSM

0

5

10

15

20

# Contexts After Teleport

0

70

140

210

Generation MSE

Four Rooms (301 | 200)

0

5

10

15

20

# Contexts After Teleport

120

240

360

480

Ten Rooms (601 | 400)

TSSM-XL

(

m

=128)

TSSM-XL

(

m

=256)

TSSM-XL

(

m

=512)

t =15

30

45

60

75

90

105

120

135

150

Truth

(m=512)(m=256)(m=128)

TSSM

TSSM

TSSM

B.2 Cache Length of TSSM-XL

In our main experiments, we have used the cache length m = 128 for TSSM-XL, because it is close to S4WM in terms of computational cost. Here we provide a more thorough investigation with larger cache lengths.

We report results on the non-teleport Four Rooms and Ten Rooms environments in Table 5 and Figures 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-dependent recall than S4WMonthe Ten Rooms teleport environment, consistent with the findings in previous work [24, 15] that Transformers are better than S4 at performing context-dependent operations.

C Additional Experiment Figures

t=15

30

45

60

75

t=30

60

90

120

150

RSSM

TSSM

S4WM

TwoRooms

Four Rooms

t=35

70

105

140

175

210

245

280

315

350

385

420

455

490

525

560

595

RSSM

TSSM

S4WM

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-XL, and S4WM, we use concat [ h t , z t ] as the input to the probing network, where h t is the output of the final Transformer/S4 block for TSSM-XL and S4WM. The probing network is an MLP with 4 hidden layers, each consisting of 1024 hidden units and followed by layer normalization [1] and SiLU [53] nonlinearity.

Both TSSM-XL and S4WM underperform RSSM-TBTT. We conjecture that this is because in RSSM-TBTT, h t captures more global information, while in TSSM-XL and S4WM, h t captures more local information for each t .

Comparing RSSM-TBTT and S4WM, we find that the S4 hidden state s t plays a similar role as the h t in RSSM-TBTT by carrying over information from previous time steps. Therefore, we conjecture that s t will capture more global information and concat [ s t , z t ] can be a more suitable embedding for probing. However, the S4 hidden state s t has a much higher (typically 64 × ) dimension than h t , making the above solution impractical.

Fortunately, S4WM is a general framework that works with many S4 variants. In particular, S5 [57] uses a hidden state s t that is of similar dimension as h t . Hence, we 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 in each S5 block, and concatenate them with z t to form the input to the probing network. Our results show that this model outperforms RSSM-TBTT by a large margin on the 9 × 9 maze. Nevertheless, the larger 15 × 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 offline-trained world model. 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 (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 in 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 to unlock every door in all three environments, while RSSM-TBTT and TSSM-XL succeed roughly half of the time.

Start

Goal

(FirstPerson View)

(Top-downView)

F Comparison with TECO

We provide a 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 a pretrained VQGAN [13] encoder/decoder, while S4WM uses a 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-based model), and is 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 ( x 0 , a 1 , x 1 , a 2 , x 2 , . . . , a T , x T ) . Namely, the agent takes action a t +1 after observing x t , and receives the next observation x t +1 . We omit the reward for simplicity.

H.1 RSSM

RSSM [28] models the observations and state transitions through the following generative process:

where z 0: T are the stochastic latent states. The approximate posterior is defined as:

The conditioning on previous states z <t and actions a ≤ t appears multiple times. RSSM uses a shared GRU [7] to compress z <t and a ≤ t into a deterministic encoding h t :

This is then used to compute the sufficient statistics of the prior, likelihood, and posterior:

The training objective is to maximize the evidence lower bound (ELBO):

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 is:

where z 0: T are the stochastic latent states. The approximate posterior is defined as:

The training objective is to maximize the evidence lower bound (ELBO):

The main difference from S4WM is that in TSSM the prior p ( z t | z <t , a ≤ t ) is computed by a stack of Transformer [61] blocks. Specifically, the Transformer blocks output an embedding vector h t through self-attention over the history:

The h t is then used for predicting the next latent state z t and decoding the latent state into image ˆ x t :

I Speed and Memory Usage Evaluation Details

All results in Figure 3 are obtained on a single NVIDIA RTX A6000 GPU. We make sure that the models have comparable number of parameters. For training, we use a batch size of 8 , sequence length of 1000 , and image size of 64 × 64 . We report the number of episodes per second processed by each model, averaged over 100 batches, and also the peak memory usage. For imagination, we use a batch size of 64 , context length of 500 , generation length of 500 , and image size of 64 × 64 . We report the number of frames per second, averaged over 8 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 of DreamerV3-S. We use a 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-XL and RSSM-TBTT, we use constant learning rates following the original papers. To ensure a fair comparison, all models use the same CNN encoder and decoder,

with convolutional and transposed convolutional layers of kernel size 4 and stride 2 , layer normalization [1], and SiLU [53] nonlinearity. The latent states z 0: T are vectors of categorical variables [30] optimized by straight-through gradients [3]. To facilitate stable training, we parameterize the categorical distributions as mixtures of 1% uniform and 99% neural network output [31]. We use KL balancing [30] to scale the gradient of the KL loss L KL ( q, p ) with respect to the posterior q and prior p :

Here, stop_grad denotes the stop-gradient operator, and we set α = 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 are fundamentally deep generative models. Therefore, they inherit the potential negative social impacts that deep generative models may have, such as generating fake images and videos that can contribute to digital misinformation and deception. Caution must be exercised in the application of these models, adhering to ethical guidelines and regulations to mitigate the risks.

<latexi sh

1_b

64="Q

uUFk

Y

Xjz

7q

T

H+

Lc

>A

C

VD

SgM

3

9Z

W

BE

o

yIr2

p

/

J

G

f

N

n

P

5

Od

w

0

R

m

K

v

8

<latexi sh

1_b

64="Q

uUFk

Y

Xjz

7q

T

H+

Lc

>A

C

VD

SgM

3

9Z

W

BE

o

yIr2

p

/

J

G

f

N

n

P

5

Od

w

0

R

m

K

v

8

LayerNorm

Linear

+

Linear

GLU

Dropout

+

LayerNorm

GELU

Dropout

GELU

Dropout

Linear

Dropout

x2

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe limitation of the benchmark environment simplicity was acknowledged.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nModel-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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

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

STRENGTHS:
1. Introduction of S4 into the field of world models and MBRL is appreciated
2. Carefully designed memory-demanding tasks and extensive experiments
3. 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. 
2. This work does not evaluate MBRL equipped with S4WM on memory-demanding tasks, such as DMLab.

QUESTIONS:
Only one minor question:

1. 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.

There does not seem to be any negative social impact of this paper that should be discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-04

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

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

QUESTIONS:
- As mentioned in the weakness section, the impact of the improved world model on MBRL algorithms.
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}