Behaviour Distillation
Authors: Andrei Lupu
Chris Lu
Jarek Luca Liesen
Robert Tjarko Lange
Jakob Nicolaus Foerster
Venue: ICLR 2024 poster (2024)
Keywords: knowledge distillation
evolutionary strategies
reinforcement learning
dataset distillation
Primary area: reinforcement learning
TL;DR: We train RL policies by evolving a small synthetic dataset of state-action pairs and doing behaviour cloning on the dataset.

Abstract: Dataset distillation aims to condense large datasets into a small number of synthetic examples that can be used as drop-in replacements when training new models. It has applications to interpretability, neural architecture search, privacy, and continual learning. Despite strong successes in supervised domains, such methods have not yet been extended to reinforcement learning, where the lack of a fixed dataset renders most distillation methods unusable.
Filling the gap, we formalize $\textit{behaviour distillation}$, a setting that aims to discover and then condense the information required for training an expert policy into a synthetic dataset of state-action pairs, $\textit{without access to expert data}$. 
We then introduce Hallucinating Datasets with Evolution Strategies (HaDES), a method for behaviour distillation that can discover datasets of $\textit{just four}$ state-action pairs which, under supervised learning, train agents to competitive performance levels in continuous control tasks.
We show that these datasets generalize out of distribution to training policies with a wide range of architectures and hyperparameters. We also demonstrate application to a downstream task, namely training multi-task agents in a zero-shot fashion.
Beyond behaviour distillation, HaDES provides significant improvements in neuroevolution for RL over previous approaches and achieves SoTA results on one standard supervised dataset distillation task. Finally, we show that visualizing the synthetic datasets can provide human-interpretable task insights.

Published as a conference paper at ICLR 2024
BEHAVIOUR DISTILLATION
Andrei Lupu'?, Chris Lu', Jarek Liesen', Robert Tjarko Lange' & Jakob Foerster' 'University of Oxford 2Meta AI 'Technical University Berlin Correspondence at a wupummeta........... com
ABSTRACT
Dataset distillation aims to condense large datasets into a small number of syn- thetic examples that can be used drop-in replacements when training new mod- els. It has applications to interpretability, neural architecture search, privacy, and continual learning. Despite strong successes in supervised domains, such methods have not yet been extended to reinforcement learning, where the lack of a fixed dataset renders most distillation methods unusable. Filling the gap, we formalize behaviour distillation, a setting that aims to discover and then condense the infor- mation required for training an expert policy into synthetic dataset of state-action pairs, without access to expert data. We then introduce Hallucinating Datasets with Evolution Strategies (HaDES), method for behaviour distillation that can discover datasets of just four state-action pairs which, under supervised learning, train agents to competitive performance levels continuous control tasks. We show that these datasets generalize out of distribution to training policies with wide range of architectures and hyperparameters. We also demonstrate applica- tion downstream task, namely training multi-task agents in zero-shot fash- ion. Beyond behaviour distillation, HaDES provides significant improvements neurocvolution for RL over previous approaches and achieves SoTA results on one standard supervised dataset distillation task. Finally, we show that visualizing the synthetic datasets can provide uuman-interpretable task insights.
INTRODUCTION
Dataset distillation (Wang et al. 2018) is the task of synthesizing a small number of datapoints that can replace training on large real datasets for downstream tasks. Not only scientific cu- riosity, distilled datasets have seen applications to core research endeavours such as interpretability.
architecture search, privacy, and continual learning (Lei & Tao 2023) (Sachdeva McAuley 2023). Despite series of successes vision tasks, and more recently in graph learning (Jin et al. (021)) and recommender systems (Sachdeva et al. 20222). ,, distillation methods have not yet been extended to reinforcement learning (RL). This is because they generally make strong assumptions about prior availability of an expert (or ground truth) dataset.
To address this gap in the literature, introduce aims to discover and condense the information required for training an expert policy into a syn- thetic dataset of state-action pairs, without access to an expert. Unlike dataset distillation, which simply replaces a hard supervised learning task by an easier one, behaviour distillation solves two challenges at once: the exploration problem (discovering trajectories with high expected return) and the representation learning problem (learning to represent policy that produces those trajectories), both of which are fundamental to deep reinforcement learning.
setting called behaviour distillation that
Thus, behaviour distillation aims to produce a dataset that obviates the need for exploration, essen- tially "pre-solving" the environment. As such, a behaviour distillation dataset does not encode summary of the full environment, but only summary of an expert policy in that environment. In other words, it reduces the joint problems of data collection (i.e. exploration) and sequential learning on a large amount of non-stationary data to one of supervised learning tiny amount of stationary non-sequential synthetic data, such as the example datasets Fig.
'The was also recently used by Furuta al. near-synonym for policy distillation.

---

