[論文レビュー] In-context Reinforcement Learning with Algorithm Distillation

Michael Laskin, Luyu Wang+12 more|arXiv (Cornell University)|Oct 25, 2022

Reinforcement Learning in Robotics被引用数 10

ひとこと要約

本稿では、強化学習（RL）エージェントの複数エピソードにわたる学習履歴を用いて、パrameterの更新なしに文脈内での方策改善を可能にする、アルゴリズム蒸留（AD）を提案する。ADは、全訓練軌道を対象とした逐次予測問題として方策改善をモデル化することで、元のRLアルゴリズムよりもデータ効率の良いRLアルゴリズムを学習する。

ABSTRACT

We propose Algorithm Distillation (AD), a method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training histories with a causal sequence model. Algorithm Distillation treats learning to reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then a causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning or expert sequences, AD is able to improve its policy entirely in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in a variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns a more data-efficient RL algorithm than the one that generated the source data.

研究の動機と目的

• パrameterの微調整なしに、試行錯誤によってモデルが改善する文脈内強化学習を可能にすること。
• 長期間の文脈を用いて、RLアルゴリズムの学習プロセスを逐次予測問題としてモデル化すること。
• 固定方策だけでなく、RL学習履歴から得られる「方策改善作用素」そのものを蒸留すること。
• 文脈内強化学習が、オフラインRL学習履歴に対する模倣学習によって達成可能であることを示すこと。
• 蒸留されたアルゴリズムが、元のRLアルゴリズムよりもデータ効率に優れていることを示すこと。

提案手法

• 多数の個別タスクにおいて、元のRLアルゴリズムから得た全学習履歴の多タスク大規模データセットを収集する。
• 各学習履歴を、時系列的かつエピソード的構造を保持した状態-行動-報酬トークンの系列として表現する。
• 全前方学習履歴を文脈として用いて、自己回帰的に行動を予測する因果的トランスフォーマーを訓練する。
• 将来の報酬ではなく、過去の状態、行動、報酬にのみ条件づけられる系列モデリング損失を用いる。
• 複数エピソードにわたる方策改善を捉えるために、十分に長い文脈長を確保する。
• 推論時において訓練済みトランスフォーマーを文脈内で利用する：重みの更新なしに、過去のエピソード履歴をプロンプトとして与え、行動を予測する。

実験結果

リサーチクエスチョン

• RQ1 系列モデルは、全学習履歴に注目することで、文脈内で方策を改善する能力を学習できるか？
• RQ2 RL学習を逐次予測問題としてモデル化することで、文脈内での探索行動と時間的報酬割り当てが可能になるか？
• RQ3 蒸留されたモデルは、新しいタスクに一般化し、元のRLアルゴリズムよりもデータ効率に優れるか？
• RQ4 模倣学習を用いて、RLアルゴリズムの「方策改善作用素」を1つのニューラルネットワークに蒸留することは可能か？
• RQ5 蒸留されたアルゴリズムの性能は、サンプル効率という観点で、元のアルゴリズムと比べてどうなるか？

主な発見

• ADは、パrameterの更新を一切行わず、過去のエピソードからの文脈のみに依存して文脈内強化学習を達成する。
• ADは、DMLabのWatermazeのような部分的に観測可能な環境でも、文脈内での探索と時間的報酬割り当てを示す。
• ADは、学習データの生成に用いられた元のRLアルゴリズムよりも、よりデータ効率の良いRLアルゴリズムを学習する。
• 本手法は、報酬がスパarsityなタスク、組み合わせ的構造、ピクセルベースの観測に対しても一般化可能である。
• 微調整やエキスパートのデモンストレーションを必要とする、従来のオフライン方策蒸留手法よりも優れた性能を示す。
• ADの成功は、複数エピソードにわたる方策改善を捉えるために、十分に長い文脈窓が必須であることに強く依存する。