Robust Adversarial Imitation Learning via Adaptively-Selected Demonstrations

1 Introduction

Imitation learning (IL) [Hussein et al., 2017] , which aims to let the agent imitate the behavior of the expert, has achieved impressive performance in many applications [Englert et al., 2013; Codevilla et al., 2018]. Compared to the intractable reward engineering [Amodei et al., 2016] in reinforcement learning (RL) [Sutton and Barto, 2018], IL provides a much easier way to directly infer a reward function from expert demonstrations, in such a way that the agent can greatly improve on the demonstrated behavior.

Generative adversarial imitation learning (GAIL) [Ho and Ermon, 2016] is one of the state-of-the-art IL methods. Based on the framework of GAN [Goodfellow et al., 2014], GAIL

∗Corresponding author

regards imitation learning as a distribution matching problem between the state-action distribution of expert policy and that of the agent policy. Many variants of GAIL have been proposed in recent years to further improve GAIL's performance from different aspects. Notice that GAIL does not directly recover a reward function since the discriminator will output 0.5 for each state-action pair at the saddle point, AIRL [Fu et al., 2017] separates the reward function from the discriminator, and can learn the reward function during the agent training. VAIL [Peng et al., 2018] uses variational information bottleneck (VIB) to constrain the ability of the discriminator since the discriminator may easily overwhelm the policy in GAIL training, which could result in the uninformative reward. So VAIL could improve upon GAIL by stabilizing the GAIL training. InfoGAIL [Li et al., 2017] follows the idea of InfoGAN [Chen et al., 2016], which aims to learn a multimodal policy with latent variable.

GAIL and its variants have shown great success in benchmark tasks with the underlying assumption that expert demonstrations are optimal. However, optimal demonstrations are often hard to acquire in real-world tasks. For example, since the number of most qualified experts is limited, the demonstrations collected in the real world are more likely to be a combination of demonstrations from both qualified experts and inferior experts. Furthermore, a qualified expert would make mistakes due to the limited energy and presence of distraction, resulting in imperfect demonstrations. As stated above, the imperfect demonstrations issue occurs and it could seriously mislead the training procedure of IL.

Existing solutions proposed to address the imperfect demonstrations issue rely heavily on the prior information or assumptions of demonstrations. In 2IWIL and IC-GAIL [Wu et al., 2019], a fraction of demonstrations labeled with confidence score need to be provided first in training. Preferencebased inverse reinforcement learning methods such as T-REX [Brown et al., 2019] need ranked trajectories to learn a relevant reward function and then conduct RL with the new learned reward. Multi-modal imitation learning method InfoGAIL can be used to recover all the demonstrators within demonstrations, however a prior about the number of demonstrators is needed first and we also cannot find the best demonstrator before we evaluate all the modal.

To address the weakness of existing solutions, we propose a new approach based on GAIL called Selective Adversarial Imitation Learning (SAIL) to solve imperfect demonstrations issue without auxiliary information in this paper. In SAIL, we first present a weighted adversarial imitation learning framework based on the Wasserstein variant of GAIL and a simple binary weight is assigned to each training demonstration to indicate whether to choose it into the training set. The weight itself can be also interpreted as the quality of the demonstration from another perspective. Instead of asking for auxiliary information about the weight such as 2IWIL, SAIL considers weight can be exactly determined by the feedback from GAIL model. Motivated by that the good demonstration is always related to a higher reward while the bad demonstration is always related to a lower reward, we ask the reward function in Wasserstein GAIL to predict weight for each demonstration. After we obtain this relation between weight prediction and model training, both of them can be jointly optimized as a whole. Besides hard binary weighting, we also propose a soft weighting scheme for SAIL. We provide a theoretical analysis to show that SAIL can converge to the saddle point. The experiment result shows the advantage of SAIL compared to baseline methods when dealing with imperfect demonstrations and the idea of SAIL can be easily extended to other GAIL-based methods such as InfoGAIL.

