A Pragmatic Look at Deep Imitation Learning

The goal of IL is to train a policy³³3In our case, a neural network with parameters $\theta$., $\hat{\pi}(a|s;\theta)$, mapping states $s$ to a distribution over actions $a$, to mimic an expert policy $\pi^{*}(a|s)$, given either the expert policy itself, or more commonly, a fixed dataset $\xi^{*}=\{\tau_{1},\ldots,\tau_{N}\}$, of trajectories $\tau=\{s_{0},a_{0},s_{1},a_{1},\ldots s_{T},a_{T}\}$ generated by the expert, where $N$ denotes the number of expert trajectories provided.

A common assumption within IL is that both the expert and our agent inhabit a Markov decision process (MDP), defined by the tuple $(\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R},p_{0},\gamma)$: $\mathcal{S}$ and $\mathcal{A}$ are the state and action spaces, ${\mathcal{T}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}}$ is the state transition dynamics, ${\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}}$ is the reward function, $p_{0}(s)$ is the initial state distribution, and $\gamma\in[0,1]$ is the discount factor (used to weight immediate vs. future rewards). The expert policy is optimal in the sense that ${\pi^{*}=\arg\!\max_{\pi\in\Pi}\mathbb{E}_{\tau\sim\pi}[R_{0}]}$, where the return at timestep $t$, $R_{t}$, is the discounted sum of rewards following a policy from state $s_{t}$ until the end of the episode at timestep $T$: ${R_{t}=\sum_{k=0}^{T-t}\gamma^{k}r_{t+k+1}}$. While in RL the goal is to interact with the environment in order to find $\pi^{*}$ (Sutton and Barto, 2018), in IL we do not have access to $\mathcal{R}$, and must instead find $\pi^{*}$ assuming that we have access to optimal trajectories.⁴⁴4In this work we do not consider the more complex settings that include suboptimal demonstrations and/or noisy observations. All following methods, unless specified otherwise, can be implemented using neural networks, providing flexible function approximation that can scale to large state and/or action spaces.

BC is very simple, and benefits from a fixed objective over a stationary data distribution. However, as $\hat{\pi}$ is only trained on $s\sim\xi^{*}$, it can fail catastrophically when it diverges from the states covered by $\pi^{*}$. In order to mitigate this, $\hat{\pi}$ must be evaluated on the environment in order to correct for discrepancies between $s,a\sim\hat{\pi}$ and $s,a\sim\pi^{*}$.

Interactive IL methods solve this issue of compounding errors (Ross et al., 2011) by iterating over running $\hat{\pi}$ in the environment, calculating $\pi^{*}(a|s)$ on $\hat{\pi}$’s state distribution, and using supervised learning on the new data (Daumé et al., 2009; Ross and Bagnell, 2010; Ross et al., 2011). While these approaches solve the data distribution issue, they require access to an interactive expert during training, which may not be available in many scenarios.

IRL instead overcomes this distribution shift by using RL to train $\hat{\pi}$ to mimic $\pi^{*}$ in the environment. The procedure consists of iterating between the following two steps:

1. 1.

   Construct a reward function⁵⁵5In the parametric case, parameterised by $\phi$, but potentially nonparametric. $\hat{\mathcal{R}}(s,a;\phi)$ using $\xi^{*}$, and optionally $\tau\sim\hat{\pi}$

2. 2.

   Train $\hat{\pi}$ using RL

RL is more complicated than the typical supervised learning setting. In particular, as the policy evolves, the data distribution changes. In the case of IRL, $\hat{\mathcal{R}}$ changing over time can introduce further non-stationarity.

