Paper: Behaviour Distillation (ICLR-2024)

Reviewer comment (Novelty):
I think the neuroevolution part is quite distinct from the dataset distillation part. It should be explained clearly why you need to use a neuroevolution technique instead of a more classical technique.

Author response:
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.

Reviewer comment (Presentation):
Maybe the naming is a bit confusing, if the method proposed is meant to distill the dataset only (and not train on it later), then it would be clearer to name the experiments as Method + HADES (whenever training a method on top of hallucinated dataset).

Author response:
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".

Reviewer comment (Evaluation):
If I understood correctly, one advantage of HADES is that it can distill behaviour by looking at some subpart of the episode (like 1/10th), and condense the gathered state-action pairs into much smaller dataset. However, from the rest of the paper, it's not clear whether they are some computational advantages of training on top of these condensed dataset. They should be made clearer. Indeed, by looking at the RL plots, in most of the cases, it's not clear if the red curve trains way faster than the green curve (apart from a few environments like humanoid and breakout-miniAtari).

Author response:
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.

Reviewer comment (Experiments):
If I understand correctly, HADES first distills a dataset into a condensed state-action pair and then trains a classical policy on it (as in classical dataset distillation). What I don't understand is the link with the neuroevolution part. Is this part required 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.

Author response:
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).

Reviewer comment (Evaluation):
In the results plot of reinforcement learning, does the number of generation directly correlate to the training time? Or are some generations quicker to run for methods trained on condensed dataset (that implies the inner loop would be faster I guess)? I think the computational advantage of training on top of condensed dataset (if any) should be made clearer.

Author response:
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.

Reviewer comment (Novelty):
Sec. 1 states that this work is "motivated by the challenge of behaviour distillation", while this (challenge) is not a clear motivation for developing the behaviour distillation. What is the advantages of using the distilled dataset in RL except for fast training? In my opinion, the networks in RL are often small and do not require long time training.

Author response:
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].

Reviewer comment (Experiments):
While I am not an expert in RL, I note that the authors employ different networks to run HaDES and the vanilla ES baseline due to "memory constraints" so that HaDES outperforms vanilla ES, while synthetic datasets often largely underperform the real data in dataset distillation. It will be more reasonable to compare HaDES and vanilla ES under the same experimental settings. Moreover, it will make this work more convincing to add the results of other RL methods instead of only ES.

Author response:
**A3:**

Those experiments are not missing; we are unable to run ES with `population_size=2048` and `width=512` on our hardware (8x Tesla V100-SXM2-32GB). As we explain in Section 5.1, we get an Out Of Memory error because the entire population must be located on a single GPU at once before running each inner loop. We tried ES with `(population_size=1024, width=512)` and `(population_size=2048, width=256)`, and reported the latter because it performs better.

In contrast, HaDES generates a population of datasets, which have much fewer parameters than policy networks. It then distributes the datasets and runs the inner loops on multiple GPUs. The policies are initialised in the inner loops and are never colocated on a single GPU, so HaDES has a much lower peak memory usage. This means we can run bigger populations and bigger networks using the same hardware. 

**A4:**

This is indeed interesting, and explaining it fully would require an in-depth analysis of learning dynamics that goes beyond the scope of the paper. However, we hyposize two contributing factors:
1. ES tends to scale poorly with the number of parameters being optimized (See section 3.2 of [3], and theoretical result in [4]). By evolving datasets rather than networks, HaDES optimises a much smaller number of parameters. For instance, in `halfcheetah`, each network has $396834$ parameters, while the dataset of size $64$ has only $16704$ parameters.
2. HaDES performs optimization in dataset-space, not policy-space. It is possible that the optimization landscape in the dataset-space is easier, or leads to solutions not easily discovered by RL, e.g. the exploit that HaDES discovers in `humanoid-standup`.

Reviewer comment (Experiments):
There lack of experiments investigating the influence of the distill budget (the size of the distlled dataset) on the final performance.

Author response:
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.

Reviewer comment (Evaluation):
It will make this work more comprehensive if there is an analysis of efficiency. In detail, the author can list the GPU time of behaviour distillation, training on distilled dataset, vanilla ES and other RL methods, which can further highlight the importance of behaviour distillation.

Author response:
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.  

**A6:**

Looking at Table 4, performing BD using HaDES is indeed slower than training a policy with ES. However, once we are done, the time required to train a policy using the final distilled dataset is equal to just 1 HaDES generation. In short:
- Avg. time to train 1 policy with ES: `num_gen * 4.46s` (8 GPUs)
- Avg. time to perform behaviour distillation with HaDES-R: `num_gen * 7.98s` (8 GPUs)
- Avg. time to train 1 policy on the final dataset:  `7.98s` (1 GPU)

Consider `halfcheetah`. ES took 2000gen * 10.6s/gen = 5.94h to train 1 policy. HaDES-R took 2000gen * 16.9s/gen = 9.39h to do BD and produce a synthetic dataset. But then training additional policies using the synthetic dataset takes only as much as a single HaDES-R generation, i.e. 16.9s. This represents a **1266x speedup** with respect to ES.

This also holds for dataset distillation: training a classifier on MNIST or CIFAR-10 data takes a few minutes. Doing dataset distillation requires 1 to 4 hours on 1 to 4 V100 GPUs [6], but training a classifier on the distilled dataset takes seconds.

Reviewer comment (Presentation):
The algorithm 1 should be placed in the main text for clear illustration.

Author response:
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.

Reviewer comment (Novelty):
Although it is interesting, the proposed behaviour distillation seems to not have a clear motivation on why it can be useful (what's the motivation for proposing this problem and its potential application, besides DD hasn't been applied to RL), and why not directly formulating the problem on expert dataset. Distilling directly from scratch can make the problem a lot harder.

Author response:
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.

Reviewer comment (Novelty):
One missing citation on BPTT [2]. It would be nice to add some discussion on optimization algorithm in DD, for example BPTT (the original one[1] and momentum-based[2]).

Author response:
> _"One missing citation on BPTT."_

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

> _"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.