Discriminator-Weighted Offline Imitation Learning from Suboptimal Demonstrations

1 INTRODUCTION . The recent success of reinforcement learning ( RL ) in many domains showcases the great potential of applying this family of learning methods to real-world applications . A key prerequisite for RL is to design a reward function that specifies what kind of agent behavior is preferred . However , in many real-world applications , designing a reward function is prohibitively difficult ( Ng et al. , 1999 ; Irpan , 2018 ) . By contrast , imitation learning ( IL ) provides a much easier way to leverage the reward function implicitly from the collected demonstrations and has achieved great success in many sequential decision making problems ( Pomerleau , 1989 ; Ng et al. , 2000 ; Ho & Ermon , 2016 ) . However , popular IL methods such as behavioral cloning ( BC ) and generative adversarial imitation learning ( GAIL ) ( Ho & Ermon , 2016 ) , assume the expert demonstration is optimal . Unfortunately , it is often difficult to obtain optimal demonstrations for many real-world tasks , because human experts often make mistakes due to various reasons , such as the difficulty of the task , partial observability of the environment , or the presence of distraction . Given such noisy expert demonstrations , which contain records of both optimal and non-optimal behaviors , BC and GAIL all fail to imitate the optimal policy ( Wu et al. , 2019b ; Ma , 2020 ) . Current methods that deal with suboptimal demonstrations either require additional labels , which can be done explicitly by annotating each demonstration with confidence scores by human experts ( Wu et al. , 2019b ) , or implicitly by ranking noisy demonstrations according to their relative performance through interacting with the environment ( Brown et al. , 2019 ; 2020 ) . However , human annotation and environment interaction are laborious and expensive in real-world settings , such as in medicine , healthcare , and industrial processes . In this work , we investigate a pure offline learning setting where the agent has access to neither the expert nor the environment for additional information . The agent , instead , has only access to a small pre-collected dataset of state-action-next state transition triples sampled from the expert and a large batch offline dataset sampled from one or multiple behavior policies that could be highly sub-optimal . Prior works that address the problem are based on variants of BC or inverse RL , Sasaki & Yamashina ( 2021 ) reuse another policy learned by BC as the weight of original BC objective , however , this requires that expert data occupy the majority proportion of the offline dataset , otherwise the policy will be misguided to imitate the suboptimal data . Zolna et al . ( 2020a ) first learns a reward function that prioritizes expert data over others and then performs offline RL based on this reward function . This algorithm is extremely expensive to run , requiring solving offline RL in an inner loop , which itself is a challenging problem and prone to training instability ( Kumar et al. , 2019 ) and hyperparameter sensitivity ( Wu et al. , 2019a ) . In this paper , we propose an offline imitation learning algorithm to learn from demonstrations that ( perhaps ) contain a large proportion of suboptimal data without additional steps of reward learning and offline RL training . Built upon the task of behavioral cloning ( BC ) , we introduce an additional task to learn a discriminator to distinguish expert and non-expert data , we propose a cooperation strategy to boost the performance of both tasks . This results in a new policy learning objective and surprisingly , we find its equivalence to a generalized BC objective , where the outputs of the discriminator serve as the weights of the BC loss function . We thus term our resulting algorithm Discriminator-Weighted Behavioral Cloning ( DWBC ) . Experimental results show that DWBC achieves higher returns and faster training speed compared to baseline algorithms , under different scenarios . 2 PRELIMINARY . 2.1 PROBLEM SETTING . We consider the standard fully observed Markov Decision Process ( MDP ) setting ( Sutton et al. , 1998 ) , M = { S , A , P , r , γ , d0 } , where S is the state space , A is the action space , P : S × A → ∆ ( S ) is the MDP ’ s transition probability , r is the reward function , γ ∈ [ 0 , 1 ) is the discount factor for future reward and d0 is the initial distribution . A policy π : S → ∆ ( A ) maps from state to distribution over actions . We denote dπ ∈ ∆ ( S × A ) as the discounted state-action distribution of π under transition kernel P , that is , dπ = ( 1− γ ) ∑∞ t=0 γ tdπt , where d π t ∈ ∆ ( S × A ) is the distribution of ( s ( t ) , a ( t ) ) under π at step t. Following the standard IL setting , the ground truth reward function r is unknown . Instead , we have the demonstrations by the expert specified by πe : S → ∆ ( A ) ( potentially stochastic and not necessarily optimal ) . Concretely , we have an expert dataset in the form of i.i.d tuples De = { si , ai , s′i } ne i=1 where ( s , a ) is sampled from distribution d πe and s′ is sampled from P ( s , a ) . In our problem setting , we also have an offline static dataset consisting of i.i.d tuples Do = { si , ai , s′i } no i=1 s.t . ( s , a ) ∼ ρ ( s , a ) , s′ ∼ P ( s , a ) , where ρ ∈ ∆ ( S × A ) is an offline state-action distribution resulting from some other behavior policies . Note that these behavior policies could be much worse than the expert πe . Our goal is to only leverage the offline batch data Db = De ∪ Do to learn an optimal policy π with regard to optimizing the ground truth reward r , without any interaction with the environment or expert . 2.2 A GENERALIZED BEHAVIORAL CLONING OBJECTIVE . In order to discard low-quality demonstrations and only clone the best behavior available , previous works ( Sasaki & Yamashina , 2021 ; Zolna et al. , 2020a ) consider a generalized behavioral cloning objective to imitate demonstrations unequally , that is , min π E ( s , a ) ∼Db [ − log π ( a|s ) · f ( s , a ) ] ( 1 ) where f : S ×A → [ 0 , 1 ] denotes an arbitrary weight function . • If f ( s , a ) = 1 for ∀ ( s , a ) ∈ S ×A , the objective ( 1 ) corresponds to the vanilla BC objective . • If f ( s , a ) = π′ ( a|s ) , where π′ is an old policy which was previously optimized with Db , the objective ( 1 ) corresponds to the objective of Behavioral Cloning from Noisy Demonstrations ( Sasaki & Yamashina , 2021 ) . Since ∑ a π ′ ( a|s ) = 1 for ∀s ∈ S is satisfied , π′ ( a|s ) can be interpreted as the weights for the weighted action sampling . • If f ( s , a ) = 1 [ Aπ ( s , a ) ] , where 1 is the indicator function which creates a boolean mask that elim- inates samples which are thought to be worse than the current policy , the objective ( 1 ) corresponds to the objective of Offline Reinforced Imitation Learning ( Zolna et al. , 2020a ) . The objective ( 1 ) can also be deemed as the objective of Soft Q Imitation Learning ( Reddy et al. , 2019 ) with f ( s , a ) = 1 for ( s , a ) ∈ De and f ( s , a ) = 0 for ( s , a ) ∈ Do in online IL literature ; or the objective of off-policy actor-critic ( Off-PAC ) algorithm ( Degris et al. , 2012 ) with f ( s , a ) = Qπ ( s , a ) · π ( a|s ) /πb ( a|s ) in online RL literature . 3 DISCRIMINATOR-WEIGHTED BEHAVIORAL CLONING . We now continue to describe our approach for offline imitation learning from demonstrations that ( perhaps ) contain large-proportional suboptimal data without additional steps of reward learning and offline RL training . Built upon the task of BC , we introduce an additional task to learn a discriminator to distinguish expert and non-expert data , we propose a cooperation strategy to boost the performance of both tasks , this will result in a new generalized BC objective . We then provide the interpretation of weights in our generalized BC objective , this gives the intuition about why our method can work . 3.1 LEARN THE POLICY AND DISCRIMINATOR SEPARATELY . It is obvious that we can avoid the negative impact of suboptimal demonstrations presented in Do by only imitating De , which can be written as min π E ( s , a ) ∼De [ − log π ( a|s ) ] ( 2 ) We call the task of learning a policy using objective ( 2 ) as the BC task . The drawback of BC task is that the learned policy may not be able to generalize due to the potential limited size and state coverage of De . It does not fully utilize the information from Do . If we can select those high-reward transitions from Do and combine them with De , we are expected to get a better policy . Now let ’ s consider another different task , which aims to learn a discriminator by contrasting expert and non-expert transitions , given by min d E ( s , a ) ∼De [ − log d ( s , a ) ] + E ( s , a ) ∼Do [ − log ( 1− d ( s , a ) ) ] ( 3 ) We call this the discriminating task , which is similar to how the discriminator is trained in GAIL ( Ho & Ermon , 2016 ) and GAN ( Goodfellow et al. , 2014a ) , except that the second term is sampled from a fixed dataset instead of new samples obtained from interacting with the environment . However , optimizing objective ( 3 ) will make the discriminator learned to assign 1 to all transitions from De and 0 to all transitions from Do . This limiting behavior is unsatisfactory because Do can contain some successful ( high-reward ) transitions . This bears similarity to the positive-unlabeled classification problem ( Elkan & Noto , 2008 ) , where both positive and negative samples exist in the unlabeled data . To solve this problem , we adopt the approach from positive-unlabeled ( PU ) learning ( Du Plessis et al. , 2015 ; Xu & Denil , 2019 ; Zolna et al. , 2020b ) . The main idea is to re-weight different losses for positive and unlabeled data , in order to obtain an estimate of model loss on negative samples that is not directly available . Applying PU learning to objective ( 3 ) yields the following objective : min d η E ( s , a ) ∼De [ − log d ( s , a ) ] + E ( s , a ) ∼Do [ − log ( 1− d ( s , a ) ) ] − η E ( s , a ) ∼De [ − log ( 1− d ( s , a ) ) ] ( 4 ) where η is a hyperparameter , corresponds to the proportion of positive samples to unlabeled samples . Intuitively , the second term in ( 4 ) could make d ( s , a ) of state-action pairs from De become 0 if similar state-action pairs are included in Do , and the third term in ( 4 ) balances the impact of the second term , i.e. , avoids d ( s , a ) of state-action pairs from De becoming 0 . To summarize , the BC task aims to imitate the expert behavior from De , but ignores the valuable information in Do ; the discriminating task aims to contrast expert and non-expert transitions from De and Do , by only using state-action information as input . Both tasks lack enough information to improve their own performance , which , however , can be obtained from the other task . It then seems natural to find a scheme to incorporate the policy into the training of the discriminator and effectively use the discriminator to help the training of the policy .

Summary: The paper deals with the following setup: offline imitation learning in the presence of both an expert dataset and a non-expert dataset. More precisely, the goal is to learn a policy as close as possible to the one(s) that generated the samples in a dataset $D_e$, while making the most of samples in a non-expert dataset $D_o$. The “reward” information is not present/used in the dataset. The authors draw inspiration from the positive-unlabeled classification as well as the adversarial imitation learning literature to propose a new algorithm to tackle this problem. They interleave the training of a discriminator and a policy. The discriminator is trained to discriminate between expert and non-expert dataset (using a positive/unlabeled loss) and takes as input the state, the action and the logit of the policy $\pi(a | s)$. The policy is trained to imitate the expert on $D_e$ and to “fool” the discriminator.   The authors present results on four environments from the Gym Mujoco suite with datasets extracted from the D4RL datasets.