Behaviour Distillation

We thank reviewer ZN7t for their comments.

0. "HaDES [...] is able to distill the initial state-actions pairs into a dataset of just four state-action pair"

We formulate Behaviour Distillation without access to expert data, as mentioned in the abstract, introduction and section 4.1. This means there are no "initial state-action pairs". Also, the size $N$ of the distilled dataset is a hyperparameter we choose. Figures 1 shows datasets with $N=2$ and $N=4$, and in section 5 we report results with $N=64$ and $N=16$.

Response to Weaknesses:

1. "explain [...] why you need to use a neuroevolution technique instead of a more classical technique"

Neuroevolution covers all methods optimising a neural network using Evolutionary Strategies (ES) [1]. We do not use neuroevolution. HaDES, and in particular HaDES-F, can act as a neuroevolution method, as we explain in section 4.2 and demonstrate in section 5.1. We use ES instead of a gradient-based approach such as Backpropagation Through Time (BPTT) because BPTT is expensive when training the inner loop policy for many steps. For context, we use 400 inner loop training steps, while [3] only uses 10 to 30. We have now clarified this in section 2.1.

2. "...it would be clearer to name the experiments as Method + HADES"

HaDES is a method that takes in an environment M and outputs a tuple $(D_\phi, \pi_{\theta^*})$, where $D_\phi$ and $\pi_{\theta^*}$ are the solutions to the optimization problem in eq. 3 and 4. HaDES uses ES in the outer loop and behaviour cloning (supervised learning) in the inner loop. Training a policy on the synthetic dataset is integral to HaDES, so we cannot split the name into "method + HaDES".

3. "looking at some subpart of the episode (like 1/10th)..."

HaDES does not divide episodes at all. It evolves the synthetic datasets from scratch. In section 1, we say that the final dataset contains "1/10th [...] of a single episode worth of data." This is about the final magnitude of the dataset, not about how HaDES works.

"... computational advantages of training on top of these condensed dataset."

The usual episode is 1000 transitions, and our datasets are at most 64 states (<10% of an episode). In contrast, training a policy using RL would require ~50k episodes [2]. Training on the distilled datasets uses between 6 and 8 orders of magnitude less data that standard RL training, making training policies extremely fast.

We have also updated Section B.2 of the manuscript with details of the GPU time required to run each method.

"looking at the RL plots [...] it's not clear if the red curve trains way faster than the green curve"

We are not sure what the reviewer is trying to say. HaDES-F (red curve) reaches statistically higher return than Vanilla ES (green curve) in 7/12 environments reported, reaches similar returns in 4/12 environments and reaches statistically lower returns in 1/12 environments. If the reviewer could please clarify what they mean, we would be happy to respond further.

Answer to Questions:

1. "HADES first distills a dataset into a condensed state-action pair and then trains a classical policy on it (as in classical dataset distillation)."

That is incorrect. As stated previously, HaDES does not take a dataset as input and does not perform classical dataset distillation. We have also expanded section 4.1 with details on why we do not rely on an expert dataset.

"...neuroevolution part. Is this part requires only for distillation or is it also required for training on top of the distilled dataset ? In the latter case, it would be interesting to see how other methods that do not use neuroevolution perform when trained on the condensed dataset."

Hopefully, our explanation under W1 clarified that we do not use neuroevolution. We use ES (which is different from neuroevolution) to optimise the synthetic dataset during the distillation process. When we train the policy (in the inner loop or on the final distilled dataset), we simply use behaviour cloning (i.e. supervised learning).

2. "In the results plot of reinforcement learning, does the number of generation directly correlates to the training time ?"

For each method and each environment, the time/generation is constant. Generally, HaDES takes a bit longer per generation since it trains a policy every inner loop, as opposed to ES, which must only evaluate the policy. We have added a time/generation breakdown to section B.2 of the appendix.

We hope our response has clarified any misunderstanding and that we have addressed the reviewer's feedback. If so, we would appreciate if they would consider increasing their score.

[1] https://link.springer.com/referenceworkentry/10.1007/978-0-387-30164-8_589

[2] ReLU to the rescue, https://arxiv.org/abs/2306.01460

[3] Dataset Distillation, https://arxiv.org/pdf/1811.10959.pdf

[4] Efficient Reductions for Imitation Learning, https://proceedings.mlr.press/v9/ross10a.html

We thank Reviewer 7Zz2 for their valuable comments. In particular, we welcome the helpful suggestions on how to improve our experiments section. We invite the reviewer to read our response below.

1. "In my opinion, the networks in RL are often small and do not require long time training."