2 RelatedWork

Several relevant researches to address imperfect demonstration issue are given below, however most of methods require prior information or assumption about demonstrations, which constrains them to be applied to a more universal setting.

Preference learning [Christiano et al., 2017] based IRL [Ziebart et al., 2008] methods have shown to be effective when handling imperfect demonstrations, especially in Atari games. With a set of ranked trajectories, T-REX [Brown et al., 2019] aims to learn a reward function that can well fit the rankings. However, these ranked trajectories may not be available in the standard IL setting. Beyond T-REX, D-REX [Brown et al., 2020] is proposed to acquire ranked trajectories automatically. Specifically, D-REX firstly learns a base policy by behavioral cloning (BC) [Bain and Sammut, 1995] over imperfect demonstrations and then derives noisy policies to generate ranked trajectories. Multi-modal imitation learning methods can be used to handle imperfect demonstrations if imperfect demonstrations are from multiple demonstrators with diverse quality. InfoGAIL [Li et al., 2017] learns a multi-modal policy where each mode of the policy represents a single demonstrator and good policies can be then obtained by selecting good modes of the multi-modal policy. VILD [Tangkaratt et al., 2019] models the distribution of multimodal demonstrations with the assumption that each demonstrator is formulated by adding different ratios of Gaussian noise to the optimal demonstrator. Thus it may suffer a performance degradation under more universal settings. From the perspective of weighting samples, imperfect demonstrations issue can be solved by assigning proper weight for different demonstrations. With a fraction of labeled demonstrations, SSIRL [Valko et al., 2013] uses semi-supervised support vector machines [Bennett and Demiriz, 1999] to separates good or bad demonstrations when learning a policy, thus improving the performance of the agent. [Wu et al., 2019] introduces weight into GAIL framework by considering a fraction of demonstrations labeled with confidence are available for training. [Wang et al., 2021] also conducts a weighted GAIL framework and learns to weight imperfect demonstrations during training without exposing too much prior.

3 Preliminaries

The framework of reinforcement learning is generally based on the Markov Decision Process (MDP) [Puterman, 1994]. An MDP consists of five elements 〈S, A, P, R, γ〉, where S is a set of states, A is a set of actions, P : S ×A× S → R is the transition probability distribution, R : S × A → R is the reward function, and γ ∈ [0, 1] is the discounting factor for future rewards. The return in the MDP is calculated as the discounted sum of rewards obtained by the agent over all episodes, the goal of RL is thus to learn a policy π : S → A that can maximize the expected return over all episodes. For any policy π, there is a one-to-one correspondence between the policy and its occupancy measure ρπ : S × A → R.

Generative Adversarial Imitation Learning. The general framework of Generative Adversarial Networks (GANs) had been applied in the imitation learning problem, which results in GAIL algorithm. In GAIL, the classical imitation learning problem is treated as an occupancy measure matching between the expert policy and the agent policy via JensenShannon Divergence. A discriminator Dψ is introduced to distinguish expert transitions from agent transitions, while the agent is to "fool" the discriminator into taking agent transitions as those expert transitions. Formally, the objective function of GAIL is written as

where ρπθ and ρπE denote the occupancy measures of agent policies πθ and πE, respectively.

The plain GAN is known to easily suffer unstable training issues (i.e. difficult to balance the performance of generator and discriminator), and GAIL may inevitably inherit these properties. WGAN [Arjovsky et al., 2017], which minimizes the Wasserstein distance instead of Jensen-Shannon divergence between distributions, is a well-known method to solve unstable training problem in GAN. While only few modifications need to be made from original GAN to WGAN, we can easily extend this idea into GAIL. The objective function of Wasserstein GAIL (WGAIL) [Xiao et al., 2019] can be defined as

where the critic rϕ is the reward function, Ψ(rϕ) = -E(s,a)∼ρˆπ (||∇rϕ(ˆ s, ˆa)||2 - 1)2 is a regularizer term to satisfy the Lipschitz constraint on reward rϕ and λ is the hyperparameter. When training to convergence, Wasserstein GAIL may learn the expert policy as well as recover the corresponding reward function.

