Quality Diversity Imitation Learning

## Abstract

Abstract Learning diverse and high-performance behaviors from a limited set of demonstrations is a grand challenge. Traditional imitation learning methods usually fail in this task because most of them are designed to learn one specific behavior even with multiple demonstrations. Therefore, novel techniques for quality diversity imitation learning are needed to solve the above challenge. This work introduces Wasserstein Quality Diversity Imitation Learning (WQDIL), which 1) improves the stability of imitation learning in the quality diversity setting with latent adversarial training based on a Wasserstein Auto-Encoder (WAE), and 2) mitigates a behavior-overfitting issue using a measure-conditioned reward function with a single-step archive exploration bonus.
Empirically, our method significantly outperforms state-of-the-art IL methods, achieving near-expert or beyond-expert QD performance on the challenging continuous control tasks derived from MuJoCo environments.

## 1 Introduction

Imitation Learning (IL) aims to mimic an expert’s behavior by learning from the demonstrations. IL has achieved great success in many real-world applications, such as robotics [], autonomous driving [], and drone control []. However, most of the traditional imitation learning methods were designed to learn one specific behavior, even given multiple demonstrations.

Learning diverse and high-quality policies is the ultimate goal of many real-world applications like robotic locomotion tasks []. Recent literature has demonstrated that quality diversity reinforcement learning (QDRL) [, ] is a promising direction for achieving this goal. Prior methods that combine Differentiable Quality Diversity (DQD) [] with off-policy RL achieved diverse and relatively high-performance policies. However, the performance gap between standard RL and QDRL still exists []. More recently, Batra et al. [] mitigate this gap by leveraging the on-policy RL method PPO [] with DQD. PPO estimates the gradients of diversity and performance objectives from the online collected data. Then, the estimated gradients are used by DQD methods like CMA-MAEGA [] that maintain a single search point and move through the behavior space by filling new regions. The synergy between PPO and DQD results in a state-of-the-art QDRL method, Proximal Policy Gradient Arborescence (PPGA) [], that achieves the ultimate goal of robotic locomotion. Generally, this kind of QDRL method finds a diverse archive of high-performing locomotion behaviors for an agent by combining PPO gradient approximations with Differentiable Quality Diversity algorithms.

However, the success of QDRL heavily relies on high-quality reward functions, which can be intractable in practice. Quality Diversity Imitation Learning (QDIL) offers a more flexible strategy for learning diverse and high-quality policies from demonstrations with diverse behaviors. In the literature, adversarial IL methods such as GAIL [] have achieved great success in learning specific behaviors for robotic locomotion tasks. Therefore, a naive solution for QDIL is to apply adversarial IL to estimate rewards from the given demonstrations and then leverage the estimated rewards for learning diverse and high-quality policies. Unfortunately, adversarial IL methods suffer from the training instability issue, which usually results in worse-than-demonstrator performance []. Moreover, when the demonstrations only contain a few behaviors, the rewards learned by adversarial IL techniques will be behavior-overfitted and unable to guide the agent to learn more diverse behaviors beyond the demonstrations.

Figure: Figure 1: Illustrations of the two issues of the Adversarial QDIL (i.e., training instability and behavior-overfitted reward) and their corresponding solutions (i.e. WQDIL and Single-Step Archive Exploration). $\delta(s)$ means the Markovian Measure Proxy of state $s$, a.k.a. the single-step measure.
Refer to caption: x1.png

The training instability issue and the behavior-overfitted reward issue heavily limit the adversarial QDIL to learning diverse and high-quality policies with limited demonstrations. In this work, we propose two synergic strategies to overcome these two challenging issues. The first strategy aims to stabilize the training of the reward model, and the second strategy focuses on encouraging the agent to do behavior space exploration. To develop the first strategy, we propose to stabilize the reward learning by applying Wasserstein adversarial training within the latent space of the Wasserstein Auto-Encoder (WAE) []. Similar to VAE [], WAE keeps the good properties of stable training and a nice latent manifold structure while generating higher-quality images than GAN []. Therefore, we propose to apply WAE to enable a more stable training of reward model in adversarial QDIL. In addition, we propose latent Wasserstein adversarial training to further improve the consistency of the reward training stability.
For the second strategy, we introduce a bonus that enables the agent to collect data with more diverse behaviors via single-step archive exploration. We call the resulting method Wasserstein Quality Diversity Imitation Learning (WQDIL) with Single-Step Archive Exploration (SSAE).
Figure 1 illustrates the two issues of the Adversarial QDIL (i.e., training instability and behavior-overfitted reward) and the corresponding solutions (i.e. WQDIL and Single-Step Archive Exploration).
The synergy between WQDIL and SSAE adherently address the two issues of adversarial QDIL, resulting in diverse and high-quality policies when learning with limited demonstrations. Figure 2 illustrates diverse behaviors induced by our learned policies.

Figure: Figure 2: Illustration of diverse behaviors learned by our Quality Diversity Imitation Learning framework on Humanoid and Walker2d, where each column represents one behavior. The “left” and “‘right” means the proportion of time the left leg or right leg contacting the ground.
Refer to caption: x2.png

We summarize our contributions as follows:

- •
First, we indicate the two main issues of the naive QDIL solution, i.e., the training instability and the behavior-overfitted reward issues of the adversarial QDIL approach.
- •
Second, we propose Wasserstein Quality Diversity Imitation Learning (WQDIL) to address the training instability issue by applying Wasserstein adversarial training within the latent space of the Wasserstein Auto-Encoder (WAE).
- •
Third, to alleviate the behavior-overfitted reward issue, we design a measure-conditioned reward, which makes the reward function sensitive to the local measure space, together with a measure-based bonus, which encourages the agent to collect data with more diverse behaviors.
- •
Empirically, we demonstrated that our method significantly outperforms state-of-the-art imitation learning methods in learning diverse and high-quality policies from limited demonstrations.

## 2 Background and Related Work

This section provides background and related works for Imitation Learning (Section 2.1), Quality Diversity Reinforcement Learning (Section 2.2), Quality Diversity Imitation Learning (Section 2.3), Training Stability (Section 2.4), and Exploration (Section 2.5).

### 2.1 Reinforcement Learning and Imitation Learning

Reinforcement Learning (RL) searches for policies that maximize cumulative reward in an environment, typically assuming the discrete-time Markov Decision Process (MDP) formalism $(S,A,r,P,\gamma)$. Here, $S$ and $A$ are the state and action spaces, $r(s,a)$ is the reward function, $P(s^{\prime}|s,a)$ defines state transition probabilities, and $\gamma$ is the discount factor. The traditional RL objective is to maximize the discounted episodic return of a policy $\mathbb{E}\left[\sum_{k=0}^{T-1}\gamma^{k}r(s_{k},a_{k})\right]$ where $T$ is the episode length.

Imitation learning (IL) trains an agent to mimic expert behaviors from demonstrations []. Behavior Cloning (BC) uses supervised learning to imitate expert behavior but suffers from severe error accumulation [].
Inverse Reinforcement Learning (IRL) seeks to recover a reward function from the demonstrations, and use reinforcement learning (RL) to train a policy that best mimics the expert behaviors []. Early IRL methods estimate rewards in the principle of maximum entropy [, , ]. Adversarial Inverse Reinforcement Learning (AIRL) [] learns a robust reward function by training the discriminator via logistic regression to classify expert data from policy data.

