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).