How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning
Authors: Max Weltevrede
Moritz Akiya Zanger
Matthijs T. J. Spaan
Wendelin Boehmer
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Generalization
Policy Distillation
Symmetries
Primary area: reinforcement_learning
TL;DR: We prove an ensemble of policies distilled on a diverse dataset improves generalisation in reinforcement learning.

Abstract: In the zero-shot policy transfer setting in reinforcement learning, the goal is to train an agent on a fixed set of training environments so that it can generalise to similar, but unseen, testing environments. Previous work has shown that policy distillation after training can sometimes produce a policy that outperforms the original in the testing environments. However, it is not yet entirely clear why that is, or what data should be used to distil the policy. In this paper, we prove, under certain assumptions, a generalisation bound for policy distillation after training. The theory provides two practical insights: for improved generalisation, you should 1) train an ensemble of distilled policies, and 2) distil it on as much data from the training environments as possible. We empirically verify that these insights hold in more general settings, when the assumptions required for the theory no longer hold. Finally, we demonstrate that an ensemble of policies distilled on a diverse dataset can generalise significantly better than the original agent.

How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning
Max Weltevrede
Moritz Zanger
Delft University of Technology
Delft University of Technology Delft, The Netherlands
Delft, The Netherlands
w.w.welleereeee@udel....
zanger@tudelfe
Matthijs T. J. Spaan Delft University of Technology
Wendelin Böhmer
Delft University of Technology
Delft, The Netherlands m. wt..jsppaa@ttuellf.....
Delft, The Netherlands
j w. Sohmert
Abstract
In the zero-shot policy transfer setting in reinforcement learning, the goal is train an agent fixed set of training environments so that it can generalise to similar, but unseen, testing environments. Previous work has shown that policy distillation after training can sometimes produce policy that outperforms the original in the testing environments. However, is not yet entirely clear why that is, what data should be used to distil the policy. In this paper, prove, under certain assumptions, generalisation bound for policy distillation after training. The theory provides two practical insights: for improved generalisation, you should 1) train an ensemble of distilled policies, and 2) distil on as much data from the training environments possible. We empirically verify that these insights hold in more general settings, when the assumptions required for the theory no longer hold. Finally, demonstrate that ensemble of policies distilled diverse dataset can generalise significantly better than the original agent.
Introduction
major challenge for developing reliable reinforcement learning (RL) agents is their ability to generalise to new scenarios they did not encounter during training. The zero-shot policy transfer setting (ZSPT, Kirk et al., 2023) tests for this ability by having an agent train on fixed set training environments, referred to as training contexts, and measuring the agent's performance set of similar, but different, testing contexts. Previous work has identified that policy distillation after training, the act of transferring knowledge from the agent's policy into freshly initialised neural network, can be used as a tool for generalisation. In particular, it has been shown that the distilled policy sometimes achieves higher test performance than the original policy (Lyle al., 2022).
held-out
However, is not yet entirely clear how policy distillation after training can improve generalisation performance in RI. Lyle et al. (2022) theoretically show that the temporal difference (TD) loss negatively affects the smoothness of the learned value function, which only indirectly explains why policy distillation after training (without TD loss) can improve generalisation. They also partially attribute the observed generalisation benefits to the stationarity of the distillation targets, which avoids negative effects induced by the non-stationary RL targets during training (Igl al., 2021), but this lacks solid theoretical justification. Moreover, recent work has shown that not only the stationarity of the RL training distribution, but also its overall diversity affect generalisation to unseen contexts (Jiang et al., 2023; Suau et al., 2024; Weltevrede et al., 2025). This additionally
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

raises the question whether we can increase generalisation performance by changing the distribution of data on which the policy distilled.
In this paper, we theoretically analyse the act of distilling a policy after training, and try to answer how the policy should be distilled and what data. Our analysis is based the idea that many real-world data distributions exhibit symmetries, and that generalising to novel inputs will require being invariant those symmetries. Although there is a lot of empirical evidence that demonstrates neural networks can learn invariances from data in a wide variety of settings and applications (Shorten and Khoshgoftaar, 2019; Feng et al., 2021; Zhang et al., 2021), proving this often requires stricter assumptions. Therefore, we analyse policy distillation in a generalisation through invariance ZSPT (GTI-ZSPT) setting, in which the agent has to learn invariance observing a subgroup of those symmetries during training. For this setting, we prove a generalisa- tion bound for a distilled policy, and deduce insights that should translate beyond the strict group theoretical framework required for the theory.
symmetry group, whilst only
Specifically, our theoretical results lead two practical insights: generalisation performance can be improved by 1) training an ensemble of distilled policies, and 2) distilling on a diverse set of states. Training an ensemble can be very costly. However, we demonstrate that generalisation can be improved (at only a fraction of the sample cost required for training the RL agent), by instead creating an ensemble after training by distilling the agent several times and averaging the resulting policy. Finally, related to our work on policy distillation, recent work has suggested that the generalisation performance of behaviour cloning (BC) is competitive with state-of-the-ar offline RL in the ZSPT setting (Mediratta et al., 2024). We demonstrate the insights for policy distillation also transfer to the BC setting and produce better generalising behaviour cloned policies. Our contributions are:
Given a policy (for example, an RL agent after training), we prove a bound on the test perfor- mance for distilled policy in the GTI-ZSPT setting. This bound is improved by 1) distilling a larger ensemble of policies, and 2) distillation over a more diverse set of states.
Inspired by the theoretical results, we empirically show that the insights gained from the theory improve generalisation of behaviour cloned and distilled policies in more general settings, when the strict assumptions required for the theory no longer hold. Furthermore, we demonstrate that an ensemble of policies distilled on diverse dataset can generalise significantly better than the original RL agent.
2 Background
The goal in reinforcement learning is to optimise decision-making process, usually formalised as a Markov decision-making process (MDP) defined by the 6 tuple M = (S, A, R.po. po.)). 7). ). In this tuple, S denotes the state space, A the action space, T : S x A P(S) the transition model, R : S x A + R the reward function, Po P(S) the initial state distribution and Y E [0.1] the discount factor, where P(S) denotes the probability function over state space S. Optimising an MDP corresponds to finding the policy æ : S - (P(A) that maximises the return (the expected discounted sum of rewards) J TEx 1E "rr)). The expectation here is over the Markov chain {so, ag, ro. $1, a1, ......} induced by following policy in MDP M (Akshay et al., 2013). The optimal policy π²" = arg max, E 12>r0 2 is the policy that maximises this return. The on-policy distribution PM= PP5) is the distribution over the states that a policy would visit in an MDP M.
A contextual Markov deciisionmaking process (CMDP) MIc (Hallak al., 2015) is MDP where the state space 5 = S' x C' can be structured as an outer product of a context space C and underlying state space S" A context E C' is sampled at the start of an episode and does not change thereafter. The context is part of the state and can influence the transitions and rewards. As such, it can be thought of as defining task specific environment that the agent has to solve in that episode. In the zero-shot policy transfer (ZSPT) setting (Kirk et al., 2023), an agent gets to train on a fixed subset of contexts Cirain C and has to generalise to a distinct set of testing contexts Ctest C, Ctrain n Ctest Ø. In other words, the agent gets to interact and train in the CMDP M M|Cerain (the CMDP induced by the training contexts Ctrain), ) but has to maximise return in the testing CMDP M Mccau.
2.1 Policy distillation
In policy distillation, a knowledge transfer occurs by distilling a policy from a teacher network into a newly initialised student network. There are many different ways the policy can be distilled
2

---

(Czarnecki et al., 2019), but in this paper we consider a student network that is distilled on fixed dataset, that is collected after training, and usually (but not necessarily) consists on-policy data collected by the teacher. For analysis, we simplify the setting by assuming deterministic, scalar student and teacher policy To : S - IR Tg : S -> TR. The distillation loss we consider is simply the mean squared error (MSE) between the output of the two policies:
LD(0, (.).).TT) = (=o(s) - TT(s))²
(I)
RED
where D = [s11...,,A Sn is the set of states we distil on. This simplified distillation setting is only used for the theoretical results, our experiments in Section 5 consider more general settings. Note, we consider behaviour cloning (BC) as specific instance of policy distillation, where the student network only has access to a fixed dataset of the teacher's behaviour (state-action tuples). For more on behaviour cloning, distillation and their differences, we refer to Appendix A.I.
If we assume a certain smoothness of the transitions, rewards and policies (in particular, Lipschitz continuous MDP and policies), it is possible to bound the performance difference between the student and an optimal policy (Maran et al., 2023, Theorem 3):
Theorem 3. Let T' be the optimal policy and Te be the student policy. If the MDP is (LT, LR)- Lipschitz continuous and the optimal and student policies are Lipsschhiitz continuous, and we have that (1 + Ls) < 1, then it holds that:
L
Jd - J 1 (1 - 7) - -Lr (1+ + ))
E Bandr (-1s), =o(-1s))
where d= -) (1)) Y PP(s
s/m=", Po) the "-discounted visitation distribution.
Proof. See Appendix E. for the proof and exact definitions of all the terms.
In other words, under these conditions the return of a student policy can be bounded by the distance between the student and optimal policies along the states visited by the optimal policy.
2.2 Symmetry groups
To formalise the notion of symmetries and invariance of function f fX + Y, is useful to define symmetrry group G. A group is a non-empty set G together with a binary operation that satisfies certain requirements such as closure, associativity and always containing an inverse and identity element.' ! A group and its elements are abstract mathematical notions. In order apply elements from group to a vector space X, we need to define group representation v/x that maps group elements to invertible matrices. In this paper, always assume the representations are orthogonal (i.e. x (1 e'x (g)T), ). Note that we only need to define the representations Φ for the analysis, as they are part of the generalisation bound but are not explicitly defined for the experiments.
We can define the invariance of a function f as
((bx(g)))) = ((x) Vo E X,9 G.
A subset B of G' is called subgroup of G (notation B S G) if the members of B form group themselves. Any group G has at least one subgroup, the trivial subgroup, consisting of only the identity element e g g o c g. Vg E G. A subgroup B < G is finite if it has a finite number of elements. For more on group theory, we refer to Appendix A.2.
Example The group SO(2) consists of the set of all rotations in two dimensions. If the input to function ʃ consists of Euclidean coordinates (Ir, y), the group representation v/x maps a rotation of a degrees to the 2D rotation matrix associated with an degree rotation. The function f would be considered rotationally invariant if (()x(a)x) = f fxx, Va E RR S0().. An example of a finite subgroup of S2) is the group C4 consisting of all 90" rotations, the subgroup consisting of only the identity element (09 rotation).
In this paper, we abuse notation slightly by denoting both the group and the non-empty set with G. depending context.
3

---