Adversarial IL methods treat IRL as a distribution-matching problem. Generative Adversarial Imitation Learning (GAIL) [] trains a discriminator to differentiate between the state-action distribution of the demonstrations and the state-action distribution induced by the agent’s policy, and output a reward to guide policy improvement.
More recently, Primal Wasserstein Imitation Learning (PWIL) [] introduces an offline reward function based on an upper bound of the Wasserstein distance between the expert and agent’s state-action distributions, avoiding the instability of adversarial IL methods. This paper includes max-entropy IRL (MaxEntIRL), the classic adversarial IL (i.e., GAIL) and the state-of-the-art PWIL as our baselines in QDIL with limited demonstrations.

### 2.2 Quality Diversity Reinforcement Learning

Quality Diversity Optimization.
Distinct from traditional optimization which aims to find a single solution to maximize the objective, Quality Diversity (QD) optimization aims to find a set of high-quality and diverse solutions in an $n$-dimensional continuous space $\mathbb{R}^{n}$. Given an objective function $f:\mathbb{R}^{n}\rightarrow\mathbb{R}$ and $k$-dimensional measure function $m:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}$, the goal is to find solutions $\theta\in\mathbb{R}^{n}$ for each local region in the behavior space $B=m(\mathbb{R}^{n})$.
QD algorithms discretize $B$ into $M$ cells, forming an archive $\mathcal{A}$. Formally, the objective is to find a set of solutions $\{\theta_{i}\}_{i=1}^{M}$ which maximises $f(\theta_{i})$ for each $i=1,\dots,M$. Each solution $\theta_{i}$ corresponds to a cell in $\mathcal{A}$ via its measure $m(\theta_{i})$, forming an archive of high-quality and diverse solutions [, ].

Previous Quality Diversity optimization methods integrate Evolution Strategies (ES) with MAP-Elites [], such as Covariance Matrix Adaptation MAP-Elites (CMA-ME) []. CMA-ME uses CMA-ES [] as an ES algorithm generating new solutions inserted into the archive, and uses MAP-Elites to retain the highest-performing solution in each cell. CMA-ES adapts its sampling distribution based on archive improvements from offspring solutions. However, traditional ES faces low sample efficiency, especially for high-dimensional parameters such as neural networks. Differentiable Quality Diversity (DQD) enhances exploration and optimization by leveraging gradient information from both the objective and measure functions. The state-of-the-art DQD algorithm, CMA-MAEGA [], utilizes CMA-ES to maintain a distribution over coefficient sets $c_{i}$, rather than directly generating solutions $\theta$. The objective function is $g(\theta)=|c_{0}|f(\theta)+\sum_{j=1}^{k}c_{j}m_{j}(\theta)$, where coefficients $c_{i}$ weight the gradients of the objective $f$ and measures $m_{j}$. CMA-MAEGA maintains a search policy and samples multiple coefficient sets after computing gradients of $f$ and $m_{i}$ to generate branched offspring solutions $\pi_{\theta_{1}},\dots,\pi_{\theta_{\lambda}}$. The archive improvement from each offspring guides the update of the coefficient distribution, which is then used to refine the search policy.

Built on the QD optimization paradigm, the Quality Diversity Reinforcement Learning (QDRL) problem can be viewed as maximizing $f(\theta)=\mathbb{E}_{\pi_{\theta}}\left[\sum_{k=0}^{T-1}\gamma^{k}r(s_{k},a_{
k})\right]$ with respect to diverse $\theta$ in a policy archive defined by measure $m$ []. In QDRL, both the objective and measure are non-differentiable, requiring approximations by DQD approaches. The state-of-the-art QDRL algorithm, Proximal Policy Gradient Arborescence (PPGA), employs a vectorized PPO (VPPO) architecture to approximate the gradients of the objective and measure functions []. PPGA introduced the Markovian Measure Proxy (MMP), a surrogate measure function that correlates strongly with the original measure and allows gradient approximation via policy gradient by treating it as a reward function. Specifically, MMP decomposes trajectory-based measures $m$ into individual steps by computing:

$$ $m_{i}(\theta)=\frac{1}{T}\sum_{t=0}^{T}\delta_{i}(s_{t}),$ (1) $$

where $\delta_{i}(s_{t})$ represents the single-step measure dependent on the state $s_{t}$. By breaking down the measure in this way, it becomes state-dependent and adheres to the Markov property. Then PPGA uses $k+1$ parallel environments with distinct reward functions – one for the original reward and $k$ for the surrogate measures. It approximates the gradients of both the objective and the $k$ measure functions by comparing the policy parameters before and after multiple PPO updates. These gradients are then passed to the modified CMA-MAEGA to update the policy archive.

### 2.3 Quality Diversity Imitation Learning

###### Definition 1 (Quality-Diversity Imitation Learning (QDIL))

Given expert demonstrations $\mathcal{D}=\{(s_{i},a_{i},\delta(s_{i}))\}_{i=1}^{n}$, where $s$,$a$,$\delta(s)$ represent state and action and the Markovian measure proxy of state $s$. QDIL aims to learn an archive of diverse policies $\{\pi_{\theta_{i}}\}_{i=1}^{M}$ that collectively maximizes $f(\theta)$ (i.e., cumulative reward) without access to the true reward. The archive is defined by a $k$-dimensional measure function $m(\theta)$, representing behavior patterns. After dividing the archive into $M$ cells, the objective of QDIL is to find M solutions, each occupying one cell, to maximize:

$$ $\max_{\{\theta_{i}\}}\sum_{i=1}^{M}f(\theta_{i}).$ (2) $$

In this paper, we focus on Adversarial QDIL, in which we apply adversarial IL, based on GAIL [] in particular, to estimate rewards from the given demonstrations $\mathcal{D}$. This adversarial training scheme considers the policy as an adversary and an additional neural network, the disciminator $D$, which discriminates between samples of the policy and samples of the demonstrations. The resulting reward model is based on the discriminator, and indicates the likelihood of corresponding to demonstration behaviours. This reward model then is used for the objective of QDIL.

###### Definition 1 (Quality-Diversity Imitation Learning (QDIL))

### 2.4 Improving Stability of Adversarial IL

The methodology of adversarial IL comes from Generative Adversarial Networks (GAN) []. Although adversarial IL methods like GAIL have achieved great success in low-dimensional environments, they also inherit the training instability issue of GAN that limits its generalization to more diverse tasks [].

Wasserstein GAN (WGAN) [] was proposed to stabilize the training of GAN by minimizing the Wasserstein distance instead of the Jensen-Shannon divergence between true data distribution and generated data distribution. However, the original WGAN with weight clipping can lead to undesired behaviors []. Gradient penalty (GP) was later proposed to address this issue, and has become a widely-used strategy for improving the stability of GAN and WGAN [, , , , , ].
We have applied these stabilization methods in adversarial QDIL, and conduct preliminary investigation experiments on continuous control tasks derived from MuJoCo. However, we only observed negligible effects of these methods on stabilizing adversarial QDIL.

