arXiv:2410.08893v2 更新类型: replace-cross 摘要：基于模型的强化学习（RL）为大多数无模型RL算法所困扰的数据效率低下提供了解决方案。然而，学习一个稳健的世界模型通常需要复杂且深度的架构，这些架构的计算和训练成本高昂。在世界模型内部，动力学模型对于准确预测尤为重要，已经探索了各种动力学模型架构，每种架构都有其自身的挑战。目前，基于递归神经网络（RNN）的世界模型面临诸如梯度消失和难以有效捕捉长期依赖性的问题。相比之下，使用变压器的方法则受到众所周知的自我注意力机制问题的影响，其中内存和计算复杂性按 \(O(n^2)\) 逐级扩展，\(n\) 表示序列长度。 为了解决这些挑战，我们提出了一个基于状态空间模型（SSM）的世界模型，特别基于Mamba，该模型实现了 \(O(n)\) 的内存和计算复杂性，同时能够有效地捕捉长期依赖性，并高效地使用较长的训练序列。我们还引入了一种新的采样方法，以缓解在训练早期由于世界模型不准确而导致的次优性问题，并结合上述技术，仅使用一个具有700万可训练参数的世界模型，实现了与其他基于模型的RL算法相当的标准化得分。该模型易于访问，并可以在现成的笔记本电脑上进行训练。代码可在 https://github.com/realwenlongwang/Drama.git 获取。