4 Methodology

Despite GAIL has achieved great success in benchmark tasks, its performance relies highly on the quality of expert demonstrations. However in some real-world imitation learning tasks, the number of most qualified experts is often limited, which means more demonstrations might come from the inferior experts. The imperfect demonstration issue therefore occurs in imitation learning, which seriously influences the performance of the learned agent.

A natural idea to exploit these imperfect demonstrations in imitation learning is to weight their influence in the objective function. In particular, we introduce a weight w(s, a) for each expert state-action pair, and then obtain the following weighted Wasserstein GAIL

However Eq. (3) could be intractable for the following reasons. In classical imitation learning, state and action pairs are the only information we collect from experts, while the ground-truth weights of the demonstrations are unknown. We therefore expect an appropriate estimation of the weights, so that they can well indicate the "goodness" of expert demonstrations.

For simplicity, we begin with the analysis on the binary weight, i.e. w(s, a) = {0, 1}. In particular, we choose a demonstration for training by setting w(s, a) = 1; on the other hand, w(s, a) = 0 implies that the corresponding demonstration will not be chosen as it does not participate in the calculation in Eq. (3). This extreme case thus provides an intuitive explanation on the use of the weight, i.e., setting the weights of good demonstration to 1 while removing those bad demonstrations via zeroing their weights. By doing this, the negative influence of the imperfect demonstrations could be largely alleviated and the performance of the agent learning can be preserved. However, the remaining question is a criterion to determine whether the weight of a given state-action pair is 1 or 0.

Recall that the aim of reinforcement learning (including imitation learning) is to maximize the reward of the agent. If a state-action pair produces a larger reward, we are then confident in taking it as a good demonstration. Given this fact, we tend to estimate the "goodness" of the expert demonstration by investigating its reward. According to Eq. (3), the reward of the state action pair (s, a) can be calculated with the help of the critic rϕ, i.e. rϕ(s, a). Given a threshold K, the demonstration with reward larger than K can be regarded as good demonstrations with weight equal to 1, otherwise the weight can be set to 0, i.e.

We find that the optimal weight w∗ (s, a) could be dynamically determined by the critic function rϕ in Wasserstein GAIL. Along with the optimization of rϕ, we proceed to rewrite Eq. (4) to the following maximization problem,

Hence, by integrating the objective function of w into the weighted Wasserstein GAIL, we achieve

The weight learning and Wasserstein GAIL training can thus be optimized alternately. For weight learning, the reward function rϕ in Wasserstein GAIL can act as a critic to weight each state-action pair. The state-action pair of a larger estimated reward will get a larger weight (e.g. w = 1) otherwise it will get a smaller weight (e.g. w = 0) in iterations. For Wasserstein GAIL training, agent policy πθ can be optimized with the solved weight variable w(s, a) from the last iteration. K is the threshold to distinguish expert demonstrations. By gradually decreasing K during the optimization, we would lower the standard to accept more expert demonstrations to be good and allow their participation in the optimization of agent policy.

Remark. The algorithm stemming fromEq. (6) is related to self-paced learning, since it will dynamically determine the pace to include "good" demonstrations into the optimization of the agent policy. However compared with the classical self-paced learning [Kumar et al., 2010], we highlight the following novelties of the proposed algorithm.

• Self-paced learning has been widely studied in various problems, including event detection [Jiang et al., 2015], object tracking [Supancic and Ramanan, 2013] and multi-view clustering [Xu et al., 2015]. But to the best of our knowledge, this is the first time to investigate the capability of self-paced learning in dealing with imperfect demonstrations in imitation learning.
• The weight w contributes to setting up a three-player adversarial game in Eq. (6), where w and the critic rϕ are optimized to maximize the Wasserstein distance of occupancy measure between agent policy πθ and expert policy πE, while the agent policy stood by θ is optimized to minimized the Wasserstein distance. In contrast, model parameters and the weight in classical self-paced learning often lead to a simple minimization problem (e.g. minimizing classification loss), and there is no complex adversarial process.
• At last, the motivations of the weighting scheme are different. Classical self-paced learning tends to progressively include examples whose loss is smaller than a threshold, while we expect to use the demonstrations whose reward is larger than a threshold.