Wasserstein Auto-Encoder (WAE) [] minimizes the optimal transport cost $W_{c}(P_{X},P_{G})$ based on the novel auto-encoder formulation.
In the resulting optimization problem, the decoder tries to accurately reconstruct the encoded training examples as measured by the cost function $f_{c}$.
Similar to VAE [], the encoder tries to simultaneously achieve two conflicting goals: it tries to match the encoded distribution of training examples $Q_{Z}:=\mathbb{E}_{P_{X}}[Q(Z|X)]$ to the prior $P_{Z}$ as measured by any specified divergence $D_{z}(Q_{Z},P_{Z})$, while making sure that the latent codes provided to the decoder are informative enough to reconstruct the encoded training examples. Using the squared cost $f_{c}(x,y)=\|x-y\|_{2}^{2}$, WAE keeps the good properties of VAEs (stable training, and a nice latent manifold structure) while generating better-quality images than GAN []. This observation inspired us to apply WAE in improving the stability of adversarial QDIL. The resulting framework is called Wasserstein Quality Diversity Imitation Learning (WQDIL).

### 2.5 Encouraging exploration in RL and IL

Exploration bonuses have a long history in bandits and reinforcement learning as part of the “optimism in face of uncertainty” principle, and include but are not limited to visitation count based approaches. In bandits, upper confidence bound (UCB, e.g. []) computes a bonus based on $\sqrt{\frac{2\ln(t)}{n(a_{t})}}$ for action $a_{t}$ at time $t$. In tabular reinforcement learning, UCRL techniques (e.g. []) provide upper confidence bounds to the reward based on a bonus proportional to $\sqrt{\frac{1}{n(s_{t},a_{t})}}$. R-MAX [] and E3 [] explicitly distinguish between known and unknown states, where in unknown states R-MAX assigns the maximal reward of the MDP and E3 performs balanced wandering (i.e. taking the action with lowest visitation count for the current state). Related concepts have also been applied in deep reinforcement learning [, , , ]. Related to exploration bonuses, recent techniques in imitation learning also considered curiosity-based exploration [, , ], which relates to discrepancy between the transition model’s prediction and the true transition. For instance, GIRIL [] computes rewards offline by pretraining a reward model using a conditional VAE [].

Our measure bonus differs from these approaches in that it is designed to explore the behavior space rather than the (state-)action space. Further, contrasting to UCRL type bonuses, due to dividing by the proportion of visits, rather than the count, our measure bonus does not shrink to zero such that underexplored behaviors will continue to receive a bonus. This is somewhat comparable in aim to the balanced wandering of E3 but differs in that it is a reward bonus rather than a policy definition (since we do not have control over the measures) and that it is maintained throughout the entire optimization process as we are interested in diversity.

## 3 Wasserstein Quality Diversity Imitation Learning (WQDIL)

This section introduces our proposed framework-WQDIL. We first begin with two solutions for the behavior overfitting problem which are called Single-Step Archive Exploration Bonus and Measure Conditioning (Section 3.1). We then introduce Wasserstein Auto-encoders (Section 3.2), and propose a WQDIL instance WAE-WGAIL that applies latent Wasserstein adversarial training for WAE optimization (Section 3.3).

### 3.1 Mitigating behavior overfitting

Compared to the vast behavior space, expert behaviors tend to be limited in diversity. This poses a problem for learning diverse behaviors, as the reward model in the traditional IL paradigm focuses solely on shaping rewards around specific expert behaviors while neglecting other behavior patterns. We refer to this type of reward as a “behavior-overfitted reward.” Consequently, this issue significantly inhibits the learning of diverse behaviors, since any RL method inherently converges to high-reward behaviors. To address this problem, we introduce a single-step archive-exploration reward bonus and measure-conditioning, which together ensure exploration across the behavior space as well as the sensitivity of the reward function to the local measure information.

##### Single-Step Exploration Bonus and Measure Conditioning

In both the bonus and the measure conditioning, we are primarily interested in single-step measures building on the Markovian Measure Proxy mentioned in Section 2.2 to make the measure information more fine-grained. This is not merely for convenience purposes but also as this will yield richer data from the demonstrations, which would otherwise be sparse (e.g. only a few episodic measures since we focus on limited-expert-demonstrations in our setting). Meanwhile, Focusing on single step measure will allow exploring behaviours with similar episodic measure but different single-step measure for potentially yielding higher fitness within a particular behavioral region. This is due to the episodic measure is the average (or sum) of each single-step measure and is less fine-grained.

For the exploration bonus, we establish the single-step archive $\mathcal{A}_{single}$, which corresponds to the state-dependent measure $\delta(s)$. Similar to the behavior archive $\mathcal{A}$, we partition $\mathcal{A}_{single}$ into numerous cells for discretization. Notably, instead of merely recording whether a cell is occupied, we track the visitation count $n_{i}$ for each cell $C_{i}$ in $\mathcal{A}_{single}$. The exploration reward bonus is defined as:

$$ $r_{exp}(s,a,\delta(s))=\frac{1}{1+p(\delta(s))},$ (3) $$

where

$$ $p(\delta(s))=\frac{n(C(\delta(s)))}{\sum_{i}n(C_{i})}.$ (4) $$

In these equations, $C_{i}$ denotes the $i$-th cell in $\mathcal{A}_{single}$, $C(\delta(s))$ represents the cell corresponding to the single-step measure proxy $\delta(s)$, and $n$ signifies the visitation count. Each time a state $s$ activates a cell in $\mathcal{A}_{single}$, the visitation count of that cell is incremented by one. This mechanism allows the single-step archive $\mathcal{A}_{single}$ to be dynamically updated during training.

The exploration bonus assigns higher rewards to regions in $\mathcal{A}_{single}$ that are less frequently visited, thereby promoting the agent to explore unseen behavior patterns. Additionally, once a region within the single-step behavior space has been sufficiently explored, the bonus decreases, facilitating the exploitation of that region to discover high-performing policies. However, note that the bonus is defined relative to the exploration of other measures such that the bonus never shrinks to zero for a particular measure. The form in Eq. 3 also avoids extreme values due to the +1 offset in the denominator. With these features together, the reward bonus can effectively mitigate the “behavior-overfitted reward’ issue by always encouraging new behavior patterns, thus facilitating diverse behaviors.

In addition to the exploration bonus, the reward model is also conditioned on the measure, according to $\mathcal{R}(s,a,\delta(s))$. This contributes to a solution to behavior overfitting: since the goal is to form behaviorally diverse policies, different behaviors require different policies to be imitated – and hence correspond to a different reward function.

### 3.2 Wasserstein auto-encoders

The Wasserstein distance is a formulation that occurs within the context of optimal transport []. Kantorovich’s formulation of the problem is given by

$$ $\mathcal{W}_{c}(P_{x},P_{G}):=\inf_{\Gamma\in\mathcal{P}(x\sim P_{X},y\sim P_{ G})}\mathbb{E}_{(x,y)\sim\Gamma}[f_{c}(x,y)]\,,$ (5) $$

where $f_{c}(x,y)\colon\mathcal{X}\times\mathcal{X}\to\mathbb{R}_{+}$ is the cost function
and $\mathcal{P}(x,y)$ is the set of joint distributions of $(x,y)$ with marginals $P_{x}$ and $P_{G}$ respectively. We are interested in the set $\mathcal{X}=S\times A\times M$, where $M$ is the measure space.

