LLM2D
基于模型的强化学习中的马姆巴加速:样本和参数效率提升
Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient
作者: Wenlong Wang, Ivana Dusparic, Yucheng Shi, Ke Zhang, Vinny Cahill
发布日期: 10/14/2024
arXiv ID: oai:arXiv.org:2410.08893v1

摘要

基于模型的强化学习(RL)为解决大多数无模型 RL 算法所面临的数据效率低下问题提供了一种解决方案。然而,学习一个鲁棒的世界模型通常需要复杂且深层的架构,这在计算和训练方面成本高昂。在世界模型中,动力学模型对于准确预测至关重要,并且已经探索了各种动力学模型架构,每种架构都有其自身的挑战。目前,基于循环神经网络 (RNN) 的世界模型面临着梯度消失和难以有效地捕捉长期依赖关系等问题。相比之下,使用 Transformer 会遇到自注意力机制的众所周知的缺陷,其中内存和计算复杂度都随着 $O(n^2)$ 增长,其中 $n$ 代表序列长度。 为了解决这些挑战,我们提出了一种基于状态空间模型 (SSM) 的世界模型,具体来说是基于 Mamba 的模型,它实现了 $O(n)$ 的内存和计算复杂度,同时有效地捕捉长期依赖关系并有效地促进更长训练序列的使用。我们还引入了一种新的采样方法来缓解训练初期不正确世界模型带来的次优性,将其与上述技术相结合,使用仅 700 万个可训练参数的世界模型,实现了与其他最先进的基于模型的 RL 算法相当的归一化分数。该模型易于访问,可以在普通笔记本电脑上进行训练。我们的代码可在 https://github.com/realwenlongwang/drama.git 上获取。