SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 16: Comparison of TSSM-XL with different cache lengths  $m$  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 as additional contexts. Increasing the cache length improves context-dependent recall. With a sufficiently large cache length, TSSM-XL without observing additional contexts can outperform S4WM.
Figure 17: Context-dependent recall from TSSM-XL with different cache lengths  $m$  in the Four Rooms teleport environment. 20 observations after the teleport are provided as additional contexts. TSSM-XL ( $m = 256$ ) and TSSM-XL ( $m = 512$ ) perform similarly in this environment, and are both better than TSSM-XL ( $m = 128$ ).
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 S4WM on the 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
Figure 18: Context-dependent recall in the teleport environments. 20 observations after the teleport are provided as additional contexts. TSSM-XL and S4WM perform similarly in the Two Rooms environment, while only S4WM is able to maintain good recall quality in the Four Rooms environment.
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-XL, and S4WM, we use concat  $\left[\pmb{h}_t,\mathbf{z}_t\right]$  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.
Table 6: Memory Maze offline probing benchmark results.
Memory 9×9
Walls (Acc. ↑)	Memory 15×15
Walls (Acc. ↑)	Memory 9×9
Objects (MSE ↓)	Memory 15×15
Objects (MSE ↓)
Constant Baseline	80.8%	78.3%	23.9	64.8
RSSM-TBTT	95.0%	81.7%	5.4	32.6
TSSM-XL	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
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  $\text{concat}[s_t, \mathbf{z}_t]$  can be a more suitable embedding for probing. However, the S4 hidden state  $s_t$  has a much higher (typically  $64 \times$ ) 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  $\mathbf{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 \times 9$  maze. Nevertheless, the larger  $15 \times 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.
Figure 20: Illustration of the task solved by the skill-level MPC agent. The task is to unlock the door specified by the goal image on the left. The agent is allowed to execute two skills: pick up one key and then attempt to 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%	42.86%
TSSM-XL	100%	40%	57.14%
S4WM	100%	100%	100%
Table 8: Comparison of generation quality on DMLab following TECO evaluation protocol.
DMLab
PSNR (↑)	SSIM (↑)	LPIPS (↓)	#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 (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)
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 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  $(\mathbf{x}_0, \mathbf{a}_1, \mathbf{x}_1, \mathbf{a}_2, \mathbf{x}_2, \dots, \mathbf{a}_T, \mathbf{x}_T)$ . Namely, the agent takes action  $\mathbf{a}_{t+1}$  after observing  $\mathbf{x}_t$ , and receives the next observation  $\mathbf{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:
$$
p \left(\mathbf {x} _ {0: T} \mid \boldsymbol {a} _ {1: T}\right) = \int \prod_ {t = 0} ^ {T} p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) d \mathbf {z} _ {0: T}, \tag {13}
$$
where  $\mathbf{z}_{0:T}$  are the stochastic latent states. The approximate posterior is defined as:
$$
q \left(\mathbf {z} _ {0: T} \mid \mathbf {x} _ {0: T}, \boldsymbol {a} _ {1: T}\right) = \prod_ {t = 0} ^ {T} q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}, \mathbf {x} _ {t}\right). \tag {14}
$$
The conditioning on previous states  $\mathbf{z}_{< t}$  and actions  $a_{\leq t}$  appears multiple times. RSSM uses a shared GRU [7] to compress  $\mathbf{z}_{< t}$  and  $a_{\leq t}$  into a deterministic encoding  $h_t$ :
$$
\boldsymbol {h} _ {t} = \operatorname {G R U} \left(\boldsymbol {h} _ {t - 1}, \operatorname {M L P} \left(\operatorname {c o n c a t} \left[ \mathbf {z} _ {t - 1}, \boldsymbol {a} _ {t} \right]\right)\right). \tag {15}
$$
This is then used to compute the sufficient statistics of the prior, likelihood, and posterior:
$$
p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) = \operatorname {M L P} \left(\boldsymbol {h} _ {t}\right), \tag {16}
$$
$$
p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \mathbf {a} _ {\leq t}\right) = \mathcal {N} \left(\hat {\mathbf {x}} _ {t}, \mathbf {1}\right), \quad \hat {\mathbf {x}} _ {t} = \operatorname {D e c o d e r} \left(\operatorname {c o n c a t} \left[ \mathbf {h} _ {t}, \mathbf {z} _ {t} \right]\right), \tag {17}
$$
$$
q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}, \mathbf {x} _ {t}\right) = \operatorname {M L P} \left(\operatorname {c o n c a t} \left[ \boldsymbol {h} _ {t}, \boldsymbol {e} _ {t} \right]\right), \quad \boldsymbol {e} _ {t} = \operatorname {E n c o d e r} \left(\mathbf {x} _ {t}\right). \tag {18}
$$
$$
\log p \left(\mathbf {x} _ {0: T} \mid \boldsymbol {a} _ {1: T}\right) \geq \mathbb {E} _ {q} \left[ \sum_ {t = 0} ^ {T} \log p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) - \mathcal {L} _ {\mathrm {K L}} \left(q \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}, \mathbf {x} _ {t}\right), p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right)\right) \right]. \tag {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 is:
$$
p \left(\mathbf {x} _ {1: T} \mid \mathbf {x} _ {0}, \boldsymbol {a} _ {1: T}\right) = \int p \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right) \prod_ {t = 1} ^ {T} p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) d \mathbf {z} _ {0: T}, \tag {20}
$$
$$
q \left(\mathbf {z} _ {0: T} \mid \mathbf {x} _ {0: T}, \boldsymbol {a} _ {1: T}\right) = \prod_ {t = 0} ^ {T} q \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {t}\right), \quad \text {w h e r e} q \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right) = p \left(\mathbf {z} _ {0} \mid \mathbf {x} _ {0}\right). \tag {21}
$$
$$
\left. \right. \log p \left(\mathbf {x} _ {1: T} \mid \mathbf {x} _ {0}, \boldsymbol {a} _ {1: T}\right) \geq \mathbb {E} _ {q} \left[ \sum_ {t = 1} ^ {T} \log p \left(\mathbf {x} _ {t} \mid \mathbf {z} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) - \mathcal {L} _ {\mathrm {K L}} \left(q \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {t}\right), p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right)\right)\right] . \tag {22}
$$
The main difference from S4WM is that in TSSM the prior  $p(\mathbf{z}_t \mid \mathbf{z}_{<t}, \mathbf{a}_{\leq t})$  is computed by a stack of Transformer [61] blocks. Specifically, the Transformer blocks output an embedding vector  $\mathbf{h}_t$  through self-attention over the history:
$$
\boldsymbol {h} _ {t} = \text {T r a n s f o r m e r B l o c k s} \left(\boldsymbol {g} _ {1: t}\right), \quad \boldsymbol {g} _ {t} = \operatorname {M L P} \left(\text {c o n c a t} \left[ \mathbf {z} _ {t - 1}, \boldsymbol {a} _ {t} \right]\right). \tag {23}
$$
The  $h_t$  is then used for predicting the next latent state  $\mathbf{z}_t$  and decoding the latent state into image  $\hat{\mathbf{x}}_t$ :
$$
p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {<   t}, \boldsymbol {a} _ {\leq t}\right) = \operatorname {M L P} \left(\boldsymbol {h} _ {t}\right), \quad \hat {\mathbf {x}} _ {t} = \operatorname {D e c o d e r} \left(\operatorname {c o n c a t} \left[ \boldsymbol {h} _ {t}, \mathbf {z} _ {t} \right]\right). \tag {24}
$$
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 \times 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 \times 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  $\mathbf{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  $\mathcal{L}_{\mathrm{KL}}(q,p)$  with respect to the posterior  $q$  and prior  $p$ :
$$
\mathcal {L} _ {\mathrm {K L}} (q, p) = \alpha \cdot \mathrm {K L} [ \operatorname {s t o p} _ {-} \operatorname {g r a d} (q) \| p ] + (1 - \alpha) \cdot \mathrm {K L} [ q \| \operatorname {s t o p} _ {-} \operatorname {g r a d} (p) ]. \tag {25}
$$
Here, stop_grad denotes the stop-gradient operator, and we set  $\alpha = 0.8$  to put more emphasis on learning the prior toward the posterior than the other way around.
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.
Figure 21: Detailed architecture of an S4 block.
Table 10: Hyperparameters for 3D environments.
S4WM	TSSM-XL	RSSM-TBTT
Training
Optimizer	AdamW	AdamW	AdamW
Batch size	8	8	32
Epochs	57	57	57
Learning rate	1 × 10-3	1 × 10-4	3 × 10-4
Weight decay	1 × 10-2	1 × 10-2	1 × 10-2
Gradient clipping	1000	1000	200
TBTT steps	-	-	48
Model
Stochastic discrete latent size	32 × 32	32 × 32	32 × 32
History encoding blocks	6	12	1
History encoding dmodel	512	512	2048
History encoding df	2048	512	-
History encoding cache length m	-	128 / 256 / 512	-
CNN layers	4	4	4
CNN multiplier	48	48	48
MLP hidden units	1000	1000	1000
Total parameters	41.1 M	43.8 M	54.1 M
Objective
KL balancing α	0.8	0.8	0.8
Table 11: Hyperparameters for 2D environments.
S4WM	TSSM-XL	RSSM-TBTT
Training
Optimizer	AdamW	AdamW	AdamW
Batch size	8	8	32
Epochs	100	100	100
Learning rate	\( 1 \times 10^{-3} \)	\( 1 \times 10^{-4} \)	\( 3 \times 10^{-4} \)
Weight decay	\( 1 \times 10^{-2} \)	\( 1 \times 10^{-2} \)	\( 1 \times 10^{-2} \)
Gradient clipping	1000	1000	200
TBTT steps	-	-	50
Model
Stochastic discrete latent size	\( 32 \times 32 \)	\( 32 \times 32 \)	\( 32 \times 32 \)
History encoding blocks	6	10	1
History encoding \( d_{\text{model}} \)	512	512	2048
History encoding \( d_{\text{ff}} \)	2048	512	-
History encoding cache length m	-	128	-
CNN layers	3	3	3
CNN multiplier	32	32	32
MLP hidden units	512	512	512
Total parameters	28.1 M	27.1 M	31.2 M
Objective
KL balancing \( \alpha \)	0.8	0.8	0.8