[論文レビュー] StARformer: Transformer with State-Action-Reward Representations for Visual Reinforcement Learning

Jinghuan Shang, Kumara Kahatapitiya+2 more|arXiv (Cornell University)|Oct 12, 2021

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

StARformerは、短い時間窓内で状態-行動-報酬（StAR）表現を明示的にモデル化することで、マルコフ的インダクティブバイアスを注入し、長期的なシーケンスモデリングを改善する、視覚的強化学習のための新しいTransformerアーキテクチャを導入する。ViT風の画像パッチと畳み込み特徴を組み合わせ、ステップとシーケンスのTransformerをインタラーリーブして融合することで、オフラインRLおよびアシスト学習設定において、AtariおよびDM Control Suiteベンチマークで最先端の手法を上回る性能を達成する。

ABSTRACT

Reinforcement Learning (RL) can be considered as a sequence modeling task: given a sequence of past state-action-reward experiences, an agent predicts a sequence of next actions. In this work, we propose State-Action-Reward Transformer (StARformer) for visual RL, which explicitly models short-term state-action-reward representations (StAR-representations), essentially introducing a Markovian-like inductive bias to improve long-term modeling. Our approach first extracts StAR-representations by self-attending image state patches, action, and reward tokens within a short temporal window. These are then combined with pure image state representations -- extracted as convolutional features, to perform self-attention over the whole sequence. Our experiments show that StARformer outperforms the state-of-the-art Transformer-based method on image-based Atari and DeepMind Control Suite benchmarks, in both offline-RL and imitation learning settings. StARformer is also more compliant with longer sequences of inputs. Our code is available at https://github.com/elicassion/StARformer.

研究の動機と目的

• 局所的な状態-行動-報酬遷移を明示的にモデル化することで、視覚的強化学習における長期的シーケンスモデリングを改善すること。
• 局所的なStAR表現を通じてマルコフ的インダクティブバイアスを注入することで、自己注意機構の負担を軽減すること。
• 全体像のコンテキストを保持するため、画像状態の細粒度な空間情報を維持するために、画像全体のCNN特徴ではなくViT風のパッチ埋め込みを用いること。
• 短期的遷移モデリングと長距離コンテキスト学習を分離することで、より長い入力シーケンスへの一般化性とスケーラビリティを向上させること。
• リターン・トゥ・ゴー、ステップワイズ報酬、および報酬なし（アシスト学習）の多様な報酬設定において、StARformerの有効性を評価すること。

提案手法

• モデルは、1つの時間ステップ内で画像状態パッチ、行動トークン、報酬トークンを自己注意するステップTransformerを採用し、局所的なStAR表現を学習する。
• 画像状態は、細粒度な空間的詳細を保持するため、ViT風のパッチにエンコードされる。一方、グローバルコンテキストを抽出するために、畳み込みネットワークを介して純粋な画像特徴が抽出される。
• シーケンスTransformerは、全トラジェクトリにわたるStAR表現とグローバル畳み込み特徴を組み合わせ、因果的自己注意を用いて行動予測を生成する。
• レイヤー単位の融合メカニズムにより、各TransformerレイヤーでStAR表現とグローバル画像特徴が統合され、モデルの深さを増さずに表現能力が向上する。
• アーキテクチャは、リターン・トゥ・ゴー、ステップワイズ報酬、および報酬なし（アシスト学習）の複数の報酬設定をサポートしており、異なる学習パラダイムにおけるアブレーションが可能である。
• 異なる統合戦略（例：P+C、C+P、StAR Fusion、StAR Stack）を用いたモデル変種を評価し、最適な特徴統合方法を特定する。

実験結果

リサーチクエスチョン

• RQ1 局所的な状態-行動-報酬遷移の明示的モデリングが、視覚的強化学習における長期的シーケンスモデリングを改善できるか？
• RQ2 ViT風のパッチ埋め込みと畳み込み特徴を融合させることで、単一の表現タイプのみを用いる場合と比較して性能が向上するか？
• RQ3 報酬信号の選択（リターン・トゥ・ゴー、ステップワイズ報酬、または報酬なし）が、モデルの性能と一般化性に与える影響は何か？
• RQ4 提示されたインタラーリーブ型ステップ-シーケンスTransformer設計は、スタック型や和集合型の変種よりも効果的か？
• RQ5 報酬信号が存在しない状況において、StARformerはアシスト学習で優れた性能を発揮できるか。これは、優れたシーケンスモデリング能力を示唆するか？

主な発見

• StARformerは、オフラインRLおよびアシスト学習設定において、AtariおよびDeepMind Control Suiteベンチマークで最先端のDecision Transformerを上回る性能を達成する。
• P+C統合戦略（ステップTransformerでパッチ埋め込み、シーケンスTransformerで畳み込み特徴）が優れた性能を発揮し、細粒度から粗粒度への特徴階層の恩恵が裏付けられる。
• StARformerは、報酬なしのアシスト学習を含むすべての報酬設定で強力な性能を維持しており、リターン・トゥ・ゴーを用いたDecision Transformerを上回る。
• レイヤー単位の融合メカニズム（オリジナルのStARモデル）は、要素ごとの和集合（StAR Fusion）やスタック型アーキテクチャ（StAR Stack）を上回る性能を示しており、多段階の抽象化が重要であることを示している。
• 注意マップの結果、行動トークンが関連する空間領域（例：Pongにおけるパドル）に意味的に注目していることが確認され、モデルが意味的に根拠のある表現を学習していることが裏付けられる。
• 長期間の入力シーケンスに対しても良好な適合性を示しており、視覚的強化学習における従来のTransformerベースのアプローチに比べ、スケーラビリティが向上していることが示唆される。