What Matters for Adversarial Imitation Learning?

1 Introduction . Reinforcement Learning ( RL ) has shown its ability to perform complex tasks in contexts where clear reward functions can be set-up ( e.g . +1 for winning a chess game ) [ 15 , 37 , 40 , 43 ] but for many real-world applications , designing a correct reward function is either tedious or impossible [ 20 ] , while demonstrating a correct behavior is often easy and cheap . Therefore , imitation learning ( IL , [ 4 , 7 ] ) might be the key to unlock the resolution of more complex tasks , such as autonomous driving , for which reward functions are much harder to design . The simplest approach to IL is Behavioral Cloning ( BC , [ 2 ] ) which uses supervised learning to predict the expert ’ s action for any given state . However , BC is often unreliable as prediction errors compound in the course of an episode . Adversarial Imitation Learning ( AIL , [ 14 ] ) aims to remedy this using inspiration from Generative Adversarial Networks ( GANs , [ 9 ] ) and Inverse RL [ 3 , 5 , 6 ] : the policy is trained to generate trajectories that are indistinguishable from the expert ’ s ones . As in GANs , this is formalized as a two-player game where a discriminator is co-trained to distinguish between the policy and expert trajectories ( or states ) . See App . C for a brief introduction to AIL . A myriad of improvements over the original AIL algorithm were proposed over the years [ 17 , 31 , 41 , 44 , 46 ] , from changing the discriminator ’ s loss function [ 17 ] to switching from on-policy to off-policy agents [ 31 ] . However , their relative performance is rarely studied in a controlled setting , and never these changes have never been compared simultaneously . The performance of these high-level choices may also depend on low-level implementation details which might be silenced in the original publications [ 19 , 29 , 36 , 42 ] , as well as the hyperparameters ( HPs ) used . Thus , assessing whether the proposed changes are the reason for the presented improvements becomes extremely ∗Equal contribution . †Univ . de Lille , CNRS , Inria Scool , UMR 9189 CRIStAL . ‡Corresponding author . E-mail : marcina @ google.com . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . difficult . This lack of proper comparisons slows down the overall research in imitation learning and the industrial applicability of these methods . We investigate such high- and low-level choices in depth and study their impact on the algorithm performance . Hence , as our key contributions , we ( 1 ) implement a highly-configurable generic AIL algorithm , with various axes of variation ( > 50 HPs ) , including 4 different RL algorithms and 7 regularization schemes for the discriminator , ( 2 ) conduct a large-scale ( > 500k trained agents ) experimental study on 10 continuous-control tasks and ( 3 ) analyze the experimental results to provide practical insights and recommendations for designing novel and using existing AIL algorithms . We release this generic AIL agent , implemented in JAX [ 25 ] as part of the Acme [ 49 ] framework : https : //github.com/deepmind/acme/blob/master/acme/agents/jax/ail . Most surprising finding # 1 : regularizers . While many of our findings confirm common practices in AIL research , some of them are surprising or even contradict prior work . In particular , we find that standard regularizers from Supervised Learning — dropout [ 10 ] and weight decay [ 1 ] often perform similarly to the regularizers designed specifically for adversarial learning like gradient penalty [ 18 ] . Moreover , for easier environments ( which were often the only ones used in prior work ) , we find that it is possible to achieve excellent results without using any explicit discriminator regularization . Most surprising finding # 2 : human demonstrations . Not only does the performance of AIL heavily depend on whether the demonstrations were collected from a human operator or generated by an RL algorithm , but the relative performance of algorithmic choices also depends on the demonstration source . Our results suggest that artificial demonstrations are not a good proxy for human data and that the very common practice of evaluating IL algorithms only with synthetic demonstrations may lead to algorithms which perform poorly in the more realistic scenarios with human demonstrations . 2 Experimental design . Environments . We focus on continuous-control tasks as robotics appears as one of the main potential applications of IL and a vast majority of the IL literature thus focuses on it . In particular , we run experiments with five widely used environments from OpenAI Gym [ 13 ] : HalfCheetah-v2 , Hopper-v2 , Walker2d-v2 , Ant-v2 , and Humanoid-v2 and three manipulation environments from Adroit [ 21 ] : pen-v0 , door-v0 , and hammer-v0 . The Adroit tasks consist in aligning a pen with a target orientation , opening a door and hammering a nail with a 5-fingered hand . Demonstrations . For the Gym tasks , we generate demonstrations with a SAC [ 28 ] agent trained on the environment reward . For the Adroit environments , we use the “ expert ” and “ human ” datasets from D4RL [ 45 ] , which are , respectively , generated by an RL agent and collected from a human operator . As far as we know , our work is the first to solve these tasks with human datasets in the imitation setup ( most of the prior work concentrated on Offline RL ) . For all environments , we use 11 demonstration trajectories . Following prior work [ 14 , 31 , 46 ] , we subsample expert demonstrations by only using every 20th state-action pair to make the tasks harder . Adversarial Imitation Learning algorithms . We researched prior work on AIL algorithms and made a list of commonly used design decisions like policy objectives or discriminator regularization techniques . We also included a number of natural options which we have not encountered in literature ( e.g . dropout [ 10 ] in the discriminator or clipping rewards bigger than a threshold ) . All choices are listed and explained in App . D. Then , we implemented a single highly-configurable AIL agent which exposes all these choices as configuration options in the Acme framework [ 49 ] using JAX [ 25 ] for automatic differentiation and Flax [ 47 ] for neural networks computation . The configuration space is so wide that it covers the whole family of AIL algorithms , in particular , it mostly covers the setups from AIRL [ 17 ] and DAC [ 31 ] . We plan to open source the agent implementation . Experimental design . We created a large HP sweep ( 57 HPs swept , > 120k agents trained ) in which each HP is sampled uniformly at random from a discrete set and independently from the other HPs . We manually ensured that the sampling ranges of all HPs are appropriate and cover the optimal values . Then , we analyzed the results of this initial experiment ( called wide , detailed description and results in App . G ) , removed clearly suboptimal options and ran another experiment with the pruned sampling ranges ( called main , 43 HPs swept , > 250k agents trained , detailed description and results in App . H ) . The latter experiment serves as the basis for most of the conclusions drawn in this paper but we also run a few additional experiments to investigate some additional questions ( App . I and App . J ) . This pruning of the HP space guarantees that we draw conclusions based on training configurations which are highly competitive ( training curves can be found in Fig . 24 ) while using a large HP sweep ( including , for example , multiple different RL algorithms ) ensures that our conclusions are robust and valid not only for a single RL algorithm and specific values of HPs , but are more generally applicable . Moreover , many choices may have strong interactions with other related choices , for example we find a surprisingly strong interaction between the discriminator regularization scheme and the discriminator learning rate ( Sec . 4 ) . This means that such choices need to be tuned together ( as it is the case in our study ) and experiments where only a single choice is varied but the interacting choices are kept fixed may lead to misleading conclusions . Performance measure . For each HP configuration and each of the 10 environment-dataset pairs we train a policy and evaluate it 10 times through the training by running it for 50 episodes and computing the average undiscounted return using the environment reward . We then average these scores to obtain a single performance score which approximates the area under the learning curve . This ensures we assign higher scores to HP configurations that learn quickly . Analysis . We consider two different analyses for each choice4 : Conditional 95th percentile : For each potential value of that choice ( e.g. , RL Algorithm = PPO ) , we look at the performance distribution of sampled configurations with that value . We report the 95th percentile of the performance as well as error bars based on bootstrapping.5 This corresponds to an estimate of the performance one can expect if all other choices were tuned with random search and a limited budget of roughly 13 HP configurations6 . All scores are normalized so that 0 corresponds to a random policy and 1 to the expert performance ( expert scores can be found in App . F ) . Distribution of choice within top 5 % configurations . We further consider for each choice the distribution of values among the top 5 % HP configurations . In particular , we measure the ratio of the frequency of the given value in the top 5 % of HP configurations with the best performance to the frequency of this value among all HP configurations . If certain values are over-represented in the top models ( ratio higher than 1 ) , this indicates that the specific choice is important for good performance . We release the raw results of our experiments7 along with a Notebook allowing to load and study it8 . 3 What matters for the agent training ? . Summary of key findings . The AIRL reward function perform best for synthetic demonstrations while − ln ( 1 − D ) is better for human demonstrations . Using explicit absorbing state is crucial in environments with variable length episodes . Observation normalization strongly affects the performance . Using an off-policy RL algorithm is necessary for good sample complexity while replaying expert data and pretraining with BC improves the performance only slightly . Implicit reward function . In this section , we investigate choices related to agent training with AIL , the most salient of which is probably the choice of the implicit reward function . Let D ( s , a ) be the probability of classifying the given state-action pair as expert by the discriminator . In particular , we run experiments with the following reward functions : r ( s , a ) = − log ( 1 − D ( s , a ) ) ( used in the original GAIL paper [ 14 ] ) , r ( s , a ) = logD ( s , a ) − log ( 1−D ( s , a ) ) ( called the AIRL reward [ 17 ] ) , r ( s , a ) = logD ( s , a ) ( a natural choice we have not encountered in literature ) , and the FAIRL [ 46 ] reward function r ( s , a ) = −h ( s , a ) · eh ( s , a ) , where h ( s , a ) is the discriminator logit . It can be 4This analysis is based on a similar type of study focused on on-policy RL algorithms [ 42 ] . 5We compute each metric 20 times based on a randomly selected half of all training runs , and then report the mean of these 20 measurements while the error bars show mean-std and mean+std . 6The probability that all 13 configurations score worse than the 95th percentile is equal 0.9513 ≈ 50 % . 7https : //storage.googleapis.com/what-matters-in-imitation-learning/data.json 8https : //storage.googleapis.com/what-matters-in-imitation-learning/analysis_ colab.ipynb shown that , under the assumption that all episodes have the same length , maximizing these reward functions corresponds to the minimization of different divergences between the marginal state-action distribution of the expert and the policy . See [ 46 ] for an in-depth discussion on this topic . We also consider clipping the rewards with absolute values bigger than a threshold which is a HP . The FAIRL reward performed much worse than all others in the initial wide experiment ( Fig . 1a ) and therefore was not included in our main experiment . This is mostly caused by its inferior performance with off-policy RL algorithms ( Fig . 22 ) . Moreover , reward clipping significantly helps the FAIRL reward ( Fig . 23 ) while it does not help the other reward functions apart from some small gains for − ln ( 1−D ) ( Fig . 82 ) . Therefore , we suspect that the poor performance of the FAIRL reward function may be caused by its exponential term which may have very high magnitudes . Moreover , the FAIRL paper [ 46 ] mentions that the FAIRL reward is more sensitive to HPs than other reward functions which could also explain its poor performance in our experiments . Fig . 25 shows that the ln ( D ) reward functions performs a bit worse than the other two reward functions in the main experiment . Five out of the ten tasks used in our experiments have variable length episodes with longer episodes correlated with better behaviour ( Hopper , Walker2d , Ant , Humanoid , pen ) — on these tasks we can notice that r ( s , a ) = − ln ( 1 −D ( s , a ) ) often performs best and r ( s , a ) = lnD ( s , a ) worst . This can be explained by the fact that − ln ( 1−D ( s , a ) ) > 0 and lnD ( s , a ) < 0 which means that the former reward encourages longer episodes and the latter one shorter ones [ 31 ] . Absorbing state ( described in App . D.2 ) is a technique introduced in the DAC paper [ 31 ] to mitigate the mentioned bias and encourage the policy to generate episodes of similar length to demonstrations . In Fig . 1b-c we show how the performance of different reward functions compares in the environments with variable length episodes depending on whether the absorbing state is used . We can notice that without the absorbing state r ( s , a ) = − ln ( 1 − D ( s , a ) ) > 0 performs much better in the environments with variable episode length which suggests that the learning is driven to a large extent by the reward bias and not actual imitation of the expert behaviour [ 31 ] . This effect disappears when the absorbing state is enabled ( Fig . 1c ) . Fig . 72 shows the performance of different reward functions in all environments conditioned on whether the absorbing state is used . If the absorbing state is used , the AIRL reward function performs best in all the environments with RL-generated demonstrations , and ln ( D ) performs only marginally worse . The − ln ( 1 − D ) reward function underperforms on the Humanoid and pen tasks while performing best with human datasets . We provide some hypothesis for this behaviour in Sec . 5 , where we discuss human demonstrations in more details . Observation normalization . We consider observation normalization which is applied to the inputs of all neural networks involved in AIL ( policy , critic and discriminator ) . The normalization aims to transform the observations so that that each observation coordinate has mean 0 and standard deviation 1 . In particular , we consider computing the normalization statistics either using only the expert demonstrations so that the normalization is fixed throughout the training , or using data from the policy being trained ( called online ) . See App . D.6 for more details . Fig . 26 shows that input normalization significantly influences the performance with the effects on performance being often much larger than those of algorithmic choices like the reward function or RL algorithm used . Surprisingly , normalizing observations can either significantly improve or diminish performance and whether the fixed or online normalization performs better is also environment dependent . Replaying expert data . When demonstrations as well as external rewards are available , it is common for RL algorithms to sample batches for off-policy updates from the demonstrations in addition to the replay buffer [ 30 , 39 ] . We varied the ratio of the policy to expert data being replayed but found only very minor gains ( Fig . 83 ) . Moreover , in the cases when we see some benefits , it is usually best to replay 16–64 times more policy than expert data . On some tasks ( Humanoid ) replaying even a single expert transitions every 256 agent ones significantly hurts performance . We suspect that , in contrast to RL with demonstrations , we see little benefit from replaying expert data in the setup with learned rewards because ( 1 ) replaying expert data mostly helps when the reward signal is sparse ( not the case for discriminator-based rewards ) , and ( 2 ) discriminator may overfit to the expert demonstrations which could result in incorrectly high rewards being assigned to expert transitions . Pretraining with BC . We also experiment with pretraining a policy with Behavioral Cloning ( BC , [ 2 ] ) at the beginning of training . Despite starting from a much better policy than a random one , we usually observe that the policy quality deteriorates quickly at the beginning of training ( see the pen task in Fig . 3 ) due to being updated using randomly initialized critic and discriminator networks , and the overall gain from pretraining is very small in most environments ( Fig . 27 ) . RL algorithms . We run experiments with four different RL algorithms , three of which are off-policy algorithms ( SAC [ 28 ] , TD3 [ 26 ] and D4PG [ 24 ] ) , as well as PPO [ 22 ] which is nearly on-policy . Fig . 7 shows that the sample complexity of PPO is significantly worse than that of the off-policy algorithms while all off-policy algorithms perform overall similarly . RL algorithms HPs . Fig . 8 shows that the discount factor is one of the most important HPs with the values of 0.97− 0.99 performing well on all tasks . Fig . 29 shows that in most environments it is better not to erase any data from the RL replay buffer and always sample from all the experience encountered so far . It is common in RL to use a noise-free version of the policy during evaluation and we observe that it indeed improves the performance ( Fig . 30 ) . The policy MLP size does not matter much ( Figs . 31-32 ) while bigger critic networks perform significantly better ( Figs . 9-10 ) . Regarding activation functions , relu performs on par or better than tanh in all environments apart from door in which tanh is significantly better ( Fig . 33 ) . Our implementation of TD3 optionally applies gradient clipping but it does not affect the performance much ( Fig . 34 ) . D4PG can use n-step returns , this improves the performance on the Adroit tasks but hurts on the Gym suite ( Fig . 35 ) .

Summary: The work conducts a detailed empirical analysis of a range of design decisions in the domain of adversarial imitation learning. The imitation learning framework is sufficiently general to encompass prior work like GAIL, AIRL, various extensions to these algorithms (e.g. regularizers), different RL algorithms used for the generator training. These decisions are tested in a comprehensive hyperparameter sweep involving over 500k trained agents across eight continuous control tasks. The results include many observations that will be useful to practitioners, such as the importance of observation normalization. It also in some cases casts doubt on received wisdom, such as the various AIL-specific methods for discriminator regularization, which seem to be little better than standard SL approaches like dropout.