Policy Distillation

Policy distillation compresses one or more big teacher policies into a small fast student by training the student to copy the teachers soft action probabilities instead of learning from rewards alone.

Why Does This Exist?

Big policies are a pain to ship. You train one expert per game, per robot task, per level, and now you own ten large networks that each need memory, compute, and upkeep. Training one small network on all tasks directly sounds better, but joint training from rewards is jumpy: tasks fight over shared weights and the small net never reaches expert quality.

Policy distillation breaks the job in two. First, train big teachers the normal way, one per task, with no size pressure. Then freeze them and teach one compact student to copy how they act. The student never touches the environment reward during this step. It just watches states, asks each teacher what it would do, and matches that answer.

This page covers that second step: the teacher to student KL loss, the temperature trick that keeps ranking detail, and multi task fusion. For the policies being copied, start with actor critic policies. For one strong teacher family, read Soft Actor Critic. For the bigger picture of deep policies, see deep reinforcement learning.

Think of It Like This

The apprentice with three tasting spoons

Picture a kitchen with three master cooks. One owns pastry, one owns grill, one owns salads. Each carries years of habit in their hands.

Now send in one apprentice with a small notebook. You do not ask the apprentice to redo their ten years of burnt pans and lucky saves. You ask each master to narrate while they work: not just "serve this dish" but "this one first, that one close second, never the third." The apprentice writes down the full ranking for every situation and practices until their own hands agree.

Temperature is how chatty the masters are. A curt master only names the winner, so the apprentice learns little. A chatty master ranks every option, so the apprentice learns which runner up is almost as good and which choice is poison.

Where the analogy stops: real teachers are frozen networks, not people, and they never adjust to the student. The student also copies mistakes faithfully, which is why the last section of this page matters so much.

How It Actually Works

A policy is a rule for picking actions. Write the teacher as $\pi_T(a \mid s)$ , the chance it picks action $a$ in state $s$ . Write the student as $\pi_S(a \mid s)$ with far fewer weights. Both output a probability over the same actions.

Distillation is supervised learning on teacher labels. Sample a state $s$ , often by rolling the teacher or a replay store, read the teacher distribution $q(a) = \pi_T(a \mid s)$ , read the student distribution $p(a) = \pi_S(a \mid s)$ , and push $p$ toward $q$ with KL divergence:

$L(s) = D_{KL}(q \parallel p) = \sum_a q(a) \log \frac{q(a)}{p(a)}$

That order matters. $D_{KL}(q \parallel p)$ punishes the student hardest where the teacher puts real weight, so the student must match the teacher top choice first. Cross entropy on the same soft targets gives the same gradients up to a constant, so you will see both names for one idea.

Why soft targets beat hard labels

An argmax label says "action 1 wins" and deletes the rest. A soft target says "action 1 gets 0.70, action 2 gets 0.20, action 3 gets 0.10," which tells the student that action 2 beats action 3 by two to one. That ranking is the dark knowledge. It guides the student through states it has barely seen.

Temperature $\tau$ controls how much ranking survives. Soften logits $z$ before the softmax:

$q_i = \frac{\exp(z_i / \tau)}{\sum_j \exp(z_j / \tau)}$

With logits $[2.0, 1.0, 0.0]$ , $\tau = 1$ gives $[0.6652, 0.2447, 0.0900]$ . Raise it to $\tau = 2$ and you get $[0.5065, 0.3072, 0.1863]$ . The winner still wins, but the third action doubles from 0.09 to 0.19, so the student keeps learning the full order. Use the same $\tau$ on both sides while training, then set $\tau = 1$ when you ship.

Multi task fusion by cycling tasks

One student can absorb many teachers. Keep a roster of frozen experts, pick a task, sample states from it, distill one batch, then switch tasks and repeat. The shared student trunk learns common structure while small task heads keep the action spaces apart. Actor Mimic follows this pattern and adds a feature match so the student copies middle layer activity too, not just final probabilities.

Worked example: KL on three actions

Take one state with three actions. The softened teacher says:

$q = [0.70, 0.20, 0.10]$

Your current student says:

$p = [0.50, 0.30, 0.20]$

Compute term by term with natural logs:

Top action: $0.70 \times \log(0.70 / 0.50) = 0.70 \times \log(1.40) = 0.70 \times 0.3365 = 0.2355$
Runner up: $0.20 \times \log(0.20 / 0.30) = 0.20 \times \log(0.6667) = 0.20 \times (-0.4055) = -0.0811$
Last action: $0.10 \times \log(0.10 / 0.20) = 0.10 \times \log(0.50) = 0.10 \times (-0.6931) = -0.0693$

Add them:

$D_{KL} = 0.2355 - 0.0811 - 0.0693 = 0.0851 \text{ nats}$

That is about $0.1228$ bits. Notice where the pain sits. The positive 0.2355 term dwarfs the rest because the student gives only 0.50 to an action the teacher loves at 0.70. Fixing that one gap cuts the loss fastest, which is exactly the behavior you want: match the teacher peak first, then polish the tail.

Code

This snippet computes the same KL and the temperature effect. It is plain Python with no extra packages.

1from math import exp, log2
3
4def softmax(logits: list[float], tau: float = 1.0) -> list[float]:5    """Softmax with temperature tau."""6    scaled = [z / tau for z in logits]7    peak = max(scaled)8    exps = [exp(v - peak) for v in scaled]9    total = sum(exps)10    return [e / total for e in exps]11
12
13def kl(q: list[float], p: list[float]) -> float:14    """KL divergence D(q || p) in nats."""15    return sum(qi * log(qi / pi) for qi, pi in zip(q, p))16
17
18teacher = [0.70, 0.20, 0.10]19student = [0.50, 0.30, 0.20]20print(f"{kl(teacher, student):.4f}")21# -> 0.085122print([f"{v:.4f}" for v in softmax([2.0, 1.0, 0.0], tau=2.0)])23# -> ['0.5065', '0.3072', '0.1863']

1# -> expected output:20.08513['0.5065', '0.3072', '0.1863']

Watch Out For

Student capacity collapse

A tiny student cannot hold five experts. The symptom looks confusing: the KL loss keeps falling on the training batches, yet reward on the hardest task falls off a cliff while easy tasks look fine. You optimized the average and starved the hard task.

Fix it by checking per task reward, not just pooled loss. Grow the student trunk, give each task its own output head, or sample hard tasks more often. If one task still lags after that, it needs its own student.

Teacher suboptimality lock-in

The student copies flaws as faithfully as skills. If the teacher brakes too late in one corner, the student brakes too late there too, and no amount of extra distillation batches will fix it because the target itself is wrong. The symptom: KL near zero, teacher and student failing the same new variation together.

Fix it with a short reward fine tune after distillation. Freeze nothing, let the student interact with the environment, and mix a little teacher KL with real reward so the student can outgrow its master where the master was weak.

The Quick Version

Policy distillation trains a small student to copy frozen teachers by minimizing KL between soft action distributions.
Soft targets carry ranking detail that hard argmax labels delete, and temperature controls how much of that detail survives.
Multi task fusion works by cycling tasks through one shared student while teachers stay frozen.
The worked KL of 0.0851 nats comes almost entirely from the student underweighting the teacher top action.
A small student can collapse on hard tasks, and a student can never beat a flawed teacher without reward fine tuning.