[論文レビュー] Transformers are Meta-Reinforcement Learners

Luckeciano C. Melo|arXiv (Cornell University)|Jun 14, 2022

Reinforcement Learning in Robotics被引用数 8

ひとこと要約

本論文では、トランスフォーマー・アーキテクチャを用いて、マルチヘッド自己注意機構を通じてエピソードメモリを再帰的に構築するメタ強化学習エージェント、TrMRLを提案する。これにより、高速な適応と頑健な一般化が可能となる。TrMRLは、HalfCheetahVel やロボット操作タスクを含む連続制御環境において、優れたもしくは同等のサンプル効率と分布外一般化性能を達成している。

ABSTRACT

The transformer architecture and variants presented remarkable success across many machine learning tasks in recent years. This success is intrinsically related to the capability of handling long sequences and the presence of context-dependent weights from the attention mechanism. We argue that these capabilities suit the central role of a Meta-Reinforcement Learning algorithm. Indeed, a meta-RL agent needs to infer the task from a sequence of trajectories. Furthermore, it requires a fast adaptation strategy to adapt its policy for a new task -- which can be achieved using the self-attention mechanism. In this work, we present TrMRL (Transformers for Meta-Reinforcement Learning), a meta-RL agent that mimics the memory reinstatement mechanism using the transformer architecture. It associates the recent past of working memories to build an episodic memory recursively through the transformer layers. We show that the self-attention computes a consensus representation that minimizes the Bayes Risk at each layer and provides meaningful features to compute the best actions. We conducted experiments in high-dimensional continuous control environments for locomotion and dexterous manipulation. Results show that TrMRL presents comparable or superior asymptotic performance, sample efficiency, and out-of-distribution generalization compared to the baselines in these environments.

研究の動機と目的

• 注意機構の文脈依存的パラメータ化を活用することで、メタ強化学習における高速適応の課題に取り組むこと。
• トランスフォーマー層を用いてエピソードメモリを再帰的に精錬することで、記憶の再活性化を模倣する記憶ベースのメタ強化学習エージェントを設計すること。
• 安定したトランスフォーマー基盤アーキテクチャを用いて、連続制御環境におけるサンプル効率と分布外一般化性能を向上させること。
• T-Fixup重み初期化が、トランスフォーマー基盤メタ強化学習エージェントの訓練を安定化させる有効性を調査すること。
• 自己注意機構がベイズリスクを最小化する一貫性のある表現を計算することにより、より良い行動予測が達成されることを検証すること。

提案手法

• TrMRLは、因果的マルチヘッド自己注意機構を用いて、トランスフォーマー層を横断して最近の作業メモリをタスクに適したエピソードメモリに結びつける。
• 各トランスフォーマー層は、ベイズリスクを最小化する一貫性のある表現を計算することで、エピソードメモリを精錬し、行動選択に適した特徴の質を向上させる。
• 最後の層からの最終エピソードメモリがポリシー・ヘッドの入力として使用され、文脈依存的行動予測が可能になる。
• T-Fixup重み初期化が、強化学習における初期探索段階の不安定性を軽減し、訓練を安定化させるために適用される。
• エピソードメモリは層を介して再帰的に構築され、各層が最適な行動を予測するのに適した一貫性のある表現へと精錬されていく。
• 各層における表現の質を評価するために、線形回帰モデルとコサイン類似度が用いられ、精錬度と予測性能が測定される。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーの自己注意機構は、文脈依存的ポリシーを形成することで、メタ強化学習における効果的な高速適応戦略として機能するか？
• RQ2 トランスフォーマー層によるエピソードメモリの再帰的精錬は、ベイズリスクを最小化する一貫性のある表現を生成し、行動予測を改善するか？
• RQ3 TrMRLは、既存のメタ強化学習ベースラインと比較して、サンプル効率および分布外一般化性能においてどのように性能を発揮するか？
• RQ4 T-Fixup初期化は、トランスフォーマー基盤メタ強化学習エージェントの訓練をどの程度安定化させるか？
• RQ5 中間層におけるエピソードメモリ表現は、専門家の行動を予測するために意味的に有用であると示せるか？これは、段階的精錬が進行していることを示唆する。

主な発見

• TrMRLは、HalfCheetahVelで最初のエピソードから高い性能を達成し、最適性能に到達するまでに約20ステップのタイムステップしか必要としない。これは、優れたオンライン適応能力を示している。
• ターゲット速度が[3.0, 4.0]に設定されたHalfCheetahVelにおける分布外一般化において、TrMRLはPEARL や VariBAD を含むすべてのベースラインを顕著に上回り、頑健なタスク表現学習能力を示している。
• 中間層のエピソードメモリ表現と最終層の表現との間の表現誤差が、層を経て減少する傾向を示しており、一貫性のある表現への段階的精錬が確認された。
• 各層からのエピソードメモリ特徴を用いた線形回帰モデルでは、テストセットにおける平均二乗誤差が減少しており、深さが進むにつれて表現が専門家の行動をより予測可能にするようになっていることが示された。
• アブレーションスタディにより、T-Fixup初期化が訓練の安定性と性能に不可欠であることが確認され、ネットワークの深さ、シーケンス長、アテンションヘッド数の変動に感受することが判明した。
• エピソードメモリの精錬プロセスは、定理4.1と整合しており、自己注意機構が各層でベイズリスクを最小化する一貫性のある表現を計算していることが実証された。