近年来，基于模型的强化学习（MBRL）在解决各种任务中展现出强大的能力。然而，现有的MBRL代理在处理长期依赖关系方面存在困难，这限制了它们在涉及行动和结果之间长时间间隔的任务，或需要回忆远距离观察结果来指导当前行动的任务中的表现。为了克服这一挑战，本文介绍了一种名为“回忆想象”（Recall to Imagine，R2I）的新方法，它将一种新型状态空间模型（SSM）整合到MBRL代理的世界模型中，以增强长期记忆和长范围信用分配能力。
长期记忆的困境：世界模型的挑战
在强化学习中，世界模型通过捕捉环境的动态变化，赋予代理感知、模拟和规划的能力。MBRL代理通过学习过去经验中的世界模型，能够“想象”其行动的未来后果，并做出明智的决策。
然而，MBRL面临着两个关键挑战：
- 长期依赖关系：世界模型需要准确地模拟环境的演化和未来奖励，并整合代理在长时间范围内的行动。由于世界模型的网络架构限制，难以学习长距离依赖关系，导致长期记忆和信用分配问题。
- 信用分配问题：需要评估行动对未来奖励的影响，这在长时间间隔的任务中尤其困难。
S4模型：突破长期依赖关系的利器
近年来，状态空间模型（SSM）在监督学习和自监督学习任务中展现出捕捉超长序列依赖关系的能力。其中，S4模型凭借其优异的性能，重新定义了长距离序列建模的研究方向。S4模型源于一个时间不变的线性动力系统，它学习状态矩阵，并能够有效地捕捉高达16K长度的依赖关系，超越了先前方法的局限性。
回忆想象：赋予世界模型记忆能力
R2I方法首次将S4模型的变体应用于MBRL，为代理提供了强大的长期记忆能力。R2I是一种通用且计算效率高的方法，在各种记忆领域展现出最先进的性能。
R2I的核心思想：
- S3M（Structured State-Space Model）：R2I将S4模型整合到DreamerV3的世界模型中，构建了S3M。S3M通过并行计算和循环推理模式，实现了快速学习和高效的轨迹模拟。
- 非循环表示模型：为了实现并行计算，R2I将表示模型从循环模型转变为非循环模型，使每个时间步的推断能够独立进行。
- SSM的并行计算：R2I采用并行扫描方法，而非卷积方法，来计算SSM的隐藏状态，这使得模型能够更有效地处理长序列数据，并支持隐藏状态的重置。
- 演员-评论家结构：R2I在演员-评论家结构中，将SSM的隐藏状态传递给策略网络，而不是使用GRU的隐藏状态。这对于解决记忆密集型环境中的长期依赖关系至关重要。
R2I的优异表现
R2I在各种记忆领域展现出优异的性能，包括：
- POPGym：R2I在POPGym中取得了最先进的性能，这是一个专门设计用来测试长期记忆能力的强化学习环境。
- 行为套件（BSuite）：R2I在BSuite中超越了其他基准方法，BSuite是一套用于评估强化学习代理能力的标准化任务。
- 记忆迷宫：R2I在记忆迷宫中超越了人类，这是一个需要极强的长期记忆能力才能解决的复杂3D环境。
此外，R2I在经典的强化学习任务中也展现出良好的性能，例如：
- Atari：R2I在Atari游戏中取得了与其他最先进方法相当的性能。
- DMC：R2I在DMC中也展现出良好的性能，DMC是一套用于评估强化学习代理在连续控制任务中的能力的标准化环境。
总结
R2I方法通过将S4模型整合到世界模型中，赋予了MBRL代理强大的长期记忆能力。R2I在各种记忆领域展现出最先进的性能，并超越了人类在记忆迷宫中的表现。R2I的通用性和计算效率使其成为解决长期记忆和信用分配问题的强大工具。
参考文献
- Hafner, D., Lillicrap, T. P., Norouzi, M., & Schmidhuber, J. (2023). DreamerV3: Learning and planning in
large, stochastic worlds. arXiv preprint arXiv:2302.04023. - Gu, S., Hoffman, M., Edwards, H., Roberts, A., Brunskill, E., & Lake, B. M. (2021a). S4: Deep
learning on symbolic sequences. arXiv preprint arXiv:2106.03107. - Smith, S. L., Gupta, A., Gu, S., & Lake, B. M. (2023). Efficiently modeling long sequences with
state space models. arXiv preprint arXiv:2302.10871.