Auto-encoders aim to represent input data in a low-dimensional latent space. Modern generative models like variational auto-encoders (VAEs) [] and generative adversarial networks (GANs) [] do so by minimizing a discrepancy measure between the data distribution $P_{x}$ and the generative model $P_{G}$. To formulate a Wasserstein based generative model, the form in Eq. 5 is intractable; however, following [], one can equate the Wasserstein distance to the following WAE objective []:

$$ $\displaystyle D_{\mathrm{WAE}}(P_{x},P_{G}):=$ $\displaystyle\inf_{Q(z|x)\in\mathcal{Q}}\mathbb{E}_{P_{x}}\mathbb{E}_{Q(z|x)} \bigl{[}f_{c}\bigl{(}x,G(z)\bigr{)}\bigr{]}$ (6) $\displaystyle+\lambda\cdot D_{z}(Q_{z},P_{z}),$ $$

where $P_{z}$ is the latent distribution, $D_{z}$ is a divergence measure, and $P_{G}$ is the decoder (i.e. generative model). We restrict our cost function to the squared Euclidian distance, i.e. $f_{c}(x,y)=||x-y||^{2}_{2}$, resulting in the $2$-Wasserstein distance.

WAE-GAN [] is an instance of WAE when choosing $D_{z}(Q_{z},P_{z})=D_{\mathrm{JS}}(Q_{z},P_{z})$ and using adversarial training to estimate $D_{z}$. Specifically, WAE-GAN uses an adversary and discriminator in the latent space $\mathcal{Z}$ trying to separate “true” points sampled from $P_{z}$ and “fake” ones sampled from $Q_{z}$ []. In the imitation learning setting, $P_{z}$ corresponds to the distribution of latent data obtained from the encoded demonstrations while $Q_{z}$ corresponds to the distribution of latent data obtained from the encoded trajectory data from the policy.

Analogously, we propose WAE-WGAN, which is equivalent to WAE-GAN except that it sets the divergence measure to the $1$-Wasserstein distance, i.e. $D_{z}(Q_{z},P_{z})=\mathcal{W}_{1}(Q_{z},P_{z})$. We choose this option based on results on the improved stability during adversarial training []. Due to strong duality with $p=1$, the Kantorovic-Rubinstein duality can be applied, i.e.

$$ $\mathcal{W}_{1}(Q_{z},P_{z})=\sup_{||f||_{L}\leq 1}\mathbb{E}_{z\sim Q_{z}}[f( z)]-\mathbb{E}_{z\sim P_{z}}[f(z)],$ (7) $$

where $f$ is a $1$-Lipschitz function $f:\mathcal{X}\to\mathbb{R}$.

### 3.3 WQDIL: WAE Meets Adversarial QDIL

This section introduces our Wasserstein Quality Diversity Imitation Learning (WQDIL) framework. We illustrate the processes of WQDIL in Algorithm 1, which is built on top of the Proximal Policy Gradient Arborescence (PPGA) algorithm []. The blue color indicates the WQDIL’s operations that are distinct from the PPGA algorithm.

Figure: Algorithm 1 Wasserstein Quality Diversity Imitation Learning (WQDIL)

Figure: Algorithm 2 Reward Estimation

In WQDIL, we input the initial reward model $\mathcal{R}$ and the expert demonstrations $\mathcal{D}$. In each iteration, we first approximate the gradients $\nabla f$ and $\nabla m$ (in step 4) using the rewards estimated from $\mathcal{R}$ (Algorithm 2). Steps 5-15 are QD optimization options used in PPGA. In step 16, we update the search policy using the reward model $\mathcal{R}$. Step 17 is a key step, in which we update the reward model $\mathcal{R}$ using the input expert demonstrations $\mathcal{D}$. We have explored a few algorithms in the following for updating the reward model $\mathcal{R}$ in WQDIL.

WAE-GAIL.
The first reward updating algorithm is WAE-GAIL, which is a naive adaptation of WAE-GAN in the QDIL framework. Algorithm 5 in Appendix A summarizes the details of WAE-GAIL.

WAE-WGAIL. Although WAE-GAIL achieves better stability and higher QD performance than GAIL in some tasks, the improvements are not consistent. To further improve WAE-GAIL, we propose to apply Wasserstein adversarial training in the latent space of WAE, which is analogous to the WAE-WGAN proposed in Section 3.2. The loss function of the auto-encoder in WAE-WGAIL is given by

$$ $\displaystyle\mathcal{L}(\phi,\psi)$ $\displaystyle=\mathcal{L}_{\text{recon}}^{e}+\mathcal{L}_{\text{recon}}^{\pi}+ \lambda D_{z}(z^{e},z^{\pi})$ $\displaystyle=\mathbb{E}_{x^{e}\sim\mathcal{D},z^{e}\sim Q_{\phi}(z^{e}|x^{e}) }[f_{c}(x^{e},G_{\psi}(z^{e}))]$ $\displaystyle+\mathbb{E}_{x^{\pi}\sim\pi,z^{\pi}\sim Q_{\phi}(z^{\pi}|x^{\pi}) }[f_{c}(x^{\pi},G_{\psi}(z^{\pi}))]$ $\displaystyle+\lambda\mathbb{E}_{x^{e}\sim\mathcal{D},x^{\pi}\sim\pi}\left[D_{ \mathcal{W}}(G_{\psi}(x^{e}))-D_{\mathcal{W}}(G_{\psi}(x^{\pi}))\right],$ (8) $$

where $\phi$ and $\psi$ are the parameters of the encoder and decoder, respectively, $x=(s,a)$ is the state-action pair, $Q_{\phi}$ is the encoder, $G_{\psi}$ is the decoder, and $D_{\mathcal{W}}$ is the Wasserstein discriminator.
The resulting algorithm is called WAE-WGAIL.
Algorithm 6 in Appendix A summarizes the details of WAE-WGAIL.

Figure: Algorithm 3 mCWAE-WGAIL.

mCWAE-WGAIL. Compared to GAIL, WAE-WGAIL has significantly better stability and consistency across different tasks. However, the latent space of WAE-WGAIL is learned from the state-action pairs. Since the demonstrations only contain very limited behaviors, the learned latent space is not adaptive to more diverse behaviors that the agent may encounter. To improve the adaptivity of the latent space and the reward model $\mathcal{R}$, we propose measure-conditioned WAE-WGAIL (mCWAE-WGAIL) that learns latent space by reconstructing state-action-measure triplet, i.e. formulating $x=(s,a,\delta(s))$ in Eq. 3.3.
Algorithm 3 summarizes the details of mCWAE-WGAIL,
our proposed algorithm for WQDIL.

## 4 Experiments

### 4.1 Experiment Setup

We evaluate our framework on three popular MuJoCo [] environments: Halfcheetah, Humanoid, and Walker2d. The measure function is the number of times a leg contacts with the ground, divided by the trajectory length.
We implement all the methods based on the PPGA implementation, which utilizes the Brax simulator [] with QDax wrappers for measure calculation [].
All Experiments were conducted on A40 GPUs for three runs per task.

### 4.2 Demonstrations

We use a policy archive obtained by PPGA to generate expert demonstrations. To follow real-world scenario with limited demonstrations, we first sample the top 500 high-performance elites from the archive as candidate pool. Then from this pool, we select a few demonstrations such that they are as diverse as possible. This process results in 4 diverse demonstrations (episodes) per environment. Figure 3 visualizes the demonstrators in the policy archive, and Table 6 in Appendix C provides the statistical properties of the demonstrations.

