Research map / Reinforcement learning / Visual and representation RL
How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves a generalisation bound for policy distillation after training in a group-symmetry ZSPT setting, and shows empirically that distilling an ensemble of policies on diverse on-policy/mixed data beats the original PPO+Explore-Go agent in Four Rooms.
Problem
Zero-shot policy transfer in RL: it is unclear why policy distillation after training can improve test generalisation over the original agent, and which data should be used for distillation.
Approach
Analyses distillation after training in a generalisation-through-invariance ZSPT setting, where training uses full data augmentation over a subgroup of a symmetry group; the bound shows test performance improves with larger distilled ensembles (N) and larger subgroup coverage (smaller kappa), which is translated into practical advice to ensemble distilled policies and distil on more diverse states
Key results
- In Four Rooms, an ensemble of N=10 distilled policies on the Explore-Go dataset reached 0.88 test return versus 0.74 for the PPO+Explore-Go teacher.
- In the Reacher illustrative CMDP, test performance rose from 0.75 (N=1) to 1.05 (N=100) for distilled ensembles trained on subgroup C4.
- Behaviour cloning ensembles (N=10) on the Explore-Go dataset reached 0.75 test return versus 0.26 for single BC on the Teacher dataset (0.37 for N=10).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase distillation ensemble size from N=1 to N=100 | Reacher policy distillation · Reacher with rotational symmetry test | return: 0.75 ± 0.147 -> 1.05 ± 0.117 | Helped |
| increase the training subgroup from C2 to C8 | Reacher policy distillation · Reacher with rotational symmetry test | return: 0.39 ± 0.0805 -> 1.11 ± 0.072 | Helped |
| add C4-rotated starting-pose trajectories | Reacher policy distillation · Reacher with rotational symmetry test | return: 0.39 ± 0.051 -> 0.48 ± 0.080 | Helped |
| add trajectories from random starting poses | Reacher policy distillation · Reacher with rotational symmetry test | return: 0.39 ± 0.051 -> 0.49 ± 0.077 | Helped |
| use a more diverse Mixed or Explore-Go dataset instead of Teacher trajectories | Four Rooms policy distillation · Four Rooms test | return: 0.56 ± 0.049 -> 0.72 ± 0.040; 0.56 ± 0.049 -> 0.78 ± 0.041 | Helped |
| increase ensemble size from N=1 to N=10 | Four Rooms policy distillation · Four Rooms test | return: 0.56 ± 0.049 -> 0.67 ± 0.054; 0.72 ± 0.040 -> 0.84 ± 0.034; 0.78 ± 0.041 -> 0.88 ± 0.036 | Helped |
| distill an N=10 ensemble on Explore-Go data rather than use the PPO+Explore-Go teacher | Four Rooms policy distillation · Four Rooms test | return: 0.74 ± 0.040 -> 0.88 ± 0.036 | Helped |
| use Explore-Go rather than Teacher data for behaviour cloning | Four Rooms behaviour cloning, N=1 · Four Rooms test | return: 0.26 ± 0.046 -> 0.56 ± 0.060 | Helped |
| use Mixed data rather than Teacher data for behaviour cloning | Four Rooms behaviour cloning, N=1 · Four Rooms test | return: 0.26 ± 0.046 -> 0.15 ± 0.024 | Hurt |
| increase behaviour-cloning ensemble size from N=1 to N=10 | Four Rooms behaviour cloning · Four Rooms test | return: 0.26 ± 0.046 -> 0.37 ± 0.054; 0.56 ± 0.060 -> 0.75 ± 0.045 | Helped |
Surprising findings
- Distilling on the Mixed dataset (containing off-policy pure exploration data) gave a better student than the Teacher-only dataset, and behaviour cloning on the Mixed dataset performed worse than on the Teacher dataset.
- Training on Random extra starting poses gave as much test benefit as the symmetry-based Training Contexts + C4 dataset, suggesting the benefit is not specific to group invariance.
Limitations
- Theory requires infinitely wide networks, scalar policies and Lipschitz assumptions that do not hold in experiments.
- Generalisation bounds for neural networks are notoriously loose; the bound is only qualitatively matched.
- Ensemble training and inference cost scales linearly with ensemble size when not parallelised.
Experimental setup
Datasets: Four Rooms (Minigrid) teacher trajectories Explore-Go trajectories Mixed teacher and pure exploration trajectories Reacher with rotational symmetry rollouts
Benchmarks: Four Rooms (Minigrid) Reacher with rotational symmetry
Metrics: mean return test performance train performance
Baselines: PPO+Explore-Go behaviour cloning single distilled student (N=1)
Builds on: Explore-Go Proximal Policy Optimization neural tangent kernel policy distillation after training
Closest papers by meaning
- Distributional Successor Features Enable Zero-Shot Policy Optimization (NeurIPS 2024 · similarity 0.88)
- On Generalization Across Environments In Multi-Objective Reinforcement Learning (ICLR 2025 · similarity 0.88)
- Generalizable Policy Improvement via Reinforcement Sampling (Student Abstract) (AAAI 2024 · similarity 0.88)
- Behaviour Distillation (ICLR 2024 · similarity 0.87)
- DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design (ICML 2024 · similarity 0.87)
- Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.