Firstly, one can match the state-action distribution under $\pi^{*}$, known as the expert’s occupancy measure ${\rho_{\pi^{*}}=\mathbb{E}_{\tau\sim\pi^{*}}\left[\sum_{t=0}^{T}\gamma^{t}\mathbbm{1}_{s,a}\right]}$ (Syed et al., 2008; Ho and Ermon, 2016), or, alternatively, feature expectations (Ng et al., 2000). This is achieved using the learned reward function, and is hence dependent on the expressivity of $\hat{\mathcal{R}}$. In particular, the constant function is underspecified and allows an infinite set of solutions. Secondly, one can “penalise” following trajectories taken by (previous iterations of) $\hat{\pi}$ (Ng et al., 2000). This allows $\hat{\mathcal{R}}$ to focus on relevant parts of the state-action space, but implicitly assumes that current/past versions of $\hat{\pi}$ are suboptimal. Thirdly, one can use the maximum entropy principle (Jaynes, 1957) to find a unique best solution out of the set of solutions that match the expert’s occupancy measure/feature expectations (Ziebart et al., 2008). Using the Lagrangian multiplier $\lambda$, and denoting $H$ as the entropy, this results in the following modified RL objective: ${\arg\!\max_{\theta}\mathbb{E}_{\tau\sim\pi(s,a;\theta)}[R_{0}]+\lambda H[\pi(s,a;\theta)]}$. Entropy regularisation is a classic technique in RL (Williams and Peng, 1991).

Adversarial IL methods instead learn a reward function online using adversarial training, motivated by maximum entropy occupancy measure matching (Ho and Ermon, 2016). In generative adversarial network training (Goodfellow et al., 2014), the “generator” is trained to output samples that fool the “discriminator” $D:\mathcal{S}\times\mathcal{A}\rightarrow(0,1)$, whilst the discriminator is trained to discriminate between samples from the generator and the data distribution. This is a minimax game, in which the equilibrium solution corresponds to minimising the Jensen-Shannon divergence between the generated and real distributions. In GAIL, $\hat{\pi}$ plays the role of the generator, and the discriminator is trained on state-action pairs from $\hat{\pi}$ and $\pi^{*}$: ${\min_{G}\max_{D}\mathbb{E}_{s,a\sim\pi^{*}}[\log(D(s,a))]+\mathbb{E}_{s,a\sim\hat{\pi}}[\log(1-D(s,a))]}$. Under this formulation, higher values indicate how “expert” $D$ believes its input to be.

There are several options for constructing $\hat{\mathcal{R}}$ from $D$. Prominent examples include those introduced in GAIL, adversarial inverse reinforcement learning (AIRL; Fu et al., 2018) (corresponding to the reverse Kullback-Leibler (KL) divergence $D_{\text{KL}}(\rho_{\hat{\pi}}\|\rho_{\pi^{*}})$), and forward KL AIRL (FAIRL; Ghasemipour et al., 2020) (Table 1). As discussed by Kostrikov et al. (2019) and empirically investigated by Jena et al. (2020), there is a potential reward bias in these functions. They note that positive $\hat{\mathcal{R}}$, i.e., $-\log(1-D(s,a))$, biases agents towards survival, whereas negative $\hat{\mathcal{R}}$, i.e., $\log(D(s,a))$ biases agents towards early termination. This bias means that even constant reward functions can outperform either of these depending on the type of the environment. We recommend the original works for discussions on the properties of various reward functions (Kostrikov et al., 2019; Jena et al., 2020; Ghasemipour et al., 2020). Kostrikov et al. (2019) also make the observation that many IL algorithms do not correctly handle terminal states, and propose appending an absorbing state indicator to states, which allows IRL algorithms to properly estimate values for terminal states. This requires processing complete trajectories from both the expert and the agent, and allowing both the RL agent and the discriminator to learn from the indicator feature.

While GAIL implicitly returns a reward function, if trained to optimality then $D$ will return 0.5 for state-action pairs from both $\hat{\pi}$ and $\pi^{*}$. Finn et al. (2016) propose changing the form of the $D$ to $\frac{\exp(f(\tau))}{\exp(f(\tau))+\hat{\pi}(\tau)}$, allowing the optimal reward function to be recovered as $f(+\text{const})$. AIRL makes a practical algorithm from this by changing $D$ to operate over state-action pairs, as in GAIL, and also further disentangling the recovered reward function $f$ as the sum of a reward approximator $g(s,a)$ and a reward shaping term (Ng et al., 1999) $h(s)$: ${f(s,a,s^{\prime})=g(s,a)+\gamma h(s^{\prime})-h(s)}$, where $s^{\prime}$ is the successor state.

One of the most significant improvements to adversarial IL methods came from moving to more sample-efficient off-policy RL algorithms (Kostrikov et al., 2019; Blondé and Kalousis, 2019), which perform updates on batches of data stored in an experience replay memory (Lin, 1992). The discriminator can similarly be more efficiently trained on replay data, and although this should include an importance weighting term to account for the change in data distribution, in practice this is not needed (Kostrikov et al., 2019).

