We appreciate the reviewer's insightful questions, which help to clarify important aspects of our work. Please find our responses to each point below.

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

References
[1] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020
[2] Data Augmentation as Feature Manipulation. Shen et al. 2022
[3] A survey on image data augmentation for deep learning. Shorten and Khoshgoftaar 2019
[4] A survey of data augmentation approaches for NLP. Feng et al. 2021
[5] Generalization in Reinforcement Learning by Soft Data Augmentation. Hansen et al. 2021
[6] Automatic data augmentation for generalization in reinforcement learning. Raileanu et al. 2021
[7] In What Ways Are Deep Neural Networks Invariant and How Should We Measure This? Kvinge et al. 2022
[8] UniformAugment: A Search-free Probabilistic Data Augmentation Approach. Lozhkov et al. 2020
[9] TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation. Müller and Hutter 2021

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.

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

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

References
[1] Fantastic Generalization Measures and Where to Find Them. Jiang et al. 2020
[2] Fantastic Generalization Measures are Nowhere to be Found. Gastpar et al. 2023
[3] Learning Dynamics and Generalization in Reinforcement Learning. Lyle et al. 2022
[4] On the Importance of Exploration for Generalization in Reinforcement Learning. Jiang et al. 2023
[5] Exploration Implies Data Augmentation: Reachability and Generalisation in Contextual MDPs. Weltevrede et al. 2025

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.

• 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:
  • Image recognition heavily relies on these symmetries, where objects are identifiable regardless of their orientation, position or scale.
  • 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.

Additionally, 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.

• 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].
• 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.
• 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.

Overall, 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.

References
[1] UniformAugment: A Search-free Probabilistic Data Augmentation Approach. Lozhkov et al. 2020
[2] TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation. Müller and Hutter 2021
[3] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020
[4] A Group-Theoretic Framework For Data Augmentation. Chen et al. 2020
[5] On the Generalization Effects of Linear Transformations in Data Augmentation. Wu et al. 2020 \ [6] Generalization in Reinforcement Learning by Soft Data Augmentation. Hansen et al. 2021
[7] Automatic data augmentation for generalization in reinforcement learning. Raileanu et al. 2021
[8] Data Augmentation as Feature Manipulation. Shen et al. 2022
[9] Learning Instance-Specific Augmentations by Capturing Local Invariances. Miao et al. 2023

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.

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

Overall, 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.

References
[1] On the Benefits of Invariance in Neural Networks. Lyle et al. 2020
[2] A Group-Theoretic Framework For Data Augmentation. Chen et al. 2020
[3] The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective. Lin et al. 2024
[4] How Much Data Are Augmentations Worth? An Investigation Into Scaling Laws, Invariance, and Implicit Regularization. Geiping et al. 2023
[5] On the Effects of Artificial Data Modification. Marcu and Prugel-Bennet 2022
[6] On the Generalization Effects of Linear Transformations in Data Augmentation. Wu et al. 2020