One approach to induce a function that is invariant to the symmetry group G is to train it with data augmentation. For groups of finite size, is possible to perform full data augmentation, which consists of applying every transformation in G, to each element of an original dataset T = (x = {(x, v) E X, J.Y).. The function is then trained on the augmented dataset To {(ex (g)x,, VZ2 ()) y) E T, e G). In general, training a function with data augmentation does not guarantee it becomes invariant (Flinth and Ohlsson, 2023), it instead can become approximately invariant invariant only on the distribution of data on which it was trained (Kvinge et al., 2022; Lyle et al., 2020; Azulay and Weiss, 2019). However, under certain conditions, the average of an infinitely large ensemble can have that guarantee (Gerken and Kessel, 2024; Nordenfors and Flinth, 2024).
2.3 Ensembles and invariance
Formally, ensemble consists of multiple neural networks So X + R with parameters 0 ~ j initialised from some distribution je and trained on the same dataset T = (x = {zz E X, YY The output of an infinitely large ensemble fc(x) at training time is given by the average over the ensemble members fo: fa(x)) Bø- [x)], where La denotes map from initial parameters 0 to the corresponding parameters after steps of gradient descent. In practice, the infinite ensemble is approximated with finite Monte Carlo estimate of the expectation fc: fe z fi N E fe.o, (z),
where 0g ~ j and N is the size of the ensemble.
2.3.1 Infinite width limit
Although there does not yet exist single comprehensive theoretical framework for how neural networks work, significant progress has been made in the field of deep learning theory in the limit of infinite layer width. In this limit, an infinite ensemble (x) trained with MSE loss follows simple Gaussian distribution that depends on the network architecture and initialisation (Jacot et al., 2018; 8; Lee et al., 2019). Gerken and Kessel (2024) prove that the infinite ensemble f(()) trained on the augmented dataset To for some group G, satisfies the definition of invariance in equation (2), for any I and any In other words, an infinitely large ensemble of infinitely wide neural networks, trained with full data augmentation for group G, is invariant under transformations from G for any input and any point during the training process. In our analysis, use Lemma 6.2 from Gerken and Kessel (2024) that bounds the invariance of an infinite ensemble of infinitely wide networks trained with full data augmentation on a finite subgroup B S G:
Lemma 6.2. Let fe(x)= Eo Ef (x)) be an infinite ensemble of neural networks with Lipschitz continuous derivatives with respect the parameters. Define the error measure of discrepancy between representations from the group G and its finite subgroup B:
K
max min gEG BEB
||xx - (b)looo
where 1opp denotes the operator norm. The prediction of an infinite ensemble trained with full data augmentation B < G deviates from invariance by
Fz(z) ((x(()))=) SxC((),
EgEG
(4)
for any time t. Here C is a function of x independent of
Proof. See Appendix E.2 for the proof and exact definitions of all the terms.
This lemma bounds the deviation from invariance of the infinite ensemble by a factor k, which is a measure of how well the subgroup B covers the full group G, in the space of representations thx. For more background on infinite ensembles in the infinitely wide limit, we refer to Appendix A.3.
3 Related work
The CMDP framework captures many RL settings focused on zero-shot generalisation (Kirk et al., 2023). Some approaches to improve generalisation focus on learning generalisable functions through inductive biases (Kansky et al., 2017; Wang et al., 2021) or by applying regularisation techniques from supervised learning (Tishby and Zaslavsky, 2015; Cobbe et al., 2019). These approaches improve

---

generalisation by changing the RL training process, whereas we distil a teacher policy after training. which in principle is agnostic how that teacher was trained. Other work improves generalisation by increasing the diversity of the data on which the agent trains, for example by increasing the diversity of the training contexts using domain randomisation (Tobin et al., 2017; Sadeghi and Levine, 2017), or creating artificial data using data augmentation (Lee et al., 2020; Raileanu et al., 2021). Our work focusses on sampling additional data from a fixed set of training contexts, but differs from data augmentation in that we do not require explicitly designed augmentations. For broader survey on zero-shot generalisation in reinforcement learning, see Kirk et al. (2023).
Policy distillation in RL has been used to compress policies. speed up learning, or train multi-task agents by transferring knowledge from teacher policies to student networks (Rusu al., 2016; Schmitt et al., 2018; Czarnecki et al., 2019). Various methods of distillation exist, balancing factors such as teacher quality, access to online data, and availability of teacher value functions or rewards (Czarnecki et al., 2019). Some studies have used distillation to improve generalisation, either by mitigating RL--peccffic con-stationarity through periodic distillation (Igl et al., 2021) 1)21)eee or by distilling from policies trained with privileged information or weak augmentations (Fan et al., 2021; Walsman et al., 2023). Most similar to our work, Lyle et al. (2022) show a policy distilled after training can sometimes generalise better than the original RL agent. But, their theory only indirectly covers policy distillation and they do not investigate how the distillation data affects generalisation.
4 Generalisation through invariance
In this section, we introduce specific ZSPT setting that allows to prove generalisation bound for a distilled policy. The main idea is that many real- world data distributions exhibit symmetries, and that generalising to novel inputs sampled from this distribution requires (at least partially) being invariant to those symmetries. Moreover, any training dataset sampled IID from this distribution will likely observe some of these symmetries.
Proving a neural network learns invariances from data is not straightforward, and usually requires assumptions on the mathematical structure of the symmetries. For this reason, we consider specific setting in which an agent has to become invariant to a symmetry group G, but trains with full data augmentation under only subgroup B I G. Even though this setting requires strict assumptions, expect the insights to apply more broadly, as there a lot of empirical evidence that data augmentation improves generalisation performance in wide variety of settings and applications (Shorten and Khoshgoftaar, 2019; Feng et al., 2021; Zhang et al., 2021; Miao et al., 2023). We formalise the idea in generalisation through invariance ZSPT (GTI-ZSPT)
Definition 1 (Generalisation through invariance ZSPT), Let Mlc be a CMDP and lei Crrains Crest C C be set of training and testing contexts that define ZSPT problem. Additionally, let at" be the optimal policy in Mlc. Mic {s E 5lp (s) >> 0} denote the set of states with non-zero support under the on-policy distribution PMil in CMDP Mlc. In the generalisation through invariance ZSPT (GTI-ZSPT), the sets STA and MMlera admit a symmetric structure:
SMiC = (bb((g)sgg E G S Miccrain
(bbs())s(b E B, B, BSG
where S ST a is a proper subset of ST Giaain and G a non-trivial symmetry group (and
B < G a finite subgroup) that leaves the optimal policy invariant: T T²(s) = T" ())(())s), Vs e S.
To quantify the discrepancy between the group and its subgroup, the following measure is defined (Gerken and Kessel, 2024):
Definition 2. For the group G and its finite subgroup B S G that define the symmetric structure of GT1-ZSPT (Definition 1), R. is measure of discrepancy between the representations of these groups:
K = max min ||vs (g) 4/s (b) lop
gEG bEB
where ||
-llop denotes the operator norm.
5

---

The constant measures how much the subgroup B < G deviates from the full group G. The bigger the subgroup is, the smaller x will become (with & 0 in the limit of G).
Example The Reacche with rotational symmetry' ZSPT in Figure satisfies the conditions of the GT1-ZSPT: This is continuous control environment where the agent has move robot arm (blue) in such a way that its hand (black circle) reaches the goal location (green circle). The four training contexts have shoulder locations that are rotated 90, 180 and 270 degrees around the goal location. In testing, the shoulder can be rotated any amount. As an example, the measure k for the subgroup C₁ of 90º rotations (as depicted in the figure), would be larger than for the bigger subgroup Cs of 45' rotations. See Appendix C.1 for more on this example and how it satisfies the assumptions.
Training
Testing
Context:
Figure 1: A 'Reacher with rotational symmetry' CMDP with four training contexts, differing in the location of the shoulder (red), positioned along a circle (dotted line). All contexts share the relative pose of the robot arm (blue). The goal is for the hand (black circle) reach the goal location (green circle) in the middle. The training contexts can be generated by applying the group of 90" rotations to context 1, and the testing contexts be generated with the full group of rotations (2O(2)).
4.1 Bounding the performance
For the GTI-ZSPT setting, can bound the performance of distilled policy in the testing CMDP with the following theorem:
Theorem 1. Consider policy distillation for deterministic, scalar teacher policy kg S - R (Equation (1) in Section 2.1) in a LT, LR-Lipschitt continuous CMDP in the GTI-ZSPT setting. Let the student policy to be an ensemble of N infinitely wide neural networks me S - R with Lipschitz continuous derivatives with respect to its parameters, distilled on an on-policy dataset Mileerain (b)sb E B,s 5} consisting of all the states in the training contexts
D DSS encountered by the teacher in the GTI-ZSPT setting. Furthermore, let the student policy be Li=" Lipschitz continuous and assume -Ir I=+,.) <1.
If the teacher is optimal in the training tasks Cirain (but arbitrarily bad anywhere else), the per- formance of the student in the testing CMDP Mleuse is bounded with probability least 1 e, by:
J - 100 (1 (1-7)(1-~<y')) Lit
LR
RCo+
1xr (E) VN
(5)
where i is the measure of discrepancy between subgroup B < G and full group G (see definition 2) and Co, Cr.,, are constants that depend on the "ddicoouneed visitation distribution of the optimal policy in MMC the network architecture, and the dataset D. Additionally, Cr.y, also depends on the network intialisation and the confidence level
Proof. Thanks to the symmetric structure of the GTI-ZSPT, we can bound the output of an infinite ensemble of distilled policies To.. when evaluated on testing states in Mc using the bound on the deviation from invariance from Section 1. This can be combined with a probabilistic bound for Monte Carlo estimators to bound the output of a finite ensemble 7 on the testing states. With this bound on the output of the student policy, we can use the performance bound for Lipschitz continuous MDPs from Section 2.1 to get our final result above. See Appendix B for the full proof.
The theorem above offers two insights:
1. The bigger the ensemble size N, the smaller the bound performance.
2. The bigger the subgroup B, the smaller the measure the smaller the bound performance.

---

As we mentioned before, even though Theorem requires strict assumptions, we believe the insights apply more broadly. Essentially, the theorem relies on the generalisation benefits induced by training on additional samples generated by performing data augmentation. In practice, it often doesn't 't matter the augmentations form a group, are consistent with the original data distribution, or are applied to all classes equally (Bishop, 1995; Wu et al., 2020; Hansen and Wang, 2021; Lin et al., 2022; Geiping et al., 2023; Miao et al., 2023). As such, we believe that in many settings, the benefits of training on a bigger subgroup B, can also be realised by simply training more diverse data, which we clarify with some examples in our experiments.
5 Experiments
In this section, we demonstrate that the insights provided by the theory translate to practical and workable principles that can improve the generalisation performance of distilled policy, beyond the performance of the original agent. In Section 5.1, we establish that bigger ensembles indeed improve generalisation and show what means to train on bigger subgroup B < G in the illustrative CMDP from Figure 1. This experiment satisfies the assumptions for the GT1-ZSPT setting, but does not strictly satisfy some of the non-practical assumptions required for Theorem 1. In Section 5.2, we demonstrate that the insights also apply to the more complex Minigrid Four Rooms environment (Chevalier- Boisverr et al., 2023) that breaks most of the assumptions required for the proof in Section 4. For experimental details, see Appendix C.
5.1 Reacher with rotational symmetry
Table 1: Performance of distilled policies in the Illustrative CMDP from Figure 1 for different ensemble sizes N (trained under subgroup B = Cc) and different subgroups B < SO(2) (for N 1). Shown are the mean and standard deviation for 20 seeds, and in bold the best returns including those with overlapping 95% confidence intervals.
Ensemble Size N:
N=1
N=10
N=100
Train Performance
1.17 0.004
1.17 + 0.004
1.17 = 0.003
Test Performance
0.75 + 0.147
0.89 + 0.107
1.05 0.117
Subgroup B < SO(2):
B C2
= C4
B Css
Train Performance
1.17 0.003
1.17 + 0.004
1.16 = 0.002
Test Performance
0.39 + 0.0805
0.75 + 0.147
1.11 0.072
The theory proves that we can reduce an upper bound on the difference to optimal performance when we increase the ensemble size N and train on a bigger subgroup < G. However, that does not always guarantee strict performance improvements (for example, if the upper bound were so large it meaningless). Furthermore, the theory requires some assumptions that are not always practical, such infinitely wide networks, scalar-valued policies, or Lipschitz-continuous reward function, that we do not expect to affect the overall result in practice. Therefore, we investigate whether the insights from Section 4 hold without these assumptions, and whether they lead to actual generalisation improvements. Table I we show that increasing the size of the ensemble, consisting of networks of finite width, does actually lead to higher test performance in the CMDP from Figure 1.
Additionally, in Table we show that generalisation performance is affected by the size of the sub- group B < S2()) we train on. Only training on two training contexts, corresponding to the subgroup C2 S (2( of 180 rotations, performs worse than training on four contexts, corresponding to the subgroup C2 S 50(2) of 90" rotations (as shown in Figure 1). Furthermore, training on eight contexts (subgroup Cs < S0(2) of all 45 rotations) is even better. In this illustrative CMDP, training on larger subgroups requires training in new contexts, but this is not always the case.
5.1.1 Improving generalisation with diverse data from the same contexts
In sufficiently complex CMDPs, there are several dimensions of variation between different contexts. For example, we can add different starting poses to the contexts the CMDP from Figure 1, such that a context is now defined by rotation of the shoulder and the relative pose of the robot arm
7

---

Training
Testing
Context:
2
Figure 2: The base context set in the illustrative reacher CMDP with varying shoulder location (red) and robot arm pose (blue), see Figure for details.
(see Figure 2). This CMDP does not strictly satisfy the symmetry conditions in Definition 1, but invariance to rotations is still a major component for generalisation. Since the training contexts now also differ in the starting pose, the dataset generated from the training contexts no longer corresponds to performing full data augmentation with respect to rotational symmetry. However, in this example, this can be fixed by training on additional data from the given set of training contexts.
To illustrate this, we compare three distillation datasets:
1. Training Contexts: This dataset consists of the teacher's trajectories in the training contexts (contexts through in Figure 2).
2. Training Contexts Cs: This dataset consists of the Training Contexts dataset, but with additional trajectories sampled from different starting poses in the same contexts. In particular, for each context, it includes trajectories starting in the rotated poses from the other contexts. This dataset corresponds to performing full data augmentation for the 90 rotations subgroup C4 on the Training Contexts dataset (see Appendix C. for a visual representation this).
3. Training Contexts Random: Like Training Contexts C4, this dataset includes additional trajectories from different starting poses. However, for this dataset the new starting poses are sampled uniformly at random.
The Training Contexts + C4 dataset illustrates how in this CMDP the subgroup B < SO(2) can be increased by sampling additional trajectories from the same training contexts (technically, the Training Contexts dataset corresponds to the trivial subgroup {e} S S0(2) consisting of only the identity element e, which is smaller than C4 S SO(2)). In Table 2, see that training on this dataset indeed produces higher test performance than the Training Contexts dataset. However, the same generalisation benefits are also observed for the Training Contexts Random dataset.
Table 2: Performance of distilled policies (for N = 1) in the Illustrative CMDP from Figure 2 for different datasets. The datasets consist the teacher's trajectories sampled for several starting states. Shown are the mean and standard deviation for 20 seeds, and in bold are the best returns including those with overlapping 95% confidence intervals.
Distillation Dataset
Train
Test
Training Contexts
1.20 + 0.157
0.39 + 0.051
Training Contexts + C1
1.11 + 0.099
0.48 0.080
Training Contexts Random
1.14 + 0.072
0.49 0.077
The Training Contexts Random dataset illustrates that the generalisation benefits of data augmenta- tion go far beyond the "training to be invariant under a group symmetry" paradigm. Some studies suggest that the benefits are simply due to the regularising effect that data augmentation can provide (Bishop. 1995; Wu et al., 2020; Hansen and Wang, 2021), or by making it more difficult to overfit to spurious correlations (Raileanu et al., 2021; Shen et al., 2022). In this sense, we expect the insight of training with full data augmentation on a bigger subgroup & G from Theorem 1, to translate in practice simply training on more diverse data, data that is sampled from the same contexts.
5.2 Four Rooms
In this section, we demonstrate that increasing ensemble size and data diversity can significantly increase the generalisation performance of distilled policy, even when most of the assumptions for Theorem no longer hold. The Four Rooms grid world environment from the Minigrid benchmark

---

does not appear to have an invariant symmetry that plays a core part in generalising to new contexts, as required for the definition of a GTI-ZSPT. The teacher is an agent trained with Proximal Policy Optimisation (PPO Schulman et al., 2017) and is therefore not necessarily optimal in the training contexts. Additionally, the teacher a stochastic policy that is distilled by regressing on the vector of probabilities or behaviour cloned using logarithmic loss (see Appendix A. for more background on these losses).
5.2.1 Policy distillation improves generalisation
For the experiments in the Four Rooms environment the teacher is a policy trained with the PPO+Explore- -o algorithm for million environment steps. The Explore-Go Go approach was intro- duced by Weltevrede et al. (2025) to increase generalisation by generating a more diverse training distribution for the RL agent. It leverages a separately trained pure exploration agent, rolled out at the beginning of each episode, Since this teacher trains on more diverse state distribution than a normal PPO agent, it provides good teaching targets for our distillation datasets. We compare the following three datasets:
artificially increase the starting state distribution for the PPO agent.
1. Teacher: This dataset consists of the teacher's trajectories in the (original) training contexts.
2. Explore-Go: This dataset mimics the training distribution for the Explore-Go approach by sampling teacher trajectories from additional starting states, generated by pure ex exploration policy rolled out at the start of each episode. This dataset has the property that all the data is on-policy for our teacher, yet more diverse than the Teacher dataset.
3. Mixed: This dataset is 50/50 mix of Teacher and trajectories collected by separately trained pure exploration policy. This dataset is diverse, but does not solely consist of states encountered by the teacher.
In Table 3 the Teacher dataset and that the ensemble of size NN = 10 outperforms the single student N = 1 for each dataset type. Moreover, the ensemble, distilled on the Explore- Go dataset, generalises signifi- cantly better than the original PPO agent, whilst only requiring around 12% additional environment steps (compared the teacher's training budget).
that the more diverse datasets (Mixed and Explore-Go) significantly outperform
Table 3: Performance of an ensemble (of size N) of policy distillation or hehaviour cloning policies on various datasets compared to the PPO+Explore-CG Shown are mean and standard deviation over 20 seeds, and in bold are the best returns including those with overlapping 95% confidence intervals (within the same category).
teacher in the Four Rooms environment.
Dataset
Train (N=1) Train (N=10)
Test (N=1)
Test (N=10)
PoO+Explore-Go
0.92 + 0.020
0.74 + 0.040
Distillation
Teacher
0.92 = 0.020
0.92 + 0.020
0.56 + 0.049
0.67 = 0.054
Mixed
0.92 + 0.020
0.92 + 0.020
0.72 + 0.040
0.84 = 0.034
Explore-Go
0.92 + 0.020
0.92 + 0.019
0.78 4 0.041
0.88 = 0.036
Behaviour Cloning
Teacher
0.91 + 0.022
0.92 + 0.020
0.26 +: 0.046
0.37 + 0.054
Mixed
0.86 + 0.031
0.91 + 0.025
0.15
0.024
0.20 + 0.026
Explore-Go
0.87 + 0.028
0.92 + 0.021
0.56 + 0.060 0.75 + 0.045
Lastly, we demonstrate in this section that the same insights also hold for a logarithmic behaviour cloning loss for stochastic policies that is widely used in practice (Foster et al., 2024). At the bottom of Table 3, we show that the an ensemble (of size N = 10), distilled on the Explore-Go dataset, generalises significantly better than single behaviour cloning agent on the Teacher dataset. Note that behaviour cloning achieves lower performances than distillation, and that BC performs considerably worse on the Mixed dataset. In our definition of behaviour cloning, the student policy learns to imitate whatever policy collected the dataset, by only observing the actions that were actually sampled during collection. Therefore, the BC agent performs worse than the distillation agent, since the latter has access to more information (all the action probabilities of the teacher). On the Mixed dataset, the BC agent clones the behaviour policy that consists of a 50/50 mix of the (optimal) Teacher policy and (suboptimal) pure exploration policy. The resulting cloned behaviour performs even worse than
2 For pure exploration, the objective focuses solely on exploring new parts of the state space, ignoring rewards.

---

the BC agent trained on the Teacher dataset. In contrast, the policy distillation agent on the Mixed dataset regresses on the action probabilities of the Teacher, on the states encountered by the 50/50 mixture of policies, and therefore has a much better learning target.
6 Discussion and limitations
The experiments in the Four Rooms environment in Section 5.2.1 serve to empirically demonstrate how insights can be leveraged to significantly enhance the generalisation performance of reinforcement learning agent through policy distillation. clear example of this is seen in our ensemble N 10, distilled the Explore-Go dataset, which achieves substantially higher test performance than the original PPO+Explore-GG teacher policy (see Table 3). The potential of policy distillation after training tool improve generalisation was initially identified in Lyle al. (2022), but we believe the results of this paper provide a more compelling argument and empirical evidence for this phenomenon.
Whether the benefits of performing data augmentation with respect to some symmetry group actually stem from induced invariance reduced overfitting and other forms of regularisation, is still an ongoing topic of discussion in the literature (Lyle et al., 2020; Shen et al., 2022). To add to this discussion, in Appendix D.1 we measure the invariance of our trained models on the 'Reacher with rotational symmetry' experiments from Table and plot it against the ensemble size N and subgroup B < S2(.). We find that in this particular experiment, the distilled policies do become more invariant as ensemble and subgroup size increase, just as our theory predicts.
Obtaining tight generalisation bounds for neural networks is notoriously challenging (Jiang et al., 2020; Gastpar et al., 2024). Moreover, some of the assumptions for Theorem such as infinitely wide neural networks, are hard to meet in practice. Therefore, we believe the true strength of our theory lies in its ability to identify crucial properties of the dataset distribution and distilled ensemble that capable of improving generalisation performance. Nonetheless, in Appendix D.2, analyse how well our results fit the + vN agree with the shape of the theoretical upper bound, suggesting that bound is not completely vacuous.
b relation identified by our theory. We find that our results reasonably
Finally, all ensemble members trained independently, and during inference, can also be evalu- ated independently (an independent forward pass with an average over the output of the ensemble afterwards). This inherent independence means that both the training and inference processes of the ensemble are parallelisable. If parallelisation not feasible, the runtime for both training and inference would increase linearly with the ensemble size. It is important note that all ensemble members are distilled on the same dataset. This means the small number of additional environment steps required to sample this dataset is independent of ensemble size.
7 Conclusion
In this paper, we investigate the advantage of policy distillation for improving zero-shot policy transfer (ZSPT) in reinforcement learning. We introduce the generalisation through invariance ZSPT setting, to prove a generalisation bound for a policy distilled after training. Our analysis highlights two practical insights: to distil an ensemble of policies, and 2) distil it on diverse set of states from the training contexts. We empirically evaluate that the insights hold in the Four Rooms environment from the Minigrid benchmark, even though it does not satisfy all the assumptions required for the theory, and that they also translate to the behaviour cloning setting. Moreover, we show that distilling an ensemble of policies on diverse set of states can produce policy that generalises significantly better than the original RL agent, thus demonstrating that policy distillation can be powerful tool to increase generalisation performance of reinforcement learning agents.
10

---

Acknowledgments and Disclosure of Funding
We thank Caroline Horsch, Laurens Engwegen and Oussama Azizi for fruitful discussions and feedback. The project has received funding from the EU Horizon 2020 programme under grant number 964505 (Epistemic AI) and was also partially funded by the Dutch Research Council (NWO) project Reliable number OCENW. 21.23.. The computational resources for empirical work were provided by the Delft AT Cluster (DAIC) (2024) and the Delft High Performance Computing Centre (DHPC) (2024).
Generalization in Deep Reinforcement Learning with project
References
S. Akshay, Nathalie Bertrand, Serge Haddad, and Loic Hélouët. The Steady-State Control Problem for Markov Decision Processes. In Kaustubh R. Joshi, Markus Siegle, Mariëlle Stoclinga, and Pedro R. D' Argenio, editors, Quantitative Evaluation of Systems- 10th International Conference, QEST 2013, Buenos Aires, Argentina, August 27-30, 2013. Proceedings, volume 8054 Lecture Notes in Computer Science, pages 290-304. Springer, 2013. doi: 10. 10.10007/978-3-64-401999 1_26.
URL https://dol
Himani Arora, Rajath Kumar, Jason Krone, and Chong Li. Multi-task Learning for Continuous Control. CoRR, abs/1802.01034, 2018. URL http://arxiv.. wrg.//orrg/aaaaa waagg/abs//n http//882..111 01034. arXiv:
1802.01034.
Aharon Azulay and Yair Weiss. Why do deep convolutional networks generalize so poorly to small image transformations? J. Mach. Learn. Res., 20:1 10:184:11-14:25, 2019. URL https:
//jmlr.or oaegrr.rr//oo org wwwgg/apprs/vv////// 11000019-1199 519. html.
Christopher M. Bishop. Training with Noise is Equivalent to Tikhonov Regularization. Neural Comput., 7(1):108-11) 116, 1995. doi: 10.116/NECC..1995.....00 URL https://doo/ o00//0.............
1162/neco. 1995. 7.1.
Maxime Chevalier- Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid && Miniworld: Modu- lar & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks. CoRR, abs/2306. 13831, 2023. URL https: //mnngrrid farama. org.
Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem
Karl Cobbe, Oleg Klimov, Christopher Hesse, Taehoon Kim, and John Schulman. Quantifying Generalization in Reinforcement Learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learn- ing Research, pages 1282-1289. PMLR, 2019, URL http:::///pooceeengg..oto wttp:iiismmm...ppe/s/9333 cobbe19t html.
Wojciech M. Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant M. Jayakumar, Grzegorz Swirszez, and Max Jaderberg. Distilling Policy Distillation. In Kamalika Chaudhuri and Masashi Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pages 1331-1340. PMLR, 2019. URI. http:tttt////ceeeeieeeee mll .press/ v89/czarnecki htmll.
Delft AI Cluster (DAIC). The Delft AI Cluster (DAIC), RRID:SCR_02500) 2024. URL https: //doc. dibo..dai.ttun wtteeettdellltl.....
Delft High Performance Computing Centre (DHPC). DelftBlue Supercomputer (Phase 2), 2024.
URL https: : h/tp://wwww/cc// tudelft. nl/ httcc/hpaaarraae (0144443//DeftBBuueeee
Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Animashree Anandkumar. SECANT: Self-Expert Cloning for Zero-Shot Generalization of Visual Policies. In Marina Meila and Tong Zhang, editors, Proceedings the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 3088-301 3099. PMLR, 2021. URL http: / http///poooeeeeeeee/eee/e/eessee mlr. press/vi3s wavvvv11aaaaaa faw./aannnc.....ma/... html.

---

Steven Y. Feng, Varun Gangal, Jason Wei, Sarath Chandar, Soroush Vosoughi, Teruko Mitamura, and Eduard H. Hovy. A Survey of Data Augmentation Approaches for NLP. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli, editors, Findings of the Association for Computational Linguisties: ACI/JCNNLP 2021, Online Event, August 1-6, 2021, volume ACL//ICNLP 2021 of Findings ACL., pages 968-988. Association for Computational Linguistics, 2021. doi: 10.18653/ V1/2021 WII.LIIINNNNNLLLLL. URL https://dol: wrt.o.......//v....... 84.
Axel Flinth and Fredrik Ohlsson. Optimization Dynamics of Equivariant and Augmented Neural Networks. CoRR, abs/2303. 13458, 2023. doi: (0.48550/ARX112333...8833 13453. URL https: wtt.:oo.g//....///........./
//doi.org/10 00..50/5/5//a.. arxiv. 2303. 13458. arXiv: 2303.13458.
Dylan J. Foster, Adam Block, and Dipendra Misra. Is Behavior Cloning All You Need? Un- derstanding Horizon in Imitation Learning. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, UIrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, Ad- vances in Neural Information Processing Systems 38: Annual Conference on Neural In- formation Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December mips. ct/ccc///ppo cace/appr_fflessss patps//ppp///22///////// Awwtrabt-tttaaoaeeeeeaeeae Confeeenceoeeeeeeeeeeeeeeee .ttll.
10 15, 2024, 2024. URL
Michael Gastpar, Ido Nachum, Jonathan Shafer, and Thomas Weinberger. Fantastic Generalization Measures are Nowhere to be Found. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL
https://openreviix..
hwtp:/iwwweeeeeeeeeeeeeeeeeeeeee wwt. wewt://oruu..ddnnnooott....
Jonas Geiping, Micah Goldblum, Gowthami Somepalli, Ravid Shwartz-Zivy Tom Goldstein, and Andrew Gordon Wilson. How Much Data Are Augmentations Worth? An Investigation into Scaling Laws, Invariance, and Implicit Regularization. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https: tttpps://operreeee wtt.:/wwwww.eeeeettrdddeedd---ee
Jan E. Gerken and Pan Kessel. Emergent Equivariance Deep Ensembles. Forty-first International Conference Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: wwt.//wennreeeeeeeeeeeeee www.eee//oooooo oetp//orrmmmiipppXXX
Dibya Ghosh, Avi Singh, Aravind Rajeswaran, Vikash Kumar, and Sergey Levine. Divide-and- Conquer Reinforcement Learning. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018. Conference Track Proceedings. OpenReview.net, 2018. URL htttt::/ wwwpeerrreee.eet/oouuurdddeeddde mett/forrumiid=xwelllet
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor -ctor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssson Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1856-1865. PMLR, 2018. URL http: //proceedings mlr. ttress/v80/haarnjjaaa
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual Markov Decision Processes. CoRR, abs/1502.00259, 2015, URL http://arriiv.rr/aass//Ste
hwtp:///////bbbb/ee/eoo 1502. 02259. arXiv: 1502.02259.
Nicklas Hansen and Xiaolong Wang. Generalization in Reinforcement Learning by Soft Data Augmentation. In IEEE International Conference on Robotics and Automation, ICRA 2021, Xi' an. China, May 30 June 5, 2021, pages 13611-13611 IEEE, 2021. doi: 10. 10.1109/CRRA44556.20011 9561 103.
Abdolhossein Hoorfar and Mehdi Hassani. Inequalities on the Lambert W func- tion and hyperpower function. J. Inequal. Pure and Appl. Math, 9(2):5- 9, 2008 q=info: 1wd.:37/Crrww::scoollaa...
URL htapp///sccholar.ggoggeeeeeerroo wwwp:t/ww...oo.c.cooooo....aaa.. bib?
AFWwaeY Alteehahhhnaaaaaarnnnooooooorrer 1e11
12

---

Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Boehmer, and Shimon Whiteson. Transient Non-stationarity and Generalisation Deep Reinforcement Learning, In 9th Interna- tional Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL. wtt.hhttt://ppeeeeeeieeeee
Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò
and Roman Gamett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 8580-8589, 2018. URL. https: pttp:////cccccco....... www.ees..c.pappeeee.eeee htcppappe//00///////// 5cc00118866220aaaeccb10022222223 Seteeeeefaa--bsraaataaaee html.
Yiding Jiang. Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic Generalization Measures and Where to Find Them. In 8th International Conference on Learning Representations, ICLR 2020, Addis Abuba, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
Yiding Jiang. J. Zico Kolter, and Roberta Raileanu. On the Importance of Exploration for Generaliza- tion in Reinforcement Learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Or- leans, LA, USA, December 10 16, 2023, 2023. URL http:: nips. cc/paper_file Absttrattbbttattttt..omooo Conf erence. html.
paper/ 2w22/hash/aeee 2ttaa////2/aaa//3334444442222662
Sham M. Kakade and John Langford. Approximately Optimal Approximate Reinforcement Learning. In Claude Sammut and Achim G. Hoffmaann, editors, Machine Learning, Proceedings of the Nine- teenth International Conference (ICML 2002), University of New South Wales, Sydney, Australia, July 8-12, 2002, pages 267-274. Morgan Kaufmann, 2002.
Ken Kansky, Tom Silver, David A. Mély, Mohamed Eldawy, Miguel Lázaro-Grreiilla, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, D. Scott Phoenix, and Dileep George. Schema Networks: Zero-shot Transfer with a Generative Causal Model of Intuitive Physics. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 1809-1818. PMLR, 2017. URL http: /ttp:::////oceednggss mlr alr.press/ 770/kanskyy76 17a. html.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A Survey of Zero-shot Generalisation in Deep Reinforcement Learning. J. Artif. Intell. Res., 76:201-2000 2023, doi: (LL1613/AAIR..144177 URL https://doi. org/10. 10. 1613/jair. 1. 14174.
Henry Kvinge, Tegan Emerson, Grayson Jorgenson, Scott Vasquez. Tim Doster, and Jesse D. Lew. In What Ways Are Deep Neural Networks Invariant and How Should We Measure This? In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Infor- mation Processing Systems 2022, NeurIPS 2022, New Orleans. LA, USA, November 28- De- wwtppppr.ppppppceeeeee /tppppaperrraaaa wwwlle//ppprr/0/// Abst:aatbbbbattttt---a-a--- Confeeectttoooeeeeeeeeeee .ttll.
cember 9, 2022, 2022. URL http:://appprrree nips cc/ 3B66fffcb1447388866111222222
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 8570-8581, 2019. URL https: sstts://procee http:://///oco......./...a
wewpoedediinngnnurripipeeee meurips.e. wwturiss.ccppppeeeee htcc//aper/2019//aeeeeee /atp///sh/////////// 118.:/001111111111888448884 c84528bd4- 1ee044522bb--bbtrracctt html.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Abuba, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
URL https: /ttp::::///peeevieeee wwt.:eeeetttteeeeeeee nett/foruum/ccctggccr...
13

---

Chi-Heng Lin, Chiraag Kaushik, Eva L. Dyer, and Vidya Muthukumar. The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective. CoRR, abs/2210.05021, 2022. doi: CE5S500/RXXI.222000055 05021. arXiv: 2210.05021.
URL. https://doii /0/ 48550/arXiv http:///...........444444/5///./ .210
Kaixiang Lin, Shu Wang, and Jiayu Zhou. Collaborative Deep Reinforcement Learning. CoRR, abs/1702.05796, 2017. URL http:/ /tt///rai..orgbb///r/aaaeaaaea/e 05796. arXiv: 1702.05796.
Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the Benefits of Invariance in Neural Networks. CoRR, abs/2005.001 000707 2020. URL. https:
//arxiv. carxii.org/ao abs/ abs/2005.00178.
arXiv: 2005.00178.
Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning Dynamics and Generalization in Deep Reinforeement Learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022. 17-23 July 2022, Baltimore, Maryland. USA, volume 162 of Proceedings of Machine Learning Research, pages 14560-14581. PMLR, 2022. URL https: //proceedings. mlr pwwps/vv6ee lylel2al..... html.
Davide Maran, Alberto Maria Metelli, and Marcello Restelli. Tight Performance Guarantees of Imitator Policies with Continuous Actions. In Brian Williams, Yiling Chen, and Jennifer Neville, editors, Thirty-Sevents AAAI Conference on Artificial Intelligence, AAAI 2023, Thirty-Fifth Con- ference on Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth Symposium Educational Advances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, Febru- ary 7-14, 2023, pages 9073-9080. AAAI Press, 2023. doi: 10.1609/AAALV3378822660 btttss://dois org/ 10. 1609/ /aa00//aa..... v37.8.26089.
URL
Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The Generalization Gap in Offline Reinforcement Learning. In The Twelfth International Conference on Learning Rep- resentations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL
https://openrevier net/ ftt 3w6xxXDOYY.
Ning Miao, Tom Rainforth, Emile Mathieu, Yann Dubois, Yee Whye Teh, Adam Foster, and Hyunjik Kim. Learning Speecif Augmentations by Capturing Local Invariances. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 o 24736. PMLR, 2023. URL https http://ppoeeenn..... amm.eee.dinnnm............m..m.. mlr
Proceedings of Machine Learning Research, pages 24720- html.
htt.g/g........./e/////////aaaaa arress/v222/miao00
Oskar Nordenfors and Axel Flinth. Ensembles provably learn equivariance through data augmentation. CoRR, abs/2410.01452, 2024. doi: 10444550AARXX1.2240000555 URL htt.h:tp/////l org/10. arXiv: 2410.01452.
48550/ arXiv. 2410. 01452.1222
Emilio Parisotto, Lei Jimmy Ba, and Ruslan Salakhutdinov. Actor-Mimic: Deep Multitask and Transfer Reinforeement Learning. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http:// /ttp////xxvvv.../ org/ abs/ argg/bbs//11.... 1111111111......
Emmanuel Rachelson and Michail G. Lagoudakis. On the locality of action domination in sequential decision making. In International Symposium on Artificial Intelligence and Mathematics, ISAIM 2010, Fort Lauderdale, Florida, USA. January 6-8, 2010, 2010. URL http://gauss.... ececs. uc. edu/Workshops/ http://aaammnnmmppmrraeaaa
pdf.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dor- mann. Stable-Baselines3: Reliable Reinforcement Learning Implementations. Journal of Machine Learning Research, 22(268):1-8. 2021. URL //jmtt .rg wwtgg/appers/v2//222 1364. .ttll.
Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic Data Augmentation for Generalization in Reinforcement Learning. In Mare' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 5402-5415, 2021. URL https: httttttss:/pprceee http///pppoec....///ee////e/seee prrooceeingssneuriiss nwwrinuu...... wttp:ipp.ccppppe///0/// 20tp///2222hhh////// /hashh/ 2b38222ff644997777766611181
Abstract html.
14

---

Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning. In Geoffrey Gordon, David B. Dunson. and Miroslav Dudík, editors, Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, AISTAT'S 2011, For Lauderdale, USA, April 11-1,, 2011, volume of mlr.press..
JMLR Proceedings, pages 627-635. JMLR.org, 2011. htR.htttp//pprcceeeeee v15/ross 1ateessssl///a 11a/rossia.a.a... pdf.
Andrei A. Rusu, Sergio Gomez Colmenarejo, Çaglar Gülçehre, Guitlaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy Distillation. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http: / //tpxv/aaivvv/aaa///a//aa/.aaaa. org/ abs/ abs/1511. 15111.691........
Fereshteh Sadeghi and Sergey Levine. CAD2RL: Real Single-Image Flight Without a Single Real Image. In Nancy M. Amato, Siddhartha S. Srinivasa, Nora Ayanian, and Scott Kuindersma, editors, Robotics: Science and Systems XIII, Massachusetts Institute of Technology, Cambridge, Massachusetts, USA. July 12-16, 2017, 2017. doi: 10.15607/RSS.20111 XII.04.. URL http: html.
//www. http://cprroeeedding.rgrrgggaeae
Simon Schmitt, Jonathan J. Hudson, Augustin Zidek, Simon Osindero, Carl Doersch, Wojciech M. Czarnecki, Joel Z. Leibo, Heinrich Küttler, Andrew Zisserman, Karen Simonyan, and S. M. Ali Eslami. Kickstarting Deep Reinforcement Learning. CoRR, abs/1803.03335, 2018. URL http:
//arxiv. orggo/ggggo. abs/1803. 03835.
arXiv: 1803.03835.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. CoRR, abs/1707.0634 0634777,7 2017. URL 06347. arxiv: 1707.06344
Ruoqi Shen, Sébastien Bubeck, and Suriya Gunasekar. Data Augmentation as Feature Manipulation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 19773-19808. PMLR, 2022. URL https: pwwpeedings.....mo htt.:tttp///rrrcciigg..... mlr. wpw.pps/v1////eeennnene shen22a. html.
Connor Shorten and Taghi M. Khoshgoftaar. A survey on Image Data Augmentation for Deep Learning. J. Big Data, 6:60, 2019. doi: 10.1186$40337-019-019790 URL https://doi org/ 10. 1011166644037-0099 019-0197-0
Miguel Suau, Matthijs T. Spaan, and Frans A. Oliehoek. Bad Habits: Policy Confounding and Trajeo--raaectoryy Generalization in RL. Reinforcement Learning Journal, 4:1711-1732, 2024. URL https: cs. umass edu/ 2////22244///// http//pppps//aaaa//////// htm1..111...1....
Yee Whye Teh, Victor Bapst, Wojeiech M. Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4496-4506, 2017. URL Abstract...t....e....... html. hash/ Hash/Oabeee
wtwpiiinsnnuuurss.cooooooo /2/1p
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop, ITW 2015, Jerusalem, Israel, April 26 May 1, 2015, pages 1-5, IEEE, 2015. doi: 10. 1109/TT.2015.77333666
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IFFF//SJ International Conference on Intelligent Robots and Systems, IROS 2017, Vancouver, BC, Canada, September 24-28, 2017, pages 23-30. IEEE, 2017. doi: 10.1 0009/IROS.2017..2220
Aaron Walsman, Muru Zhang, Sanjiban Choudhury, Dieter Fox, and Ali Farhadi. Impossibly Good Experts and How to Follow Them. In The Eleventh International Conference on Learning Representutionss ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net 2023. URL. net www.:eeenneettttotnneeeeeeeee orromiiidssciA.ggooof
15

---

Xudong Wang, Long Lian, and Stella X. Yu. Unsupervised Visual Attention and Invari- ance for Reinforcement Learning. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pages 6677-6687. Computer Vi-
sion Foundation IEEE, 2021. doi: 10.1 (0.1109/CVPR44437.220 1.0611. URL. https:
//openaccess thecvf. com/ content//CPR2022///mm/ /Wang Unsupervised heteviiiiediiissaleeeeaaaaaaaale html.
Max Weltevrede, Caroline Horsch, Matthijs T. J. Spaan, and Wendelin Böhmer. Exploration Implies Data Augmentation: Reachability and Generalisation in Contextual MDPs, March 2025. URL
attt::/arxiv..e http://oggbbbbbbbbababb
03565. arXiv:2410.03563 [es].
Sen Wu, Hongyang R. Zhang, Gregory Valiant, and Christopher Ré. On the Generalization Effects of Linear Transformations in Data Augmentation. In Proceedings the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 10410-10420. PMLR, 2020. URL http: //proceedings. mlr press//119/mu20g.. html.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun, ACM, 64(3):107-115, 2021. doi: 10.1 145/3446776.
16

---

A Extended background
A.1 Policy distillation & behaviour cloning
In policy distillation, a knowledge transfer occurs by distilling a policy from a teacher network into a newly initialised student network. Depending on the objective of the knowledge transfer, the student network can be smaller, the same size, or bigger than the teacher network. Moreover, there are many different ways the policy can be distilled (Czarnecki et al., 2019), depending on the specific loss function used (Ghosh et al., 2018; Teh et al., 2017), whether the student can collect additional data during distillation (Lin et al., 2017; Parisotto et al., 2016; Ross al., 2011), 1), or has access to additional information like rewards or a teacher's ' value function (Czarnecki et al., 2019).
We consider student network with the same architecture and size the teacher that is distilled on fixed dataset (so without allowing additional interactions the student with the environment). This fixed dataset is collected after training, and usually consists of on-policy data collected by the teacher itself. this paper, we analyse a simplified setting where both the student and teacher policy are assumed to be deterministie and scalar: To S + R, T S R. A simple distillation loss in this setting is the mean squared error (MSE) between the output of the two policies:
(.0.DDTT)) = Σ (To(s) - =
.ED
where D = (s1,..., 8n) is the set of states we distil on.
More generally, distillation can be performed between deterministic, vector valued student and teacher policies Te S + Fd T : S + Rd with the loss
1(0, D, m3)
Σ (s))))) T (s)
(6)
For stochastic policies, it is more common to minimise the Kullback-Leibler (KL) divergence between the student and teacher policies (Arora et al., 2018), sometimes including an entropy regularisation term (Teh et al., 2017; Lyle al., 2022)
1
1(0 D, m3) n
Σ
DKL (T)(S)) + (m))))
where H-) denotes the entropy of the policy. An alternative approach for discrete, stochastic policies is to regress towards the logits or probabilities over actions from the teacher:
1(0, D, nT)
Σ
where T-Iss) now indicates the vector (of dimension |A)) of probabilities logits that policy produces in state s.
As mentioned above, usually the policy is distilled on-policy data collected by the teacher. However, in general, a policy can in principle be distilled on any distribution over states, since the targets produced by the teacher (i.e., (s(s) or TT(+ss) can be trained off- off-policy, independently of how the state s was reached, or which action was taken s during collection.
A.1.1 Behaviour cloning
We consider behaviour cloning (BC) as specific instance of policy distillation, where the student network only has access fixed dataset of the teacher's 's behaviour (state- acttonaction tuples) and not additional information like the teacher's policy, value function or environment rewards. The goal in behaviour cloning is learn to imitate the behaviour policy (e.g., the teacher) that collected the dataset. In this sense, it differs from the general distillation setting, in that the learning targets are always on- polico with respect to the policy (or the mixture of policies) that collected the data.
Just as for distillation, the BC loss can differ depending on whether the student policy is deterministic or stochastic. For deterministic policies, the loss usually the MSE between the student and the action observed in the dataset:
1
1(0, D3)
no (8:)
7
17

---

where D3 {(11 ....... (sn, an) is dataset of behaviour of size n. For stochastic policies, it is more common to use a logarithmic loss (Foster et al., 2024)
48, (B) = Σ In (a:ls))
(8)
Note that these losses mainly differ from the distillation losses in that the learning targets are the actions ag taken by the policy that collected the dataset, rather than the (potentially off- oifl-polcyy teacher policy (s:).
A.2
Group Symmetry
group is non-empty set together with binary operation that satisfies the following require- ments:
a.beC, be
(Closure)
(a b) c a- (b.c), Va,b,cEG
(Associativity)
Be G, e.a=a a VaEG Va G, G, a-a = a a=
(Identity)
(Inverse)
We will abuse notation slightly by denoting both the group and the non-empty set with G, depending on context.
We can define group representation ex acting X, map v G - GL(X) from G to the general linear group GL(X) of a vector space X, where the general linear group defined as the set of x invertible matrices (for finite dimensional vector space X with dimension n) with matrix multiplication as operator and where the map & is a group homomorphism, i.c (b) (a b), Va, b E G. With these definitions, invariance of function is defined as
[(bx(1))) f(x) Ya E E
useful property when performing full data augmentation with a group G. is that applying a transformation from G to any of the training samples in Te, is equivalent applying permutation Py the augmented training dataset indices:
vx (g)x; = (p)
where E (1, [Ta]]
A.3 The infinite width limit
In the limit of infinite layer width, ensemble of neural networks from random initialization follows a Gaussian process that is characterised by the neural tangent kernel (NTK Jacot et al., 2018) defined
O x') = Σ Baww ( 3000
(ofole)
T ( (otele")
20
where we assumed the network fo has L layers and (()) denotes the parameters at layers 1 E [1, L.
respectively. The Gaussian process time has mean m, and covariance E, (Lee et al., 2019):
m²(x) -x,,)))(x,,)))) 2:) 'Tiles
Ex x') = K[2,, x') + (1)1 (2, (x,
+ h.c.)
where we use the Einstein notation convention to indicate implicit sums over the dataset indices i, h.c. indicates the Hermitian conjugate of the preceding term, T, = exp(-n/Ot)), K(#,, x²) = Eg (x)(=) (x) (ffo(()((())))))) fo())) is the neural network Gaussian process (NNGP) kernel, and 2 and x are defined as follows:
(2, O(2, 0(s,,)))0-'',,,,0
(x, 0(z,,;)(0-TTi, K(,,,,
18

---

We use shorthand notation Ex(r,x) (x) for the NNGP variance.
An infinite ensemble fe equals the mean my of the Gaussian process: [x) m4(x). Note that for t + 0o, the output of the infinite ensemble for on the training inputs X' converges to the targets )
(2) = mao (x) = @ x) = y
19

---

B Proof of Theorem 1
In this section, we will go through the steps for the proof of the main theorem of section 4. We first repeat the definition of the GTI-ZSPT and associated discrepency measure K
Definition 1 (Generalisation through invariance ZSPT). Let Mlc be a CMDP and let Civains Crest C C' be set of training and testing contexts that define a ZSPT problem. Additionally, let n² be the optimal policy in Mlc. Sii = {s E Sl (s) > 0} denote the set of states with non-zero support under the on-policy distribution PMMle in CMDP Mlc. In the generalisation through invariance ZSPT (GTI-ZSPT), the sets M and S Micraai admit symmetric structure:
SMile = (
S Milcirain
= (bb((b))bb B, BEG
where S ST Gyrain
is proper subset of SM, Civain
and G a non-trivial symmetry group (and Cirain
B S a finite subgroup) that leaves the optimal policy invariant: n T²()) T (bs(g)s), S.
Definition 2. For the group G and its finite subgroup B < G that define the symmetric structure of GTI-ZSPT (Definition 1), K is measure of discrepancy between the representations of these groups:
max min |1vb +bs(b)) lop gEG bEB
where ||
!lap denotes the operator norm.
B.1 Invariance of an ensemble
In order to prove Theorem 1, we first repeat Lemma 6.2 from Gerken and Kessel (2024) that bounds the invariance of an infinitely large ensemble of infinitely wide neural networks trained with full data augmentation on some finite subgroup B S G:
Lemma 6.2. Let : S + IK be an infinitely wide neural network with parameters 0 and with Lipschitz continuous derivatives with respect to the parameters. Furthermore, let ūe be an infinite ensemble (s) = Eg (T c.o(s)), where the initial weights 0 are sampled from distribution and the operator L, maps 0 its corresponding value after steps of gradient descent with respect MSE loss function. Define the error as measure of discrepancy between representations from the group G and its finite subgroup B:
max min (1))(() v's
(10)
gEG beB
The prediction of an infinite ensemble trained with full data augmentation on B S G deviates from invariance
(T) - ne (x)(()))) s Ce(s).
4g E G
(11)
for any time t. Here s E can by any state and Ce is independent of g.
Proof. For completeness repeat the proof in own notation in F.2.
Next, we prove a lemma that bounds the prediction error between finite ensemble and an infinite ensemble:
Lemma 1. The difference between the infinite ensemble Re and its finite Monte Carlo estimate ी of size N, is bounded by
[T)( - #e(s)) VN (x,
(12)
with probability least 11-c. c. Here Σ, the variance of the NNGP at time and Cr, (8, depends on Σg, the state s and confidence level
Proof. We start with Lemma B.4 from Gerken and Kessel (2024) that holds for any Monte-Carlo estimator:
20

---

Lemma B.4. The probability that the deep ensemble गe and its Monte-Caue estimate कe differ by more than a given threshold 6 is bounded by
120 exp
82
[[| (s)) 6|
S
π
20²
where we have defined
a : (ar(tt) Es(s)
where (x) is the NNGP variance and N is the finite ensemble size.
Proof. For completeness we repeat the proof in our own notation in E.3.
We this lemma bound the probability of the deviation between finite and infinite ensemble to be smaller than a threshold 6:
PP ([( (x)) S = P[(
#.(+) > 8]
2 -
where =
= / exp ( 207) Next, we rewrite ő in terms of a given confidence level
2
8²
exp
8
20
2
8
8²
exp
mT2
0
8²
Wo
o²
2
8 0o Wo TT
where Wo is the principal branch of the Lambert W function and the second to last step holds because 52
o3 2 F and 2 > 0 for given probability If we know the value for e, Wo + can be solved for numerically. However, in general, the principal branch of the Lambert W function has no closed-form solution, but was upper bounded by Hoorfar and Hassani (2008)
Wo(x) s In (1+ In(a + +1)
2x +1 I+
for x 2 -//.. Which means we can upper bound 8 with:
(s(s)
<
In
(13)
N
TH + + 2 + πC2) ) + NT In(
Cr. Cr, (3,6)
(14)
We can now prove an intermediate lemma that bounds the deviation from the optimal policy for finite ensemble (rather than an infinite one, as in Lemma 6.2)
Lemma 2. Let the student policy o. be an ensemble of infinitely wide neural networks Te with Lipschitz continuous derivatives with respect its parameters. distilled on-policy dataset
D 8 consisting of all the states encountered by the teacher in M McCaaa
Ml
the teacher is optimal in the training lasks Ctrain (but arbitrarily bad anywhere else), the deviation is bounded with probability at least e, by:
from the optimal policy for any test state s' M Cicst
1x²()') - #co (8') S RCo(s')
VN CO
21

---

where R. is the measure of discrepancy between subgroup B S G and full group G (see definition 2) and Ce Or depend on the state s' E Mi the NTK @ (i.e. network architecture), and the dataset D. Additionally, CT. also depends on the NNGP kernel K. (i.e. network initialisation) and the confidence level
Proof. Because we assume the teacher optimal in the training tasks, our training dataset is actually D
ST MlCerain S Mla Furthermore, by definition of the GTI-ZSPT setting, we have for the states encountered by the optimal policy in MIc: SMc (15(g))lgg E G,8 5}. Furthermore, we have for the states encountered by the optimal policy in MlCras?i Mlc: S MMl = (b)s( E B, E 5) for B & G. This means that for any state 8 E Mlo' there
exists symmetry transformation g E G from s to state & E S C DD in the training dataset that leaves the policy invariant:
Vs E SMa
=g- E G, s.t.
"s( 3
= (3),
for some 8 E D (15)
Since this holds for any state Sii it also holds for any state in S
C CSMI
Grest
Now, Lemma 6.2 holds for any state E S and any g G. So, if we choose 5 and g such that ( s' for testing state s' E S MMlli
we have
me 8
tee (v(ss))))) S K Ce(
(5) Te (5() S x Ca(s')
where we write that Ce is now function of s' E S SMlcca instead of 8 E D, which can do
because there exists a one-to-one mapping between the two: 8 " The above bound holds for any time I. If we choose + oo, we have that the infinite ensemble of infinitely wide neural networks म. trained on D, will converge 7oo (5) (5) (5), V3 D. Furthermore, due to our choice of g E G, we have that To (5) = x (3) x (s'), and the bound becomes
x (x'( aoz (s') S K Co(s'), Vs' E S
This bounds the output of the infinite ensemble after training Toc, evaluated in a testing state E Micc.., to the optimal policy in that state. We can now combine this bound with our Lemma I above to bind the policy of finite ensemble the optimal policy in any testing state:
(s") (s))l = | (s') - मे (s') + (s') - Ato (s')l
)-))
< (s') - #oo (s')) + (s') #oc (s')) < Ce(s') + (s') - 5xx (5))
s Ca(s)) + 1
Cza (s'.c) with probability 21-c, As' E S
VN
Mc
B.2 Performance during testing
We can now use Theorem 3 from Maran et al. (2023) to prove performance bound for our student policy o (s') in the testing CMDP M Ccnat in terms of the Wasserstein distance between the student and optimal policy in this testing CMDP:
Theorem 3. Let xt" be the optimal policy and it be the student policy. If the CMDP (Lr, Ln)- Lipschitz continuous and the optimal and student policies are L-piscchhtzz continuous, and we have that LL + Ltx) < 1, then holds that:
Ln
F - jax < (1-7)(1 LL Lx ))
Benda (-V(= (1os),
where d= (5) (1-= Eie P((s = s/m", .po) ) the yddiccuuntdd visitation distribution and the the facttonffactoo
22

---

Proof. For completeness we repeat the proof in our notation in Appendix E.1.
With this, we can finally prove the main theorem:
Theorem 1. Consider policy distillation for a deterministic, scalar teacher policy Rg : - R (Equation (1) Section 2.1) in Lr, LR-Lipschitz continuous CMDP in the GT1-ZSPT setting. Let the student policy & be an ensemble of N infinitely wide neural networks ne S - R with Lipschitz continuous derivatives with respect to its parameters, distilled on an on-policy dataset {b)(b E B,s 5} consisting of all the states in the training contexts encountered by the teacher in the GT1-ZSPT setting. Furthermore, let the student policy be La=- Lipschitz continuous and assume -I- I=a_. <1
D S
Ceraim
If the teacher is optimal in the training tasks Cirain (but arbitrarily bad anywhere else), the per- formance of the student in the testing CMDP MlCuus is bounded with probability at least 1 by:
(1-7)(1- - -Ir(1 + L#w)) 1) (CCo ( / (r.0(())) - 11
LR
1 VN
(S)
where K. is the measure of discrepancy between subgroup B S G and full group G (see definition 2) and Ce. CT.,. are constants that depend on the d-dicoounted visitation distribution of the optimal policy in M Cr..s the network architecture, and the dataset 'D Additionally, Cr.,. also depends on the network initialisation and the confidence level e.
Proof. We have for deterministic polices that the Wasserstein distance reduces
W-r"(+ss), oo (+|s) | (s) (s)
So, if we invoke Theorem from Maran et al. (2023) on the testing CMDP M Cruse we can use Lemma 2 show
LR (1 (1-7) (1-7)(1 -LLT(1 + Lxxz))
5 J R s
Bender (M(( (-(s= ao (-[s))
LR
5
B Bende ||= (8) (x)
(1 7)(1 -IL((1 Litww)
(1 7) - yL11111rr11110 + Lxxe)) (NOE+ +
LR
1 1 C (0)
VN
where Cg -Bangr (s) depends the NTK O (i.e. network architecture) and Cs (c)
Exnde' [Oz (8,e) depends on the NNGP kerne K. (i.e. network initialisation).
23

---

C
Experimental details
The code for all the experiments in the main text can be found at https://github.com htw.ee//eiisiilllaattttttttterrr
C.1 Reaahee with rotational symmetry' CMDP
In the RReacher with rotational symmetry' CMDP from Figure I, the state s= (I's, Us. Te. Be, 2/. Vh consists of the 2D Euclidean coordinates the shoulder (Ie: Vs), elbow (æes ye) and hand (2), yo. centred around the target location, and the continuous 2D action space consists of the torque to rotate the shoulder and elbow joints. The episode terminates and the agent receives a reward of the robot arm is within small area around the target location. Elsewhere, the reward function equals 1-0,5d
if the hand
.0.5T Ödgarget <dmin where dtarget is the distance between the target location and hand, T = 200 is the maximum number of steps before timeout, and ddtarge=<ddmnn is only when the current dearget smaller than the minimal distance dmii target achieved that episode, and otherwise. In the experiments from Section 5.1, policies are distilled on datasets collected by rolling out trajectories from a teacher agent fixed set of training contexts. Ensembles are created by independently distilling N policies (with different seeds) and afterwards evaluating by averaging over the output of the polices.
C.1.1 Satisfying the assumptions for the GTI-ZSPT setting
The 'Reacher with rotational symmetry' CMDP from Figure satisfies the symmetric structure assumed in the GTI-ZSPT setting. To illustrate this, we could define the states encountered by the optimal policy in context 1 in Figure as the subset of states S in the GTI-ZSPT definition. With that definition we can see that subgroup B = C4 would generate all the states in S MlC group G' = S0(2) would generate all the states in SMile*
and full
For the states (Ie: Mai Tes Ue: The vh), the representation (a) for rotation with angle is the block diagonal matrix:
cos
sin
0
0
0
0
sin a
COS Ck
0
(a)
0
COS
sin
0
0
0
sin a
cos or
0
0
cOs
sin
0
0
sin o
COS aX
which is orthogonal.
Additionally, the CMDP is ---Lisschitz continuous since there collisions causing non-smooth transitions. Moreover, with a proper choice of reward function (for example, R deanger ), it is also
LR-Lipschitz continuous. Note that for our experiments, we choose non smooth reward function since it helped with training the teacher.
C.1.2 Figure & Table
In the setting from Figure I and Table 1, the contexts only differ in the location of the shoulder. The robot arm pose always starts at a 45" degree angle for the shoulder joint (counter-clockwise with respect to an axis drawn from the shoulder to the target), and 90 degree angle for the elbow joint (clockwise with respect to an axis drawn from the shoulder to elbow). In the testing distribution, the shoulder can be located anywherr along circle around the target location, but the starting pose is always the same.
For the training contexts, the shoulder is located at different, evenly spaced, intervals around the 360 circle. the bottom half of Table we train on three different data sets denoted by the corresponding subgroups of SO(2) that generate the training contexts: C2, C2 and Cg. For each of the datasets, we always context Figure as the base context, and apply various rotations to generate the other training contexts. The C2 set consists of context together with the subgroup of 180" rotations (the 0 and 880 rotations, resulting in context and context Figure 1). The set Ca consists of the 90" rotations and the corresponding four training contexts depicted in Figure 1. Lastly, the Cs set

---

consists of the 45 rotations, half of which are the 90" rotations from Figure 1, and the other half are the 45 rotations in between those. Note that for the results of varying ensemble size the top half of Table 1, we used the C4 dataset.
The teacher policy is handcrafted policy (a (-2, for 12 steps and a (2, 2) afterwards) that is optimal for the starting pose considered in this setting. The neural network consists of three fully connected hidden layers of size [64, 64, 32) with ReLU activation functions. The policy is distilled with the MSE loss in (6) (which is the same loss as (1) but for vector-valued actions instead of scalar). The exact hyperparameters can be found in Table 4.
Table Hyper-paramelers used for the 'Reacher with rotational symmetry' CMDP experiments 'Reacher with rotational symmetry'
Hyper-parameter
Value
Epochs
500
Batch size
6
Learning rate
1 x 10 10-4
C.1.3 Figure Table
In the setting from Figure 2 and Table 2, the contexts not only differ in the shoulder location (as described in the subsection above), but also in starting pose of the robot arm. For testing, random shoulder location and starting pose sampled for each episode.
There are four training contexts in this setting, whose shoulder location correspond the C4 dataset described above, but whose initial arm poses are sampled randomly by sampling two angles between and 360 degrees for the shoulder and elbow joint. Each seed has its own set of random training poses (but the same shoulder location). The dataset created from these four training contexts is referred as the Training Contexts dataset in Table 2. The Training Contexts C2 dataset essentially consists of 16 training contexts, four of which are the ones from Training Contexts, and the other 12 are the random poses from the Training Contexts contexts, duplicated for each of the other shoulder locations. Figure 3 illustrates this for an example set of four Training Contexts. The Training Contexts + Random dataset the same as the Training Contexts C4 set, except that instead duplicating the random poses from the four Training Contexts, 12 new random poses are sampled.
For this more complicated version of the 'Reacher with rotational symmetry' CMDP, is much more convoluted to handcraft an optimal policy. Instead, we train an soft actor-critic agent (SAC Haarojja al., 2018) with the distribution, to get close to an optimal policy for any context. The network for the SAC teacher consists of two fully connected hidden layers of size [400, 300) with ReLU activation functions. The other hyperparameters for the SAC agent can be found in S. Note that for the results in Table 2, we evaluate single distilled policies (N = 1) and we use the same network and hyperparameters for distillation the experiments for Table 1.
(Raffin et al., 2021) implementation on the full context
25

---

Table Hyper-parameters used for the SAC teacher agent in the 'Reacher with rotational symmetry'
CMDP
SAC Teacher
Hyper-parameter
Value
Total timesteps
500 000
Buffer size
300 000
Batch size
256
Discount factor
0.99
Gradient steps
64
Train frequency (steps)
64
Target update interval (steps)
Target soft update coefficient
0.02
Warmup phase
10 000
Share feature extractor
False
Target entropy
auto
Entropy coeff
auto
Use State Dependent Exploration (gSDE)
True
Adam
Learning rate
5 10-4
Training
Training Contexts
Training Contexts C4
7
Figure 3: The Training Contexts and Training Contexts + C4 context sets in the 'Reacher with rotational symmetry' reacher CMDP with varying shoulder location (red) and robot arm pose (blue), see Figure for details.
C.2 Four Rooms
In the Four Rooms environment (Figure 4), an agent (red triangle) starts in a random location and facing a random direction, and has to move to the goal location (green square) whilst navigating the doorways connecting the four rooms. We modify original Minigrid implementation little bit, by reducing the action space from the default seven (turn left, turn right, move forward, pick up an object,
26

---

drop an object, toggle/activate an object, end episode) to only the first three (turn left, turn right, move forward). Moreover, we use a reward function that gives reward of when the goal reached step count
and zero elsewhere, which differs from slightly from the default one that gives - 0.9 max steps for reaching the goal. Lastly, our implementation of the Four Rooms environment allows for more control over context dimensions, allowing for the construction of distinct training and testing sets. Our version of the Four Room environment can be found at <redacted for review>.
Training
Testing
Figure 4: Example of Four Rooms training and testing contexts.
In this environment, the contexts differ in the topology of the doorways connecting the four rooms, the initial location of the agent, the initial direction the agent is facing, and the goal location. The teacher is the PPO+Explore- Go agent from Weltevrede et al. (2025), that is trained on 200 training contexts. We used separate validation and testing sets consisting of unseen contexts of size 40 and 200 respectively. The validation set was used for algorithm development and hyperparameter tuning, and the test set was only used as a final evaluation (and is reported Table 3). Since the optimal policy in Four Rooms is deterministic, we also evaluate performance of our policies deterministically by always taking the action with maximum probability in a given state.
The Teacher dataset used for distillation and behaviour cloning, simply consists of the states en- countered by rolling out the stochastic teacher policy in the 200 training contexts until the desired dataset size was reached. We use a dataset size of 500.000, since this was the replay buffer size of the DQN agent from Weltevrede would look like for the PPO+Explore-Gr teacher. It is created by running pure exploration agent (trained as part of the PPO+Explore-Gr teacher) for steps at the beginning of each episode, and afterwards rolling out the stochastic teacher policy until termination of the episode. The number of pure exploration steps is sampled uniformly from range [0. K), where K used in Weltevrede et al. (2025)). The pure exploration experience is not added to the dataset, only the states encountered by the teacher. The Explore-Go dataset also has size 500.000, but requires additional interactions with the environment (on average 2 K 25 steps per episode) that are not added to the dataset. Since the average episode length of the teacher is of similar size, the Explore-Go dataset requires roughly twice the dataset size of additional environment steps create. Lastly, the Teacher dataset size 250.000, and a dataset created by rolling
(2025). The Explore-Gr dataset mimics what rollout buffer
(the same as was
mixed dataset a 50/50 mixture out the pure exploration policy for 250.000 steps. We generate 20 PPO+Explore- teachers, and generate one dataset of each type per teacher (for a total of 20 datasets).
An important thing to note, is that although the states for the distillation and behaviour cloning experiments are same, the learning targets are not. This has the biggest effect on the Mixed dataset, where the learning targets for behaviour cloning are the actions that were taken to create the dataset, and for distillation, the targets are the PPO+Explore-OO teacher's probabilities the given state (independent of what action taken in that state during the creation of the dataset).
27

---

The Four Room experiments were executed on a computer with an NVIDIA RTX 3070 GPU, Intel Core 17 12700 CPU and 32 GB of memory. Training the teacher (PPO agent) would take approximately hours, and a single distillation run would take approximately 10 minutes. The code for our experiments can be found at <redacted review>.
C.2.1 Implementation details
For the distillation experiments (top of Table 3), we used the same architecture as the teacher in Weltevrede et al. (2025). We distil the stochastic teacher policy by regressing on the probabilities (as in Equation 7). We found this to work significantly better than alternative loss functions, since normalised range for the targets helps the averaging in the ensemble. We tune the distillation hyperparameters by performing grid search over the following values
Learning rate: {1 x 10-4, 1 10-3,
10-
Batch Size: {64, 256, 512, 1024, 2048)
Epochs: {10, 20, 30, 50, 60, 70, 80, 90, 100}
We performed the tuning for single teacher seed by, for each dataset type, splitting the dataset into a training set (sampled from the first 150 training contexts) and validation set (sampled from the other 50 training contexts), and choosing the combination of hyperparameters that minimised the distillation loss on the validation set. The final results are distilled on the full datasets and evaluated in the testing contexts. The final hyperparameters can be found in Table 6.
Table 6: Hyperparameters used for policy distillation in the Four Rooms environment.
Four Rooms Distillation
Hyper-parameter
Value
Teacher
Epochs
100
Batch size
64
Learning rate
x 10
Explore-Go
Epochs
50
Batch size
512
Learning rate
x 10
Mixed
Epochs
50
Batch size
256
Learning rate
x 10
For the behaviour cloning experiments (bottom of Table 3), we used the same architecture as for the distillation experiments. We trained using the logarithmic BC loss for stochastic policies (Equation (8)). We also tuned the behaviour cloning in the same way as the distillation policies above, but with a smaller range for the epochs: 21,2,3, 33,4,,5 4,5, 5, 6, 7, 8, 9, 10). since we found would overfit much sooner than the distillation experiments. The final hyperparameters can be found in Table 7.
28

---

Table 7: Hyperparameters used for behaviour cloning the Four Rooms environment.
Four Rooms Behaviour Cloning
Hyper-parameter
Value
Teacher
Epochs
1
Batch size
64
Learning rate
1 10-3
Explore-Go
Epochs
1
Batch size
64
Learning rate
1 10-3
Mixed
Epochs
2
Batch size
256
Learning rate
1 10-3
29

---

D Additional experiments
D.1 Measure of invariance
order to identify invariance as key factor for the increased performance the 'Reacher with rotational symmetry' environment from Table I, we will measure how invariant the policy becomes when increasing ensemble or subgroup size. a measure invariance we use variance of the network's output across the group orbit (K (Kvinge et al., 2022) (for completely invariant network, this should be zero). In practical terms, we evaluate each policy on all (360) integer degree rotations of the starting state (the orbit) and compute the total variation (trace of the covariance matrix) over the produced outputs. Figure 5, we see that for both increased ensemble and subgroup size, as the test performance increases, does the measure of invariance decrease.
1.1
Total Variation Over Orbits
1.0
Total Variation Over Orbits
Measure
10
Test Return
Test Return
0.9
1003
0.8-
-
Performance
100
*
Measure Invariance
10
10F
Ensemble Size (N)
Subgroup (B SO(2))
(b)
Figure 5: Test return (left axis) compared with the total variation (trace of the covariance matrix) over orbits of the SO(2) group of rotations (right axis) for (a) different ensemble sizes and (b) subgroups B 3 SO(2). The total variation a measure how invariant the agent has become with respect to rotations, zero total variation would correspond perfect invariance. Shown are the mean and 95% confidence intervals over 20 seeds.
D.2
V generalisation bound
vN form the generalisation bound from Theory 1. We evaluate different ensemble sizes on the 'Reacher with rotational symmetry' environment from Figure 1, trained on subgroup B C2 Table This table includes the results from Table 1, as well as additional results for ensemble sizes N = 1000 and N = 10.000. In Table 9, we repeat the results for different ensemble sizes distilled on the Explore-Go dataset the Four Rooms environment from Table 3, with the addition of ensemble size N 100.
In this section, we investigate how well our results the + 4
In Figure 6, compare the difference to optimal performance - with ensemble size N. and plot the best fit to the V. relation predicted by theory. The exact optimal performance JF not necessarily known, but we estimate by taking the average train performance instead (which seems have converged both environments). The best h fit is obtained by computing the /
optimal linear fit between _== J and x' NN (and then transforming the solution back N).
The results in Figure seem to follow the a + upper bound to some extend. The upper bound is likely not very tight, but the results seem to indicate the bound is also not vacuous.
30

---

0.5
Best Fit
0.150
Best: Fit
0.4
Data
0.125
Data
0.3
0.100
0.2
0.075
0.1-
0.030
0.025-
10a
10
10a
10"
10
10F
10
Ensemble Size (N)
Ensemble Size (N)
(a) Reacher with rotational symmetry
(b) Four Rooms
Figure 6: Difference optimal performance as a function of ensemble size in the testing contexts for the (a) Reacher with rotational symmetry and (b) Four Rooms environments. Shown are the mean and 95% confidence intervals over 20 seeds from (a) Table and (b) Table 9, together with the best
possible fit for the relation + VN as predicted by our theory.
Table 8: Performance of distilled policies in the Illustrative CMDP from Figure for different ensemble sizes N (trained under subgroup = Ca). Shown are the mean and standard deviation for 20 seeds, and in bold are the best returns including those with overlapping 95% confidence intervals.
Ensemble Size N:
N=1
N=10
N=100
N=1000
N=10.000
Train Performance
1.17 + 0.004
1.17 + 0.004
1.17 +: 0.003
1.17 + 0.003
1.17 0.002
Test Performance
0.75 0.147
0.89
0.107
1.05 = 0.117
1.13 0.072
1.17 0.038
Table 9: Performance of an ensemble (of size N) of policy distillation behaviour cloning policies on Explore-Go dataset in the Four Rooms environment. Shown are mean and standard deviation over 20 seeds, and in bold are the best returns including those with overlapping 95% confidence intervals.
Ensemble Size N: N=1
N=10
N=100
Train Performance
0.92 0.020
0.92 + 0.019
0.92
0.021
Test Performance
0.78 + 0.041
0.88 + 0.036
0.89 + 0.033
31

---

E Repeated proofs
E.1 Theorem 3 from Maran et al. (2023)
Here we repeat the proof for Theorem in Maran al. (2023). We mostly change the notation to be consistent with our paper. Note that in Maran et al. (2023) they consider MDPs rather than CMDPs, but since any CMDP (including the testing CMDP Mlc.... in a ZSPT setting) is just special instance of an MDP, the theorem readily applies.
We first introduce number of definitions and assumptions used in our derivations
Definition 3. We will call a function f L-Lipschitz continuous for two metric sets (X, dz), (Y, dy). where das dy are distance metrics, we have
V2 X, (z((() & (21
Definition 4. We define [| to the Lipschitz semi-norm of J, with
d.((()))
llfllz
21 22),
Definition 5. We introduce the Wasserstein distance between probability distributions p and q as
/ [dp
W(r.)) = sup II 1251
fdq
We assume an (Lr, Lii-hiict continuous MDP with a metric state and action space and associated distances d=(= ds) and da d,), for which we have
Assumption 1.
W(T(-+s, T(188, a)) s (ds(((s, 8) da(a, a)), R(s,a) R(s, a)l LR(ds(s, 8) da(a,ā));
Vss (8. E ((sa)), (8, E S A.
Assumption 2. We assume L, Lipschitt continuuus policies, which satisfy
(1(((1s),((((())) L,dd(ssss) S.
Note that this definition subsumes deterministic policies.
With this, we can state the theorem in question (Maran et al., 2023)
Theorem 3. Let x² be the optimal policy and o be the student policy. If the CMDP is (Lr. LR)- Lipschitz continuous and the optimal and student policies are LipLisschitz continuous, and we have that LT (1 + Ltx)) < 1, then holds that:
LR
J J
Bande (M(=" (())
(1 (1- TTr(1 + l(x)
where d= (s)= (1- Eieo of P = = Po) the y-discounted visitation distribution and ?
the the discount factor.
Proof. From assumptions and 2, it follows that if 7I, < 1 the value function Q" associated with i LQ- -ipschitt continuous (Rachelson and Lagoudakis, 2010) with
LR
Ler <
1- -I- L.))
The first part of our proof derives performance differences under Lipschitz value functions. We begin with the performance difference Theorem (Kakade and Langford, 2002) stating
J
Eandr. (TTT) T (s)]
32

---

where V is the state value function. Focusing on the inner expectation we have
/ (T (s. - V (s))
E (1 (s)
/ Q* (s) - (*)
A
/ (s, (dals) - m2(dals)).
Now, let L. = be the Lipschitz semi-norm of Q=a w.r.t and define ga(a)
(s, Lo with the property I1alle = 1. This yields
(+)) = / Q (s.a - =2(da(s))
/ 9.(a)Lo T2(da) (dda)s))
L, / (o(a)(m T2(dals))
By definition of the Wasserstein distance
/ ((a)(s) - ( 1alll (((()
sup
such that we have
|Ba (s))
L. g (((()a)((dass)) - =2(dals))
A
/ g(a)(T)(dals) - (d(ass))
S L. sup
L. (2( (11)).
Now, we recall L., 1sQ**(s, and by our assumptions Q is - Lipschitz that
continuuuuu continuous such
L, $ sup (s. S
< 1Q ||ا
< LQ
Putting these results together, we can obtain
LQ Es-aro
J s
After setting - and n2 = t and using that Liqea S 1-
LR
we have
In
Jx - A (1-) - 7L7(1 +
Band Band" (W(x (-o (1)
Lzw))) +
E.2
Lemma 6.2 from Gerken and Kessel (2024)
Here we repeat the proof for Lemma 6.2 in Gerken and Kessel (2024) the notation used this paper.
Lemma 6.2. Let Te : S + TE an infinitely wide neural network with parameters 0 and with Lipschitz continuous derivatives with respect to the parameters. Furthermore, let 7e be an infinite ensemble ihe (s) Eg (iT (x),,)), where the initial weights 0 are sampled from distribution p and the operator L, maps 8 its corresponding value after t steps of gradient descent with respect
33

---

MSE loss function. Define the error as measure of discrepancy between representations from the group G and its finite subgroup B:
K
max min ll (g) - (b)( lop
(10)
gEG B
The prediction of an infinite ensemble trained with full data augmentation on B S G deviates from invariance by
Tee (s) - te (1))())) S Ce(s), Vg E G
(11)
for any time L. Here 8 E can by any state and Ce is independent of g.
Proof. Lets denote a set of states with D = [st] and a training dataset
{(st, ys [Vs, E
D, Yi E ) where Yi E 3/ indicates the target for sample s; (for example, y/- (T()s) ) for distillation with respect to a teacher πT). Using the definition of the measure of discrepancy is and the property of full data augmentation with a finite group from (9), we can write for any training sample (*,, Ms)
TB = (Dp,Vn) = {(bs (b)s, v) v) T, b E B} and any g E G and E B:
[g)5( - ())l = [g)) - + (6)s & ([)s(g) - 4hs (b)llpp))ss)) < x||sy|| Additionally, we can use the definition of the mean of the NNGP my write for any 8 $ (s) - Fe (g)) = (m,(s) - m. ( (v> (g)=))
= [(((s, DB) ((bs(g)), Da)) 'I- ex)((
We use Lemma 5.2 from Gerken and Kessel (2024), made specific for scalar- and vector-valued functions:
Lemma 5.2. For scalar- and vector-valued functions, functioon data augmentation implies that the permutation 11 commutes with any matrix-valued analytical function F involving the NNGP kernel K., the NTK θ and their inverses:
funntiicaauunction
= where lg(g) denotes the permutation matrix applying the permutation Pg associated with to each training point.
to show that:
O (I) (exp--ex(((00]]1 O(,, 0 (I (00T-cxppnn0t)_eee
(s, 0,, (I exp((n)@t))
= =ss,, (4) 10, (I - exp(-
where we also used invariance of the labels: TTb))a = yB. Plugging this into the expression above: (s) T. (1)())))
(00ss, (p)(()) (g),, 0))) (T -xx(
[(O(s, (0) O 1/ (0)s (II exp( - 54.8
Now, we bound the following:
A 3) = 10 (3) -O(s'
1
(Orelo)
@mT(s)
(m)
)
ag(1)
a0())
1=1
O (0(1
< ||s 8|
Baww
= ||a =((((s
where L((()) is the Lipschitz constant of dges Ty. Finally, using the triangle inequality:
(T())) - म. (vs(g)s)] s C(+),
1|sp (6) - #s(g)s|1
exp((n/Ot))saun)
Σ ||se|12
Σ
1 (T -eppp--p----pp (3p(-100)))su)11 = =C)(s)
/
j.k
34

---

E.3 Lemma B.4 from Gerken and Kessel (2024)
Here we repeat the proof for Lemma B.4 in Gerken and Kessel (2024) in the notation used in this paper.
Lemma B.4. The probability that the deep ensemble ne and its Monte-Carle estimate tn differ by more than a given threshold ó is bounded by
20
8²
[[|ma(s) ( 8] S
exp
π
20
where we have defined
Ex(s)
0 = (tr(
N
where (s) is the NNGP variance and N is the finite ensemble size.
Proof. In the infinite width limit, the ensemble members for our Monte-Carlo estimator 7 are i.i.d. random variables drawn from a Gaussian distribution with mean म; and variance Σ. Therefore, the probability of deviation a given threshold given by
x22 20
IP ((() (s)) 8)
exp
dr
V2TO,
where of [x(s) N
the variance of the Monte-Carlo estimator ät. With change of integration 2n for u 2 min (u), we get:
variable 2.v2
and using the fact that 1 v
2min(x]
[[] #c(s)| 8)
2 VT
exp
(-
u²
du 5
v20,,
00
(2u) exp (22² du
π
Vina
Vioa
Finally, this integral evaluates to
2 o
P (T)) - +c())]
> 8 S
exp
20
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims supported by the theoretical and experimental contributions in sections 4 and 5.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: The limitations of the theoretical assumptions are investigated and discussed
in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well- vell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The assumptions are clearly listed in Section 4 and the proof can be found in Appendix B.
Guidelines:
36

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
The experimental details can be found in Appendix C including a link to the
code.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived
well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky some cases, in which case authors are welcome describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code can be found at https: http::////tttttt.too cwn wow.e/bb..c.///eeeeereeeeeee distillation toto---ttt--trrnninnn
Guidelines:
37

---

The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines fattps://nips.. cc/ Sewdddee/oooeuumiiieseeeeeeeeeee
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results, See the NeurIPS code and data submission guidelines (https:
/ /nips. wttppccccpccaceeccccccccccccccct hubblii/guiddee/dddeubbiiiii /tteee///ooeeuuiiiiiiiiooooooo sminiioiiopyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental details can be found in Appendix C
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For all results, standard deviation is reported and their significance is evaluated by checking for overlapping 95% confidence intervals.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
38

---

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The amount of compute required is mentioned in Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't It make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https htt////eurrpppp wwtppaalic/tccccaliiiieeeeeeeeee
Answer: [Yes]
Justification:
We conform to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper concerns fundamental research for which is difficult to reason about any societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
39

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do have data or models that have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with
necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Any original creators or owners are properly credited and licences are re- spected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that souree should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Code for the experiments is released with the camera ready version. Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
40

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: We do not have any experiments with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
This paper does not involve LLMs as any important, original, or non-standard components.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve L.MMs any important, original, non-standard components. Please refer to our LLM policy ptttps://neurppee cct.oo/feeeeeccce/222255.. for what should or should not be described.
41


Reviews:
[{"review_id": "O9ztenvu1L", "reviewer": "Reviewer_3N2K", "summary": "The paper provides a group-theoretical understanding of policy distillation in RL when training a policy on a diverse set of contexts. It is shown that it is better to train an ensemble of distilled policies rather than a single policy. Further, when the data distribution exhibits symmetries, invariances introduced into the training data distribution tighten a generalisation bound on the distilled policy performance.  A reacher task with rotational invariances and a four-room navigation problem are provided as benchmarks, showing that larger ensembles and larger set of invariances boost performance.", "questions": "1. How does the runtime increase in the size of the ensemble?\n\n2. For more complex tasks that do not possess obvious invariances as the reacher experiment, how should the subgroup B be selected? Can this be learned from the given data?\n\n3. How does the method perform on large-dimensional problems?", "limitations": "I think more discussion should be provided about possible failure modes of the method in real-world settings where good invariances might be difficult to determine. I think a more detailed discussion of runtime costs, especially pertaining to ensemble size and how to mitigate it, could be useful.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\n1. The generalisation bound for distilled ensembles under subgroup/symmetry coverage is novel and mathematically elegant.\n\n2. The paper provides clear takeways from the theory that can be translated into practice: distill into an ensemble and use diverse states. While seemingly simple concepts, the theoretical understanding of why these things matter for policy distillation performance are less trivial.\n\n3. Ablations on ensemble size and distillation-data diversity convincingly support the theory.\n\nWeaknesses and comments:\n\n1. It seems that runtime should also increase significantly as a function of N, the size of the ensemble. It would be good to also show the runtimes associated with N = 1, N = 10 and N = 100 in Section 5.1. \n\n2. While the theory suggests that larger subgroups result in a lower performance gap, based on my understanding, the current experiments show it suffices to generate maximally diverse data from many invariances to improve performance. However, for more complex tasks that do not possess such obvious invariances as the reacher experiment, it remains largely unclear to me how the subgroup B should be constructed in practice. Do you suggest to do data augmentation with random invariances, hoping to generate data that sufficiently covers the testing distribution? What happens if you apply the wrong invariance for the problem or an invariance that fails to adequately capture the variation of the test instances? Does the method fail? Are there more principled ways to extract a (small) set of invariances that provide maximum coverage/generalization?  I think a section discussing the limitations and strategies to address these limitations would be helpful. Ideally, an experiment showing sensitivity of the method to poorly chosen invariances would be even more insightful so we can better understand what to focus on in the future.\n\n3. The existing experiments are somewhat limiting and focused on 2-D problems. I think some additional experiments on higher dimensional problems (i.e. Mujoco/D4RL) where invariances may be less trivial to design would be helpful.", "was_revised": true, "final_justification": "I am satisfied with the rebuttals which answered all the main concerns of all the reviewers. Overall I really enjoyed reviewing this paper, and I believe the theoretical insights into how and why data augmentation works for imitation and policy distillation will be important in moving RL forward. I will therefore maintain my positive score.", "rebuttal": "We thank the reviewer for their insightful comments and questions, which have helped us to further clarify and strengthen our manuscript. We have addressed each point below.\n\n- **Regarding the runtime increase of the ensemble:** All ensemble members are designed to be trained independently, and during inference, they can also be evaluated independently (an independent forward pass with an average over the output of the ensemble afterwards). This inherent independence means that both the training and inference processes of the ensemble are fully parallelisable. If parallelisation is not feasible, the runtime for both training and inference would increase linearly with the ensemble size. Furthermore, it is important to note that all policies are distilled using the same distillation dataset. This means the small number of additional environment steps required to sample a more diverse dataset is independent of ensemble size.\n- **Regarding more complex tasks that do not possess obvious invariances:** This is an important question, and we agree that for more complex tasks, learning the invariant subgroup B from the data could be a potential solution. However, at this stage, we are not entirely convinced that this approach will offer significant benefits. In our experiments with the Reacher environment, we observed no substantial advantage from training on a specific subgroup B compared to simply training on random transformations (see Table 2). While the Reacher environment is relatively small and conclusions should not be overly generalised, this finding aligns with our understanding of how data augmentation operates in practice. Although performing data augmentation with group symmetrical assumptions is necessary for certain theoretical guarantees [1,2], there seems to be a gap between the theoretically proven and the empirically observed benefits of data augmentation. In fact, a large body of empirical evidence suggests that data augmentation improves generalisation outside of such strict assumptions on the structure of the augmentations and underlying data distribution. For example, generalisation benefits are observed _even when_ the augmentations result in distorted data outside the support of the true data distribution (e.g., [3-6]). As such, we believe that collecting a more diverse dataset for distillation has the potential to increase generalisation performance, even if the additional data does not perfectly align with the testing distribution. However, we acknowledge that a more principled approach to creating such a diverse distillation dataset could more efficiently lead to increased generalisation in more complex domains (beyond our experiment in the rotational Reacher environment). To accomplish this, we believe the further analysis of the existing and future literature on the theoretical benefits of data augmentation is an exciting direction for future work.\n- **Regarding performance on large-dimensional problems:** This is an excellent question that points to a critical area for future research. It is indeed still an open question how well our method performs on more complex, higher-dimensional problems where the invariance may be less obvious. We believe our experiments on the Four Rooms environment (Table 3) represent a crucial first step in addressing this. In this environment, the exact symmetries governing the data distribution are not immediately obvious to us, yet we still observed significant benefits from distilling an ensemble ($N=10$) on more diverse data. This was evident with two distinct approaches to attaining this diversity: the Explore-Go dataset and the Mixed dataset. We are confident that further analysis into what properties of environments and distillation datasets could lead to improved performance will be a very interesting and fruitful direction for future work.\n\nOverall, we agree more discussion on the limitations of our work would strengthen the paper even more, and we will add a dedicated limitations section to achieve this. We hope these clarifications adequately address the reviewer's concerns and demonstrate the potential and significance of our approach. We are open to discuss any further questions during the discussion phase.\n\n**References** \\\n[1] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020 \\\n[2] A Group-Theoretic Framework For Data Augmentation. Chen et al. 2020 \\\n[3] The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective. Lin et al. 2024 \\\n[4] How Much Data Are Augmentations Worth? An Investigation Into Scaling Laws, Invariance, and Implicit Regularization. Geiping et al. 2023 \\\n[5] On the Effects of Artificial Data Modification. Marcu and Prugel-Bennet 2022 \\\n[6] On the Generalization Effects of Linear Transformations in Data Augmentation. Wu et al. 2020"}, {"review_id": "egfFRovaWe", "reviewer": "Reviewer_q7PZ", "summary": "This paper aims to investigate, under a zero-shot policy transfer setting (ZSPT) in reinforcement learning (RL), why policy distillation leads to more generalizable policies and what data should be used for distilling the policy. Specifically, the authors prove a generalization bound for a distilled policy in a generalization through invariance ZSPT (GTI-ZSPT) setting.  The bound provides two insights: \n1. If the policy distillation is performed on an ensemble of N infinitely wide neural networks, a larger N leads to a lower generalization bound. \n2. The generalization bound can be lowered by distilling the policy with more data from the environments. \n\nThe experiments validate that increasing the number of ensemble policies or training on a larger subgroup yields better results. In addition, the experiments also demonstrate that some of the assumptions can be relieved while the insights provided by the theory still hold.", "questions": "My questions are mostly mentioned in the weaknesses/suggestions section. \n1. What types of real-world data satisfy or do not satisfy the symmetric property? \n2. How to choose the smallest subgroup that can achieve good enough generalizability? \n3. What is the relationship between the number of ensemble policies and performance improvement? Is it linear or non-linear? \n4. Do you expect the insights provided by your theory to hold in more complex environments, e.g., Playground (in Minigrid) or robotics tasks?", "limitations": "It is difficult to tell which part of Section 5 mentions limitations. Please try to have a separate section discussing the limitations.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths\nThis paper is theoretically sound in its discussion of the key factors that facilitate generalizable policy transfer in RL through policy distillation. It provides mathematical validation of the effects of ensemble policies and the diversity of the data trained on, which are intuitive results, but opens the door to more advanced algorithm development. \n\nIn the experiments, the paper also demonstrates that even if the assumptions required by the theory do not hold, we can still observe the effects caused by the number of ensemble policies and diverse training data. This improves the applicability of the theoretical results.\n\n# Weaknesses/Suggestions\n1. The theory primarily relies on the symmetries present in many real-world data distributions. However, it is not straightforward to determine which types of real-world data satisfy this assumption. In addition to the Lie group SO(2) example provided in the paper, there can be more discussion on other data types that have symmetric properties. \n2. The theory suggests that if the training data covers a larger subgroup, the generalizability will be improved. I understand that in practice, it is challenging to train the policy on the entire group G; therefore, the experiments also suggest that random data can enhance generalizability. Yet, based on the theoretical results, I am interested in seeing how to wisely choose a subgroup that can achieve good enough generalizability. This discussion brings insights into balancing sample efficiency and performance. I am also interested in learning how to determine the number and diversity of ensemble policies. \n3. There can be more experiments on environments that do not have invariance symmetry, as validating in one environment might not be sufficient to validate your conclusions.", "was_revised": true, "final_justification": "I have read through all the reviews and authors’ responses. After careful consideration, I decided to raise my rating accordingly, as the proposed theoretical findings are crucial to gaining a deeper understanding of how commonly applied techniques, e.g., data augmentation, work or can be improved.", "rebuttal": "We sincerely thank the reviewer for their thoughtful comments and insightful questions, which have helped us clarify and strengthen our manuscript. We've carefully considered each point and provide the following responses.\n- **Regarding what types of real-world data satisfy or do not satisfy the symmetric property:** The real world is rich with symmetries, including mirror, translational, and scaling symmetries, or simply invariances to noise or irrelevant background details. For example:\n\t- **Image recognition** heavily relies on these symmetries, where objects are identifiable regardless of their orientation, position or scale. \n\t- **Robotics** often deals with actions and environments exhibiting symmetries in physics or optimal paths. For example, rotational and translational symmetries when moving from A to B. \n\nAdditionally, there might be many more symmetries existing in the real world that we are not immediately aware off since they are not obvious from our current viewpoint. We believe the strength of our approach lies in identifying insights that can hold even when the symmetries are no longer obvious to us, as demonstrated in the Four Rooms experiments in Table 3. \n- **Regarding the smallest subgroup that can achieve good enough generalisability:** Our theory suggests minimising the discrepancy measure $\\kappa$. For any fixed subgroup size $|B|$, this measure is minimised when the subgroup has maximal coverage over the full group $G$ (as measured by the operator norm in representation space $\\psi_S$). This aligns with some prior empirical work on data augmentation, where distributing augmentations uniformly leads to good performance [1,2]. \n- **Regarding the relationship between the number of ensemble policies and performance improvement:** Our theory indicates that the performance of a distilled ensemble is bounded by optimal performance with a term of the form $b + \\frac{c}{\\sqrt{N}}$, where $N$ is the ensemble size. This shows a non-linear relationship, with diminishing returns as $N$ increases. This relationship provides an upper bound on the performance difference, so it might not perfectly predict actual realised performance improvement. However, we are careful not to overstate the predictive power of the current theory, as it relies on simplifying assumptions (like infinitely wide neural networks) that are hard to meet in practice. We believe the true power of our theory lies in helping identify relevant dataset properties and distillation ensemble characteristics for improved generalisation, rather than providing exact quantitative predictions for ensemble size versus performance. However, we believe our work is an important first step that could lead to more precise theoretical and empirical developments in future work.\n- **Regarding more complex environments:** While our theory focuses on specific symmetric structures, we believe its insights could hold much more generally in complex environments like Minigrid's Playground or robotics tasks. For example, our rotational Reacher experiment (Table 2 in Section 5.1.1) saw no significant difference between training on an exact subgroup $B=C_4$ (as per theory) versus random augmentations. From our understanding, this aligns with the known gap between theoretical and empirical benefits of data augmentation. Often theoretical results require global group symmetry [3,4], yet empirical evidence shows data augmentation improves generalisation even when augmentations aren't strictly symmetric [5-7] or aren't global [8,9]. Therefore, we believe our insights have the potential to boost generalisation in even more complex settings. Further analysis of what environmental properties are required for our insights to hold is a very interesting direction for future work.\n\nOverall, we agree more discussion on the limitations of our work would strengthen the paper, and we will add a dedicated limitations section to the to accomplish this. We hope these clarifications adequately address the reviewer's concerns and demonstrate the potential of our approach. We are open to discuss any further questions during the discussion phase.\n\n**References** \\\n[1] UniformAugment: A Search-free Probabilistic Data Augmentation Approach. Lozhkov et al. 2020 \\\n[2] TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation. Müller and Hutter 2021 \\\n[3] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020 \\\n[4] A Group-Theoretic Framework For Data Augmentation. Chen et al. 2020 \\\n[5] On the Generalization Effects of Linear Transformations in Data Augmentation. Wu et al. 2020 \\ \n[6] Generalization in Reinforcement Learning by Soft Data Augmentation. Hansen et al. 2021 \\\n[7] Automatic data augmentation for generalization in reinforcement learning. Raileanu et al. 2021 \\\n[8] Data Augmentation as Feature Manipulation. Shen et al. 2022 \\\n[9] Learning Instance-Specific Augmentations by Capturing Local Invariances. Miao et al. 2023"}, {"review_id": "I4fITmESLm", "reviewer": "Reviewer_QfQd", "summary": "This paper explores the generalization ability of policy distililation both theoretically and empirically, showing that ensemble of policy distilation on augmented data improves generalization.", "questions": "1.Is the ensemble used during training or during testing? Which contributes more?\n\n2.Could the gains stem primarily from regularization (e.g., via ensemble averaging or noise injection) rather than group invariance? How do experiments isolate invariance as the key factor?\n\n3. Is there a principled way to determine the minimal sufficient diversity/size of the distillation dataset for a task? Does the theory offer guidelines beyond \"more is better\"?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n1. The paper proposes a novel theoretical framework (GTI-ZSPT) that formalizes how group invariance and ensemble distillation contribute to policy generalization.\n2. Theoretical results are non-trivial and connect deeply with recent findings in deep learning theory and group invariance.\n3. The exposition is clear, with thoughtful examples and visual aids.\n\n**Weaknesses**\n1. There is a large gap between theory and practice: the theory focuses on the size of the *symmetry group* of the dataset, but in practice, collecting diverse states will simply benefit, as suggested in Section 5.\n2. Using 10-policy ensembles incurs significant training/inference overhead. No analysis is provided for smaller ensembles (e.g., 3–5 policies) or efficiency trade-offs, limiting practical adoption.", "was_revised": true, "final_justification": "I will maintain my original rating and recommend accept. All my concerns have been addressed.", "rebuttal": "We appreciate the reviewer's insightful questions, which help to clarify important aspects of our work. Please find our responses to each point below.\n- **Regarding how and when the ensemble is used:** The ensemble members are independently distilled on a single distillation dataset and the forward pass through each ensemble member is performed independently (and the output is then averaged) at each step during testing. Therefore, the ensemble is not used during training; its role is solely during testing to average over the ensemble members' outputs at each step. This approach allows for high parallelisability during both training and inference, keeping computational overhead relatively small. If parallel computation is not possible, the computational overhead would scale linearly with ensemble size. We also note that since all ensemble members are distilled on the same dataset, the number of additional environment interactions required to create this dataset (which is typically considered the main bottleneck in RL) is independent of the ensemble size.\n- **Regarding regularisation versus group invariance:** This is indeed a very important and actively debated point in the data augmentation literature. There is ongoing academic discussion regarding how data augmentation precisely improves generalisation. Some studies suggest that data augmentation, even when performed with respect to a clearly defined group invariance, does not guarantee the learning of an invariant function [1]. Instead, some work suggests that the real underlying mechanism has more to do with reducing overfitting to spurious correlations, than inducing symmetry invariance in the model [2]. As such, whether the benefits of performing data augmentation with respect to some symmetry group actually stem from induced invariance or reduced overfitting and other forms of regularisation, is still an open topic of discussion. Nonetheless, the empirical benefit of data augmentation in various domains and settings is widely acknowledged [3-6], and we demonstrate in Section 5.2.1 that the insights from our theory translate to environments in which the group symmetries are no longer obviously present. \n- **Regarding how experiments isolate invariance as the key factor:** As discussed above, whether the generalisation benefits from data augmentation stem from learned invariance or other mechanisms (e.g., reduced overfitting and regularisation), is still an open topic of research in the data augmentation literature. However, there are attempts made to try to analyse the induced invariance of the model. If the symmetries are known, it is possible to measure the variance of the network output across the group orbit [7]: if the model is invariant, this variance should be zero. Unfortunately, we cannot perform this analysis for the Four Room experiment in Section 5.2, since we do not know the exact symmetries that govern this environment. However, we are committed to perform this analysis for the rotational Reacher experiments in Section 5.1 and add it to the paper. This will not provide a definitive answer, but can give an indication of the level of invariance induced in the model for the Reacher environment that most closely aligns with our theory. \n- **Regarding a principled way to determine the minimal sufficient diversity/size of the distillation dataset for a task:** Following our theoretical framework, the relevant metric to minimise would be the discrepancy measure $\\kappa$. For any fixed subgroup size $|B|$, this measure is minimised if the subgroup has maximal coverage over the full group $G$ (as measured by the operator norm in representation space $\\psi_S$ ). This aligns with some prior empirical work on data augmentation, where distributing augmentations uniformly leads to good performance [8,9]. However, we are careful not to overstate the predictive power of the current theory, as it relies on assumptions that are difficult to hold in practice (such as the infinitely wide neural networks assumption). We believe the primary strength of our theory lies in its ability to help identify relevant properties of the data distribution and the distillation ensemble that can be leveraged to improve generalisation performance. It serves as a valuable first step that could potentially facilitate further theoretical developments in future work, allowing these types of questions to be answered with greater certainty.\n\n**References** \\\n[1] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020 \\\n[2] Data Augmentation as Feature Manipulation. Shen et al. 2022 \\\n[3]  A survey on image data augmentation for deep learning. Shorten and Khoshgoftaar 2019 \\\n[4] A survey of data augmentation approaches for NLP. Feng et al. 2021 \\\n[5] Generalization in Reinforcement Learning by Soft Data Augmentation. Hansen et al. 2021 \\\n[6] Automatic data augmentation for generalization in reinforcement learning. Raileanu et al. 2021 \\\n[7] In What Ways Are Deep Neural Networks Invariant and How Should We Measure This? Kvinge et al. 2022 \\\n[8] UniformAugment: A Search-free Probabilistic Data Augmentation Approach. Lozhkov et al. 2020 \\\n[9] TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation. Müller and Hutter 2021"}, {"review_id": "R5kw1jjoWS", "reviewer": "Reviewer_3YDA", "summary": "The paper studies policy transfer in RL environments. Assume we have a contextual MDP, where we are allowed to observe training contexts $C_{train}$ but wish to generalize to test contexts $C_{test}$. It has been shown in practice that sometimes, distilling a policy can lead to a policy with better generalization behavior than the original policy. This paper is primarily a theory paper, deriving a generalization bound on distillation, assuming finite and infinite ensembles of policies. Note that the paper views BC as a specific instance of policy distillation (distillation of the optimal policy).\n\nThe proof focuses on a specific subcase where the contexts of the contextual MDPs are elements of a symmetry group $G$, and the training contexts are elements of a subgroup $B$. This assumption allows the authors to define a constant $\\kappa$, which measures the\ndiscrepancy between subgroup $B$ and the entire group $G$. (Intuitively, a measure for how \"close\" the training contexts $B$ are to test time contexts in $G - B$.\n\nThe paper further assumes an infinitely wide neural network, following the NTK line of work, and an infinite number of training samples from the training contexts $C_{train}$. This allows the authors to prove a bound between the optimal return $J^{\\pi^*}$ and return of the ensemble of $N$ policies $J^{\\pi_\\infty}$, that depends on:\n\n* A constant term that depends on the discount factor, and Lipschitz constants for transitions, rewards, and policies.\n* An additive constant based on $\\kappa$ and the network architecture, which is irreducible error even if number of ensembled models $N$ approach infinity.\n* And one term that scales with $O(1/\\sqrt{N})$ on the number of models in the ensemble.\n\nIn short, the bound grows tighter as subgroup $B$ covers more of the group $G$, and as number of models in the ensemble increases.\n\nTo evaluate how well this appears practically, the authors study a 2D reacher environment, where the symmetry group is the SO(2) group of rotations, and the subgroups are $C_2, C_4, C_8$ corresponding to 180 / 90 / 45 degree rotations around the circle.\n\nThe authors finally demonstrate that widening the initialization (context for the MDP) causes better test time generalization, as does changes to the makeup of the trajectory dataset to include more exploratory data.", "questions": "Could the authors explain why Distillation (Teacher dataset) performs so differently than Behavior Cloning (Teacher dataset)? Shouldn't these two be identical? As the author state earlier, they consider behavior cloning to be a special case of distillation.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Most of the results in this work are intuitively understood, but this paper proves more formal bounds on expected generalization of policies trained via distillation. The paper is helpful on providing references to background theoretical material, and although I did not read the proofs very closely, at a high level nothing about them seemed wrong / I expect them to hold if checked more closely.\n\nThe subgroup + symmetry group assumption in particular is fairly restrictive however. Given that the purpose of this assumption is solely to define $\\kappa$, it's not clear to me that this actually adds much to the proof either? I would assume that instead of requiring a symmetry group, you could instead directly define $\\kappa$ as an unobserved constant measuring the distance between $C_{train}$ and $C_{test}$, along with the appropriate Lipschitz assumptions, and get an equivalent theorem that's more general.\n\nI think it is good that the authors show that test performance improves as the ensemble grows in side, even while train-time performance has saturated (Section 5.1). However I was a bit disappointed that there was no attempt made to check how closely this follows the $1/\\sqrt{N}$ derived earlier. If you abstract Eqn (5) to $a(b + \\frac{c}{\\sqrt{N}})$ for unknown constants $a,b,c$ and try to fit lines to it, surely that would be a better argument for how good the bound in (5) is? As it stands there is no intuition for whether it's tight or loose. Similar things could be said about measuring $\\kappa$.\n\nFinally I'm not sure Section 5.2.1 does anything for the paper, since as far as I can tell it has little bearing on the theoretical result proven by the paper. Taken on its own as an empirical result, it argues that using a wider distribution of training data leads to better test time performance in environments where we fully saturate the training set, but this is already a well-known result from the rest of deep learning. So it seems irrelevant to the paper itself.", "was_revised": true, "final_justification": "Did not have many issues with the original paper, believe paper is technically solid and authors have addressed my concerns.", "rebuttal": "We thank the reviewer for their insightful comments and questions, which have helped us to improve the clarity and impact of our paper. We address each point below.\n- **Regarding distillation vs. behaviour cloning:** This indeed raises an important point about the difference between distillation and behaviour cloning as defined in our paper. The key difference lies in their optimisation objectives and learning targets. Both aim to imitate the same behaviour policy, but behaviour cloning can only observe the sampled actions from that policy, whereas distillation has access to all action probabilities. This implies different losses: behaviour cloning uses a cross-entropy loss, whereas distillation regresses the probabilities with a MSE loss (see Appendix A.1.1). As a result, distillation *should* be much more sample efficient, as it has access to more information than behaviour cloning.\n- **Regarding the $\\frac{1}{\\sqrt{N}}$ generalisation bound:** We appreciate the suggestion to analyse how closely our results align with the $\\frac{1}{\\sqrt{N}}$ theoretical bound. This is indeed an interesting analysis that could provide valuable intuition. While we believe this would be a valuable addition, conducting the necessary experiments would unfortunately before the current discussion period ends. However, we are committed to performing this analysis and will include it in the final version of the paper. That being said, we do want to manage expectations regarding the tightness of this bound. Obtaining tight generalisation bounds for neural networks is notoriously challenging [1,2]. Furthermore, some of our theoretical assumptions, such as infinitely wide neural networks, are hard to meet in practice. We believe the true strength of our theory lies in its ability to identify crucial properties of the dataset distribution and distilled ensemble that are capable of improving generalisation performance. Moreover, our paper introduces a novel analytical approach to policy distillation, paving the way for more advanced theoretical developments in the future.\n- **Regarding section 5.2.1:**  While it is true that these experiments break a significant number of the theoretical assumptions, we firmly believe this section is relevant to the paper's overall message. Section 5.2.1 serves to empirically demonstrate how our insights can be leveraged to significantly enhance the generalisation performance of a reinforcement learning agent through policy distillation. A clear example of this is seen in our ensemble $N=10$, distilled on the Explore-Go dataset, which achieves substantially higher test performance than the original PPO+Explore-Go teacher policy (see Table 3). The potential of policy distillation to improve generalisation was initially identified in [3], but we believe our paper provides a more compelling argument and empirical evidence for this phenomenon. Thus, rather than aiming to confirm theoretical predictions, Section 5.2.1 showcases a significant empirical finding for which our theory provides the most convincing justification and motivation to date. Furthermore, the idea of broadening the training distribution by sampling more extensively from existing training environments, as opposed to generating entirely new ones, is an emerging concept in the RL literature, recently introduced in [4,5]. Our paper demonstrates that this concept also applies effectively to policy distillation, and to the best of our knowledge, we are the first to highlight these specific generalisation benefits within this context.\n\nTo improve the clarity of the paper, we will add a dedicated limitations section that will discuss all the limitations and points raised during review. We hope the clarifications outlined here adequately address the reviewer's concerns and demonstrate the potential of our approach. If not, we are of course open to discuss any further questions during this discussion phase.\n\n**References** \\\n[1] Fantastic Generalization Measures and Where to Find Them. Jiang et al. 2020 \\\n[2] Fantastic Generalization Measures are Nowhere to be Found. Gastpar et al. 2023 \\\n[3] Learning Dynamics and Generalization in Reinforcement Learning. Lyle et al. 2022 \\\n[4] On the Importance of Exploration for Generalization in Reinforcement Learning. Jiang et al. 2023 \\\n[5] Exploration Implies Data Augmentation: Reachability and Generalisation in Contextual MDPs. Weltevrede et al. 2025"}]

Author rebuttal:

Meta-review: (a) Scientific Claims and Findings

The central claim of this paper is that the generalization of a distilled policy in RL can be provably improved by using an ensemble of student policies. The authors theoretically “analyse the act of distilling a policy after training, and try to answer 35 how the policy should be distilled and on what data”. They formalise this by proving a generalization bound in a "Generalisation Through Invariance" setting, where the performance gap to an optimal policy is shown to decrease with the ensemble size (N) and with better coverage of the underlying symmetry group (a smaller discrepancy measure, k). The empirical findings confirm these two principles hold in practice, both in a setting that matches the theory and in a more general setting that violates the theory's strict assumptions.

(b) Strengths

The paper's primary strength is its novel theoretical contribution as it provides a formal answer to the important question of why 1) policy distillation can improve generalisation, moving beyond prior empirical observations and indirect explanations, 2) data augmentation via invariances. The resulting insights—use ensembles and diverse data—are simple, intuitive, and directly actionable. I found the connection made between generalization in RL, group theory, and the neural tangent kernel (NTK) literature both elegant and insightful. The experiments are well-designed to specifically test the predictions of the theory and demonstrate their robustness when assumptions are relaxed.

(c) Weaknesses

The main weakness of the paper lies in the strong assumptions required to prove the theoretical result, including Lipschitz continuity of the MDP, the existence of a clear symmetry group, and the use of infinitely wide neural networks. This creates a noticeable gap between formal theory and messy real-world applications. The small experiments demonstrate empirically that the qualitative insights from the theory are still predictive and useful in a more practical setting (Minigrid's Four Rooms) where these assumptions do not hold, however the method was not tested in more challenging, less artificial settings.

(d) Reasons for Decision

The paper is recommended for acceptance because the paper is technically solid and offers clear takeaways that are backed by well-executed experiments. It also brings together various results in DL and empirical observation, thus it advances our understanding of why certain techniques work, which is a valuable contribution to the field. The reviewers were unanimously positive, and their initial concerns were effectively addressed during the rebuttal period, leading to a strong consensus.

(e) Discussion and Rebuttal

The author-reviewer discussion was highly constructive. Authors addressed all major points, clarifying distinctions and agreeing to add an empirical analysis. Their effective engagement strengthened the paper and resolved reviewer doubts.
Decision: Accept (poster)