Figure: (a) HalfCheetah
Refer to caption: x3.png

### 4.3 Performance

Figure: Figure 4: Visualization of the policy archive of Expert, GAIL, WAE-GAIL, WAE-WGAIL, WAE-WGAIL-Bonus and mCWAE-WGAIL-Bonus on Humanoid. The color indicates the cumulative rewards of best performing policy in the archive cells.
Refer to caption: x6.png

Figure: Figure 5: Learning curve comparison of our mCWAE-WGAIL-Bonus against WAE-WGAIL and state-of-the-art IL methods.
Refer to caption: x7.png

To validate the effectiveness of WQIL with Single-Step Archive Exploration, we compare our mCVAE-WAIL-Bonus against state-of-the-art IL methods within QDIL framework. These methods includes our WAE-WGAIL (Algorithm 6, Appendix A), GAIL [], PWIL [], AIRL [], MaxEntIRL [] and GIRIL []. We apply a gradient penalty to all the methods that require the online update of the reward model. Appendix B summarizes the hyperparameters.

We compare the performance under four QD metrics []: 1) QD-Score, the sum of scores of all nonempty cells in the archive. QD-score is the most important metric in QDIL as it aligns with the objective of QDIL as in equation (2); 2) Coverage, the percentage of nonempty cells, indicating the algorithm’s ability to discover diverse behaviors; 3) Best Reward, the highest score found by the algorithm; and 4) Average Reward, the mean score of all nonempty cells, reflecting the ability to discover both diverse and high-performing policies. We use the true reward functions to calculate these metrics.

The heatmaps in Figure 4 visualize the performance across the policy archive for Expert, GAIL, WAE-GAIL, WAE-WGAIL, WAE-WGAIL-Bonus and mCWAE-WGAIL-Bonus on Humanoid. Even using true rewards, the expert explores less than half of the archive in the Humanoid environment. The high-performance policies are more restricted, covering around a quarter of the full archive. GAIL nearly explores the full archive space, however, most of the learned policies are low-performance (in dark color). Due to the training instability and behavior-overfitting issues, GAIL only learns high-performance policies within a small region of the behavior space. WAE-GAIL learns more high-performance policies than GAIL since it maintains better stability than GAIL. However, the explored region is still limited. By applying latent Wasserstein adversarial training, WAE-WGAIL further enlarges the explored archive region and coverage of high-performance policies. WAE-WGAIL-Bonus updates the archive using the rewards learned by WAE-WGAIL and the estimated single-step archive exploration bonus. The explored region of WAE-WGAIL-Bonus is clearly larger than WAE-WGAIL, which we attribute to the exploration bonus yielding improved coverage of the measure space. Thanks to the measure conditioning, the explored archive space of mCWAE-WGAIL-Bonus is slightly smaller but is more focused in the sense that it yields the largest set of high-performing policies among the five QDIL methods.

**Table 1: Ablation study results on the effect of latent Wasserstein adversarial training (-WGAIL).**
|  | HalfCheetah | Walker2d | Humanoid |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg |
| PPGA-trueReward | $6.75\times 10^{6}$ | 94.08 | 8,942 | 2,871 | $3.64\times 10^{6}$ | 77.04 | 5,588 | 1,891 | $5.71\times 10^{6}$ | 49.96 | 9,691 | 4,570 |
| mCWAE-WGAIL-Bonus | $\mathbf{3.80\times 10^{6}}$ | 98.28 | 4,553 | 1,547 | $\mathbf{3.07\times 10^{6}}$ | 87.39 | 4,142 | 1,407 | $\mathbf{6.40\times 10^{6}}$ | 75.05 | 6,875 | 3,447 |
| mCWAE-GAIL-Bonus | $3.07\times 10^{7}$ | 99.32 | 3,350 | 1,235 | $1.89\times 10^{6}$ | 56.88 | 1,860 | 1,337 | $3.26\times 10^{6}$ | 55.95 | 5,182 | 2,530 |
| WAE-WGAIL | $3.74\times 10^{6}$ | 93.77 | 5,463 | 1,615 | $2.11\times 10^{6}$ | 73.20 | 3,327 | 1,148 | $3.98\times 10^{6}$ | 68.08 | 8,553 | 2,359 |
| WAE-GAIL | $2.91\times 10^{6}$ | 92.85 | 3,792 | 1,178 | $1.21\times 10^{6}$ | 43.59 | 1,357 | 1,097 | $4.01\times 10^{6}$ | 66.45 | 5,721 | 2,381 |

**Table 2: Ablation study results on the effects of single-step archive exploration (-Bonus) and measure conditioning (mC).**
|  | HalfCheetah | Walker2d | Humanoid |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg |
| PPGA-trueReward | $6.75\times 10^{6}$ | 94.08 | 8,942 | 2,871 | $3.64\times 10^{6}$ | 77.04 | 5,588 | 1,891 | $5.71\times 10^{6}$ | 49.96 | 9,691 | 4,570 |
| mCWAE-WGAIL-Bonus | $\mathbf{3.80\times 10^{6}}$ | 98.28 | 4,553 | 1,547 | $\mathbf{3.07\times 10^{6}}$ | 87.39 | 4,142 | 1,407 | $\mathbf{6.40\times 10^{6}}$ | 75.05 | 6,875 | 3,447 |
| WAE-WGAIL-Bonus | $2.66\times 10^{7}$ | 76.89 | 4,788 | 1,246 | $2.55\times 10^{6}$ | 73.93 | 3,310 | 1,381 | $4.97\times 10^{6}$ | 67.27 | 8,620 | 3,09 |
| mCWAE-WGAL | $3.78\times 10^{7}$ | 87.76 | 5,728 | 1,623 | $2.37\times 10^{6}$ | 71.61 | 3,611 | 1,326 | $3.54\times 10^{6}$ | 66.96 | 8,485 | 2,114 |
| WAE-WGAIL | $3.74\times 10^{6}$ | 93.77 | 5,463 | 1,615 | $2.11\times 10^{6}$ | 73.20 | 3,327 | 1,148 | $3.98\times 10^{6}$ | 68.08 | 8,553 | 2,359 |

Figure 5 shows the learning curve comparison of QD performance for different QDIL methods. Among the four metrics, QD-Score is the most important because it reflects both high diversity and high quality of performance. In HalfCheetah, our method mCWAE-WGAIL-Bonus achieves one of the highest QD-Score with much better sample efficiency than the other counterparts (i.e., achieving the same QD-Score with much fewer iterations). In Walker2d and Humanoid environments, our method achieves a significantly higher QD-Score than the other methods.
PWIL achieves a slightly higher QD-Score than our method in HalfCheetah, however, it performs the worst in Humanoid. MaxEntIRL performs worse than GAIL, while AIRL outperforms GAIL slightly. Pretraining from limited demonstrations, GIRIL performs similar to GAIL in HalfCheetah and similar to WAE-WGAIL in Humanoid but much worse in Walker2d.
For Coverage, our mCWAE-WGAIL-Bonus is the best in HalfCheetah and Walker2d. Although our coverage in Humanoid is not the best, we significantly outperform the baselines in terms of AverageReward.
The quantitative results are summarized in Table 7 of Appendix D.1.

### 4.4 Ablations

