[論文レビュー] Policy Distillation

Andrei A. Rusu, Sergio Gómez Colmenarejo+7 more|arXiv (Cornell University)|Nov 19, 2015

Neural Networks and Reservoir Computing被引用数 87

ひとこと要約

この論文は、パフォーマンスに損なわれることなく、深層Qネットワーク（DQN）からより小さな効率的な学生ネットワークへ方策を転送する手法であるポリシー蒸留を導入する。モデル圧縮、マルチタスク方策の統合、オンライン蒸留を可能にし、アタリ環境における単一タスク教師および共同学習されたDQNエージェントよりも優れたパフォーマンスを達成する。

ABSTRACT

Policies for complex visual tasks have been successfully learned with deep reinforcement learning, using an approach called deep Q-networks (DQN), but relatively large (task-specific) networks and extensive training are needed to achieve good performance. In this work, we present a novel method called policy distillation that can be used to extract the policy of a reinforcement learning agent and train a new network that performs at the expert level while being dramatically smaller and more efficient. Furthermore, the same method can be used to consolidate multiple task-specific policies into a single policy. We demonstrate these claims using the Atari domain and show that the multi-task distilled agent outperforms the single-task teachers as well as a jointly-trained DQN agent.

研究の動機と目的

• DQNを用いて訓練された強化学習エージェントの高い計算コストと大きなモデルサイズに対処すること。
• 蒸留を用いて、大規模でタスク特化型のDQN教師エージェントから、より小さな効率的な学生ネットワークへ知識を転送すること。
• 複数の単一タスクDQN方策を1つのマルチタスク方策に統合し、個々の教師を上回る一般化性能を得ること。
• 学習中に最もパフォーマンスの良い方策を継続的に追跡することで、DQN学習の安定化を図るオンライン蒸留の探求。

提案手法

• 事前に訓練されたDQN教師の行動価値出力分布を模倣するように学生ネットワークを訓練する。この際、ソフトラベルを用いる。
• 温度調整付きソフトマックスを用いて行動価値分布をなめらかにし、より良い知識転送を実現する。
• 行動ギャップに基づいて行動分類の重みを付ける知識蒸留損失関数を適用する。これはCAPIフレームワークに類似している。
• 教師方策が生成する軌道に基づいて、教師の行動価値出力を回帰するように学生ネットワークを教師あり学習で訓練する。
• 学習中に現在最もパフォーマンスの良いDQN方策を定期的に更新することで、オンライン蒸留を実装する。
• 共有畳み込み特徴量とタスク固有のヘッドを持つマルチコントローラアーキテクチャを用いて、複数ゲームへの一般化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ポリシー蒸留は、パフォーマンスの低下を伴わずに、DQN方策を15倍程度に圧縮できるか？
• RQ2 複数の単一タスクDQN方策を1つのマルチタスク方策に蒸留し、個々の教師よりも一般化性能が向上するか？
• RQ3 学習中に最もパフォーマンスの良い方策をリアルタイムで追跡することで、オンライン蒸留がDQN学習の安定化に寄与するか？
• RQ4 学習中に教師方策が著しく変化する場合、蒸留はどの程度の性能を示すか？
• RQ5 非確率的かつ実数値の行動価値設定において、強化学習の蒸留において最良のパフォーマンスを達成する損失関数設計は何か？

主な発見

• ポリシー蒸留により、単一ゲームのアタリタスクにおいて、DQNモデルのサイズを最大15倍に圧縮してもパフォーマンスに劣化が生じない。
• 蒸留されたマルチタスクエージェントは、10個の単一タスクDQN教師の幾何平均パフォーマンスの89.3％を達成し、Q*bert や Seaquest などの複数のゲームで教師を上回った。
• 3ゲームのマルチタスク設定では、蒸留エージェント（Multi-Dist-KL）が単一タスクDQN教師の116.9％のパフォーマンスを達成し、共同学習されたマルチタスクDQNエージェント（83.5％）を顕著に上回った。
• オンライン蒸留により、学生エージェントはDQN教師と同等またはそれ以上のパフォーマンスを達成し、訓練の分散が著しく低減された。
• 行動ギャップに基づいて重み付けされたソフトマックス損失（CAPIと同様）が最も優れた結果をもたらし、強化学習における蒸留では損失関数の設計が極めて重要であることを示した。
• 反復的相互作用やデータ分布の制御がなくても、RLにおける蒸留は有効であることが確認され、一般化正則化手法としての役割を果たすことが示された。