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

要約

モデルベースの強化学習（RL）は、ほとんどのモデルのないRLアルゴリズムを悩ませるデータの非効率性に対するソリューションを提供します。
ただし、堅牢な世界モデルを学習するには、多くの場合、複雑で深いアーキテクチャが必要であり、計算的に高価で訓練が困難です。
世界モデル内では、シーケンスモデルが正確な予測に重要な役割を果たし、それぞれに独自の課題があるさまざまなアーキテクチャが検討されています。
現在、再発性ニューラルネットワーク（RNN）ベースの世界モデルは、勾配の消失と長期依存関係のキャプチャに苦労しています。
一方、変圧器は、$ o（n^2）$としてスケーリングし、$ n $がシーケンスの長さである、自己触媒メカニズムの2次記憶と計算の複雑さに悩まされます。
これらの課題に対処するために、$ O（n）$メモリと計算の複雑さを達成しながら、長期依存関係を効果的にキャプチャし、より長いシーケンスで効率的なトレーニングを可能にする状態空間モデル（SSM）ベースの世界モデル、ドラマ、特にMambaを活用することを提案します。
また、初期のトレーニング段階で誤った世界モデルによって引き起こされる最適性を緩和するための新しいサンプリング方法を紹介します。
これらのテクニックを組み合わせることで、ドラマは、700万パラメーターの世界モデルのみを使用して、他の最先端（SOTA）モデルベースのRLアルゴリズムと競合するAtari100Kベンチマークで正規化されたスコアを実現します。
ドラマは、標準のラップトップなど、既製のハードウェアでアクセスしやすくトレーニング可能です。
私たちのコードは、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 requires complex and deep architectures, which are computationally expensive and challenging to train. Within the world model, sequence models play a critical role in accurate predictions, and various architectures have been explored, each with its own challenges. Currently, recurrent neural network (RNN)-based world models struggle with vanishing gradients and capturing long-term dependencies. Transformers, on the other hand, suffer from the quadratic memory and computational complexity of self-attention mechanisms, scaling as $O(n^2)$, where $n$ is the sequence length. To address these challenges, we propose a state space model (SSM)-based world model, Drama, specifically leveraging Mamba, that achieves $O(n)$ memory and computational complexity while effectively capturing long-term dependencies and enabling efficient training with longer sequences. We also introduce a novel sampling method to mitigate the suboptimality caused by an incorrect world model in the early training stages. Combining these techniques, Drama achieves a normalised score on the Atari100k benchmark that is competitive with other state-of-the-art (SOTA) model-based RL algorithms, using only a 7 million-parameter world model. Drama is accessible and trainable on off-the-shelf hardware, such as a standard laptop. Our code is available at https://github.com/realwenlongwang/Drama.git.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google