Drama：Mamba 驱动的模型基强化学习在样本与参数效率方面的突破

基于模型的强化学习 (RL) 提出了解决大多数基于模型的 RL 算法数据效率低下的解决方案。然而，学习稳健的世界模型通常需要复杂且深入的架构，这在计算上昂贵且难以训练。在世界模型中，序列模型在准确预测中发挥关键作用，已探索多种架构，每种架构都有其独特挑战。当前，基于循环神经网络 (RNN) 的世界模型在梯度消失和捕捉长期依赖方面受限。而变压器则受制于自注意力机制的二次内存和计算复杂度，随序列长度呈 $O(n^2)$ 增长。为此，我们提出一种基于状态空间模型 (SSM) 的世界模型，Drama，具体利用 Mamba，实现 $O(n)$ 的内存和计算复杂度，同时有效捕捉长期依赖，实现更长序列的高效训练。我们还引入一种新颖的采样方法，以缓解早期训练阶段因错误世界模型导致的次优问题。综合这些技术，Drama 在 Atari100k 基准测试上实现了与其他最先进 (SOTA) 模型基 RL 算法竞争的归一化得分，使用仅 700 万参数的世界模型。Drama 可在标准笔记本电脑等常规硬件上访问和训练。我们的代码已公开于 https://github.com/realwenlongwang/Drama.git。