[Paper Review] Distillation Strategies for Proximal Policy Optimization

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

Reinforcement Learning in Robotics13 references4 citations

TL;DR

This paper proposes distillation strategies for Proximal Policy Optimization (PPO), transferring knowledge from a high-capacity teacher agent to a low-capacity student agent using offline rollouts. The method enables distilled students to outperform directly trained low-capacity agents and achieve parity with the teacher after fine-tuning, demonstrating that distillation significantly improves sample efficiency and performance in actor-critic reinforcement learning.

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.

Motivation & Objective

• To adapt knowledge distillation from DQN to the more advanced PPO actor-critic algorithm.
• To enable efficient training of low-capacity student agents by leveraging high-quality rollouts from a pre-trained teacher.
• To investigate whether distilled students can surpass directly trained agents of the same capacity.
• To evaluate the impact of distillation epochs and fine-tuning on student performance.
• To demonstrate the feasibility of using distillation for hardware-aware policy co-design in deep reinforcement learning.

Proposed method

• A high-capacity PPO agent (teacher) is trained in the environment and then used to collect state observations and action probabilities in a replay buffer.
• The replay buffer stores state-action pairs and teacher-provided action probabilities, forming a high-quality offline dataset.
• A low-capacity student network is trained using distillation loss, where the student matches the teacher’s action probabilities rather than just the optimal actions.
• Distillation is performed over multiple epochs using the offline replay buffer to optimize the student policy.
• After distillation, the student is further fine-tuned in the environment using standard PPO updates to improve performance.
• The method is evaluated across multiple Atari environments to compare distilled students against directly trained agents and the teacher.

Experimental results

Research questions

• RQ1 Can knowledge distillation from a high-capacity PPO teacher improve the performance of a low-capacity student agent?
• RQ2 Does distillation enable a student to outperform a directly trained agent of equal capacity?
• RQ3 How does the number of distillation epochs affect the final performance of the student?
• RQ4 Does fine-tuning the distilled student in the environment further improve performance?
• RQ5 Can distillation be used as a practical method for efficient policy co-design under hardware constraints?

Key findings

• Distilled students consistently outperform directly trained agents of the same capacity, demonstrating the effectiveness of knowledge transfer.
• In the Medium vs. Medium and Low vs. Low comparisons, distilled students matched or exceeded the performance of directly trained agents.
• For some environments like Pong and Freeway, 10 distillation epochs were sufficient to match teacher performance, while others like Breakout and Ms. Pac-Man required hundreds of epochs.
• Fine-tuning the distilled student for 20 million time steps brought its performance to parity with the teacher, especially for medium-capacity models.
• The performance plateau observed in distillation curves suggests that early stopping may not be necessary, though this requires further investigation.
• The results indicate that distillation is a robust method for reducing inference costs while maintaining or improving policy performance in actor-critic RL.