We conduct two ablations to study the effects of each component of our method mCWAE-WGAIL-Bonus. The first ablation in Section 4.4.1 studies the effects of latent Wasserstein adversarial training in the QDIL framework. On the other hand, Section 4.4.2 focuses on studying the effects of single-step archive exploration bonus and measure conditioning.

#### 4.4.1 Ablation on Latent Wasserstein Adversarial Training

To study the effect of latent Wasserstein adversarial training, we additionally introduce mCWAE-GAIL-Bonus, which reinforces WAE-GAIL with measure conditioning and a single-step archive exploration bonus.
Table 1 compares the QD performance of our mCWAE-WGAIL-Bonus with mCWAE-GAIL-Bonus and WAE-WGAIL. As we can see, applying Wasserstein adversarial training improves the QD-Score of WAE-GAIL and mCWAE-GAIL-Bonus by 27.5% and 23.7% respectively on HalfCheetah. In Walker2d, the latent Wasserserstein adversarial training enhances the QD-Score of WAE-GAIL and mCWAE-GAIL-Bonus by 74.3% and 62.4%, which is significant. Moreover in Humanoid, mCWAE-WGAIL-Bonus achieves 2x QD-Score of mCWAE-GAIL-Bonus, impressively outperforming the expert (i.e., PPGA-trueReward) by 12%. The learning curves are shown in Figure 6 in the Appendix D.2.

#### 4.4.2 Ablations on Single-Step Archive Exploration and Measure Conditioning

Table 2 compares the effects of single-step archive exploration bonus and measure conditioning. Measure conditioning improves the QD-Score for WAE-WGAIL-Bonus by 42.8%. In Walker2d, the contributions of measure conditioning to WAE-WGAL and WAE-WGAIL-Bonus are 12.3% and 20.3% respectively. In Humanoid, mCWAE-WGAIL-Bonus outperforms WAE-WGAIL-Bonus by 28.7%.
Now, we study the effect of the single-step archive exploration bonus. The bonus does not show much improvement on HalfCheetah. Comparing mCWAE-WGAIL-Bonus and mCWAE-WGAIL, we found that the contributions of bonus on the QD-Score improvement are 29.5% in Walker2d and 80% in Humanoid. Comparing WAE-WGAIL-Bonus and WAE-WGAIL, the contributions of the bonus on the QD-Score improvement are 20.8% in Walker2d and 24.8% in Humanoid.
We illustrate the corresponding learning curves of this ablation study in Figure 7 of Appendix D.2.

## 5 Conclusion

This paper focuses on addressing two key issues of adversarial QDIL, i.e., the training instability of adversarial training and the behavior-overfitting in the design of the reward function. We propose WQDIL and a measure-conditioned reward function with Single-Step Archive Exploration to address these two issues respectively. Through two thorough ablation studies, we confirm the significance of each component (latent Wasserstein training, measure conditioning, and single-step exploration bonus) and find that latent Wasserstein adversarial training contributes most significantly to the QD-Score. The collaboration between these components contributes to the superior QD performance of our method on three continuous control tasks. Our method enables the agent to stably learn diverse and high-quality behaviors from limited demonstrations.

## Appendix A Algorithms

Figure: Algorithm 4 Wasserstein Auto-Encoder with GAN-based penalty (WAE-GAN).

Figure: Algorithm 5 WAE-GAIL.

Figure: Algorithm 6 WAE-WGAIL.

## Appendix B Hyperparameter Setting

### B.1 Hyperparameters for PPGA

Table 3 summarizes a list of hyperparameters for PPGA policy updates.

**Table 3: List of relevant hyperparameters for PPGA shared across all environments.**
| Hyperparameter | Value |
| --- | --- |
| Actor Network | [128, 128, Action Dim] |
| Critic Network | [256, 256, 1] |
| $N_{1}$ | 10 |
| $N_{2}$ | 10 |
| PPO Num Minibatches | 8 |
| PPO Num Epochs | 4 |
| Observation Normalization | True |
| Reward Normalization | True |
| Rollout Length | 128 |
| Grid Size | 50 |
| Env Batch Size | 3,000 |
| Num iterations | 2,000 |

### B.2 Hyperparameters for IL

Table 4 summarizes a list of hyperparameters for AIRL, GAIL, measure-conditioned GAIL, and MConbo-GAIL.

**Table 4: List of relevant hyperparameters for GAIL and WAEs shared across all environments.**
| Hyperparameter | Value |
| --- | --- |
| Discriminator | [100, 100, 1] |
| Learning Rate | 3e-4 |
| Discriminator Num Epochs | 1 |
| Regularization coefficient $\lambda$ | 1 |

Table 5 summarizes a list of hyperparameters for GIRIL.

**Table 5: List of relevant hyperparameters for GIRIL shared across all environments.**
| Hyperparameter | Value |
| --- | --- |
| Encoder | [100, 100, Action Dim] |
| Decoder | [100, 100, Observation Dim] |
| Learning Rate | 3e-4 |
| Batch Size | 32 |
| Num Pretrain Epochs | 10,000 |

## Appendix C Demonstration Details

**Table 6: Demonstrations statistics on the three environments.**
| Tasks | Demo num | Attributes | min | max | mean | std |
| --- | --- | --- | --- | --- | --- | --- |
| HalfCheetah | 4 | Length | 1000 | 1000 | 1000.0 | 0.0 |
| Return | 3766.0 | 8405.4 | 5721.3 | 1927.6 |  |  |
| Walker2d | 4 | Length | 356.0 | 1000.0 | 625.8 | 254.4 |
| Return | 1147.9 | 3721.8 | 2372.3 | 1123.7 |  |  |
| Humanoid | 4 | Length | 1000.0 | 1000.0 | 1000.0 | 0.0 |
| Return | 7806.2 | 9722.6 | 8829.5 | 698.1 |  |  |

## Appendix D Additional Experimental Results

### D.1 Quantitative Overall Performance.

**Table 7: Compare the QD performance of our mCWAE-WGAIL-Bonus against WAE-WGAIL, GAIL, PWIL, AIRL, MaxEntIRL and GIRIL across three environments. Cov, Best, and Avg refer to Coverage, Best Reward, and Average Reward respectively.**
|  | HalfCheetah | Walker2d | Humanoid |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg | QD-Score | Cov(%) | Best | Avg |
| PPGA-trueReward | $6.75\times 10^{6}$ | 94.08 | 8,942 | 2,871 | $3.64\times 10^{6}$ | 77.04 | 5,588 | 1,891 | $5.71\times 10^{6}$ | 49.96 | 9,691 | 4,570 |
| mCWAE-WGAIL-Bonus | $3.80\times 10^{6}$ | 98.28 | 4,553 | 1,547 | $\mathbf{3.07\times 10^{6}}$ | 87.39 | 4,142 | 1,407 | $\mathbf{6.40\times 10^{6}}$ | 75.05 | 6,875 | 3,447 |
| WAE-WGAL | $3.74\times 10^{6}$ | 93.77 | 5,463 | 1,615 | $2.11\times 10^{6}$ | 73.20 | 3,327 | 1,148 | $3.98\times 10^{6}$ | 68.08 | 8,553 | 2,359 |
| GAIL | $2.02\times 10^{6}$ | 67.83 | 5,115 | 1,165 | $2.48\times 10^{6}$ | 69.29 | 4,031 | 1,429 | $1.86\times 10^{6}$ | 82.27 | 6,278 | 924 |
| PWIL | $\mathbf{4.17\times 10^{6}}$ | 99.65 | 3,626 | 1,675 | $2.27\times 10^{6}$ | 64.45 | 2,835 | 1,410 | $1.12\times 10^{6}$ | 91.73 | 841 | 492 |
| AIRL | $3.11\times 10^{6}$ | 83.57 | 5,183 | 1,410 | $2.53\times 10^{6}$ | 70.53 | 4,280 | 1,437 | $2.31\times 10^{6}$ | 71.47 | 7,661 | 1,308 |
| MaxEntIRL | $1.12\times 10^{6}$ | 85.48 | 2,570 | 523 | $1.80\times 10^{6}$ | 68.81 | 3,756 | 1,046 | $1.82\times 10^{6}$ | 83.27 | 4,658 | 882 |
| GIRIL | $2.17\times 10^{6}$ | 95.96 | 3,466 | 909 | $0.52\times 10^{6}$ | 25.08 | 1,139 | 821 | $4.33\times 10^{6}$ | 67.40 | 6,992 | 2,590 |

