[論文レビュー] Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation

Emilio Parisotto, Ruslan Salakhutdinov|arXiv (Cornell University)|Apr 4, 2021

Reinforcement Learning in Robotics被引用数 9

ひとこと要約

本論文は、推論遅延制約のある強化学習設定において、大規模で高容量なTransformerの学習者モデルから、小規模で高速なLSTMのエージェントモデルへ知識を継続的に転送するActor-Learner Distillation (ALD) を提案する。Transformerのサンプル効率を活用しながらLSTMの推論速度を維持することで、小型モデルと同等の学習時間で優れたサンプル効率を達成する。

ABSTRACT

Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on un-accelerated hardware such as CPUs, which likewise restricts model size to prevent intractable experiment run times. These "actor-latency" constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during acting, we develop an "Actor-Learner Distillation" (ALD) procedure that leverages a continual form of distillation that transfers learning progress from a large capacity learner model to a small capacity actor model. As a case study, we develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, at the cost of significantly higher computational complexity. With transformer models as the learner and LSTMs as the actor, we demonstrate in several challenging memory environments that using Actor-Learner Distillation recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.

研究の動機と目的

• 推論速度が厳しく制限されるエージェント-遅延制約付き強化学習環境において、高容量モデルの導入を困難にする課題に対処すること。
• ロボットICSのような実世界の応用において、計算リソースと電力制約が推論時のモデルサイズを制限する中で、サンプル効率の高い大規模モデル（例：Transformer）を活用できるようにすること。
• 高容量の学習者から低容量のエージェントへ知識を圧縮する継続的蒸留手順を開発し、サンプル効率を損なわずに実現すること。
• 分散型オンポリシー強化学習において、高速で小規模なエージェントモデルを用いて実験の実行時間を短くしつつ、大規模な学習者モデルの学習進捗を活用できること。
• 特に部分観測環境において、異なるアーキテクチャ間（具体的にはTransformerとLSTM）の蒸留が、帰納的バイアスとパフォーマンス向上を効果的に転送できることを示すこと。

提案手法

• ALDフレームワークは、相互作用から学習し、サンプル効率を向上させる大規模なTransformerモデルを備えた集中型学習者を用いる。
• 別個に、推論速度制約を満たす軽量なLSTMベースのエージェントモデルが、低遅延意思決定を実現する。
• エージェントの行動を模倣するように、学習者モデルの知識が継続的ポリシー蒸留を通じて転送される。
• 蒸留は学習中にオンラインで適用され、学習者のポリシー出力をソフトターゲットとして用いてエージェントのポリシー更新を監視する。
• 本手法は、エージェントが軌道を収集し、それを中央学習者に送信してポリシー更新を行う分散型オンポリシー強化学習設定で動作する。
• 本手法はLSTMエージェントの計算効率を維持しつつ、記憶容量が大きい環境におけるTransformer学習者の優れたサンプル効率を捉える。

実験結果

リサーチクエスチョン

• RQ1 大規模でサンプル効率の高いTransformerモデルを用いても、推論速度を損なわず、小規模で高速なLSTMエージェントモデルを訓練可能か？
• RQ2 高容量の学習者から低容量のエージェントへの継続的蒸留が、部分観測環境において学習者のサンプル効率の利点を保持できるか？
• RQ3 ALDは、小型モデルと同等の学習実行時間で、大規模モデルと同等のパフォーマンスを達成できるか？
• RQ4 エージェント-遅延制約付き強化学習環境において、特にTransformerとLSTM間の異なるアーキテクチャ間での知識転送はどの程度効果的か？
• RQ5 TransformerからLSTMへの蒸留が、記憶容量が大きい強化学習タスクにおける一般化性能とロバスト性を向上させるか？

主な発見

• ALDにより、LSTMエージェントは挑戦的な記憶環境において、Transformer学習者と同等のサンプル効率を達成し、初期学習とは著しく優れた性能を示す。
• 本手法は、大規模なTransformerを学習に使用しても、LSTM単体のベースラインと同等の実験実行時間を維持する。
• 部分観測環境では、ALDがTransformerのサンプル効率の利点を回復させつつ、LSTMの低遅延推論を維持する。
• 継続的蒸留プロセスは、TransformerからLSTMへの帰納的バイアスの効果的な転送を実現し、学習の安定性とパフォーマンスを向上させる。
• ALDは、異なる帰納的バイアスを持つ学生と教師間の知識蒸留が、エージェント-遅延制約付き強化学習において非常に効果的であることを示している。
• 本手法はアーキテクチャの効率化改善と直交するため、さらに性能向上を図るための効率的Transformer変種と組み合わせて利用可能である。