World models are a fundamental component in model-based reinforcement learning (MBRL) agents. 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 S4-based world model that can 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 specifically 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.

翻译：世界模型是基于模型的强化学习（MBRL）智能体中的基本组件。为了在部分可观测环境中对未来进行时间上延展且一致的仿真，世界模型需要具备长期记忆能力。然而，当前最先进的MBRL智能体（如Dreamer）主要采用记忆容量有限的循环神经网络（RNN）作为其世界模型骨干。本文旨在探索替代性的世界模型骨干网络，以改进长期记忆性能。具体而言，受Transformer和结构化状态空间序列（S4）模型在低维序列中捕捉长程依赖关系的卓越能力及互补优势启发，我们研究了它们作为世界模型骨干的有效性。我们提出了S4WM——首个基于S4的世界模型，能够通过潜在想象生成高维图像序列。此外，我们在四组专门为评估世界模型关键记忆能力（包括长期想象、上下文相关回忆、奖励预测和基于记忆的推理）而定制的环境中，对基于RNN、Transformer和S4的世界模型进行了全面比较。研究结果表明，S4WM在长期记忆方面优于基于Transformer的世界模型，同时在训练和想象阶段展现出更高的效率。这些结果为开发更强大的MBRL智能体铺平了道路。