[論文レビュー] CoBERL: Contrastive BERT for Reinforcement Learning

Andrea Banino, Adrià Puidomenech Badia+4 more|arXiv (Cornell University)|Jul 12, 2021

Reinforcement Learning in Robotics被引用数 8

ひとこと要約

CoBERLは、対照的でBERTにインspiredされた表現学習目的とゲート付きLSTM-Transformerアーキテクチャを組み合わせることで、強化学習におけるデータ効率を向上させる。マスクされた予測と手作業のオーグメンテーションを用いない時間領域の対照的損失を活用することで、Atari、コントロールスイート、3D環境において、GTrXLなどのベースラインを常に上回る優れたサンプル効率とパフォーマンスを達成する。

ABSTRACT

Many reinforcement learning (RL) agents require a large amount of experience to solve tasks. We propose Contrastive BERT for RL (CoBERL), an agent that combines a new contrastive loss and a hybrid LSTM-transformer architecture to tackle the challenge of improving data efficiency. CoBERL enables efficient, robust learning from pixels across a wide range of domains. We use bidirectional masked prediction in combination with a generalization of recent contrastive methods to learn better representations for transformers in RL, without the need of hand engineered data augmentations. We find that CoBERL consistently improves performance across the full Atari suite, a set of control tasks and a challenging 3D environment.

研究の動機と目的

• raw pixelから深層強化学習におけるデータ効率を向上させるために、より良い表現を学習すること。
• 大規模で密度の高いネットワークを用いたRL学習におけるノイズの多い勾配と強い時間的相関の課題に対処すること。
• 手作業で設計されたデータオーグメンテーションを回避する自己教師付き表現学習手法を開発すること。
• 学習可能なゲートメカニズムを用いてLSTMとTransformerの長所を統合し、動的な知識利用を可能にすること。
• 2Dおよび3Dの制御タスクを含む多様なRL環境において、強固で一般化可能な表現学習を可能にすること。

提案手法

• BERTにインspiredされた双方向マスク予測とReLICの時間領域一般化を組み合わせた、新しい対照的損失を提案し、自己教師付き表現学習を実現する。
• 因果的にマスクされたゲート付きTransformer-XL（GTrXL）を用いて観測値を処理し、マスクされた入力を予測することで、自己回帰的整合性を保証する。
• GTrXLとLSTMの出力を統合するための学習可能なゲートメカニズムを導入し、タスクの要件に応じて表現パスの動的選択を可能にする。
• ゲートをRLポリシー損失を用いて訓練することで、アーキテクチャ全体のエンドツーエンド最適化を実現する。
• 予測されたマスク入力からの正規化された埋め込みと元の観測値との間で、インスタンス識別と不変性ペナルティに基づく対照的目的を適用する。
• TransformerとLSTMのコンポONENTに供給する前に、raw pixel観測値を残差畳み込みニューラルネットワークエンコーダーで埋め込む。

実験結果

リサーチクエスチョン

• RQ1 BERTにインスパイアされた対照的表現学習目的は、手作業のオーグメンテーションを用いずに、強化学習におけるデータ効率を向上させることができるか？
• RQ2 学習可能なゲートを介してLSTMとTransformerを統合することで、多様なRL環境におけるパフォーマンスとサンプル効率にどのような影響を与えるか？
• RQ3 時間領域におけるマスク予測と対照的学習の統合は、RLにおけるより一貫性があり一般化可能な表現をもたらすか？
• RQ4 CoBERLは、複数のベンチマーク環境において、GTrXLなどの既存アーキテクチャをどの程度上回るか？
• RQ5 マスキング比が補助事前学習目的および最終的なRLパフォーマンスに与える影響は何か？

主な発見

• CoBERLは、全57種類のAtariゲームにおいて一貫したデータ効率の向上を達成し、特にサンプル効率において顕著な向上を示した。
• DeepMind Control Suiteでは、GTrXLおよび他のベースラインと比較して、より優れたサンプル効率と高い最終スコアを達成した。
• 3D DMLab-30環境では、CoBERLは優れた記憶力と推論能力を示し、AUC（t(60) = 2.616, p = 0.0011）においてGTrXLを上回った。
• アブレーションスタディにより、対照的損失とゲートアーキテクチャの両方がパフォーマンス向上に独立して寄与していることが確認され、特に表現学習において対照的目的が顕著に有効であった。
• 高いマスキング比（例：50％）はパフォーマンスを低下させ、補助タスクにおける有用な時間的文脈が減少することが示された。
• 本手法はドメインをまたいで良好に一般化し、2Dプラットフォーマーと3Dファーストパーソン/サードパーソン制御タスクの両方で強力なパフォーマンスを示した。