Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient

要約

モデルベースの強化学習（RL）は、ほとんどのモデルのないRLアルゴリズムを悩ませるデータの非効率性に対するソリューションを提供します。
ただし、堅牢な世界モデルを学ぶには、しばしば複雑で深いアーキテクチャが必要であり、それは計算と訓練に費用がかかります。
世界モデル内では、ダイナミクスモデルは正確な予測に特に重要であり、それぞれに独自の課題があるさまざまなダイナミクスモデルアーキテクチャが調査されています。
現在、再発性ニューラルネットワーク（RNN）ベースの世界モデルは、勾配の消失や長期依存関係を効果的にキャプチャする困難などの問題に直面しています。
対照的に、変圧器の使用は、メモリと計算の複雑さの両方が$ o（n^2）$として、$ n $がシーケンスの長さを表す、自己関節メカニズムのよく知られている問題に苦しんでいます。
これらの課題に対処するために、特にMAMBAに基づいた状態空間モデル（SSM）ベースの世界モデルを提案します。これは、$ O（n）$メモリと計算の複雑さを達成しながら、長期依存性を効果的にキャプチャし、より長いトレーニングシーケンスの使用を効率的に促進します。
。
また、トレーニングの初期段階で誤った世界モデルによって引き起こされる亜最適性を緩和するための新しいサンプリング方法を紹介し、前述のテクニックと組み合わせて、他の最先端のモデルベースのRLアルゴリズムに匹敵する正規化されたスコアを達成する
わずか700万のトレーニング可能なパラメーターワールドモデルを使用しています。
このモデルにはアクセス可能で、既製のラップトップでトレーニングできます。
私たちのコードは、https：//github.com/realwenlongwang/drama.gitで入手できます

要約(オリジナル)

Model-based reinforcement learning (RL) offers a solution to the data inefficiency that plagues most model-free RL algorithms. However, learning a robust world model often demands complex and deep architectures, which are expensive to compute and train. Within the world model, dynamics models are particularly crucial for accurate predictions, and various dynamics-model architectures have been explored, each with its own set of challenges. Currently, recurrent neural network (RNN) based world models face issues such as vanishing gradients and difficulty in capturing long-term dependencies effectively. In contrast, use of transformers suffers from the well-known issues of self-attention mechanisms, where both memory and computational complexity scale as $O(n^2)$, with $n$ representing the sequence length. To address these challenges we propose a state space model (SSM) based world model, specifically based on Mamba, that achieves $O(n)$ memory and computational complexity while effectively capturing long-term dependencies and facilitating the use of longer training sequences efficiently. We also introduce a novel sampling method to mitigate the suboptimality caused by an incorrect world model in the early stages of training, combining it with the aforementioned technique to achieve a normalised score comparable to other state-of-the-art model-based RL algorithms using only a 7 million trainable parameter world model. This model is accessible and can be trained on an off-the-shelf laptop. Our code is available at https://github.com/realwenlongwang/Drama.git

arxiv情報

提供元, 利用サービス

arxiv.jp, Google