[論文レビュー] Facing Off World Model Backbones: RNNs, Transformers, and S4

Fei Deng, Junyeong Park+1 more|arXiv (Cornell University)|Jul 5, 2023

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

この論文では、S4 や S5 のような並列化可能な構造的状態空間モデル（PSSM）と互換性を持つ世界モデルフレームワークである S4WM を紹介する。これにより、長期間にわたる視覚的シーケンスの効率的で高次元なモデリングが可能になる。S4WM は、長期的想像、文脈依存的想起、報酬予測、記憶に基づく推論において、RNN や Transformer を用いた世界モデルを上回り、かつトレーニングが速く、シーケンスの生成もより効率的である。

ABSTRACT

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.

研究の動機と目的

• 高次元の視覚的シーケンスを対象とする、並列化可能な SSM（PSSM）である S4 や S5 と互換性を持つ一般化された世界モデルフレームワークの開発。
• RNN、Transformer、S4 をバックボーンとして用いた視覚的世界モデルの長期間にわたる記憶を要するタスクにおける有効性の調査。
• 長期間の想像、文脈依存的想起、報酬予測、記憶に基づく推論といった主要な記憶能力について、異なるバックボーンアーキテクチャの性能を評価・比較すること。
• S4WM が RNN や Transformer を用いたベースラインと比較して、優れた長期記憶能力と推論効率を達成していることを示すこと。

提案手法

• S4 をバックボーンとして画像シーケンス生成に用いる、変分推論に基づく確率的潜在変数モデルである S4WM を提案する。
• パラメータ化された状態遷移を備えた構造的状態空間モデル（SSM）を採用し、二次未満の計算量と並列トレーニングを実現する。
• 高次元の観測を圧縮するための潜在空間ボトルネックを導入し、潜在空間における想像を通じて画像シーケンスの効率的生成を可能にする。
• 記憶能力を評価するための4つの専用環境（長期間、文脈的想起、報酬予測、複数ドアキーモード）を設計する。
• 再構成誤差と生成誤差の MSE、想像の正確性、報酬予測のパフォーマンスを評価指標として用いる。
• オフラインプローブとモデル予測制御（MPC）実験を用いて、S4WM が下流のポリシー学習における可能性を検証する。

実験結果

リサーチクエスチョン

• RQ1 S4 を用いた世界モデルは、高次元の視覚的シーケンスにおける長距離依存関係を効果的にモデル化できるか？
• RQ2 S4WM は、長期的想像と記憶保持の観点から、RNN や Transformer を用いた世界モデルと比較してどのように異なるか？
• RQ3 S4WM は、視覚的世界モデリングにおいて、RNN や Transformers よりも優れたトレーニング効率と推論速度を達成できるか？
• RQ4 S4WM は、動的で文脈依存的な推論タスクにおいて、正確な記憶更新を維持できるか？
• RQ5 S4WM は、部分的に観測可能な環境における複数ステップの計画のような複雑で記憶を多く要する推論タスクに一般化できるか？

主な発見

• 長期間タスクにおいて、S4WM は最も高い想像の正確性を達成し、RNN や Transformer は長期間にわたって一貫した状態表現を維持できなかった。
• 文脈的想起タスクでは、S4WM は全長の文脈長で高い正確性を維持するが、TSSM-XL や RSSM-TBTT は100ステップを超えると著しく性能が低下する。
• 想像内での報酬予測において、S4WM のみが正確に予測を実行しており、TSSM-XL は長期間のシーケンスで失敗し、RSSM-TBTT はほぼランダムな推測に近い結果となった。
• 複数ドアキーモードでは、S4WM が最も複雑な設定（7つの鍵、654ステップのクエリ）で生成 MSE 0.10 を達成し、TSSM-XL（9.24）と RSSM-TBTT（6.28）を上回った。
• S4WM は RNN や Transformer のベースラインより高速にトレーニングされ、RNN よりも高い想像スループットを示しており、優れた計算効率を示した。
• オフラインプローブと MPC 実験の結果、S4WM は RSSM よりも高いタスク成功率を達成しており、MBRL におけるポリシー学習の可能性を裏付けた。