### D.2 Learning Curves of Ablation Studies

Figure: Figure 6: Learning curve comparison of our mCWAE-WGAIL-Bonus against WAE-GAIL-Bonus, WAE-WGAIL and WAE-GAIL.
Refer to caption: x8.png

Figure: Figure 7: Learning curve comparison of our mCWAE-WGAIL-Bonus against the components: WAE-WGAIL-Bonus, mCWAE-WGAIL and WAE-WGAIL.
Refer to caption: x9.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Is the concept of a measure bonus constrained to QD-IL? Is there anything that would prevent adding it to QD-RL algorithms, such as PPGA, for QD-RL problems as well?
* Currently, MConbo-GAIL can outperform the PPGA expert in the Humanoid environment, even though it is trained using demonstration trajectories from experts selected from PPGA expert's archive - which is somewhat surprising. If the measure bonus were added to the PPGA expert, could this drastically improve its performance? In that case, could MConbo-GAIL still outperform it?
soundness: 3
strengths: 1. This paper is the first to explicitly integrate standard QD algorithms with imitation learning, though it is not the first to study imitation learning methods that learns a diverse set of policies.
* (There is another paper that studies imitation learning methods in a similar setting inspired by QD algorithms, although it diverges from standard QD algorithms: N. Justesen, M. González-Duque, D. Cabarcas, J. -B. Mouret and S. Risi, "Learning a Behavioral Repertoire from Demonstrations," 2020 IEEE Conference on Games (CoG), Osaka, Japan, 2020, pp. 383-390, doi: 10.1109/CoG47356.2020.9231897.)

2. The measure bonus, designed to encourage exploration of policies with diverse measures, is straightforward and intuitively sound.

3. The effectiveness of the designed MConbo-IRL is demonstrated through empirical results on several contibuous control tasks, which show that the MConbo- version of GAIL/VAIL overall outperforms PPGA + GAIL/VAIL and sometime outperforms PPGA + true reward in terms of QD metrics, including QD-Score, Coverage, Best Reward, and Average Reward.
summary: This paper proposes a novel Quality Diversity Imitation Learning (QD-IL) framework designed to learn a diverse set of high-performing policies based on varied but limited expert demonstrations. The proposed framework, MConbo-IRL, combines elements from standard QD algorithms, such as CMA-MEGA, and adversarial imitation learning methods such as GAIL and VAIL, and introduces two key modifications: (1) adding an extra reward term that encourages exploration of policies with diverse measures (referred to as the measure bonus), and (2) adding the single-step measure of the current state as an input to the discriminator (referred to as measure conditioning). Experimental results across several continuous control tasks demonstrate that the proposed framework outperforms the baseline (PPGA trained under a single reward function learned by GAIL/VAIL) in QD metrics.
weaknesses: 1. The effects of measure conditioning: While the purpose of the measure bonus is straightforward, it is less clear how the measure conditioning affects the QD-IL performance. The paper states (such as in Introduction, Section 4.2.1, Appendix E) that measure conditioning (i.e., adding the single-state measure as an input to the discriminator) "enables the agent to learn broader behavior patterns". However, this claim lacks sufficient empirical support. The ablation study on measure conditioning in Appendix E shows that it impacts QD-Score and Average Reward, but the mechanism of how it impacts the policy learning remains unclear without further analysis of the algorithm's inner components.
My confusion is as follows: if adding the single-state measure input to the discriminator improves the discriminator performance, doesn't this only makes it easier for the discriminator to distinguish policy trajectories that demonstrate single-state measures absent from expert demonstrations? It seems that this could potentially discourage the policy learning to be more diverse. Could the authors clarify if this reasoning is correct, or if measure conditioning is intended to encourage diversity policy learning in a different way?

2. In Section 5.3, the observation that MConbo-GAIL achieves much higher coverage in the Humanoid environment than the PPGA expert is attributed to the "synergy between the measure bonus and CMA-MEGA". However, there are other baselines, such as PWIL and GAIL, that also achieve much higher coverage than the PPGA expert in this environment. This suggests that other factors may contribute to coverage performance in this case.

3. The expert demonstration data contains trajectories collected from multiple experts with diverse behaviors. These experts, as optimized agents from PPGA under the true reward of the environment, each optimizes their own version of the reward function. Therefore, there is no shared reward function among them. This places the PPGA + GAIL/VAIL baseline at an intrinsic disadvantage, as it is designed to learn a single reward function. It may be worthwhile to explore additional forms of baseline methods based on standard GAIL/VAIL. For example, applying PPGA + GAIL/VAIL to create a policy archive for each demonstration expert individually, then aggregating all policy archives together, could provide another baseline.

4. The paper is not very clearly written, as it sometimes references terms (often from QD algorithms) without defining them or pointing to exact external references, making the paper less self-contained.
* For instance, the format of the policy archive and the exact procedure of archive updates, are not detailed in the paper or the pseudo-code.
* Similarly, there are other parts of the pseudo-code that are not completely defined, making it difficult to understand the exact procedure.
* What is the term $p$ in Equation 2? If $p$ is a fixed constant, can it then be ignored? If that is the case, how does the definition of measure bonus in Equation 2 differ from an indicator function (in the sense that an indicator function scaled by a value $q$ is still binary)? Are $p$ and the scaling factor hyperparameter $q$ being adjusted adaptively during the training process?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - Are there any training / wall clock time numbers of the proposed method? given that a gpu sim like brax is used I think wall time may be somewhat relevant to be aware of as the purpose of GPU sim is to improve training wall time.
- How well might this method fair with sub-optimal demonstrations? Or would some form of offline RL approach now be needed?

