SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
Table 1: Evaluation of long-term imagination. Each environment is labeled with (context steps | query steps). All models obtain good reconstruction, while S4WM is much better at long-term generation up to 500 steps. All models struggle in the Ten Rooms environment.
Two Rooms (301 | 200)	Four Rooms (501 | 500)	Ten Rooms (1101 | 900)
Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)
RSSM-TBTT	1.7	62.2	1.5	219.4	1.5	323.1
TSSM-XL	2.5	62.9	2.4	224.4	2.6	360.4
S4WM	1.8	27.3	1.7	44.0	1.8	224.4
Figure 4: Long-term imagination in the Four Rooms environment. While RSSM-TBTT and TSSM-XL make many mistakes in wall colors, object colors, and object positions, S4WM is able to generate much more accurately, with only minor errors in object positions.
TSSM-XL. TSSM [5] is the first Transformer-based world model for improving long-term memory. It was originally evaluated on sequences of length  $\sim 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 TSSM-XL. 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  $\sim 10 \times$  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  $\sim 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.
Figure 5: Generation MSE per imagination step. Each environment is labeled with (context steps | query steps). S4WM maintains a relatively good generation quality for up to 500 steps, while RSSM-TBTT and TSSM-XL make large generation errors even within 50 steps.
Figure 6: Evaluation of 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. TSSM-XL performs the best in the Two Rooms environment where the context phase is short, and is able to recall successfully without additional observations. When the context phase is longer, S4WM performs the best.
To this end, we develop three environments with increasing difficulty, namely Two Rooms, Four Rooms, and Ten Rooms, based on the 3D Memory Maze [50]. The top-down views are shown in Figure 2. In the context phase, the data collecting policy starts from 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
Table 2: Reward prediction accuracy in the Distracting Memory environments. Each environment is labeled with (context steps | query steps). S4WM succeeds in all environments. TSSM-XL has limited success when observing the full sequence, but fails to predict rewards within imagination. RSSM-TBTT completely fails.
Width = 100 (199 | 51)	Width = 200 (399 | 101)	Width = 400 (799 | 201)
Inference Accuracy (↑)	Imagination Accuracy (↑)	Inference Accuracy (↑)	Imagination Accuracy (↑)	Inference Accuracy (↑)	Imagination Accuracy (↑)
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%
Figure 7: Imagination in the Distracting Memory environment of width 100. RSSM-TBTT and TSSM-XL fail to keep track of the agent's position, leading to inaccurate reward prediction within imagination.
to recall depends on the new observations. Transformers have been shown to be better than S4 at performing such context-dependent operations in low-dimensional sequence manipulation tasks [24] and synthetic language modeling tasks [15]. We investigate this in the context of world model learning, with high-dimensional image inputs.
To help the model retrieve the correct events in history, we provide up to 20 observations after the teleport as additional contexts. The generation MSE of the recall 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.
Table 3: Memory-based reasoning in the Multi Doors Keys environments. Each environment is labeled with (context steps | query steps). S4WM performs well on all environments, while others struggle.
Three Keys (76 | 174)	Five Keys (170 | 380)	Seven Keys (296 | 654)
Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)
RSSM-TBTT	0.05	5.16	0.04	6.36	0.03	6.28
TSSM-XL	0.05	1.27	0.03	6.05	0.02	9.24
S4WM	0.01	0.04	0.02	0.27	0.02	0.10
Our results show that only S4WM is able to accurately predict rewards within imagination. TSSM-XL has limited success when observing the full sequence, but fails to imagine future rewards accurately. RSSM-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.
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.
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.
[1] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[2] Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Kuttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, et al. DeepMind Lab. arXiv preprint arXiv:1612.03801, 2016.
[3] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
[4] Guy Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Carnegie Mellon University, 1990.
[5] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with Transformer world models. In Deep RL Workshop NeurIPS 2021, 2021.
[6] Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for Gymnasium, 2018. URL https://github.com/Farama-Foundation/Minigrid.
[7] 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.
[8] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International Conference on Machine Learning, 2020.
[9] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[10] Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: Efficient sequence-based RL via state spaces layers. In International Conference on Learning Representations, 2023.
[11] Fei Deng, Ingook Jang, and Sungjin Ahn. DreamerPro: Reconstruction-free model-based reinforcement learning with prototypical representations. In International Conference on Machine Learning, 2022.
[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[13] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming Transformers for high-resolution image synthesis. In CVPR, 2021.
[14] Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charles Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. In Advances in Neural Information Processing Systems, 2019.
[15] Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. In International Conference on Learning Representations, 2023.
A Environment and Dataset Details
For each 3D environment (i.e., Two Rooms, Four Rooms, and Ten Rooms), we generate  $30\mathrm{K}$  trajectories using a scripted policy, of which  $28\mathrm{K}$  are used for training, 1K for validation, and 1K for testing. For each 2D environment (i.e., Distracting Memory and Multi Doors Keys), we generate  $10\mathrm{K}$  trajectories using a scripted policy, of which 8K are used for training, 1K for validation, and 1K for testing. All reported results are obtained from the test trajectories, using the model checkpoints that achieve the best validation loss. The image observations are of size  $64\times 64\times 3$  for 3D environments, and  $40\times 40\times 3$  for 2D environments.
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
$$
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) \tag {10}
$$
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:
$$
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} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right). \tag {11}
$$
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  $\mathbf{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:
$$
q \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {\leq t}, \boldsymbol {a} _ {\leq t}\right) = \operatorname {M L P} \left(\boldsymbol {h} _ {t}\right), \quad \boldsymbol {h} _ {0: T}, \boldsymbol {s} _ {T} = \operatorname {S 4 B l o c k s} \left(\boldsymbol {g} _ {0: T}, \boldsymbol {s} _ {- 1}\right), \quad \boldsymbol {g} _ {t} = \operatorname {M L P} \left(\text {c o n c a t} \left[ \boldsymbol {e} _ {t}, \boldsymbol {a} _ {t} \right]\right). \tag {12}
$$
We have defined the dummy action  $\pmb{a}_0 = \varnothing$  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-Full-Posterior 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.
Figure 8: Architecture of S4WM-Full-Posterior. It maintains the parallel training ability while computing the posterior from the full history.  $\varnothing$  denotes dummy actions.
Table 4: Comparison of alternative S4WM architectures on long-term imagination. Each environment is labeled with (context steps | query steps). S4WM-Full-Posterior is comparable to S4WM on the Four Rooms environment, but is better on the more challenging Ten Rooms environment. In contrast, S4WM-No-MLP performs much worse, suggesting the importance of the MLP in the S4 blocks to long-term imagination.
Four Rooms (501 | 500)	Ten Rooms (1101 | 900)
Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)
S4WM	1.7	44.0	1.8	224.4
S4WM-Full-Posterior	1.7	38.0	2.0	171.1
S4WM-No-MLP	2.3	68.6	2.5	277.0
Figure 9: Generation MSE per imagination step of alternative S4WM architectures. Each environment is labeled with (context steps | query steps). S4WM-Full-Posterior performs the best as imagination horizon increases.
Figure 10: Long-term imagination from alternative S4WM architectures in the Four Rooms environment. S4WM and S4WM-Full-Posterior have similar imagination quality in this environment, while S4WM-No-MLP makes more mistakes in wall colors and object positions.
S4WM
Figure 11: Comparison of alternative S4WM architectures on context-dependent recall in teleport environments. Each environment is labeled with (context steps | query steps). We provide up to 20 observations after the teleport as additional contexts. S4WM-Full-Posterior performs similarly as S4WM when the context phase is short, and becomes better when the context phase is longer. S4WM-No-MLP performs the worst, suggesting the importance of the MLP in the S4 blocks to context-dependent recall.
S4V
Full-Posterior
S4WM-No-MLP
Figure 12: Context-dependent recall from alternative S4WM architectures in the Four Rooms teleport environment. 20 observations after the teleport are provided as additional contexts. S4WM and S4WM-Full-Posterior perform similarly in this environment, while S4WM-No-MLP makes more mistakes in wall colors and object positions.
Figure 13: Comparison of speed and memory usage for alternative S4WM architectures. S4WM-Full-Posterior is more computationally demanding than S4WM during training, while S4WM-No-MLP is the most efficient.
S4WI
Table 5: Comparison of TSSM-XL with different cache lengths  $m$  on long-term imagination. Each environment is labeled with (context steps | query steps). Larger cache lengths show better generation quality, at the cost of more computation.
Four Rooms (501 | 500)	Ten Rooms (1101 | 900)
Recon. MSE (↓)	Gen. MSE (↓)	Recon. MSE (↓)	Gen. MSE (↓)
TSSM-XL (m = 128)	2.4	224.4	2.6	360.4
TSSM-XL (m = 256)	2.4	124.4	2.5	295.1
TSSM-XL (m = 512)	2.4	38.4	2.5	277.6
Figure 14: Generation MSE per imagination step of TSSM-XL with different cache lengths  $m$ . Each environment is labeled with (context steps | query steps). Increasing the cache length helps improving long-term imagination.
Figure 15: Long-term imagination from TSSM-XL with different cache lengths  $m$  in the Four Rooms environment. TSSM-XL ( $m = 512$ ) works well in this environment and is comparable to S4WM, while TSSM-XL ( $m = 128$ ) and TSSM-XL ( $m = 256$ ) make many mistakes.
TSSM-XL  $(m = 128)$
TSSM-XL  $(m = 256)$
TSSM-XL  $(m = 512)$