Structured State Space Models for In-Context Reinforcement Learning

要約

構造化状態空間シーケンス (S4) モデルは、最近、長距離シーケンス モデリング タスクで最先端のパフォーマンスを達成しました。
これらのモデルは、高速な推論速度と並列化可能なトレーニングも備えているため、多くの強化学習設定で役立つ可能性があります。
隠れ状態の初期化とリセットを並行して実行できるようにする S4 のバリアントへの変更を提案し、強化学習タスクに取り組むことができるようにします。
単純なメモリベースのタスクで、変更したアーキテクチャが Transformers よりも漸近的に高速に実行され、LSTM モデルよりも優れたパフォーマンスを発揮することを示します。
次に、長距離シーケンスを処理するモデルの能力を活用することにより、エージェントにランダムにサンプリングされた連続制御環境が与えられ、環境のランダムにサンプリングされた線形投影と組み合わされる、困難なメタ学習タスクで強力なパフォーマンスを達成します。
観察と行動。
さらに、結果として得られるモデルが、分散外の保留タスクに適応できることを示します。
全体として、このホワイト ペーパーで提示された結果は、S4 モデルが、コンテキスト内強化学習に使用されるデフォルト アーキテクチャの強力な候補であることを示唆しています。

要約(オリジナル)

Structured state space sequence (S4) models have recently achieved state-of-the-art performance on long-range sequence modeling tasks. These models also have fast inference speeds and parallelisable training, making them potentially useful in many reinforcement learning settings. We propose a modification to a variant of S4 that enables us to initialise and reset the hidden state in parallel, allowing us to tackle reinforcement learning tasks. We show that our modified architecture runs asymptotically faster than Transformers and performs better than LSTM models on a simple memory-based task. Then, by leveraging the model’s ability to handle long-range sequences, we achieve strong performance on a challenging meta-learning task in which the agent is given a randomly-sampled continuous control environment, combined with a randomly-sampled linear projection of the environment’s observations and actions. Furthermore, we show the resulting model can adapt to out-of-distribution held-out tasks. Overall, the results presented in this paper suggest that the S4 models are a strong contender for the default architecture used for in-context reinforcement learning

arxiv情報

提供元, 利用サービス

arxiv.jp, Google