Distilling Policy Distillation (1902.02186v1)

Abstract: The transfer of knowledge from one policy to another is an important tool in Deep Reinforcement Learning. This process, referred to as distillation, has been used to great success, for example, by enhancing the optimisation of agents, leading to stronger performance faster, on harder domains [26, 32, 5, 8]. Despite the widespread use and conceptual simplicity of distillation, many different formulations are used in practice, and the subtle variations between them can often drastically change the performance and the resulting objective that is being optimised. In this work, we rigorously explore the entire landscape of policy distillation, comparing the motivations and strengths of each variant through theoretical and empirical analysis. Our results point to three distillation techniques, that are preferred depending on specifics of the task. Specifically a newly proposed expected entropy regularised distillation allows for quicker learning in a wide range of situations, while still guaranteeing convergence.

Distilling Policy Distillation: An Analysis

This paper titled "Distilling Policy Distillation," authored by researchers at DeepMind, presents a rigorous exploration of policy distillation in reinforcement learning, evaluating different formulations both theoretically and empirically. Central to this discussion is the delineation between the various strategies employed in policy distillation—each with distinct motivations, strengths, and operational mechanisms.

Overview of Policy Distillation

Policy distillation refers to the technique of transferring knowledge from one agent, typically a well-trained or expert agent known as the "teacher," to another agent called the "student." This is achieved by training the student agent to imitate the teacher's policy, which is a distribution of actions over states. Despite its conceptual simplicity, policy distillation encompasses a broad array of methodologies, and subtle variations in these methodologies can significantly influence their efficacy and the eventual outcome of the distillation process.

Key Contributions

The authors compare these approaches by introducing a unifying framework that facilitates a holistic view of policy distillation methods, laying out specific variations and their respective implications:

1. Control Policy and Sampling Distribution:
   • Teacher-driven distillation uses trajectories sampled from the teacher policy, optimizing for policy replication at the cost of slower adaptation to new environments.
   • Student-driven distillation utilizes trajectories sampled from the student policy, potentially leading to faster convergence by reducing divergence between training and test conditions. This approach is shown to replicate the teacher's behavior more robustly across a wider range of states due to greater exploration.
   • It is noted that many common student-driven distillation techniques do not correspond to valid gradient vector fields, leading to potential non-convergent dynamics unless corrected by certain reward-based modifications.
2. Evaluation of Update Methods:
   • A detailed comparison between methods using trajectory-based cross-entropy losses and those that use rewards intrinsic to the teacher's value function is presented.
   • Expected entropy regularized distillation emerges as a preferred choice, offering advantages such as reducing variance in gradient estimations and providing more reliable outcomes across diverse environments and action spaces.
3. Utilization of Actor-Critic Framework:
   • Suggestions for leveraging value functions from the teacher in actor-critic methods are provided, offering ways to bootstrap learning and introduce intrinsic rewards based on the teacher's critic. This enables the handling of suboptimal or imperfect teacher policies more effectively.
   • The paper demonstrates how actor-critic approaches could be modified to improve student policy beyond the teacher's performance.

Empirical and Theoretical Implications

From an empirical standpoint, the paper thoroughly examines how different distillation methodologies perform across randomly generated Markov Decision Processes (MDPs), illustrating the strengths of student-driven distillation in terms of its potential for faster learning and broader behavioral replication. Theoretically, by drawing on notions from optimization and gradient vector fields, a clear narrative is constructed around when and why certain distillation techniques are effective.

Future Directions and Considerations

While the paper makes substantial strides in characterizing variations in policy distillation, the authors acknowledge potential avenues for future research. Specifically, they highlight the necessity to explore the impact of distillation techniques in conjunction with large-scale function approximators, such as deep neural networks, on real-world, complex environments. Extending the theoretical insights to incorporate the influence of neural architectures remains a critical open question in this domain.

In sum, this analysis paves the way for more nuanced decision-making in reinforcement learning tasks involving policy distillation—a crucial step towards creating more efficient and capable AI agents.