We agree that the networks used in RL tend to be smaller than the convolutional networks used in vision. However, RL regularly requires tens of millions of environment steps to learn a good policy, which can easily translate to several hours or even days of training time [1, 2]. This is why there has been a recent push to build hardware-accelerated RL environments and high-performance RL implementations [2,3,4,5,6]. The distilled datasets allow us to completely avoid environment bottlenecks by turning policy training into a very small supervised learning problem.

2. "advantages of using the distilled dataset in RL except for fast training?"

Beyond training speed, there are multiple applications to the datasets produced by HaDES. Notably:

• In section 5.2, we show that distilled datasets readily generalise to a large range of architectures and hyperparameters, which implies they can be used for effective Neural Architecture Search (NAS). NAS has notably been identified as a potential application of Dataset Distillation in the supervised domain as well [7].

• In section 5.3, we demonstrate that the datasets evolved for individual environments can be combined to train a multi-task agent, therefore enabling faster research on generalist foundational models.

• In section 5.5, we show briefly that the datasets can help explain the behaviour of the policies and make them more interpretable.

• As mentioned in our conclusion, we also hope HaDES to act as a starting point for new and improved methods for continual learning, where resetting and retraining policies repeatedly is common to fight plasticity loss [2, 9].

3. "It will be more reasonable to compare HaDES and vanilla ES under the same experimental settings"

Thank you for raising this point. We have now updated Figure 3 b) with additional runs where the network widths match those used for vanilla ES (width=256). The impact of the on the final performance is negligible in SpaceInvaders, Asterix and Freeway, but has a noticeable impact on Breakout. Nonetheless, the order of the different methods remains unchanged.

For Brax (Figure 3.a)), this was already the case (both HaDES and ES use width=512). Nonetheless, HaDES-F outperforms ES in all Brax environments tested, and HaDES-R matches or beats the baseline in 6/8 environments.

4. "add the results of other RL methods instead of only ES"

We agree that reporting an RL baseline in addition to ES would help contextualise our results better. We have updated Figure 3 with a PPO baseline, which is a very strong RL baseline for our suite of environments [10]. While ES methods still lag behind RL in terms of performance, HaDES helps narrow the gap significantly. This is in addition to the main benefit of HaDES, namely producing the distilled datasets.

5. "lack of experiments investigating the influence of the distill budget (the size of the distilled dataset) on the final performance"

We thank the reviewer for this helpful suggestion. We have added such an analysis to section C.1 of the appendix, with a reference to it at the beginning of section 5.

6. "list the GPU time of behaviour distillation, training ..."

We have now expanded section B.2 with additional details regarding the GPU time for different methods. In particular, please refer to Table 4 for a breakdown per environment.

7. Minor points:

We thank the reviewer for their thoroughness! We have added a clarification regarding the term "Behaviour Distillation" to avoid ambiguity. Regarding the algorithm, we have considered moving it from the appendix to the main paper, but have decided against it due to lack of space.

We hope that we addressed the reviewer's comments, which genuinely helped us improve our manuscript! As a result, we hope the reviewer will consider recommending our revised submission for acceptance.

[1] R2D2, https://openreview.net/pdf?id=r1lyTjAqYX

[2] Bigger, Better, Faster, https://arxiv.org/abs/2305.19452

[3] Madrona, https://madrona-engine.github.io/shacklett_siggraph23.pdf

[4] Brax, https://arxiv.org/abs/2106.13281

[5] Scaling Distributed RL, https://arxiv.org/abs/2306.16688

[6] MuJoCo 3 announcement, https://github.com/google-deepmind/mujoco/discussions/1101

[7] Data Distillation: A Survey, https://arxiv.org/pdf/2301.04272.pdf

[8] PureJaxRL, https://github.com/luchris429/purejaxrl

[9] The Primacy Bias in Deep RL, https://arxiv.org/abs/2205.07802

[10] ReLU to the rescue, https://arxiv.org/abs/2306.01460

I appreciate the authors' response, which clears part of my concerns. I still have some concerns as below:

Q1: "The distilled datasets allow us to completely avoid environment bottlenecks by turning policy training into a very small supervised learning problem." Does the vanilla policy training use supervised learning on an un-distilled large (state, action) pair dataset? If yes, I think the behavior distillation still lies in a supervised learning regime instead of reinforcement learning.

Q2: The authors show many pros of behavior distillation in NAS, interpretability, continual learning, etc. In my opinion, these pros are also inherent advantages for dataset distillation and lie in supervised learning regime. How about the exclusive advantage of behavior distillation especially in RL field.

Q3: There still lack experiments on ES neuroevolution, width=512 in Fig. 3(b)

Q4: I am curious why networks trained on small distilled data can outperform those trained on large source real data in behavior distillation, while the generalization performance of distilled largely lags behind the real data in dataset distillation.

