Facing Off World Model Backbones: RNNs, Transformers, and S4

要約

ワールド モデルは、モデルベースの強化学習 (MBRL) の基本コンポーネントです。
部分的に観測可能な環境で、時間的に拡張された一貫した将来のシミュレーションを実行するには、世界モデルが長期記憶を備えている必要があります。
ただし、Dreamer などの最先端の MBRL エージェントは、メモリ容量が限られているリカレント ニューラル ネットワーク (RNN) をワールド モデル バックボーンとして主に採用しています。
この論文では、長期記憶を改善するための代替世界モデルのバックボーンを探求します。
特に、トランスフォーマーと構造化状態空間シーケンス (S4) モデルの有効性を調査します。これは、低次元シーケンスの長距離依存関係を捕捉する顕著な能力とその補完的な強みを動機としています。
私たちは、S4 とそのバリアントを含む並列化可能な SSM と互換性のある初の世界モデルである S4WM を提案します。
S4WM は、潜在変数モデリングを組み込むことにより、潜在的な想像力を通じて高次元の画像シーケンスを効率的に生成できます。
さらに、RNN、Transformer、および S4 ベースの世界モデルを 4 つの環境セットにわたって広範に比較し、長期的な想像力、状況依存の想起、報酬予測、
そして記憶に基づいた推論。
私たちの調査結果は、S4WM が長期記憶の点で Transformer ベースのワールド モデルよりも優れており、トレーニングと想像力の際に優れた効率を示していることを示しています。
これらの結果は、より強力な MBRL エージェントの開発への道を開きます。

要約(オリジナル)

World models are a fundamental component in model-based reinforcement learning (MBRL). To perform temporally extended and consistent simulations of the future in partially observable environments, world models need to possess long-term memory. However, state-of-the-art MBRL agents, such as Dreamer, predominantly employ recurrent neural networks (RNNs) as their world model backbone, which have limited memory capacity. In this paper, we seek to explore alternative world model backbones for improving long-term memory. In particular, we investigate the effectiveness of Transformers and Structured State Space Sequence (S4) models, motivated by their remarkable ability to capture long-range dependencies in low-dimensional sequences and their complementary strengths. We propose S4WM, the first world model compatible with parallelizable SSMs including S4 and its variants. By incorporating latent variable modeling, S4WM can efficiently generate high-dimensional image sequences through latent imagination. Furthermore, we extensively compare RNN-, Transformer-, and S4-based world models across four sets of environments, which we have tailored to assess crucial memory capabilities of world models, including long-term imagination, context-dependent recall, reward prediction, and memory-based reasoning. Our findings demonstrate that S4WM outperforms Transformer-based world models in terms of long-term memory, while exhibiting greater efficiency during training and imagination. These results pave the way for the development of stronger MBRL agents.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google