Policy distillation

In distillation a student model learns to reproduce the outputs of a teacher model. With policies, the student is penalised for choosing different actions from the teacher in the same situation. The teacher can be a bigger network, an expert, or simply an older copy of the student.

For a policy that outputs a mean action $\mu(s)$ (see stochastic policy), the simplest distillation term adds the squared difference to the training loss:

$$L_{\text{teacher}} = w \cdot \frac{1}{|M|}\sum_{s \in M} \left\lVert \mu_\theta(s) - \mu_{\text{teacher}}(s) \right\rVert^2$$

Here $w$ sets how strongly the student is pulled towards the teacher, and $M$ is the set of states where the teacher applies.

Distillation as an anchor

When a trained policy is fine-tuned on a new skill, the updates for the new skill also change how it does the old ones, and the old skills can fade (catastrophic forgetting). Keeping a frozen copy of the original policy as a teacher anchors the old behaviour. The trick is the mask $M$: the teacher applies only where the old behaviour is still wanted, so the new skill is free to learn.

In this project

When the cross, hooks and uppercut were added to the jab, the first run had no anchor and the jab’s hit rate fell from 28 % to 2–7 % in 2.7 M steps. The next run used the best jab policy as a teacher, with $w = 2$, summed over the 43 motor outputs. The mask was every moment with no new punch running: the stance, the rhythm and the jab. The jab then held 5–6 m/s through 36 M steps while four new punches were learnt. Later runs freed more of the policy from the teacher: the cross, then every press with a target, then the neck and head, whose tilt the teacher had been passing on.

When the fighter learnt to walk, the best punching model was the teacher everywhere except during a step. At weight 2 the punches still drifted over about 100 M steps of walking. Raising the weight didn’t help, because of a bug: the teacher saw its inputs scaled with the student’s running statistics (observation normalisation), not its own, so its advice was slightly off. With the teacher on its own scaling and weight 10, the walker scored 107 of 200 punch presses against the punching model’s 112 (see Footwork).

Many teachers, one student

The same term can bring several specialists into one student. Each state is routed to the teacher in charge of it, and the student is pulled towards that teacher’s action there. Since one teacher owns each state, the targets never conflict, and the student has to fit them all in one network. In the fighter, the routing is by input: while a technique’s input is on, its one-technique teacher applies, and otherwise the base model does. The student learns on-policy: the teachers label the states it visits itself, and it still earns every task reward. distill1 brought both knees into guard18 this way. distill4, with $w = 4$ (from 2), brought in the body straight and both knees, and held punches at 60/120 against the base’s 59, with no falls. The pull counter came back only at the bag: 6/6 there, 1–2/6 in the air, against its teacher’s 8/8 (see Shadow clones).