# Facing Off World Model Backbones: RNNs, Transformers, and S4

arXiv: https://arxiv.org/abs/2307.02064 · HTML: https://mlsanity.com/paper/2307.02064/

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 55 of its 70 references that are in the map.
- First posted on arXiv 5 Jul 2023.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 30 of 55)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [On the Parameterization and Initialization of Diagonal State Space Models](https://mlsanity.com/paper/2206.11893/) (2022)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Flexible Diffusion Modeling of Long Videos](https://mlsanity.com/paper/2205.11495/) (2022)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Long Range Language Modeling via Gated State Spaces](https://mlsanity.com/paper/2206.13947/) (2022)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [It's Raw! Audio Generation with State-Space Models](https://mlsanity.com/paper/2202.09729/) (2022)

## Built on by (most cited first, 30 of 31)

- [World Models for Autonomous Driving: An Initial Survey](https://mlsanity.com/paper/2403.02622/) (2024)
- [Adaptive Horizon Actor-Critic for Policy Learning in Contact-Rich Differentiable Simulation](https://mlsanity.com/paper/2405.17784/) (2024)
- [StateSpaceDiffuser: Bringing Long Context to Diffusion World Models](https://mlsanity.com/paper/2505.22246/) (2025)
- [Slot State Space Models](https://mlsanity.com/paper/2406.12272/) (2024)
- [SSM Meets Video Diffusion Models: Efficient Long-Term Video Generation with Structured State Spaces](https://mlsanity.com/paper/2403.07711/) (2024)
- [Deep Continuous-Time State-Space Models for Marked Event Sequences](https://mlsanity.com/paper/2412.19634/) (2024)
- [Accelerating Model-Based Reinforcement Learning with State-Space World Models](https://mlsanity.com/paper/2502.20168/) (2025)
- [Four Principles for Physically Interpretable World Models](https://mlsanity.com/paper/2503.02143/) (2025)
- [Rethinking Transformers in Solving POMDPs](https://mlsanity.com/paper/2405.17358/) (2024)
- [Discrete Codebook World Models for Continuous Control](https://mlsanity.com/paper/2503.00653/) (2025)
- [medDreamer: Model-Based Reinforcement Learning with Latent Imagination on Complex EHRs for Clinical Decision Support](https://mlsanity.com/paper/2505.19785/) (2025)
- [How Safe Am I Given What I See? Calibrated Prediction of Safety Chances for Image-Controlled Autonomy](https://mlsanity.com/paper/2308.12252/) (2023)
- [Hieros: Hierarchical Imagination on Structured State Space Sequence World Models](https://mlsanity.com/paper/2310.05167/) (2023)
- [DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing](https://mlsanity.com/paper/2311.01450/) (2023)
- [Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?](https://mlsanity.com/paper/2405.12094/) (2024)
- [MambaPlace:Text-to-Point-Cloud Cross-Modal Place Recognition with Attention Mamba Mechanisms](https://mlsanity.com/paper/2408.15740/) (2024)
- [Learning to Play Atari in a World of Tokens](https://mlsanity.com/paper/2406.01361/) (2024)
- [Zero-Shot Reinforcement Learning Under Partial Observability](https://mlsanity.com/paper/2506.15446/) (2025)
- [Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient](https://mlsanity.com/paper/2410.08893/) (2024)
- [Counterfactual World Models via Digital Twin-conditioned Video Diffusion](https://mlsanity.com/paper/2511.17481/) (2025)
- [Composition of Memory Experts for Diffusion World Models](https://mlsanity.com/paper/2605.18813/) (2026)
- [Learning Local Causal World Models with State Space Models and Attention](https://mlsanity.com/paper/2505.02074/) (2025)
- [On Memory: A comparison of memory mechanisms in world models](https://mlsanity.com/paper/2512.06983/) (2025)
- [Physically Viable World Models: A Case for Query-Conditioned Embodied AI](https://mlsanity.com/paper/2605.30542/) (2026)
- [Learning Off-policy with Model-based Intrinsic Motivation For Active Online Exploration](https://mlsanity.com/paper/2404.00651/) (2024)
- [HiCMamba: Enhancing Hi-C Resolution and Identifying 3D Genome Structures with State Space Modeling](https://mlsanity.com/paper/2503.10713/) (2025)
- [Vector Quantization in the Brain: Grid-like Codes in World Models](https://mlsanity.com/paper/2510.16039/) (2025)
- [Real-World Robot Control by Deep Active Inference With a Temporally Hierarchical World Model](https://mlsanity.com/paper/2512.01924/) (2025)
- [Probing the Impact of Scale on Data-Efficient, Generalist Transformer World Models for Atari](https://mlsanity.com/paper/2605.08578/) (2026)
- [Predictive Training with Latent Imagination for Visual Quadruped Navigation](https://mlsanity.com/paper/2607.17574/) (2026)

## Abstract (arXiv)

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.