There are countless more advances within adversarial IL, making it difficult to know which techniques increase performance robustly. Orsini et al. (2021) performed a large-scale hyperparameter search over many of these methods. The key takeaways were that off-policy RL algorithms help improve sample efficiency, discriminator regularisation is key, and that hyperparameter choices which are optimal for AI-generated trajectories are not always the same for human-generated trajectories—a valuable distinction that lies out of the scope of this work. Their work also shows the importance of large-scale empirical evaluation, as their results overturned theoretical claims about the importance of discriminator regularisation (Blondé et al., 2022).

One disadvantage of adversarial training is the requirement for the discriminator, which is also undergoing training as part of the minimax game, to provide a useful training signal to the generator. There are several other IRL algorithms that also attempt to match the expert and agent’s state-action distributions, but use non-adversarial methods.

One solution is to replace the discriminator with a nonparametric model (Li et al., 2015; Dziugaite et al., 2015). Specifically, distribution matching can be achieved by minimising the maximum mean discrepancy (MMD; Gretton et al., 2012) defined over a reproducing kernel Hilbert space (RKHS). Given distributions, $P$ and $Q$, and a mapping ${\psi:\mathcal{X}\rightarrow\mathcal{H}}$ from features $X\in\mathcal{X}$ to an RKHS $\mathcal{H}$, the MMD is the distance between the mean embeddings of the features: ${\text{MMD}(P,Q)=\|\mathbb{E}_{x\sim P}[\psi(x)]-\mathbb{E}_{y\sim Q}[\psi(y)]\|_{\mathcal{H}}}$. Using a kernel function $k$, one can calculate $\text{MMD}^{2}(P,Q)=\mathbb{E}_{x,x^{\prime}\sim P}k(x,x^{\prime})+\mathbb{E}_{y,y^{\prime}\sim Q}k(y,y^{\prime})-2\mathbb{E}_{x\sim P,y\sim Q}k(x,y)$.

GMMIL (Kim and Park, 2018) extends this principle to the IL setting. Dropping terms that are constant with respect to $\hat{\pi}$, GMMIL has the reward function:

where $M$ and $N$ are the number of state-action pairs from $\pi^{*}$ and $\hat{\pi}$, respectively.

Two disadvantages of GMMIL are that 1) the “discriminator” cannot learn relevant features, and 2) it has $O(MN)$ complexity. RED (Wang et al., 2019) solves these issues by building upon random network distillation (RND; Burda et al., 2018). In RND, a predictor network ${f_{\phi}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}^{K}}$ is trained to minimise the mean squared error (MSE) against a fixed, randomly initialised network ${f_{\bar{\phi}}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}^{K}}$. Empirically, the MSE indicates how out-of-distribution new data is. RED uses a Gaussian function over the MSE, resulting in

where $\sigma$ is a bandwidth hyperparameter. Wang et al. (2019) interpret RND as an approximate support estimation method, and hence the RED reward function encourages the agent to have a support over its state-action distribution that matches the expert’s.

Similarly to RED, DRIL (Brantley et al., 2020) constructs a reward function based on the disagreement between models trained on the expert data, and can also be interpreted as a support estimation method. However, unlike the other methods which operate over the joint distribution of state-action pairs, DRIL builds simply upon BC, operating over $p(a|s)$. DRIL first trains an ensemble of $E$ different policies using the BC objective (Equation 1) on the expert data, and then uses a function of the (negative of the) variance between the policies to estimate a reward for the agent:

where the $q$ is a top quantile of the uncertainty cost computed over the expert dataset. While deep ensembles are known to produce reasonable uncertainty estimates (i.e., variance in outputs) on out-of-distribution data (Lakshminarayanan et al., 2017), it is also possible to approximate them using sampling with dropout (Srivastava et al., 2014). Brantley et al. (2020) showed empirically that this performed comparatively to using independent models.

