Source-linked AI summary
Policy Distillation
Andrei A. Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, Raia Hadsell
TL;DR
DQN can learn strong visual-task policies, but transferring and combining such policies is difficult because reinforcement-learning action values are unbounded, unstable, and sometimes ambiguous. Policy distillation trains an untrained student from softened teacher action distributions, and the resulting agents compress policies, combine tasks, and can outperform DQN teachers. The method is demonstrated across single-game, multi-game, and online settings, with multi-task distilled agents exceeding jointly trained DQN performance in the reported three-game comparison.
Problem
Reinforcement-learning Q-values are unbounded and unstable, while similar action values make direct Q-value or single-best-action imitation difficult.
Method
Policy distillation transfers Q-network policies by training a student to regress toward softened teacher action distributions produced from Q-values.
Results
Multi-task distilled agents outperform their DQN teachers, achieving mean scores of 105.1% and 116.9% in the three-game experiment, while KL cost gives the best-performing students across four games.
Takeaways & Limitations
Distillation supports policy compression, multi-game agents, and online policy stabilization, including a 15-times compression claim without performance degradation.
Takeaways & Limitations
Multi-task DQN learning remains highly challenging because policy interference, reward scaling, and value-function instability hinder full single-game performance.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Deep reinforcement learning can learn effective pixel-to-action policies, but DQN requires substantial training and policy distillation is introduced to transfer those policies into smaller or combined networks despite challenges in Q-value outputs.
- 1 INTRODUCTION: DQN learns pixel-to-action policies by training a convolutional network to approximate action values from observations and rewards.At each step, the agent selects the action with the highest predicted cumulative reward.
- 1 INTRODUCTION: Policy distillation transfers one or more action policies from Q-networks to an untrained target network.The paper applies this approach to single-game compression, multi-game policies, and online distillation.
- 1 INTRODUCTION: 15 times smaller networks can be produced without performance degradation, while distilled multi-task policies can outperform their original experts.The paper also investigates continual online distillation of the best policy into a target network.
- 1 INTRODUCTION: Supervised distillation trains a target network to reproduce a teacher’s output distribution, often using softened targets.The paper extends this established compression and ensemble-consolidation idea to reinforcement-learning policies.
- 1 INTRODUCTION: Distillation is difficult for reinforcement learning because action values are real-valued, unbounded, reward-scale dependent, and sometimes similar across actions.These properties differ from the class-probability outputs commonly used in traditional distillation.
2 PREVIOUS WORK
The paper connects policy distillation to model compression, reinforcement learning, multi-task learning, and imitation learning, while addressing Atari-specific differences that complicate shared-task training.
- 2 PREVIOUS WORK: Policy distillation follows earlier work compressing ensemble or deep-network knowledge into a smaller student network.These methods train the student from outputs produced by a teacher model.
- 2 PREVIOUS WORK: Imitation-learning approaches regress to teacher signals from model-based algorithms or Monte Carlo tree search.The paper’s approach is related because it also transfers policy information through regression.
- 2 PREVIOUS WORK: Multi-task learning uses related tasks as a shared inductive bias, but Atari games have diverse images without a common statistical basis.Although the games share pixel inputs, their visual distributions make multi-task learning more difficult.
3 APPROACH
The approach distills DQN policies into student networks using softened action-value targets, and extends this process to combine multiple game-specific experts into one policy. Multi-task distillation is presented as an alternative to jointly training multi-game DQN, which generally struggles to reach single-game performance.
- 3.2 SINGLE-GAME POLICY DISTILLATION: Policy distillation transfers one or more DQN action policies to an untrained student network through supervised learning.The teacher supplies gameplay observations and Q-value targets while the student is trained separately.
- 3.2 SINGLE-GAME POLICY DISTILLATION: Softmax targets are sharpened for Atari policy distillation, making the student targets easier to learn despite visual and Q-value variation across frames.Figure 1 contrasts similar Pong frames with different Q-values and different Space Invaders frames with similar Q-values.
- 3.2 SINGLE-GAME POLICY DISTILLATION: Predicting every action's Q-value is difficult because values are unbounded, unstable, and costly to evaluate, while predicting only the best action loses information when actions have similar values.These properties make direct Q-function transfer and single-action imitation problematic.
- 3.2 SINGLE-GAME POLICY DISTILLATION: The method compares NLL on the teacher's best action, MSE on the full Q-value vector, and KL divergence applied with a temperature parameter.NLL preserves only the selected action, whereas MSE preserves the complete set of action-values.
- 3.3 MULTI-TASK POLICY DISTILLATION: Multi-task distillation trains one student sequentially from separate replay buffers generated by n single-game DQN experts, using task-specific controller output layers.The game switches every episode, and the task id selects the appropriate output layer.
- 3.3 MULTI-TASK POLICY DISTILLATION: Joint multi-game DQN remains extremely challenging and generally fails to reach full single-game performance, which the authors associate with policy interference, reward scaling, and value-function instability.Policy distillation is proposed as a way to combine policies without these damaging interference and scaling problems.
4 RESULTS AND DISCUSSION
The experiments evaluate policy distillation for single-game compression, multi-task consolidation, and online tracking of evolving DQN policies. Distilled students retain or improve performance while using substantially smaller or shared networks.
- Single-game policy distillation: KL distillation produced the best-performing students and outperformed DQN teachers on most games.The evaluation compared MSE, NLL, and KL criteria using human starting points.
- Single-game policy distillation: 10 Atari games were evaluated using student networks with 25%, 7%, and 4% of the DQN teachers’ parameters.The evaluation used human-expert starting states to test generalization.
- Single-game policy distillation: Four-times-smaller distilled agents outperformed DQN, while 15-times-smaller agents performed on par with their teachers.The four-times-smaller Dist-KL-net1 contained 428,000 parameters.
- Multi-game policy distillation: The multi-task DQN reached 83.5% of single-task performance, whereas Multi-Dist-NLL and Multi-Dist-KL reached 105.1% and 116.9%, respectively.Multi-Dist-KL achieved as much as 50% higher scores on Q*bert.
- Multi-game policy distillation: A single distilled network for 10 Atari games achieved 89.3% of the geometric-mean generalization score of 10 single-task DQN teachers.The student network was four times larger than a single DQN, and three games scored much higher than their teachers.
- Online policy distillation: Online policy distillation tracked the best DQN policy during Q-learning by periodically updating the student’s teacher network.The experiment addressed policy changes during training and used two initial random seeds.
5 DISCUSSION
Policy distillation serves three roles in deep Q-learning: compressing single-game policies, combining multiple games, and stabilizing DQN through online transfer. The discussion emphasizes loss-function choice and reports that distilled policies can outperform their teachers.
- Policy distillation compresses single-game policies, builds multi-game agents, and improves DQN stability through online distillation.
- Soft-max action-gap weighting gives the best reinforcement-learning distillation results among the tested approaches.The authors connect this weighting to the CAPI framework.
- Distillation applies without iterative training or letting the student control its training-data distribution.
- Distilled policies can outperform their teacher networks, supporting distillation as a general model-regularization principle.
A EXPERIMENTAL DETAILS
The experiments generate student data from DQN teachers, train students with alternative action-value targets, and evaluate compressed and multi-task architectures under standardized gameplay conditions. Multi-task models use shared representations with task-specific controller networks.
- Training procedure: Teachers generate emulator-frame inputs and Q-value outputs for replay-memory training of student networks.Replay memory stores up to 10 hours of gameplay, and student updates follow each additional hour of teacher data.
- Distillation targets: Three distillation targets use teacher Q-values, the highest-valued action, or action distributions derived from the teacher outputs.These targets correspond to MSE, NLL, and KL-based training criteria.
- Network architectures: Student networks use linear one-unit-per-action outputs, with a final softmax for NLL and KL distillation.
- Network architectures: Compression scales down units in each layer while preserving the basic architecture, with most saved parameters coming from the top fully connected layer.
- Multi-task models: Multi-task agents share representations but use separate controller networks or output layers for each game.A four-times-larger network was trained for ten-game multi-task distillation.
- Evaluation: Agents are evaluated for up to 30 minutes using human-generated starting states and 5% random actions.Multi-game scores use a geometric mean because reward scales differ across games.
B SUPPORTING TABLES FOR POLICY DISTILLATION FIGURES
The supporting tables report performance for compressed single-task networks across ten Atari games and multi-task distilled agents across three games.
- Table B1 reports single-task compressed-network performance on 10 Atari games.Best relative scores are outlined in bold.
- Table B2 reports multi-task distilled-agent performance on 3 Atari games.Best relative scores are outlined in bold.
C ADDITIONAL RESULTS USING ONLINE POLICY DISTILLATION
Online policy distillation occurs during DQN learning by transferring the current best policy into a new network. The figure reports experiments on two games across two initial random seeds using the same learning rate.
- Online distillation transfers the current best DQN policy into a new network during DQN training.
- The online experiments cover 2 games with 2 initial random seeds and one learning rate across all runs.
D VISUALIZATION OF REPRESENTATION OVER 10 ATARI GAMES
Figure D2 visualizes activations from 10 Atari games after Multi-Dist-KL distillation, comparing representations in the first convolutional and final fully connected layers.
- The figure uses t-SNE to visualize activations from 10 Atari games at the first convolutional and final fully connected layers.
- Layer-1 embeddings suggest that low-level representation statistics may remain game-specific, potentially reflecting diverse inputs.
- Final-layer embeddings show game-specific activations with greater within-game variance, probably reflecting output statistics.