Soft Weighting. In Eq. (4), the binary weight w = {0, 1} simply classifies the demonstrations into two categories, however, for some certain demonstrations it is hard to tell that one is absolute "good" or "bad" demonstration. So the soft weighting scheme can be conducted and weight w can be regarded as the probability of the demonstration to be "good" demonstration. Beyond Eq. (4), we suggest the optimal soft weight as

Notice that Eq. (7) can be exactly adapted to the sigmoid function. The soft weight in Eq. (7) is a smooth function of hard weight in Eq. (4). Instead of directly separating the demonstrations into "good" and "bad" by the hard weight, the soft weighting scheme provides a more feasible interpretation to address the probability or likelihood of a demonstration to be "good". So we may have a more robust measurement about the quality of demonstrations, especially for those demonstrations whose reward is around the threshold K.

To further naturally embed the optimization of w into Wasserstein GAIL, we can develop the following objective function as a reformulation of Eq. (7),

whose solution is exactly w∗ in Eq. (7). By defining

we can then easily extend Eq. (6) to a soft weight version,

Optimization. Eq. (6) and (10) can be solved by ACS (Alternative Convex Search) [Bazaraa et al., 2013], which suggests that weight w and GAIL can be alternatively updated.

The threshold K should be declined to include more samples during the training process. Different from [Kumar et al., 2010], which provides a fixed way to update K, we consider using a new dynamic way to update the threshold K to ensure that more demonstrations could be included in during training. Suppose Kj is the threshold in the jth iteration, and we can calculate it as

where rmid and rmin are the median and min value within batched demonstration reward estimated by rϕ, m is the hyper-parameter, Npre and N denote the number of pretraining iteration and total iteration. In Eq. (11), the threshold K is always constrained to be in the interval [rmin, rmid] and gradually decreases. We can also easily observe that KN > rmin is satisfied. So the idea underlying Eq. (11) is to ensure more demonstrations could be gradually included during training while a few outliers will not be considered eventually. The hyper-parameter m is used to control the speed to include good demonstrations into training. For example, Ki would decrease slowly in the initial stage when m is large, which means it may slowly include more samples during the early training.

To demonstrate the advantage of our proposed dynamic decrease way to update threshold K, we compare it with two naive fixed update way in self-paced learning, i.e., exponential decrease K = K/µ, linear decrease K = K - µ, the learning curve of three decrease ways is provided in Figure 1. Convergence Analysis. We also provide theoretical analysis on the convergence of SAIL. We first show that the objective function of SAIL can converge with respect to θ by Theorem

1. Then we discuss the convergence of ϕ and w respectively since they are independent of θ. After showing all the variables can converge to a unique saddle point, we can therefore alternatively update them in each iteration.

Theorem 1. Rewrite the objective function of SAIL as minθ maxw,ϕ Jw,ϕ(θ), where

Assume that δk is the upper bound of the KL-constraint step in TRPO in kth iteration, i.e DKL(πk θ , πk+1 θ ) ≤ δk and rϕ is bounded by M, then Jw,ϕ(θ) can converge with repect to θ.

