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

TL;DR

The paper evaluates backbone architectures for visual world modeling in model-based RL, introducing S4WM—a general, parallelizable SSM-based world model that supports latent imagination. Through offline experiments on memory-centered 3D and 2D environments, S4WM outperforms RNN- and Transformer-based backbones in long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning, while enabling faster training. The study demonstrates the potential of parallelizable SSMs for scalable, long-horizon planning, and shows that S5 variants can further enhance memory capabilities. Limitations include the use of relatively simple, deterministic environments, with future work aimed at stochastic settings and policy-driven evaluation.

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.

Figures (22)

• Figure 1: We propose S4WM, the first S4-based world model for improving long-term memory. S4WM efficiently models the long-range dependencies of environment dynamics in a compact latent space, using a stack of S4 blocks. This crucially allows fully parallelized training and fast recurrent latent imagination. S4WM is a general framework that is compatible with any parallelizable SSM including S5 and other S4 variants.
• Figure 2: Partially observable 3D (Top) and 2D (Bottom) environments for evaluating memory capabilities of world models, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning.
• Figure 3: Comparison of speed and memory usage during training and imagination. S4WM is the fastest to train, while RSSM-TBTT is the most memory-efficient during training and has the highest throughput during imagination.
• 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.
• 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.