[論文レビュー] Distillation Strategies for Proximal Policy Optimization

Sam Green, Craig M. Vineyard+1 more|arXiv (Cornell University)|Jan 23, 2019

Reinforcement Learning in Robotics参考文献 13被引用数 4

ひとこと要約

本論文は、教師エージェントから生徒エージェントへの知識伝達を可能にする、Proximal Policy Optimization (PPO) における蒸留戦略を提案している。高容量の教師エージェントを用いてオフラインのロールアウトを生成し、低容量の生徒エージェントに知識を伝達することで、生徒エージェントの性能を向上させている。この手法により、直接訓練された低容量エージェントを上回る性能を示し、ファインチューニングを経て教師エージェントと同等の性能に到達することが可能となった。これは、蒸留がアクタ・クリティック強化学習におけるサンプル効率と性能を著しく向上させることを示している。

ABSTRACT

Vision-based deep reinforcement learning (RL) typically obtains performance benefit by using high capacity and relatively large convolutional neural networks (CNN). However, a large network leads to higher inference costs (power, latency, silicon area, MAC count). Many inference optimizations have been developed for CNNs. Some optimization techniques offer theoretical efficiency, such as sparsity, but designing actual hardware to support them is difficult. On the other hand, distillation is a simple general-purpose optimization technique which is broadly applicable for transferring knowledge from a trained, high capacity teacher network to an untrained, low capacity student network. DQN distillation extended the original distillation idea to transfer information stored in a high performance, high capacity teacher Q-function trained via the Deep Q-Learning (DQN) algorithm. Our work adapts the DQN distillation work to the actor-critic Proximal Policy Optimization algorithm. PPO is simple to implement and has much higher performance than the seminal DQN algorithm. We show that a distilled PPO student can attain far higher performance compared to a DQN teacher. We also show that a low capacity distilled student is generally able to outperform a low capacity agent that directly trains in the environment. Finally, we show that distillation, followed by "fine-tuning" in the environment, enables the distilled PPO student to achieve parity with teacher performance. In general, the lessons learned in this work should transfer to other modern actor-critic RL algorithms.

研究の動機と目的

• DQNからより高度なPPOアクタ・クリティックアルゴリズムへの知識蒸留の適応を目的とする。
• 事前に訓練された教師エージェントから得られる高品質なロールアウトを活用して、低容量の生徒エージェントの効率的な訓練を可能にする。
• 同じ容量の直接訓練されたエージェントを上回る性能を示せるかどうかを検証すること。
• 蒸留エポック数とファインチューニングの影響が生徒エージェントの性能に与える影響を評価すること。
• ハードウェア制約下での効率的なポリシー共同設計のための蒸留の実用性を示すこと。

提案手法

• 高容量のPPOエージェント（教師）を環境で訓練し、その後、リプレイバッファに状態観測と行動確率を収集する。
• リプレイバッファは状態-行動ペアと教師が提供する行動確率を格納し、高品質なオフラインデータセットを形成する。
• 低容量の生徒ネットワークを、教師の行動確率に一致させるように蒸留損失関数を用いて訓練する。この際、最適な行動にのみ合わせる代わりに、教師の行動確率そのものを模倣する。
• オフラインのリプレイバッファを用いて複数エポックにわたり蒸留を実行し、生徒ポリシーを最適化する。
• 蒸留終了後、標準的なPPO更新を用いて環境でさらにファインチューニングすることで、生徒エージェントの性能を向上させる。
• 複数のAtari環境で、蒸留された生徒エージェントと直接訓練されたエージェント、および教師エージェントを比較評価する。

実験結果

リサーチクエスチョン

• RQ1 高容量のPPO教師エージェントからの知識蒸留は、低容量の生徒エージェントの性能向上に寄与するか？
• RQ2 蒸留により、同じ容量の直接訓練エージェントを上回る性能を示せるか？
• RQ3 蒸留エポック数が生徒エージェントの最終的性能に与える影響は何か？
• RQ4 環境で蒸留済みの生徒エージェントをさらにファインチューニングすることで、性能が向上するか？
• RQ5 ハードウェア制約下での効率的ポリシー共同設計に、蒸留を実用的な手法として用いることができるか？

主な発見

• 蒸留された生徒エージェントは、同じ容量の直接訓練エージェントを一貫して上回り、知識伝達の有効性を示している。
• ミディアム vs. ミディアムおよびローバイローバイ比較において、蒸留された生徒エージェントは直接訓練エージェクトエージェントと同等またはそれを上回る性能を示した。
• Pong や Freeway などの一部の環境では10エポックの蒸留で教師の性能に到達したが、Breakout や Ms. Pac-Man のような他の環境では数百エポックを要した。
• 2000万時間ステップにわたるファインチューニングにより、特にミディアム容量モデルでは生徒エージェントの性能が教師エージェントと同等にまで向上した。
• 蒸留曲線に見られる性能の飽和傾向から、早期停止は必ずしも必要ではないと考えられるが、さらなる検証が必要である。
• 結果から、蒸留は推論コストを低減しつつ、アクタ・クリティックRLにおけるポリシー性能を維持または向上させる強固な手法であると示唆された。