Published as a conference paper at ICLR 2024
00301 +0
6422n -1
337
Figure 1: Entire synthetic datasets required to train optimal Cartpole policy (top) and an expert Hopper policy with behaviour cloning (bottom). The state-action pairs help interpret the learned policies. Red box contains observation features for Cartpole and action labels (torques) for Hopper.
Motivated by the challenge of behaviour distillation, we introduce Hallucinating Datasets with Evolution Strategies (HaDES), a method based on a meta-evolutionary outer loop. Specifically, HaDES optimizes the synthetic dataser using bi-level optimization structure, which uses evolu- tionary strategies (ES) to update the datasets on the outer loop and supervised learning ("behaviour cloning") on the current dataset in the inner loop. The fimess function for ES is the performance of the policy at the end of the supervised learning step. We show that the generated datasets can be used retrain policies with vastly different architectures and hyperparameters from those used to produce the datasets and achieve competitive returns to training directly on the original environment while doing behaviour cloning on less than 1/10th or in some cases less than 1/100th of a single episode worth of data. We also demonstrate the applicability of these datasets to downstream tasks and open- open-source them in the hope of accelerating future research. Furthermore, while HaDES is tailored to behaviour distillation, we show it is also competitive when applied to popular computer vision dataset distillation benchmarks.
There is recent resurgence of interest in evolutionary strategies (ES) for machine learning, fuelled by their generality and applicability to non-differentiable objectives, as well as to long-horizon tasks with delayed rewards Lu et al. (2023); Salimans et al. (2017). However, current evolutionary opti- mization methods are limited in the number of parameters that can be evolved, since a large number show
of parameters combined with large population size induce large memory footprint, in Section 5.1 This limits the use of ES for training large neural networks.
To tackle this issue, adapt HaDES into alternative parametrization and training scheme for neuroevolution by not resampling the initial weights of the inner loop policy. This parametrization has the benefit of scaling independently of the number of parameters in the evolved policy, thereby reducing the memory footprint and resulting in competitive performance across multiple environ- ments when compared vanilla ES.
Our main contributions are:
1. We formalize the setting of behaviour distillation, which extends the principles of dataset distillation to reinforcement Searning((Section 4.1):
2. We introduce HaDES, the first method for behaviour distillation (Section 4):
3. We show that a minor change to our method provides a parametrization for neuroevolution through ES that reduces its memory (Section 5.1).
4. We demonstrate empirically that HaDES can produce effective synthetic datasets for chal- lenging discrete and continuous control environments that generalize to training policies with large range of architectures and hyperparameters (Section 5.2);
5. We use the synthetic datasets for a downstream task: quickly training multi-task agent from datasets produced for individual environments (Section 5.3):
6. We achieve SoTA for common dataset distillation benchmark with HaDES (Section 5.4)
7. We open-source our code and synthetic datasets under https:: http:::///hthii com/
FLA TTAAAOoO//eeaavvor--tisttt
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORKS
2.1 DATASET DISTILLATION
Efforts to reduce the amount of training data required for machine learning can be traced back to reduced support vector machines (Lee & Mangasarian] 2001] Lee & Huang 2007] Wang et al.) 2005). In deep learning. [Bache et al. (2017) and Coleman et al. (2019) have called coreset selection the problem of selecting a small number of representative examples from a dataset that suffice to train model faster without degrading its final performance.
Wang et al. forgo this restriction to real examples in favour of producing synthetic datasets, coining the task of dataset distillation. Since then, numerous approaches have been proposed distill super- vised learning datasets (Lei & Tao (023). Most involve a bi-level optimization procedure and can be divided into four big categories (Sachdeva & McAuley 2023). Gradient matching methods (Zhao iet al. 2020) aim to minimize the difference in gradient updates that a model receives when training on the synthetic the distance between checkpoint parameters of models trained on either dataset. Unfortunately, nei- ther of these techniques is applicable to reinforcement learning without prior access to an expert policy, its checkpoints or at the very least a dataset of expert trajectories. More recently, Zhao & BBlen] and Wang al.] directly align the synthetic dataset distribution to match the real one. While highly effective, such approach is also inapplicable to reinforcement learning due to the non- stationary and policy-specific data distribution. The oldest and most closely related approach to our method is meta-model matching (Wang et al.) 2018 Nguyen et al. 2020 Loo et al. 2022), which involves fully training a model on the synthetic data in the inner loop, while updating the dataset in the outer loop to minimize the model's loss the real data. These works either compute expen- sive meta-gradients using back- sneek-roppggttion through time (Wang et al. (2018); [eng & Russakovsky (2022), BPTT), use neural tangent kernel rather than a finite width neural network in the inner loop such that they can compute the classifier loss on the target dataset in closed form (Nguyen et al., 2020 While these methods could be applied RL by choosing an appropriate loss (e.g. REIN- FORCE (Williams) 1992)), we instead replace meta- gradients by an evolutionary approach in the outer loop, making the cost of the outer loop updates independent of both the network size and the number of updates in the inner loop. This is important since we can use hundred of policy updates in the inner loop in practice, making the use of BPTT prohibitively expensive.
the real dataset, while trajectory matching (Cazenavette et al.) 2022) minimizes
A few other works have extended dataset distillation beyond image classification to graphs (Jin et al., 2021] 2022) and recommender systems (Sachdeva et al.] 2022), but to the best of our knowledge no previous work has broken away from assuming access to pre-existing target dataset. As such, our work is the first to break the data-centric paradigm and introduces the first general distillation
method applicable to distillation in reinforcement learning.
2.2 NEUROEVOLUTION AND INDIRECT ENCODINGS
Neuroevolution (Schwefel 1977) has been shown to perform comparably to reinforcement learning on several benchmarks (Such et al. 2017; Salimans et al. 2017). Part of our work can be viewed as a form of indirect encoding (Stanley et al. 2019) for neuroevolution an alternative parameteri- zation for evolving neural network weights. Rather than evolve the parameters of neural network directly, indirect encoding evolves a "genotype" in an alternative representation (which is usually compressed) that then maps the parameters. A well-known example is HyperNEAT (Stanley al. 2009) precursor to HyperNetworks (Ha et al.) 2016). which evolves smaller neural network to generate the weights of a larger one. Indirect encoding is desirable because evolution strategies can scale poorly in the number of parameters (Hansen] 2016). Our work, instead of evolving neural network, evolves small dataset on which we train larger neural network with supervised learning-
Other related work has also evolved other aspects of reinforcement learning training. For example, other works have evolved RL policy objectives (Lu et al.) 2022 Co-Reyes et al. 2021 Houthooft et al. 2018 Jackson et al. 2024) and environment features (Lu et al. (003).. Most related to work is Synthetic Environments (Ferreira et al., 2022). which evolve neural networks to replace an environment's 's state dynamics and rewards to speed up training. Instead of evolving transition and reward functions and training with RL, our work evolves supervised data for behavioural cloning (BC). This greatly aids the interpretability and simplifies the inner-loop.
3

---

Published as a conference paper at ICLR 2024
3 BACKGROUND
The goal of this paper is to discover behavioural dataset which, in combination with supervised learning, solves Markov Decision Processes. We describe and formalize the conceptual basis of our paper below.
3.1 REINFORCEMENT LEARNING
in which S,A,P, R, define the state space, action space, transition probability function (which maps from a state and action distribution over the next state), reward function (which maps from a state and action to scalar return), and discount factor, respectively. At each step t, the agent observes a state St and uses its policy Ty function from states to actions parametrized by to select an action a,- The environment then samples a new state 81+1 according to the transition function P and scalar reward r, according to the reward function R. The objective reinforcement learning discover policy that maximizes the expected discounted sum of rewards:
A Markov Decision Process (MDP) is defined by tuple (S,A,P, R,7)
J(0) EN
Σx' 1=
3.2 EVOLUTION STRATEGIES
Many reinforcement learning algorithms use the structure of the MDP to update the policy using gradient-based methods and techniques such as the Bellman equation (Bellman) (966). An alter- native approach is to treat the function (0) as a blackbox function and directly optimize 0. One popular approach to this is known as Evolution Strategies (Slimaan et al. (2017). Given an arbitrary function F()). ES optimizes the following smoothed objective:
E (F + ac)]
where NO0, 7) is the standard multivariate normal distribution and is its standard deviation. We estimate the gradient F by sampling noise from N(0, 7) and evaluating F at the resulting points. Specifically, the gradient is estimated by:
VoE. [F(ó + oe)] = E. (.-.(
[ F( +ot)
We then apply this update to our parameters and repeat the process. When applied to meta- optimization, ES allows us optimize functions that would otherwise require taking meta-gradients through hundreds or thousands of update steps, which is often intractable Metz et al. (2021).
3.3 DATASET DISTILLATION
In the context of supervised learning, dataset distillation is the task of generating selecting proxy dataset that allows training a model to a similar performance as training on the original dataset, often using several orders of magnitude fewer samples. Formally,
assume dataset D {2,, y/ / of which Draain C D is the training (sub)set, and a training algorithm alg. Also, let falg(D) 2; + y: be the classifier obtained by training D with alg. Then, the dataset distillation objective is find a synthetic dataset Dø, [Dol
<< (Dirain) such that
v) = (e (x).
where & indicates that the synthetic dataset can be parametrized and learned, rather than being sampled. In practice, [D&l | is often set to fixed number of examples, e.g. Daol = n/YY,. where |Y| is the total number of discrete classes, or to be determined by fixed number of parameters, i.e. [0] = n[[æ| + lyl). The latter formulation, being more permissive, admits factorized representations of the data, c.g. by representing (Dol with generative neural network. While this is promising avenue for future work, we focus exclusively on non-factir factoredd distillation, which allows for better interpretability and tractability of the synthetic dataset.
4

---

Published as a conference paper at ICLR 2024
4 PROBLEM SETTING AND METHOD
4.1 BEHAVIOUR DISTILLATION
We introduce behaviour distillation as a parallel to dataset distillation. Rather than optimizing proxy dataset to minimize a loss in supervised learning, behaviour distillation optimizes a proxy dataset that maximizes the discounted return in reinforcement learning, after supervised learning. More formally, the aim of behavioural distillation is to find dataset D3 E (S x AA), where N is number of points in the dataset, to solve the following bi-level optimization problem:
max J( (Dx))
(3)
D3
s.t. 0 (D arg min L(6, Do),
(4)
0
where 0 are the parameters of the neural network used in the policy wa, J is the discounted sum of returns defined in Eq. and I. is any supervised learning loss.
Crucially, this setting does not assume access expert policy dataset, for two reasons. Firstly, expert data may not always be available or may not be easily compressible, for instance if the expert erratic idiosyncratic. Secondly, standard imitation learning is plagued by cascading errors: if ı is an imitation learning policy that deviates from some expert Texpert with probability < e, then it is likely to end up off-distribution, further increasing its error rate. As such, it generally incurs a regret J (TTexp)))
(Tx 2 T where T is the episode horizon (Ross & Bagnell 2010). Given dataset distillation lossy, applying it naively would result in a large and therefore a poor performance for æ, Since we ultimately care about maximizing the expected discounted return rather than reproducing some specific expert behaviour, this is how we formalize behaviour distillation.
4.2 HADES
To tackle behaviour distillation, we introduce method "Hallucinating Datasets with Evolution Strategies" (HaDES). HaDES optimizes the inner loop objective to obtain 6 using gradient descent, and optimizes the outer loop objective (the return of m&. using ES (Salimans et al. 2017).
In the inner loop, HaDES uses the cross-entropy loss for discrete action spaces, and the negative log likelihood loss for continuous action, albeit other losses could be substituted as well. We provide pseudocode in Algorithm
4.2. POLICY INITIALIZATION
We further specify two variants of our methods, which have distinct use cases. They differ only in the way inner loop policies are initialized, which leads different inductive biases. The variants described here are visualized in Fig.
The first variant is HaDES with fixed policy initialization, or HaDES-F. In this variant, sample single policy initialization dg at the very beginning of meta-training. The policy is re-trained every lnner-loop, but always starting from this fixed fg
The second variant is HaDES with randomized policy initialization, or HaDES-R. In this variant, we use multiple (k 2) policy initializations (f)...., 0%); in the inner loop, and we resample the initializations randomly at every generation i.
HaDES-F has stronger inductive bias in that is only optimizing D& for single initialization. We expect it might be able to "overfit" on that initialization and achieve higher policy returns, but at the cost of the synthetic dataset having poor generalization properties to other initializations. HaDES- F will therefore be stronger for neuroevolution (where only the final return of the specific policy matters), but weaker for behaviour distillation.
HaDES-R has a weaker inductive bias in that it optimizes D, for range of initializations. We expect this to result in decreased policy returns, but the synthetic dataset to be a useful artifact of training that generalizes will therefore be a better choice for behaviour distillation, but a weaker one for neuroevolution. We confirm both those intuitions empirically in Section 5
unseen initializations
unseen policy architectures. HaDES-R
5

---

Published as a conference paper at ICLR 2024
Wil T6p
-T.
eval
Env
Iain
evall
Env
tan
eval
Env
J
ES
J
ES
Figure 2: Left: Standard neuroevolution. Middle: HaDES-F: Right: HaDES-R. HaDES-F uses
single fixed policy initialization. HaDES-R samples k 2 2 policy initializations every generation.
EXPERIMENTS
Next, we show that HaDES can be successfully applied to broad range of settings. More specifi- cally, we investigate the following four applications:
1. We demonstrate the effectiveness of HaDES indirect encoding for neuroevolution in RI. tasks (Section 5.1), with an analysis of the distillation budget Appendix C.I
2. We show that synthetic datasets successfully generalize across a broad range of architec- tures and hyperparameters (Section 5.2).
3. We show that the synthetic datasets be used train multi-task models (Section 5.3)
4. We show that while focus RI,, method is also competitive when applied to dataset distillation in supervised settings (Section 5.4
We refer the reader to Appendix B.2 for experimental details, including runtime comparisons.
5.1 PERFORMANCE EVALUATION OF HADES-DISCOVERES DISCOVEEED RL DATASETS
We first test the effectiveness of HaDES way of training competitive policies in 8 continuous control environments from the Brax suite. Fig. 3a shows the performance HaDES with a fixed policy initialization (HaDES-F), HaDES with two randomized initializations (HaDES-R), and direct neuroevolution through ES. HaDES-F achieves the highest return across the board, while HaDES-R R method discov-
also matches or beats the baseline in 6/8 environments. In Hlmanoidd-Sanndie ers glitch in the collision physics, propelling itself into the air to achieve extremely high return
In MinAtar, find that HaDES-F outperforms the ES baseline in Breakout, and matches it in SpaceInvaders and Freeway. We hypothesize MinAtar is a harder setting for our method due to the symbolic (rather than continuous) nature of the environment. For both settings, we also plot the performance of PPO policies after 5 x 107 training steps. While there is still a performance discrepancy between ES and RL, HaDES narrows the gap significantly.
We use policy networks with width 512 and population size of 2048 for HaDES, but are forced to cut the network widths by half for the ES baseline on MinAtar due to memory constraints. Indeed, ES requires the entire population to be allocated on single device at once when new population is generated and when estimating the meta-gradient at each generation. While distributed approaches to the outer loop computation are feasible, they involve an engineering overhead, which our method alleviates. As such, our method is drastically more memory efficient, enabling larger populations and network sizes with minimal code changes. We also run HaDES with width 256 for completeness.
5.2 HADES DATASETS GENERALIZE ACROSS ARCHITECTRRES & HYPERPARAMETERS
We now turn our attention to the datasets themselves, and use them to train new policies in zero- shot fashion, i.e. without additional environment interactions. take two synthetic datasets for Hopper
generated with HaDES-R and another with HaDES-F and use them to train new policies from scratch. For each dataset, and for each of 7 different policy network sizes, we train 50 policies, each using different learning rate and number of training epochs. In particular, the learning rates span 3 orders of magnitude and training epochs ranged uniformly between 100 and
video link

---

Published as a conference paper at ICLR 2024
higpet
walford 24
Inverted. doot pard.lum
2300
15000
2000
8280
2000
(030
4200
-530
500
2305
2030
500
1000
15:00
2000
500
1930
1500
2000
503
3200
1530
2003
523
1000
3505
2030
humanoil
sumnnoistandap
5000
6000
250000
4.000
50.00
4030
330000
2000
2000
220000
1000
SEC 1000
15/00 2000
520
10:00
1500
2000
503
3000
1530
2003
523
1666
3505
2030
Careredora
Goneration
Gane witor
Generations
FPO
25 rearsoeouusnn
14025-3
(a)
spocerwedes Minttar
Brokkott Wirdtar
Asteric MuNTE
Frewway ralaa
150
50
10
2003
2000
3303
4006
5300
2003
2000 3003 4000 5300
2382 2006 3002 4000 5300
2303
2006 3003
4000 5300
Generatiora
Generatior
Figure 3: HaDES trains competitive policies on a) Brax, using 64 state-action pairs and b) MinAtar using 16 state-action pairs. For each environment, we show the mcan return of the population at each generation for HaDES-F, HaDES-R and direct neuroevolution through ES, well as the PPO final performance after 5 x 10" steps. HaDES-F matches or outperforms direct ES on all Brax environments, outperforms ES in one out of four MinAtar environments and matches it in two others. We also observe a significant gap between HaDES-F and HaDES-R, as predicted in /Section 4.2
INOOS!
2800
254
X12 2025 2054
Figure 4: Hopper dataset transfer to other architectures and training parameters. We take synthetic dataset of 64 state-action pairs evolved for policy networks of size 512 (highlighted) and use it to train policies with varying widths and 50 hyperparameter combinations per width. We plot the top 50% within each width group. HaDES-F indicates that the dataset was trained with fixed mo- The HaDES-R dataset was trained with randomized (m)...., = and generalizes much better across all architectures and training parameters. This holds generally across environments (see Appendix C.2)
500. For each dataset and width, we discard the worst 25 policies and plot the return distribution of the remaining 25 in Fig. 4
The datasets were evolved for policies of width 512, a fixed learning rate, and a fixed number of epochs, but readily generalize out of distribution training with different settings and architectures. In particular, we see that the HaDES-R R dataset is more robust to changes in both policy architecture and training parameters than the HaDES-F dataset, which incorporates stronger inductive bias.
We hypothesize that the generalization properties of the synthetic datasets can further be improved by randomizing not only the policy initialization, as in HaDES-R, but also architectures and training
7

---

Published as a conference paper at ICLR 2024
parameters, further reducing some of the inductive biases present in our implementation. How to best navigate the trade-off between generalization, dataset size and policy performance remains an interesting question for future work.
5.3
HADES DATASETS CAN BE APPLIED TO ZERO-SHOT MULTI-TASKING
1.3
/ 1.25
Nomaltzes Rgtum
0.8
Normaltoos Rebor
1.00
0.6
wrng merged
0.75
wrong merged
0.4
0.50
0.25
halfchootah
hapser
humanaid
fumrain dstandup
Envionment
Environment
Figure 5: We use the synthetic datasets to train multi-task agents without any additional environment interaction. We plot the normalized fitness of agents trained either on the correct dataset for their environment, the wrong dataset for their environment, catenation and zero-padding to have the observation sizes match. Left: we train multi-task agents that achieve 50% normalized fitness for Halfcheetah and Hopper. Right: train agents that achieve ž 100% normalized fitness for Humanoid and Humanoidstandup. The multi-task policy architecture and training parameters were not optimized. We plot mean + stderr. across 10 seeds. This shows that synthetic datasets can accelerate future research on RL. foundation models.
combined dataset, merged through con-
We showcase an application of the distilled dataset by training multi-task Brax agents in a zero-shot fashion. To this end, we merge datasets for two different environments in the following way. Let Di = (S), and 'D2 = AS) be the two datasets. each with 64 state-action pairs. We first zero pad the states in Sj to the left and zero pad the states in S the right, such that each padded state has now the combined dimension |s1| [82). where S₁ Rl and S2 C R |2 We then do the
same for the actions, such that each padded action now has size |a1| |a2), where A1 C Rlosl and |aa)
A2 C. R I Finally, take the union of both padded datasets build merged dataset D merged with 128 state-action pairs. We then train agents on either D1, D2 or D mergedd using behaviour cloning and report the normalized performance on both environments Fig.
We perform this experiment for two pairs of environments: (Halfcheetah, Hopper) and (Humanoid, Humanoid Standup). Blue indicates the baseline performance of training on D; and evaluating on environment i. Red shows the performance of training on D; and evaluating on environment -i, and is a loose proxy for how much the data from one environment helps learn a policy for the other. Finally, green shows the performance of policies trained on Dmerged- We see that the multi-task agents achieve roughly 50% of the single-task performance in the first pair of environments, see no loss in performance in the second pair.
This shows that the synthetic datasets evolved by HaDES can vastly accelerate future research on RL foundation models. Indeed, given those datasets, training new models takes only a few seconds, and makes it possible to experiment with architectures and multi-task representation learning at a fraction of the original computational cost. Furthermore, it allows studying the properties of cross-task parameter sharing and representation learning in isolation, separately from the exploration issues inherent to reinforcement learning.
5.4 HADES CAN BE APPLIED TO SUPERVISED DATASET DISTILLATION
While the focus is on behaviour distillation RI,, our method readily applicable the standard dataset distillation setting by replacing environment return with cross-entropy loss on some target dataset. We apply our method to dataset distillation with image/class We use the cross-entropy loss on the training set as the fitness for
MNIST and FashionM-
NIST. with results in Table I HaDES and report the mean accuracy on the test set obtained by training classifiers on the final synthetic dataset. We run different seeds and train 20 classifiers for each final datasets. Similar to ZZhao et al. (2020) and Zhao & Bilen (202T), we report the mean and standard deviation across all 60 final classifiers and compare against the best method in these settings, namely RFAD Loo et al.
8

---

Published as a conference paper at ICLR 2024
MNIST
FNShionMNIST
RFAD
Ours
RFAD 78.6 1.3
Ours
img/els
94.4 + 1,5
90.1 + 0.3
80.2 + 0.4
Test set accuracy of classifiers trained on datasets composed of image per class. We com- pare to RFAD, which is the SotA for non-factorized dataset distillation on these datasets (Sachdeva) & McAuley 2023). RFAD uses a ConvNet architecture for testing, while we use smaller CNN. Despite being designed for RL, our method is also competitive for image-based dataset distillation and achieves state-of-the-ar distillation for 10-image FashionMNISTT
Table 1:
(022). We find that HaDES performs competitively in 10 image MNIST and achieves state-of-the- art results in FashionMNIST. However, we were unable to scale our methods to CIFAR-10, which have many more parameters due to being RGB rather than greyscale.
In tuning method, find that the two single most important parameters the outer learning rate and the dataset initialization. We find that initializing the dataset at the class-wise mean of the data works best. This is similar to findings by Zhao & Bilen who also find that warm starting the dataset performs better than initializing from scratch.
5.5 HADES-EXPLAINABLLE
The final benefit of the HaDES datasets is that the synthetic examples lend themselves to inter- pretability. Going back Fig. I that the resulting datasets have intuitive properties. For instance, the two state dataset for Cartpole captures that the policy should go left if the pole is leaning left and go right otherwise.
Performing an explainable RI. study lies beyond the scope of this paper, but in a critical analysis, Atrey et al. highlight the importance of taking hypothesis-drived drriven approach to explaining deep RL. policies, formulating possible explanations, then testing them rigorously with ablations and careful experiments. With that mind, we argue that our synthetic datasets an effective starting point for such hypothesis- tepptingss-testngg.. for instance by applying transformations to datasets and observing how it impacts the trained policies.
6 DISCUSSION AND CONCLUSION
In this paper, we introduced a new parametrization for policy neuroevolution by evolving small synthetic dataset and training a policy on through behavioural cloning. We showed that our method produces policies with competitive return in continuous control and discrete tasks. Our method can be used for behaviour distillation, summarizing all relevant information about optimal behaviour in an environment into small synthetic dataset. This "behaviour floppy disk" quickly train new policies parametrized by a range of different architectures. We then demonstrated the utility of the distilled dataset by training multi-task models. We finished by showing that although our focus is RL, our method also applies to vanilla dataset distillation in supervised learning, where we achieved state-of-the-arr in one settings.
The main limitation of this work is of computational nature, since evolutionary methods require a large population to be effective. Furthermore, while our alternative parameterization enables to evolve larger neural networks than standard neuroevolution, the number of parameters still grows linearly with the number of datapoints, especially in pixel-based environments which tend to be very highly dimensional. Tackling this issue, for instance by employing factorized distillation, is therefore a promising avenue for future work. Another downside of our work is the number of hyperparameters, since we need to tune both the ES parameters in the outer loop and the supervised learning ones in the inner loop. However, anecdotal evidence seems to indicate that ES can adapt to the inner loop parameters, for instance by increasing the magnitude of the dataset if the learning rate low. Understanding the interplay between parameters would allow for faster an better tuning. A related approach would also be to evolve the inner loop parameters along with the dataset. Finally, possible applications of the distilled datasets are ripe for investigation, for instance in continual or life-long learning, and regularizing datasets further promote interpretability.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
To encourage reproducibility, we described the method in detail in Section 4.2) include pseudocode Algorithm 1) and provide hyperparameters in Appendix B.1 in format that corresponds directly to the configs used by our code. We also open-source our code and our synthetic datasets at https: 7/githab. com/ Coon/FLA FLA Cwom///AAOx////eeee
riitrriddssilllooo lationi
ACKNOWLEDGMENTS
Andrei Lupu was partially funded by Fonds de recherche du Québec doctoral training scholarship.
REFERENCES
Akanksha Atrey. Kaleigh Clary, and David Jensen. Exploratory not explanatory: Counterfactual analysis of saliency maps for deep reinforcement learning. arXiv preprint arXiv:1912.05743, 2019.
Olivier Bachem, Mario Lucic, and Andreas Krause. Practical coreset constructions for machine learning. arXiv preprint
2017.
Richard Bellman. Dynamic programming. Science, 153(3731):34 1966.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. compo///ecco/goooge/////aaaa
George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A Efros, and Jun- Yan Zhu. Dataset distillation by matching training trajectories. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4750-4759, 2022.
John D Co-Reyes, Yingjie Miao, Daiyi Peng, Esteban Real, Sergey Levine, Quoe V Le, Honglak Lee, and Aleksandra Faust. Evolving reinforcement learning algorithms. arXiv preprint arXiv:2101.03955 2021.
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning. arXiv preprint a906 /829, 2019.
Zhiwei Deng and Olga Russakovsky. Remember the past: Distilling datasets into addressable mem- ories for neural networks. Advances in Neural Information Processing Systems, 35:34391-34404, 2022.
Fabio Ferreira, Thomas Nierhoff, Andreas Saelinger, and Frank Hutter. Learning synthetic environ- ments and reward networks for reinforcement learning. arXiv preprint wtx.ix:2222.0279000 2022.
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax differentiable physics engine for large scale rigid body simulation, 2021. URL http:
//github. com/ http///ooollaaoooooo
Hiroki Furuta, Yusuke Iwasawa, Yutaka Matsuo, and Shixiang Shane Gu. A system for morphology-task generalization via unified representation and behavior distillation. arXiv preprint ar 14296, 2022.
David Ha, Andrew M Dai, and Quoe V Le. Hypernetworks. In International Conference on Learn- ing Representations, 2016.
Nikolaus Hansen. The ema evolution strategy: A tutorial. arXiv preprint 1rX1v:1604.00772, 2016.
Rein Houthooft, Richard Y Chen, Phillip Isola, Bradly C Stadie, Filip Wolski, Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. arXiv preprint arXiv: 1802 0482/, 2018.
10

---

Published as a conference paper at ICLR 2024
Matthew Thomas Jackson, Chris Lu, Louis Kirsch, Robert Tjarko Lange, Shimon Whiteson, and Jakob Nicolaus Foerster. Discovering temporally-aware reinforcement learning algorithms. arXiv preprint 1aXii::24220082222 2024.
Wei Jin, Lingxiao Zhao, Shichang Zhang, Yozen Liu, Jiliang Tang, and Neil Shah. Graph conden- sation for graph neural networks. arXiv preprint arXiv:21 2021.
Wei Jin, Xianfeng Tang, Haoming Jiang, Zheng Li, Danqing Zhang, Jiliang Tang, and Bing Yin. Condensing graphs via one-step gradient matching. In Proceedings of the 28th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 720-730, 2022.
Robert Tjarko Lange. gymnax: A jax-based reinforcement learning environment library, 2022b.
URL http://github. com/RoberiTlanggegynmaa 2022.
Robert Tjarko Lange. evosax: Jax-based evolution strategies. In Proceedings of the Companion Conference on Genetic and Evolutionary Computation, pp. 659-662, 2023.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mmist/, 1998.
Yuh-Jye Lee and Su- Yun Huang. Reduced support vector machines: A statistical theory. IEEE
Transactions on neural networks, 18(1):1-13, 2007.
Yuh-Jye Lee and Olvi Mangasarian. Rsvm: Reduced support vector machines. In Proceedings of the 2001 SIAM International Conference on Data Mining, pp. 1-17. SIAM, 2001.
Shiye Lei and Dacheng Tao. A comprehensive survey dataset distillation. arXiv preprint arXiv:2301.0560) 2023.
Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus. Efficient dataset distillation using random feature approximation. Advances in Neural Information Processing Systems, 35:13877- 13891, 2022.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:16-55- 16468, 2022.
Chris Lu, Timon Willi, Alistair Letcher, and Jakob Nicolaus Foerster. Adversarial cheap talk. In International Conference on Machine Learning, pp. 22917-22941. PMLR, 2023.
Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv preprint arXiv:2111.05803 2021.
Timothy Nguyen, Zhourong Chen, and Jachoon Lee. Dataset meta-learning from kernel ridge- regression. arXiv preprint arXiv:2011.00053 2020.
Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Yee Whye Teh and Mike Titterington (eds.), Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pp. 661-668, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL https:
//proceedings mlr press/vv/ross... html
Noveen Sachdeva and Julian McAuley. arXiv:2301.04272, 2023.
Data distillation: A survey.
arXiv preprint
Noveen Sachdeva, Mehak Dhaliwal, Carole-Jean Wu, and Julian McAuley. Infinite recommendation networks: A data-centric approach. Advances in Neural Information Processing Systems, 35: 31292-31305, 2022.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as a scalable alternative reinforcement learning, arXiv preprint 003.03864000..0. 2017.
11

---

Published as a conference paper at ICLR 2024
Hans-Paul Schwefel. Evolutionsstrategienn für die numerische optimierung. In Numerische Opti- mierung von Computer-Modellee mittels der Evolutionsstrattggi,, pp. 123-176. Springer, 1977.
Kenneth O Stanley, David B D' Ambrosio, and Jason Gauci. A hypercube- based encoding for evoly- ing large-scale neural networks. Artificial life, 15(2):185-212 2009.
Kenneth Stanley, Jeff Clune, Joel Lehman, and Risto Miikkulainen. Designing neural networks through neuroevolution. Nature Machine Intelligence, 1(1):44-35 2019.
Felipe Petroski Such, Vashisht Madhavan, Edoardo Conti, Joel Lehman, Kenneth O Stanley, and Jeff Clune. Deep ncuroevolution: Genetic algorithms are a competitive alternative for training deep neural networks for reinforcement learning. arXiv preprint arXiv:17/2.06567, 2017.
Jigang Wang, Predrag Neskovic, and Leon N Cooper. Training data selection for support vector machines. In International Conference on Natural Computation, pp. 554-564. Springer, 2005.
Kai Wang, Bo Zhao, Xiangyu Peng. Zheng Zhu, Shuo Yang, Shuo Wang. Guan Huang, Hakan Bilen, Xinchao Wang, and Yang You. Cafe: Learning to condense dataset by aligning features. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 72196-12205, 2022.
Tongzhou Wang, Jun- Yan Zhu, Antonio Torralba, and Alexei Efros. Dataset distillation. arXiv preprint arXiv:1811.1090 2018
Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhuber. Natural evolution strategies. The Journal of Machine Learning Research, 15(1 (519)994,,980 2014.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229-256, 1992.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv:/708.0777 2017.
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv:1903.03176. 2019.
Bo Zhao and Hakan Bilen. Dataset condensation with differentiable siamese augmentation. In International Conference on Machine Learning, pp. 12674-12685. PMLR, 2021.
Bo Zhao and Hakan Bilen. Dataset condensation with distribution matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision pp. 6514-6523, 2023.
Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen. Dataset condensation with gradient matching. arXiv preprint arXiv:2006.05920 2020.
12

---

Published as a conference paper at ICLR 2024
A HADES ALGORITHM
Algorithm HaDES
Require: dataset size n, environment Env
Require:
number of meta steps T, pop. size P, learning rate a, std a 1: Initialize Dj {(s, (.......,,,,,)).,
e.g. randomly sampled
2: for meta_step 0, o...., do
3: Initialize { (0)
4: for = 0,... P do
5: if i is even then
6:
Perturb De with noise & & to get D; else
> antithetic noise
7:
8:
Perturb D& with noise & - to get D;
9:
Update & ~~N(( o)
10:
end if
Initialize policy me Train policy on D, using BC Unroll my and compute expected return J4 = J (TT, En, end for
12:
13:
14:
15:
Approximate J = Pr Σ. Jis Update D. Do +aV.J &J
16:
17:
end for
18:
Train policy me on final D
19: return (mg, D()
B IMPLEMENTATION DETAILS
B. HYPERPARAMETERS
Parameter
Value
LR
0.005
NUMLENVS
NUM STEPS
1024
UPDATE EPOCHS
400
MAX_GRAD GRAD_NORY
0.5
ACTIVATION
tanh
WIDTH
512
ANNEAL.LR
False
GREEDY
False
CONST NORMALIZE OBS
False
NORMALIZE OBS
True
NORMALIZE REWARD
True
popsize
2048
dataset_size
64
rollouts perppperrr.rrreeeeee candidate
1
n generations
2000
sigma_init
0.03
sigma_decay
1.0
Irate init Evo. strategy
0.05
OpenES
Table 2: Hyperparameters for HaDES in Brax. Top: inner loop parameters. Bottom: Outer loop parameters.
13

---

Published as a conference paper at ICLR 2024
Parameter
Value
NET
mlp
LR
0.03
NUM_ENVS
8
NUM JUM.STEPS
1024
UPDATE EPOCHS
64
MAX. GRAD NNOM
0.5
ACTIVATION
relu
WIDTH
512 or 256
ANNEAL_RR
True
GREEDY LACT
False
NORMAALORRA
True
NORMALIZE OBS
False
NORMALIZE REWARD
False
popsize
2048
dataset_size
16
vollouts.per.candiie
2
generations
5000
sigma init sigma_limit
0.5
0.01
sigma_decay
Irate_ine init
0.05
Irate .decay
1.0
Evo. strategy temperature
SNES
20.0
Table Hyperparameters for HaDES in MinAtar. Top: inner loop parameters. Bottom: Outer loop parameters.
B.2 EXPERIMENTAL DETAILS
For all RL tasks we use Brax (Freeman et al. 2021). ). suite of continuous control environments, and MinAtar (Young & Tian 2019) of Atari-like environments.
For dataset distillation, we report results on two image classification tasks: MNIST (LeCun 1998), which composed of handwritten digits, and FashionMNIST (Xiao al. 2017). which features different clothing items.
For the evolutionary algorithm, cation, and use SNES Wierstra et al. (014)) for MinAtar. In the inner loop, we minimize either the cross-entropy loss (discrete cases) or the negative log likelihood of the synthetic actions (continuous action cases). All of our runs use 8 Nvidia V100 GPUs and take between and 17 seconds per outer loop generation. Detailed generation times are reported Table 4 These times include outer loop operations (all methods), inner loop policy training (HaDES only), and inner loop policy evaluation (all methods). HaDES-R is slightly slower than HaDES-F since it trains two policies instead just one. Because train policies from scratch every generation, the times reported are strict upper bounds to how long takes to train a policy on the final distilled datasets.
OpenES Salimans et al. 2017 for Brax and image classifi-
In image classification and MinAtar, assign labels (i.c. classes discrete actions) uniformly, whereas Brax we evolve the dataset labels alongside the observations since the environments feature continuous actions.
We implement our algorithm in JAX (Bradbury et al. 2018 using the PureJaxRL (Lu et al. 2022), gymnax (Lange) 2022) and evosax (Lange] 2023) libraries to enable parallel training hardware accelerators. We also use virtual batch normalization (Salimans et al., 2016) to stabilize training, which was previously found to be crucial in stabilizing ES (Salimans al.) 2017)
14

---

Published as a conference paper at ICLR 2024
Environment Name
ES Neuroevolution
HaDES-F
HaDES-R
Hopper
4.0
5.6
6.3
Walker2d
3.4
7.9
8.5
Reacher
2.7
6.7
7.2
Inverted Double Pendulum
2.6
4.4
4.6
Ant
6.6
11.1
12.1
Halfcheetah
10.7
14.8
16.9
Humanoid
7.8
13.9
15.1
HumanoidStanduu
8.6
14.8
16.4
1.6
1.5
1.6
Breakout-MinAtar
1.3
1.6
1.8
Asterix-MinAtta Freeway-MinnAte
1.9
2.1
2.4
2.3
2.2
2.9
Table 4: Runtime of the different neuroevolution methods in seconds/generation. Times averaged over seeds rounded the nearest tenth of second. Standard deviation omitted, but the difference between the fastest and slowest runs for any setting is usually smaller than 0.2 seconds.
C ADDITIONAL RESULTS
C.1 IMPACT OF DISTILLATION BUDGET ON PERFORMANCE
hopper
haffchictah
humanaid
1870
4000
Maan Naturn
4000
1730
3000
2000
2000
1000
1990
236
256
64
256
Dataset Size (15,0) pains)
Dataset Size ((saa) DO 151
Dotaset Size (1s.a)
1R0
IN3ESE
Figure 6: Final return of HaDES policies as function of distillation budget (i.e. dataset size). ES neuroevolution and PPO returns also plotted for reference.
In Fig. 6
we investigate the impact the distillation budget on the final performance of policies trained with HaDES on three different environments. What observe is that dataset sizes that too small degrade performance, likely because they cannot contain all the information required to train expert policy. This is particularly noticeable in Humanoid, where for dataset of state- action pairs. the score drops as low as 1013. However, a score of 1000 corresponds to a humanoid policy that keeps its balance and stays immobile, with lower scores indicating that the policy falls and causes an early termination. This is an indication that for distillation budgets that are too low capture expert behaviour, HaDES does fail learn, and will still optimize return within the constraints of the budget.
On the opposite end of the spectrum, we also observe return dropping for large dataset sizes ( 256), despite the increased expressivity. This is possibly due to ES (and therefore HaDES) scaling poorly to large number of parameters. This problem can be alleviated by relying on better ES methods, by using a factorized approach to distillation.
C.2 DATASET GENERALIZATION ACROSS ARCHITECTURES AND HYPERPARAMETRO
Here we plot generalization plots for additional environments. As expected, HaDES-R generalizes better than HaDES-F both to new hyperparameters and new architectures.
15

---

Published as a conference paper at ICLR 2024
2300
64
120
256
512
2004
2236
of
Figure 7: Dataset transfer to hopper architecture and training parameters.
6000
HECES HNCESF
2000
2000
Pally Network witth
Figure 8: Dataset transfer to walker2d architecture and training parameters.
seacher
-503
1603
1007
2002
256
Figure Dataset transfer to reacher architecture and training parameters.
inverted, dousle_enndium
12503
INGES HICESE
10003
1502
328
256
Peliry Network
Figure 10: Dataset transfer to inverted .eribullepeduuummmmm architecture and training parameters.
HEEES
1500
Palicy Nebworc Witth
Figure 11: Dataset transfer architecture and training parameters.
16

---

Published as a conference paper at ICLR 2024
Rcheetah
.000
2002
1002
228
Policy Nctwoo's dits
Figure 12: Dataset transfer to halfcheetah architecture and training parameters.
humarolc
4003
1002
128
256
Figure 13: Dataset transfer humanoid architecture and training parameters.
3000000
INGESE
2000000
100000
Pelicy Network W dts
Figure 14: Dataset transfer to humanoidstandup architecture and training parameters.
17


Reviews:
[{"review_id": "1eOWTMihWN", "reviewer": "Reviewer_f5UH", "summary": "This paper proposes a new branch of dataset distillation - behaviour distillation. Behaviour distillation aims at distilling a set of synthetic RL dataset without having access to the expert data. The authors handles indifferentiability of the formulation using evolution strategy, and show good performance.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The introduction of behaviour distillation can enrich the literature and direction of Dataset Distillation. Different from standard DD, behaviour distillation does not require access to the expert datasets. This is quite close to a standard /basic RL setting and could be a good starting point.\n\n+ The author's writing is easy to follow and pretty clear on the technical details\n\n+ The experimental section show promising results using the proposed algorithm HADES.", "weaknesses": "- 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 it's 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.\n- It would be great if the authors can discuss the behaviour difference of a standard RL algorithm and a synthetic data-driven RL algorithm (HADES). What's the potential benefits?\n- 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]).\n\nOverall the paper is quite interesting. Looking forward to authors' response on the above questions.\n\n[1] Dataset Distillation\n\n[2] Remember the Past: Distilling Datasets into Addressable Memories for Neural Networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mQh0JXsBZR", "reviewer": "Reviewer_7Zz2", "summary": "This manuscript proposed a behavior distillation algorithm of HaDES that aims to distill a few synthetic (state, action) pairs ($\\mathcal{D}_\\phi$) in the reinforcement learning (RL) setting, and the network can fastly learn a satisfied policy by training on the small distilled dataset. Concretely, the authors firstly formulate that behavior distillation problems as a bi-level optimization where the inner loop optimizes the policy network parameters by supervised learning on distilled data, while the outer loop maximizes the cumulative reward ($J$) w.r.t. the RL task.  Then, synthetic (state, action) pairs are optimized by gradient ascent where  the gradient is estimated by evolution strategies (ES). The authors verified that distilled dataset outperform vanilla ES algorithm on multiple RL datasets including Brax and MinAtar exvironments.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper first introduces the dataset distillation into the RL setting, and the core pros of this paper are presented as below:\n\n1. Formulate the behaviour distillation by incorporating dataset distillation with an RL reward function;\n\n2. Propose a novel behaviour distillation algorithm of HaDES that learns a few (state, action) pairs to fastly train a policy network with supervised learning;\n\n3. Multiple empirical studies are conducted to verify the effectiveness and robustness of the synthetic behaviour dataset.", "weaknesses": "While the authors creatively incorporate dataset distillation into RL setting, there are some weaknesses mainly lie on the motivation and experiments, which hurt the contribution of this manuscript. \n\n**Q1:** 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. \n\n**Q2:** 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.\n\n**Q3**: There lack of experiments investigating the influence of the distill budget (the size of the distlled dataset) on the final performance.\n\n**Q4**: 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.\n\nBased on these observations, I think this manuscript is marginally below the acceptance, but I would like to increase my rating if the above questions are well addressed. Below are some minor questions and typos.\n\n---\n\nMinors and typos:\n1. The term Behavior Distillation has already been used in [1], which is borrowed from the concept of knowledge distillation instead of dataset distillation. The authors should discuss this to avoid ambiguity.\n\n2. The algorithm 1 should be placed in the main text for clear illustration.\n\n2. Page 2: \"scaling independently\" -> \"scaling independence\"\n\n3. Page 2: \"policy, reducing\" -> \"policy, thereby reducing\"\n\n4. Page 3: \"train a model faster\" -> \"train a model fastly\"\n\n5. Page 4: The formulation of dataset $\\mathcal{D}$: \"$\\mathcal{D} = \\\\{x_i, y_i\\\\}$\" -> \"$\\mathcal{D} = \\\\{x_i, y_i\\\\}_{i=1}^N$\n\n6. Page 5: this first line: \", i.e. \" -> \", i.e., \"\n\n7. Page 5: there exists an extra right bracket in Eq. (3)\n\n[1] Furuta, Hiroki, et al. \"A System for Morphology-Task Generalization via Unified Representation and Behavior Distillation.\" The Eleventh International Conference on Learning Representations. 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yjBsHWtpJL", "reviewer": "Reviewer_ZN7t", "summary": "The authors attempt to transfer the concept of dataset distillation to the realm of reinforcement learning. To that end, they introduce the concept of behaviour distillation, which is trying to condense a dataset of state-action pairs required for training a RL agent. After introducing this concept, they propose a new method named HADES, that implements this concept and is able to distill the initial state-actions pairs into a dataset of just four state-action pair that, when trained upon by an agent can make it reach competitive performance.", "questions": "- 1) 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 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.\n- 2) In the results plot of reinforcement learning, does the number of generation directly correlates 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", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Transfer the concept of dataset distillation to reinforcement learning\n- The condensed dataset can result in a quite small one\n- The results on supervised classification dataset distillation also seem good", "weaknesses": "- 1) 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\n- 2) 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).\n- 3) 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)