Happy to raise score to the accept range if there is more discussion on other aspects of the proposed algorithm. I primarily feel that there are a number of aspects that should be investigated to better understand the new approach with e.g., respect to its demonstration data which often influences resulting behaviors significantly. Especially given that only 4 demonstrations are used which can be more susceptible to high variance in results.
soundness: 3
strengths: - The proposed modified version of GAIL, MConbo-GAIL, addresses shortcomings of GAIL such as overfitting to demonstration specific behaviors and the inability of the discriminator to generalize to new states. The changes made are sound and help encourage the discriminator to focus on high-level more generalizable features.
- The improvements shown in table 1 and figure 6 show the impact of MConbo variations of GAIL and VAIL inverse RL methods.
summary: The submission combines imitation learning and quality diversity optimization to learn an archive/set of policies that all try and maximize the unknown oracle reward function of an environment. Only expert demonstrations are provided. A modified version of Proximal Policy Gradient Arborescence (PPGA) is used as the QD-RL method to get more diverse behaviors to then generate expert demonstrations and a reference archive of expert behaviors. Measure bonus and conditioning are the two main contributions to enhance existing off the shelf inverse RL algorithms to achieve better balance of exploration and exploitation and address the challenge of localized rewards.
weaknesses: - I think an ablation on the demonstrations used to verify the robustness of the proposed method could be important. It seems that a fixed set of 4 demos per environment tested are used.
- Overall I believe this paper lacks experiments investigating the method at a deeper level. There is just one ablation and then main results, however, I think it is important to understand more about several aspects including the demonstrations and perhaps how the learned behaviors might correlate with those demonstrations. E.g. a question could be "How much of the new archive is generalization beyond demonstrations?"

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Overall, I think the idea is interesting, but I am concerned about how the method can be scaled to other complex tasks. I have the following questions:

- I don’t quite understand the Figure 1(a). Can you give a simple robot example for it? Like what is the policy space here? Does it mean a different observation space or action space?
- In Figure 1(b) Multiple behavior demos, traditional behavior cloning has the mode averaging problem, but some current methods such as diffusion-based imitation learning has the ability to learn multi-modality, right?
- QD-IL relies on the definition of behavior diverse measures $m$.  In Chapter 4.2.1, it says that the measure abstracts high-level task features. Does this mean you always need to design task-specific behavior patterns in advance? I would assume this is very hard for more complex and long-term tasks.
- The experiments do not contain any manipulation tasks. Could QD-IL be adapted to complex manipulation tasks? Could the authors give a simple case for that?
- In Figure 4, what is the x and y axis?
soundness: 3
strengths: - The paper proposes the first generic QD-IL framework that integrates QD algorithms into imitation learning.
- Introducing measure bonus into GAIL and VAIL improves policy diversity while maintaining the policy’s performance. This is demonstrated using 4 QD metrics across 3 tasks.
- The paper provides ablation studies on the measure bonus which indicate the measure bonus directly encourages the exploration of new behaviors.
- The paper is well-written with the contributions explained clearly.
summary: This paper presents Quality Diversity Imitation Learning (QD-IL), a method that allows policies to acquire a diverse range of skills using limited demonstrations. Quality Diversity (QD) is leveraged to explore multiple solutions to optimization problems, ensuring each solution maximizes its fitness value. The authors propose a general imitation learning framework that incorporates QD within adversarial imitation learning methods, including GAIL and VAIL. The effectiveness of QD-IL is demonstrated across three simulation tasks, showcasing its potential for learning diverse behaviors.
weaknesses: - The measure bonus calculation requires high-level task features design, which could be challenging for long-term and complex tasks. I would assume it always needs task-specific design.
- I can’t find the details about how the authors define the behavior spaces for the 3 tasks. Maybe I missed something, but I would appreciate it if the authors could explain more about how to implement the behavior measures in the experiments.
- The experiments do not contain tasks with manipulation. It would be great to show how the QD-IL scales to other robot manipulation tasks.
- Equation 3 lacks the explanation of $\tau^{E}$.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. How sensitive is QD-IL to the choice of hyperparameters? It would be helpful to have a robustness/sensitivity test.
2. Can the framework be adapted to domains with high-dimensional input data, other than the three environments stated in the paper? It is also questionable to say that the "most challenging" humanoid environment, since a pixel-based input environment like Atari is clearly more challenging. It would be helpful to make it clearer.
3. Could the author give more analysis on the learned behavior patterns? Do they really correspond to "skills" in the semantic level? E.g. moving left leg, moving right leg, etc. It should be more convincing if these diversified policies can be visualized and validated as different behavior patterns for better interpretability.
soundness: 3
strengths: 1. The writing is clear and easy to follow.
2. The problem is well-formed and the algorithm is solid. The experimental result shows clearly the superiority of the proposed QD-IL framework.
summary: The paper proposes a novel framework called Quality Diversity Imitation Learning (QD-IL) designed to enable agents to learn a wide range of skills from limited demonstrations. The framework combines quality diversity (QD) optimization with adversarial imitation learning (AIL). It also introduces measure-based reward bonuses and measure conditioning. Experimental results show that QD-IL significantly outperforms traditional methods in Mujoco tasks.
weaknesses: 1. Missing discussion of related works: there is a large literature regarding Multi-domain Imitation Learning [1] [2] [3], skill- or option-discovery [4] [5] [6] that deal with similar problems as stated in the paper. Namely, they are all designed to "find diverse solutions" so that the policies can represent diverse behaviors. If the setting is different, I suggest the authors be clearer on that; if not, these literature should be included in the discussion, and incorporated into the baselines (other than only GAIL and VAIL). It should also be suggested the advantage of the proposed QD-IL over other related works.

   

2. Currently, the learned diverse policies can only be read by the curves (like that in Fig.1.(a)). There lacks a formal and precise theoretical interpretation of how and why it is connected to diverse behavior patterns. How strongly does a local optimal suggest a behavior pattern and how well it is guaranteed? It would be helpful to develop more theoretical analysis in this regard.



   [1] Seyed Kamyar Seyed Ghasemipour, Shixiang Shane Gu, and Richard Zemel. Smile: Scalable meta inverse reinforcement learning through context-conditional policies. *Advances in Neural Information Processing Systems*, 32, 2019.

   [2] Lantao Yu, Tianhe Yu, Chelsea Finn, and Stefano Ermon. Meta-inverse reinforcement learning with probabilistic context variables

   [3] Jiayu Chen, Dipesh Tamboli, Tian Lan, and Vaneet Aggarwal. Multi-task Hierarchical Adversarial Inverse Reinforcement Learning. In *Proceedings of the 40th International Conference on Machine Learning*, pages 4895–4920. PMLR, July 2023.

   [4] Jesse Zhang, Karl Pertsch, Jiefan Yang, and Joseph J Lim. Minimum description length skills for accelerated reinforcement learning. In *Self-Supervision for Reinforcement Learning Workshop-ICLR*, volume 2021, 2021.

   [5] Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez-Gonzalez, Edward Grefenstette, Pushmeet Kohli, and Peter Battaglia. CompILE: Compositional Imitation Learning and Execution, May 2019.

   [6] Yiding Jiang, Evan Zheran Liu, Benjamin Eysenbach, J Zico Kolter, and Chelsea Finn. Learning Options via Compression.

Meta-review: The paper presents an imitation learning based approach for obtaining diverse skills. The reviewers have raised several concerns on the novelty of the ideas and particularly on the experimental results. I'm glad to see that the authors have tried to address these concerns. However, the reviewers still had outstanding concerns on the quality of writing, algorithmic details such as the reward function design, depth of experimental insight, limitation to locomotion / navigation tasks, and an analysis of the algorithmic design decisions on performance. I would hence recommend expanding the paper in the aforementioned directions along with incorporating the discussion with reviewers in the next version of the paper.

Decision: Reject