[论文解读] Facing Off World Model Backbones: RNNs, Transformers, and S4

Fei Deng, Junyeong Park+1 more|arXiv (Cornell University)|Jul 5, 2023

Reinforcement Learning in Robotics被引用 4

一句话总结

该论文提出了S4WM，这是首个与可并行化的结构化状态空间模型（PSSM）如S4和S5兼容的世界模型框架，实现了高效的长上下文视觉序列建模。S4WM在长期想象、上下文依赖性回忆、奖励预测和基于记忆的推理方面优于基于RNN和Transformer的世界模型，同时训练速度更快，序列生成效率更高。

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.

研究动机与目标

• 开发一种与可并行化的S4和S5等S4WM兼容的一般性世界模型框架，用于高维视觉序列。
• 研究RNN、Transformer和S4作为视觉世界模型主干在网络长期记忆密集型任务中的有效性。
• 评估并比较不同主干架构在关键记忆能力上的表现：长期想象、上下文依赖性回忆、奖励预测和基于记忆的推理。
• 证明S4WM在长期记忆能力和推理效率方面优于基于RNN和Transformer的基线模型。

提出的方法

• 提出S4WM，一种基于变分推断的概率潜在变量模型，采用S4作为图像序列生成的主干网络。
• 采用具有参数化状态转移的结构化状态空间模型（SSM），实现次二次复杂度并支持并行训练。
• 引入潜在空间瓶颈以压缩高维观测，通过潜在空间想象实现图像序列的高效生成。
• 设计了四个专用环境——长时程、上下文回忆、奖励预测和多门钥匙，用于评估记忆能力。
• 使用重建与生成的MSE、想象准确性以及奖励预测性能作为评估指标。
• 通过离线探针和模型预测控制（MPC）实验，验证S4WM在下游策略学习中的潜力。

实验结果

研究问题

• RQ1 基于S4的世界模型能否有效建模高维视觉序列中的长距离依赖？
• RQ2 在长期想象和记忆保持方面，S4WM与基于RNN和Transformer的世界模型相比表现如何？
• RQ3 在视觉世界建模中，S4WM是否实现了比RNN和Transformer更优的训练效率和推理速度？
• RQ4 在动态、上下文依赖的推理任务中，S4WM能否保持准确的记忆更新？
• RQ5 S4WM能否泛化到复杂、记忆密集型推理任务，如部分可观察环境中的多步规划？

主要发现

• 在长时程任务中，S4WM实现了最高的想象准确性，优于RNN和Transformer，后者在长序列中无法保持一致的状态表示。
• 在上下文回忆任务中，S4WM在所有上下文长度下均保持高准确率，而TSSM-XL和RSSM-TBTT在超过100步后显著下降。
• S4WM是唯一能准确预测想象中奖励的模型，TSSM-XL在长序列上失败，而RSSM-TBTT的表现接近随机猜测。
• 在多门钥匙环境中，S4WM在最复杂设置（7把钥匙，654次查询步骤）下生成MSE最低（0.10），优于TSSM-XL（9.24）和RSSM-TBTT（6.28）。
• S4WM的训练速度优于RNN和Transformer基线模型，且想象吞吐量高于RNN，展现出更优的计算效率。
• 离线探针和MPC实验表明，S4WM在任务成功率上优于RSSM，证实其在MBRL中策略学习的潜力。