世界模型骨干网络对决：RNN、Transformer与S4
机器学习
2023-11-10 v2
摘要
世界模型是基于模型的强化学习（MBRL）的基本组成部分。为了在部分可观测环境中执行时间上扩展且一致的未来模拟，世界模型需要具备长期记忆。然而，最先进的MBRL智能体（如Dreamer）主要采用循环神经网络（RNN）作为世界模型骨干，其记忆容量有限。本文旨在探索替代性世界模型骨干以提升长期记忆。具体而言，我们研究了Transformer与结构化状态空间序列（S4）模型的有效性，其动机在于二者在低维序列中捕捉长程依赖的卓越能力及互补优势。我们提出S4WM，首个兼容可并行化SSM（包括S4及其变体）的世界模型。通过引入潜变量建模，S4WM可通过潜想象高效生成高维图像序列。此外，我们在四组环境上广泛比较了基于RNN、Transformer和S4的世界模型，这些环境经我们定制以评估世界模型的关键记忆能力，包括长期想象、上下文相关回忆、奖励预测与基于记忆的推理。结果表明，S4WM在长期记忆上优于基于Transformer的世界模型，同时在训练与想象中展现更高效率。这些结果为开发更强的MBRL智能体铺平了道路。
引用
@article{arxiv.2307.02064,
title = {Facing Off World Model Backbones: RNNs, Transformers, and S4},
author = {Fei Deng and Junyeong Park and Sungjin Ahn},
journal= {arXiv preprint arXiv:2307.02064},
year = {2023}
}
备注
NeurIPS 2023. Added instantiation with S5. Project page: https://fdeng18.github.io/s4wm