Source-linked AI summary
Distilling Policy Distillation
Wojciech Marian Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant M. Jayakumar, Grzegorz Swirszcz, Max Jaderberg
TL;DR
Policy distillation transfers knowledge between reinforcement-learning policies, but its many formulations can optimize different objectives and exhibit different learning dynamics. This paper unifies and compares these variants theoretically and empirically, finding expected entropy-regularized distillation especially reliable while identifying task-dependent alternatives.
Problem
Policy distillation is widely used, yet its many mathematical formulations and their effects on performance and optimized objectives have not been systematically compared.
Method
The paper provides a unified theoretical analysis and large-scale empirical comparison of distillation variants, including control-policy, actor-critic, and teacher-value approaches.
Results
Expected entropy-regularized distillation is the most reliable formulation theoretically and empirically, while other techniques are preferred depending on task specifics.
Takeaways & Limitations
Distillation methods should be selected using the paper’s comparative guidance, with expected entropy regularization as a broadly reliable choice and teacher critics useful for imperfect teachers.
Takeaways & Limitations
The study uses synthetic MDPs, leaving behavior on large-scale real-world problems with neural-network function approximators unresolved.
Abstract
from arXiv · showhide
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.
1 Introduction
Policy distillation transfers knowledge from trained policies, but its many mathematical formulations can optimize different objectives and produce different learning behavior. This paper develops a comparative roadmap through theoretical analysis, empirical evaluation, and a method-selection rule of thumb.
- Motivation: Distillation transfers a teacher’s state-dependent action distribution to a student policy, enabling knowledge transfer in reinforcement learning.The paper motivates distillation as a way to reduce the experience required for training and support stronger or otherwise untrainable agents.
- Problem: Different formulations vary in trajectory sampling and divergence choice, including teacher-, student-, or mixture-driven sampling and state- or trajectory-level KL objectives.These variations can substantially change the resulting optimization problem.
- Contributions: The paper proves that commonly used student-trajectory distillation need not form a gradient vector field and can oscillate after rewards are introduced.It also presents methods for recovering the gradient-vector-field property.
- Method selection: Figure 1 summarizes the paper’s method-selection guidance as an initial rule of thumb rather than an exhaustive decision procedure.The diagram references teacher, on-policy, entropy-regularized, N-distill, expected entropy-regularized, and teacher-value methods.
- Evaluation: The authors compare control policies, actor-critic methods, and teacher-value-based approaches mathematically and across thousands of random MDPs.The combined findings support proposed distillation variants and a decision tree for choosing among methods.
2 Preliminaries
The paper formulates distillation in finite Markov decision processes, where a teacher policy’s knowledge is transferred to a parameterized student through environment trajectories. It specifies the policy, transition, reward, and discounting notation used for subsequent update rules.
- MDP setting: The setting is a finite Markov decision process with states, actions, stochastic policies, transition dynamics, and rewards.The student policy maps states to action distributions, while trajectories are generated through policy-environment interaction.
- Distillation setup: Knowledge is transferred from teacher policy π to student policy πθ using trajectories generated by a control policy qθ.The control policy may differ from the teacher and student policies.
- Discounting: Theoretical results usually use γ = 1, whereas empirical experiments generally use γ = .99.The paper states that the proofs extend trivially to arbitrary γ ≤1.
3 Policy distillations
The paper analyzes how control policies, update rules, and reward terms shape policy distillation, showing that common on-policy updates may lack convergence guarantees while corrected and student-driven variants improve practical learning. It compares these choices theoretically and empirically, including expected entropy regularization and distillation from imperfect teachers.
- 3.1 Control policy: On-policy distillation updates generally do not form valid gradient vector fields, so they cannot always be interpreted as optimizing a loss.The paper connects this issue to incompatibility between trajectory losses and policy gradients.
- 3.1 Control policy: Tabular on-policy distillation converges when the control policy samples every teacher-reachable state with nonzero probability.A strongly stochastic student policy satisfies this condition and can converge to the teacher policy over reachable states.
- 3.1 Control policy: Adding environment rewards can produce oscillations and non-convergence, but a next-state reward correction restores the gradient-vector-field property and convergence.The correction uses the negative distillation loss at the next state; in the illustrated experiment, N-distill+R converges and minimizes the loss.
- 3.1 Control policy: Student-driven trajectories better match teacher behavior under the student’s own state distribution, while uniform control achieves broader state-space coverage but converges extremely slowly.The choice of control policy therefore trades off replication fidelity, coverage, and convergence speed.
- 3.2 Empirical evaluation of the control policy: Student-driven distillation reaches full teacher performance with about 3× fewer steps than teacher-driven distillation and about 10× fewer than uniform-driven distillation.Teacher-driven distillation performs better only for KL measured on teacher-generated trajectories, a setting the paper characterizes as rarely encountered in practice.
- 3.3 Empirical evaluation of various updates: Entropy-only reward distillation becomes much slower as irrelevant actions expand the action space, whereas expected entropy regularization recovers traditional distillation performance.The expected method splits the entropy term between a one-step loss and a reward component, reducing estimator variance.
4 Policy distillation from Actor-Critic
The paper studies how a teacher’s value function can support distillation, including bootstrapping, value-based guarantees, and intrinsic rewards. These methods can improve learning over direct policy replication, but their effectiveness depends strongly on teacher quality and whether the objective solves the original task.
- Using the teacher’s critic: Teacher value functions provide information beyond policy cloning and can help estimate how much to trust imperfect teachers.The paper introduces critic-based distillation because access to Vπ allows the student to use the teacher’s expected returns.
- Using the teacher’s critic: A value-weighted policy loss guarantees a converged student is at least as strong as its teacher in expected value.The guarantee is stated for the expected value over the initial-state distribution, under the paper’s loss condition.
- Using the teacher’s critic to bootstrap: Bootstrapping from the teacher’s value function also guarantees a student whose expected value is no lower than the teacher’s.However, this update no longer optimizes the student’s true return because the bootstrap value is defined for the teacher policy.
- Using the teacher’s critic to bootstrap: Teacher-value bootstrapping can outperform direct replication but may saturate early when the teacher is weak, because it is not solving the original problem.The paper gives TD+V0, whose teacher predicts zero value everywhere, as a poor example.
- Using the critic as an intrinsic reward: Value-based shaping rewards preserve optimal solutions without forcing complete policy cloning, while strong teachers can speed convergence and weak teachers can slow it.The shaping reward combines the teacher-value change with the true reward, so obtaining a reward is not penalized when the value decreases.
- Using the critic as an intrinsic reward: An adversarial teacher can impose penalties for actions that differ from its worst-policy behavior, significantly slowing training without changing the learning solution.This is the value-based analogue of the optimal-teacher shaping analysis.
- Using the critic as an intrinsic reward: In the corridor experiment, policy cloning quickly recovered optimal scores from an optimal teacher, whereas shaping was slower but still beat no-distillation training.With an adversarial teacher, cloning methods failed, while Teacher V reward still learned; teacher-value bootstrapping improved only over the adversarial teacher, not necessarily the task.
5 Conclusions
The paper unifies policy-distillation variants and evaluates their mathematical and empirical strengths and weaknesses. It identifies expected entropy regularised distillation as the most reliable formulation while noting that real-world deep-network behavior remains unresolved.
- Conclusions: The paper presents a unifying view of policy-distillation techniques and compares them through theoretical analysis and large-scale synthetic-MDP experiments.The synthesis is summarized in Figure 1.
- Conclusions: Many widely used distillation methods are not valid gradient vector fields and may therefore exhibit non-convergent learning dynamics.The paper proposes modifications that address some of these dynamical issues.
- Conclusions: Expected entropy regularised distillation is identified as the most reliable formulation both theoretically and empirically.The conclusion also states that a teacher critic can help handle imperfect teachers when available.
- Conclusions: The behavior of these approaches on large-scale real-world problems with deep neural-network function approximators remains an open question.The mathematical role of the function approximator is also left for future work.
A.1 MDPs
The study uses finite grid-world Markov decision processes with deterministic state rewards, stochastic movement, terminal states, and either full or local visual observations.
- Environment structure: The environments are W × H grid worlds with a finite state space, a special terminating state, and a central initial state.Agents can move left, right, up, or down, subject to transitions that remain within the state space.
- Environment structure: Each action has a desired directional effect, while invalid moves are constrained by the grid boundary.The state space also includes a special state to which the agent transitions with probability 0.01 after each action.
- Rewards and termination: Rewards are assigned to states and are deterministic, while terminal states end episodes and return the agent to the initial state.This creates finite episodes for the experiments.
- Observations: The experiments compare fully observable state-index observations with partially observable local object views.Partial observations encode nearby walls, reward values, and whether states terminate.
- Observations: Using 9×9 local visual observations instead of other tested extents did not change the paper’s qualitative results.The authors therefore treat the particular visual extent as non-crucial.
A.2 Distribution over MDPs
The paper samples synthetic grid-world MDPs by modifying a base grid with walls, rewards, and terminal states under fixed probabilities, while retaining connectivity to a positive-reward state.
- Generation procedure: The MDP-generation procedure begins by creating the state space and traversing states from the upper-left corner horizontally, then vertically.The procedure applies modifications independently as it traverses the grid.
- Example environment: The example 20×20 grid world depicts the initial state, terminal states, rewards, and walls using distinct visual encodings.The caption identifies blue for the initial state, green for terminal states, and black for walls.
- Validity constraint: The generated MDP is retained only if a path exists from the initial state to the +10 state.Otherwise, the generation process is repeated.
A.3 Actor Critic
The actor-critic experiments sample full episodes under the student policy and update using either Monte Carlo returns or TD(1) bootstrapped estimates, while fitting a baseline with the corresponding value targets.
- Full episodes are sampled under the student policy before actor-critic updates are performed.
- Updates use either single-sample Monte Carlo returns or TD(1) bootstrapped estimates.
- Experiments use γ = 0.99, with qualitatively similar results reported for γ = 0.95 and γ = 0.999.
- The baseline function is fitted after each update using the same Monte Carlo or TD value with an L2 loss.
A.4 Q-Learning
The Q-learning setup uses a standard update with zero-initialized values, then treats the learned policy as either greedy or Boltzmann depending on temperature; policies and value functions are tabular by observation.
- Q-Learning: Q-learning applies its standard update after each visited state, with Q-values initialized to zero.The learning rate is λ = 0.01, and training lasts 30k iterations.
- Teacher policy: A temperature of 0 makes the learned teacher policy greedy, while other temperatures produce a Boltzmann policy.
- Representations: Actor-critic policies are parameterized by logits for every action and unique observation.
- Representations: Value functions are represented as one scalar per observation.
- Extended figures: Figure 8 extends Figure 5 with an A2C teacher, Figure 9 extends Figure 6 with more corridor sizes, and Figure 10 extends Figure 4 with additional agents.
C Proofs for Section 3 (Policy distillations)
The proofs characterize when policy-distillation updates converge, fail to define a gradient field, or alter the learned behavior through loss direction and reward correction. They establish convergence under specific stochasticity and optimization assumptions and explain mode- versus mean-seeking behavior under constrained distributions.
- Gradient-field analysis: Student-driven distillation need not be a gradient vector field when per-trajectory loss gradients are not proportional to trajectory-probability gradients.The theorem states this under differentiability and a non-proportionality assumption.
- Convergence: With strongly stochastic tabular policies and suitable loss minimization, student-driven updates converge to the teacher on all reachable states.Strong stochasticity ensures the student visits the reachable state space, while the loss must attain its minimum when student and teacher distributions match.
- Convergence: The update direction is strictly descending for reachable states whenever teacher and student policies do not yet match.
- Reward correction: Adding the negative next-step distillation loss as a reward correction restores the gradient-vector-field property for one-step on-policy updates.The correction is a negative loss at the next state or state-action level, depending on the loss form.
- Loss direction: Cross-entropy direction changes the optimum: H×(p∥q) is minimized by q = p, whereas H×(q∥p) selects the most probable action under a unique-mode assumption.
- KL direction: For constrained distribution families, KL(p∥q) is mean-seeking while KL(q∥p) is mode-seeking and may ignore other mixture components.For mixtures of Gaussians, the former matches the mean and the latter covers one Gaussian while ignoring others.
D Proofs for Section 4 (Policy distillation from Actor-Critic)
The actor-critic proofs show that distillation using a teacher’s value function can produce a policy whose expected value is at least the teacher’s under stated finite-MDP assumptions, including when the teacher is optimal.
- Policy matching: If the optimized cross-entropy is no greater than the teacher’s entropy in every state, the student’s expected value is at least the teacher’s.
- Value-function distillation: With access to the true teacher value function, the actor-critic update converges to a policy whose expected value is no smaller than the teacher’s.
- Value-function distillation: The value-augmented update treats each action’s target as its immediate reward plus the teacher’s next-state value.The proof interprets this as a bandit problem whose continuation follows the teacher policy.
- Optimal teachers: For an optimal teacher, actions deviating from an optimal path receive an immediate value-based penalty relative to their extrinsic reward.