Another view on distribution matching, known as moment matching⁸⁸8Matching the moments of the model distribution to the empirical target distribution., can be achieved through optimising integral probability metrics (IPM; Müller, 1997). IPMs provide a distance function between two distributions, $\sup_{f\in\mathcal{F}}\mathbb{E}_{x\sim P}[f(x)]-\mathbb{E}_{y\sim Q}[f(y)]$, for a function class $\mathcal{F}$ containing functions $f:\mathcal{X}\rightarrow\mathbb{R}$. Different function classes $\mathcal{F}$ recover different IPMs; for instance, $\mathcal{F}=\{f:\|f\|_{\mathcal{H}}\leq 1\}$ with RKHS $\mathcal{H}$ gives the MMD and $\mathcal{F}=\{f:\|f\|_{L}\leq K\}$ with bounded Lipschitz constant $K$ gives the Wasserstein distance.

We evaluate all algorithms on the popular MuJoCo simulated robotics benchmarks: Ant, HalfCheetah, Hopper, and Walker2D (Todorov et al., 2012; Brockman et al., 2016). To improve reproducibility and enable fairer comparison against other reported results, we use the D4RL “expert-v2” trajectory datasets (Fu et al., 2020).⁹⁹9Although this benchmark was developed for offline RL, we use it for IL by ignoring the saved rewards. When loading the expert data we process each episode to distinguish between “true” and time-dependent terminations (Pardo et al., 2018), and provide absorbing state indicators (Kostrikov et al., 2019); these are also tracked for agent episodes. By default, we maximise available data by not subsampling expert transitions. We choose 3 different trajectory “budgets” for the IL algorithms to learn from: 5, 10 and 25 expert trajectories.

All IL algorithms use SAC (Haarnoja et al., 2018a), with automatic entropy tuning (Haarnoja et al., 2018b), as the base RL agent, as theoretically required by SQIL and AdRIL, and as was empirically shown to be performant for adversarial IL algorithms (Orsini et al., 2021). The actor applies a tanh transformation to scale actions $\in(-1,1)$, and dual critics are trained to reduce value overestimation (Fujimoto et al., 2018). We also include BC as a baseline, with the same actor architecture. All algorithms are optimised with AdamW (Loshchilov and Hutter, 2019). We use PyTorch (Paszke et al., 2019) for all of our code.

We group the IL algorithms and their variants into 6 key methods: AdRIL, DRIL, GAIL, GMMIL, PWIL, and RED. Our hyperparameter search spaces were determined based on hyperparameter ranges within the original works, a large subset of options tried by Orsini et al. (2021), and general hyperparameters such as learning rate and batch size, resulting in 7-18 hyperparameters to tune per algorithm. For each trajectory budget, we give each algorithm 30 hyperparameter evaluations using Bayesian optimisation (Balandat et al., 2020), with the minimum of the cumulative reward for each of the 4 environments used as the optimisation objective, which can be seen as minimising regret over a set of environments. We then evaluate agents over 10 seeds with the best hyperparameters found, and report performance according to best practices (Agarwal et al., 2021). We therefore train 2880 agents for the final results, not including the extensive training and testing of agents performed while replicating and augmenting IL algorithms.

DRIL, GAIL and RED include several options for their trained discriminators, including network hidden size, depth, activation function, dropout, and weight decay. The GAIL discriminator has additional options, detailed below.

AdRIL options include balanced sampling (alternating sampling expert and agent data batches vs. mixed batches, Swamy et al., 2021), and the discriminator update frequency $\geq 1$, which determines the number of “rounds”. We also include “0” in the search space, which if chosen reverts to using the SQIL reward function.

DRIL options include the quantile cutoff $\in[0,1]$. As per the original work (Brantley et al., 2020), we include an auxiliary BC loss during training. For simplicity we use the dropout ensemble. As using absorbing state indicators requires importance sampling, we adapt the BC loss used within DRIL to account for importance weights.

GAIL options include reward shaping (Fu et al., 2018), subtracting $\log\hat{\pi}(a|s)$ from the predicted reward (Fu et al., 2018), the GAIL, AIRL and FAIRL reward functions (Ho and Ermon, 2016; Fu et al., 2018; Ghasemipour et al., 2020), discriminator gradient penalty (Kostrikov et al., 2019; Blondé and Kalousis, 2019), discriminator spectral normalisation (Blondé et al., 2022), discriminator entropy bonus $\geq 0$ (Orsini et al., 2021), binary cross-entropy, Mixup, and nn-PUGAIL discriminator loss functions (Ho and Ermon, 2016; Chen et al., 2021; Xu and Denil, 2021), and 3 additional hyperparameters for these loss functions (Mixup alpha $\geq 0$, positive class prior $\geq 0,\leq 1$, and non-negative margin $\geq 0$).