Q5: Can the behavior distillation integrate with other RL algorithms instead of ES? If ES largely lags to modern mainstream RL algorithms, the contribution of behavior distillation will be greatly undermined.

Q6: Because efficiency is a main advantage for behavior distillation, a runtime analysis in necessary. However, according to results in Table 4, (1) ES has shorter runtime than HaDES, which means ES is more efficient, am I right? (2) The time unit in Table 4 is second, which means that both training with vanilla setting or behavior distillation have a satisfied speed and thus undermines the necessity to develop behavior distillation. This is why I asked Q1 in my (initial) official review.

We thank reviewer f5UH for their feedback and are glad they believe our paper can enrich the literature surrounding Dataset Distillation.

1. Behaviour Distillation as a setting is directly motivated by the success of Dataset Distillation in the supervised domain, with applications such as differential privacy, neural architecture search, continual learning, and federated learning [1]. With our work, we hope to extend the reach of those applications to RL, while also enabling new RL-specific ones, such as the training of multi-task policies without additional environment interactions (section 5.3 of our paper).

We formulated Behaviour Distillation to be independent of an expert policy for two main reasons:

• Firstly, an expert dataset is not always available, and so we do not wish to rely on this assumption. Furthermore, even if an expert dataset was available, the expert's behaviour could be erratic and hard to distill, which means that we would need to increase the size of the dataset to capture the idiosyncrasies of the expert. Instead, HaDES restricts its optimization to policies that can be distilled into datasets of size $N$, eliminating policies that are too erratic by design.
• Secondly, even when we do have access to a dataset, supervised dataset distillation is still likely to fail. In RL, a policy trained with behaviour cloning (BC) on an expert dataset incurs regret with respect to the expert that is quadratic in the length of the episode [2]. Intuitively, this happens because each timestep $\pi$ deviates from $\pi_{expert}$, it risks ending up off-distribution, which in turn makes it more likely to make additional mistakes. Since dataset distillation is often lossy, we can therefore expect $\pi$ to achieve low returns in many long horizon tasks, even if it achieves low (but non-zero) cross-entropy with the expert.

We recognize this motivation wasn't explicit in the original manuscript and have therefore expanded section 4.1 with the details above.

2. "discuss the behaviour difference of a standard RL algorithm and a synthetic data-driven RL algorithm (HADES). What's the potential benefits?"

We are not sure whether the reviewer is asking about the difference between algorithms or between the resulting policies. We answer both:

• Differences between algorithms:

  • Deep RL algorithms, in short, will repeatedly unroll the policy in the environment to collect transitions, and use those transitions to update the policy through SGD, in a way that ultimately maximises the expected discounted return.
  • Vanilla Evolutionary Strategies (ES) perform Neuroevolution: they start from a randomly initialised neural network parametrization $\theta$ of the policy $\pi_\theta$. Each generation, they sample a population of random variations of that parametrization, estimate the expected return of each member of the population (the candidates), and use that information to update the policy parameters.
  • As described in section 4.2, HaDES is similar to ES, but instead of keeping track of a policy $\pi_\theta$, it keeps track of a synthetic dataset $D_\phi$. Each generation, it then samples a population of datasets and uses BC to obtain a policy from each dataset candidate. The advantage is that $|\phi|$ < $|\theta|$ typically, which reduces the single-GPU memory of the algorithm and means there are fewer parameters to optimise. In practice, this also translates to significantly higher returns, as shown in Section 5.1. Additionally, the advantage of HaDES over both Deep RL and ES is that it produces the dataset $D_\phi$, which can be used for fast retraining and other applications.

• Differences between policies:

  • Because HaDES and Deep RL rely on very different optimization methods, they can in theory produce policies with different behaviours. For instance, HaDES only restricts optimization over relatively "simple" policies, i.e. those that can be condensed in a dataset of predetermined size. This is beneficial in scenarios where interpretability and predictability are important. We believe exploring the full impact of the training method on the resulting policies would be a promising direction for future work.

3. "One missing citation on BPTT."

Thank you for the missing reference. We have now added it to the paper.

4. "It would be nice to add some discussion on optimization algorithm in DD, for example BPTT"

We already had a brief discussion regarding the connection between BPTT/meta-gradients and HaDES in section 2.1. Following feedback from the reviewer, we have added a few more details to the revised manuscript. If this is not what the reviewer meant, we ask them to kindly clarify.

We hope this covers the reviewer's questions, and are happy to continue the conversation should the reviewer have any follow-up.

[1] Data Distillation: A Survey, https://arxiv.org/pdf/2301.04272.pdf [2] Efficient Reductions for Imitation Learning, https://proceedings.mlr.press/v9/ross10a.html