The proof of Theorem 1 is provided in supplementary material. As stated in Theorem 1, Jw,ϕ(θ) can converge with respect to θ, which means for any fixed w and ϕ we have Jw,ϕ(θk) converges uniformly to Jw,ϕ(θ∗ ). On the other hand, Jw,ϕ(θ) is a continuous and concave function of ϕ, so we have ϕ∗k = arg supϕ Jw,ϕ(θk) also converges to ϕ∗ = arg supϕ Jw,ϕ(θ∗ ). This result implies that there is a saddle point (θ∗ , ϕ∗ ) in the objective function of SAIL and both θ and rϕ can converge to their optimal solutions alternately by gradient descent. As for weight w, since it is largely determined by the rϕ dynamically, it could also converge as rϕ reaches convergence. Another intuitive explanation on convergence of weight can be stated from self-paced learning. The weight could reach convergence since all the weight is set to 1 at the end of the training.

5 Experiment

In this section, we are going to figure out two questions with the experiment. (1) Does SAIL outperform other comparison methods when dealing with imperfect demonstrations? (2) Can the idea of SAIL be extended to other GAIL-based methods?

5.1 Experiment Setting

We conduct experiments to evaluate our proposed method on four Mujoco [Todorov et al., 2012] continuous control tasks with two kinds of imperfect demonstrations, i.e. suboptimal demonstrations (stage 1) and near-optimal demonstrations (stage 2). Six checkpoints (π1, π2, π3, π4, π5, π6) with increasing quality are used to form stage 1 and 2 demonstrations. Stage 1 demonstrations are formed by trajectories from π1, π2, π3 and π4 while stage 2 demonstrations are formed by trajectories from π1, π2, π5 and π6, so the quality of stage 2 demonstrations is obviously better.

The reward function rϕ in WGAIL and SAIL is constrained into [0, 5] by a sigmoid function. As [Kumar et al., 2010] suggested, we conduct pre-training on WGAIL with about 10% of total interactions before the weight learning step in SAIL. The threshold K is initialized such that half of the demonstrations can be included. We evaluate the agent every 5,000 transitions in training and the reported result in Table 1 is the average of the last 100 evaluations. Also, we conduct our experiment with five random seeds. More details on data quality and the hyper-parameters used to reproduce the results are available in the supplementary material.

5.2 Performance

The result provided in Table 1 shows that the SAIL performs generally better than WGAIL and GAIL, which demonstrates the effect of adding weight for demonstrations. Recall that the Wasserstein variant can improve the performance over GAIL by improving the stability of GAN training empirically. We can observe that WGAIL could outperform GAIL when dealing with highly sub-optimal demonstrations (Stage 1), but it does not show quite an advantage when dealing with nearoptimal demonstrations (Stage 2). We think the reason is that plain GAIL can generally perform well given near-optimal demonstrations (Stage 2), however it may suffer serious unstable training issues with highly sub-optimal demonstrations (Stage 1), which makes it hard to perform well as the results suggested. By contrast, WGAIL can well solve the unstable training problem so it could perform better than GAIL under highly imperfect demonstrations setting, but it is not likely to substantially outperform GAIL much. As for SAIL, credit to the use of weight, it can have a significant improvement over WGAIL with both stage 1 and stage 2 demonstrations. The soft label performs better than the hard label at most time and further improves the performance of SAIL.

We also compare SAIL to other GAIL-based methods such as InfoGAIL, AIRL and VAIL. Since the performance of InfoGAIL is measured by the average of all recovered modals, it may not outperform GAIL. AIRL and VAIL are proposed to address different issues about GAIL, which may help them achieve better results. As the results suggested, AIRL and VAIL only have minor improvements compared to plain GAIL and both of them do not perform well when dealing with highly imperfect demonstrations. D-REX is an IRLbased method proposed to handle imperfect demonstrations without auxiliary information, which properly suits our setting. However, since its performance may highly rely on the initial policy trained by BC, imperfect demonstrations may therefore have a negative influence on the final result of DREX. 2IWIL requires auxiliary confidence scores as a prior for training and we can observe that SAIL reaches comparable performance with 2IWIL.