To adapt GMMIL for absorbing states, we adapt it to use the weighted MMD (Yan et al., 2017). Due to the $O(MN\cdot\dim(X))$ complexity of MMD, it is prohibitive to use the entire expert dataset per update in the off-policy setting, and hence we randomly sample a batch of expert transitions per update. For this reason we also restrict the maximum batch size of GMMIL’s hyperparameter search space.

PWIL options include the reward scale $\alpha\geq 0$ and reward bandwidth scale $\beta\geq 0$. To make PWIL more comparable to the other algorithms, we use the “nofill” variant, which does not prefill the replay buffer with expert transitions (Dadashi et al., 2021).

To adapt RED for absorbing states, we train its discriminator with a weighted MSE loss.

The hyperparameter search spaces and optimised hyperparameters can be found documented in the codebase.

All IL algorithms (except BC) are trained with $10^{6}$ environment interaction steps, and are evaluated 30 times with a deterministic policy every $10^{4}$ steps. To aggregate test returns, we calculate the interquartile mean (IQM) over the 30 evaluation episodes, and then once again over seeds, reporting the final IQM ± 95% confidence interval (CI) using 1000 stratified bootstrap samples (Agarwal et al., 2021). The final returns over all environments are reported in Table 2.¹⁰¹⁰10For reference we also include the results of our SAC implementation, trained for $3\times 10^{6}$ steps. We also show performance over time in Figure 1, where the returns are aggregated over environments, normalised by the D4RL environment min and max reference scores: $\text{normalised score}=\frac{\text{score - random score}}{\text{expert score - random score}}$, with the reference scores produced by random and expert agents.

Overall, GAIL performs robustly across all trajectory budgets, presumably due to the large amount of research that has gone into improving the performance and robustness of adversarial IL methods. However, AdRIL, which is practically much simpler (and faster), performs the same (with higher sample efficiency) with higher budgets, as near-expert data from the agent is less likely to overpower the effect of the expert data. This can also be seen in the hyperparameter optimisation, with the discriminator update frequency going from 0 (reverting to the SQIL reward function) to 12500 to 25000 as the trajectory budget increased. PWIL also performed well across budgets. Hyperparameter optimisation set the reward ($\alpha$) and reward bandwidth ($\beta$) scales to 1 for 5 trajectories, but increased the values for these for 10 and 25 trajectories. However, we note that since each environment step requires computing the reward between the current state-action pair and a subset of the expert data, starting with the entire dataset at the start of each episode, it is computationally expensive for high budgets. BC scales well as the number of trajectories increase.

Unfortunately, we were unable to successfully optimise off-policy versions of DRIL, GMMIL and RED. In an early version of our codebase we were able to optimise their original, on-policy versions successfully, so with considerable effort put into hyperparameter tuning and trying additional regularisation strategies, we believe that their is some fundamental issue caused from going from training on on-policy to off-policy returns. One would expect that for successful training the inferred rewards for the agent’s trajectories should increase over time, but this was observed for runs of these methods as well, and is therefore not predictive of successful imitation. Q-values are a function of the predicted rewards, so did not provide further diagnostic insights. We also created variants of DRIL and RED in which the discriminators were trained online, similarly to GAIL, but were unsuccessful; however, there are many ways to do so, and our attempts do not preclude a successful online discriminator variant from being developed. Finally, we note that DRIL performs better at the start of training, which we can attribute to the BC auxiliary loss; experiments with uncertainty-only DRIL (UO-DRIL; Brantley et al., 2020) did not show this trend, with scores remaining low during the entirety of training.

Some weak trends we noticed from hyperparameter optimisation were that both batch and discriminator sizes increased with the trajectory budget. We hypothesise that the former is due to added stochasticity in optimisation aiding when data is scarce, whilst the latter is due to the need to prevent overfitting in low-data regimes. However, we caution that these trends do not always hold, as, for example, the optimal batch size for GAIL decreased with the trajectory budget.