Extend to InfoGAIL. We also show that the idea of SAIL can be easily adapted to other GAIL-based methods such as InfoGAIL, as shown in Table 1. Follow the idea of SAIL, we first rewrite InfoGAIL into its Wasserstein variant (InfoWGAIL), then the weight and a regularizer on weight can be added into InfoWGAIL to get InfoSAIL. InfoGAIL performs well in near-optimal demonstrations (Stage 2), especially in HalfCheetah-v2 it nearly reaches the performance of the optimal policy. However, when given highly sub-optimal demonstrations, InfoGAIL may also inevitably suffer unstable training issues and does not perform quite well like plain GAIL. Wasserstein variant might help for improving the performance, but there is an upper bound for further improving its performance. Due to the use of weight, InfoSAIL can perform generally better than InfoWGAIL while the soft label is slightly better than the hard label generally, which is consistent with GAIL. As the result suggested, we believe the idea of SAIL can be easily adapted to other GAIL-based methods to help deal with imperfect demonstrations.

5.3 Sample Efficiency

We provide an intuitive model optimization trajectories map via a 2D weight-space slice in Figure 2, from which we can conclude that SAIL is sample efficiency. The optimization trajectories are plotted along the contours of the average reward surface for this slice. To plot this figure, we follow the instruction of [Li et al., 2018]. Let θi denotes the parameter of policy model πθi in ith iteration and θn is the final policy model parameter πθn . Given the matrix M = [θ0 - θn; · · · ; θn-1 - θn], we can apply PCA to M and then select two most explanatory directions. The parameters of the policy model in different checkpoints are then projected onto the plane defined by these two directions for plotting. The model on each point is evaluated by the average cumulative reward. We save 10 checkpoints of the policy model in every 5×105 interactions during training, then plot the optimizer trajectories (Red dots) and average reward surface along with PCA directions.

We can observe the advantages of SAIL in Figure 2 from two aspects. First, SAIL can finally achieve higher performance than WGAIL when given sub-optimal demonstrations (Stage 1). Second, both SAIL and WGAIL achieve similar performance when given near-optimal demonstrations (Stage 2), but SAIL can converge fast.

In sub-optimal demonstrations (Stage 1), SAIL's final model (see the star symbol) is always located in the higher reward zone of parameter space while the model of WGAIL is not, as shown in Figure 2(a). Therefore, we can conclude that SAIL can achieve higher performance than WGAIL when given sub-optimal demonstrations. When given near-optimal demonstrations (Stage 2), both WGAIL and SAIL can perform well. The advantage of SAIL under this condition is that it can converge fast, which means it is sample efficiency. As shown in Figure 2(b), we can observe that in Ant-v2, both SAIL and WGAIL start heading to the optimal solution directly at around 1.5 × 106 interactions (see the first large point). However, SAIL approaches the final high reward zone at around 2.5 × 106 interactions while GAIL approaches the high reward zone at around 3.5 × 106 interactions (see the second and third large points). The same thing happens in HalfCheetah-v2. SAIL approaches the high reward zone at around 2.5×106 interactions. However inWGAIL, only the final policy model reaches the high reward zone. As stated above, we can conclude that SAIL is sample efficiency.

6 Conclusion

In this paper, we propose a new method called SAIL to improve the performance of GAIL when dealing with imperfect demonstrations. We use Wasserstein GAIL as a baseline and a binary weight variable can be assigned to each expert demonstration to address their influence on the objective function. The weight can be also interpreted with whether to choose the demonstration into training or not. Different from some existing solutions which require prior information about weight, we can automatically acquire this weight. Besides hard binary weighting, we also propose a soft weighting scheme for SAIL. Experiment results on Mujoco show that the advantage of SAIL over compared methods and our idea can be also easily adapted to other GAIL-based methods.

Acknowledgements

This work was supported in part by the National Key R & D Program of China under Grant 2018AAA0101100, National Natural Science Foundation of China under Grants 61822113, 41871243, the Science and Technology Major Project of Hubei Province (Next-Generation AI Technologies) under Grant 2019AEA170, the Australian Research Council under Projects DE180101438 and DP210101859.