Disagreement-Regularized Imitation Learning

1 INTRODUCTION . Training artificial agents to perform complex tasks is essential for many applications in robotics , video games and dialogue . If success on the task can be accurately described using a reward or cost function , reinforcement learning ( RL ) methods offer an approach to learning policies which has proven to be successful in a wide variety of applications ( Mnih et al. , 2015 ; 2016 ; Lillicrap et al. , 2016 ; Hessel et al. , 2018 ) . However , in other cases the desired behavior may only be roughly specified and it is unclear how to design a reward function to characterize it . For example , training a video game agent to adopt more human-like behavior using RL would require designing a reward function which characterizes behaviors as more or less human-like , which is difficult . Imitation learning ( IL ) offers an elegant approach whereby agents are trained to mimic the demonstrations of an expert rather than optimizing a reward function . Its simplest form consists of training a policy to predict the expert ’ s actions from states in the demonstration data using supervised learning . While appealingly simple , this approach suffers from the fact that the distribution over states observed at execution time can differ from the distribution observed during training . Minor errors which initially produce small deviations become magnified as the policy encounters states further and further from its training distribution . This phenomenon , initially noted in the early work of ( Pomerleau , 1989 ) , was formalized in the work of ( Ross & Bagnell , 2010 ) who proved a quadratic O ( T 2 ) bound on the regret and showed that this bound is tight . The subsequent work of ( Ross et al. , 2011 ) showed that if the policy is allowed to further interact with the environment and make queries to the expert policy , it is possible to obtain a linear bound on the regret . However , the ability to query an expert can often be a strong assumption . In this work , we propose a new and simple algorithm called DRIL ( Disagreement-Regularized Imitation Learning ) to address the covariate shift problem in imitation learning , in the setting where the agent is allowed to interact with its environment . Importantly , the algorithm does not require any additional interaction with the expert . It operates by training an ensemble of policies on the demonstration data , and using the disagreement in their predictions as a cost which is optimized through RL together with a supervised behavioral cloning cost . The motivation is that the policies in the ensemble will tend to agree on the set of states covered by the expert , leading to low cost , but are more likely to disagree on states not covered by the expert , leading to high cost . The RL cost ∗Work done while at Microsoft Research . thus guides the agent back towards the distribution of the expert , while the supervised cost ensures that it mimics the expert within the expert ’ s distribution . Our theoretical results show that , subject to realizability and optimization oracle assumptions1 , our algorithm obtains aO ( κT ) regret bound , where κ is a measure which quantifies a tradeoff between the concentration of the demonstration data and the diversity of the ensemble outside the demonstration data . We evaluate DRIL empirically across multiple pixel-based Atari environments and continuous control tasks , and show that it matches or significantly outperforms behavioral cloning and generative adversarial imitation learning , often recovering expert performance with only a few trajectories . 2 PRELIMINARIES . We consider episodic finite horizon MDP in this work . Denote by S the state space , A the action space , and Π the class of policies the learner is considering . Let T denote the task horizon and π ? the expert policy whose behavior the learner is trying to mimic . For any policy π , let dπ denote the distribution over states induced by following π. Denote C ( s , a ) the expected immediate cost of performing action a in state s , which we assume is bounded in [ 0 , 1 ] . In the imitation learning setting , we do not necessarily know the true costs C ( s , a ) , and instead we observe expert demonstrations . Our goal is to find a policy π which minimizes an observed surrogate loss ` between its actions and the actions of the expert under its induced distribution of states , i.e . π̂ = arg minEs∼dπ [ ` ( π ( s ) , π ? ( s ) ) ] ( 1 ) For the following , we will assume ` is the total variation distance ( denoted by ‖ · ‖ ) , which is an upper bound on the 0−1 loss . Our goal is thus to minimize the following quantity , which represents the distance between the actions taken by our policy π and the expert policy π ? : Jexp ( π ) = Es∼dπ [ ‖π ( ·|s ) − π ? ( ·|s ) ‖ ] ( 2 ) Denote Qπt ( s , a ) as the standard Q-function of the policy π , which is defined as Q π t ( s , a ) = E [ ∑T τ=t C ( sτ , aτ ) | ( st , at ) = ( s , a ) , aτ ∼ π ] . The following result shows that if ` is an upper bound on the 0 − 1 loss and C satisfies certain smoothness conditions , then minimizing this loss within translates into an O ( T ) regret bound on the true task cost JC ( π ) = Es , a∼dπ [ C ( s , a ) ] : Theorem 1 . ( Ross et al. , 2011 ) If π satisfies Jexp ( π ) = , and Qπ ? T−t+1 ( s , a ) −Qπ ? T−t+1 ( s , π ? ) ≤ u for all time steps t , actions a and states s reachable by π , then JC ( π ) ≤ JC ( π ? ) + uT . Unfortunately , it is often not possible to optimize Jexp directly , since it requires evaluating the expert policy on the states induced by following the current policy . The supervised behavioral cloning cost JBC , which is computed on states induced by the expert , is often used instead : JBC ( π ) = Es∼dπ ? [ ‖π ? ( ·|s ) − π ( ·|s ) ‖ ] ( 3 ) Minimizing this loss within yields a quadratic regret bound on regret : Theorem 2 . ( Ross & Bagnell , 2010 ) Let JBC ( π ) = , then JC ( π ) ≤ JC ( π ? ) + T 2 . Furthermore , this bound is tight : as we will discuss later , there exist simple problems which match the worst-case lower bound . 3 ALGORITHM . Our algorithm is motivated by two criteria : i ) the policy should act similarly to the expert within the expert ’ s data distribution , and ii ) the policy should move towards the expert ’ s data distribution 1We assume for the analysis the action space is discrete , but the state space can be large or infinite . Algorithm 1 Disagreement-Regularized Imitation Learning ( DRIL ) 1 : Input : Expert demonstration data D = { ( si , ai ) } Ni=1 2 : Initialize policy π and policy ensemble ΠE = { π1 , ... , πE } 3 : for e = 1 , E do 4 : Sample De ∼ D with replacement , with |De| = |D| . 5 : Train πe to minimize JBC ( πe ) on De to convergence . 6 : end for 7 : for i = 1 , ... do 8 : Perform one gradient update to minimize JBC ( π ) using a minibatch from D. 9 : Perform one step of policy gradient to minimize Es∼dπ , a∼π ( ·|s ) [ C clip U ( s , a ) ] . 10 : end for if it is outside of it . These two criteria are addressed by combining two losses : a standard behavior cloning loss , and an additional loss which represents the variance over the outputs of an ensemble ΠE = { π1 , ... , πE } of policies trained on the demonstration data D. We call this the uncertainty cost , which is defined as : CU ( s , a ) = Varπ∼ΠE ( π ( a|s ) ) = 1 E E∑ i=1 ( πi ( a|s ) − 1 E E∑ i=1 πi ( a|s ) ) 2 The motivation is that the variance over plausible policies is high outside the expert ’ s distribution , since the data is sparse , but it is low inside the expert ’ s distribution , since the data there is dense . Minimizing this cost encourages the policy to return to regions of dense coverage by the expert . Intuitively , this is what we would expect the expert policy π ? to do as well . The total cost which the algorithm optimizes is given by : Jalg ( π ) = Es∼dπ ? [ ‖π ? ( ·|s ) − π ( ·|s ) ‖ ] ︸ ︷︷ ︸ JBC ( π ) +Es∼dπ , a∼π ( ·|s ) [ CU ( s , a ) ] ︸ ︷︷ ︸ JU ( π ) The first term is a behavior cloning loss and is computed over states generated by the expert policy , of which the demonstration data D is a representative sample . The second term is computed over the distribution of states generated by the current policy and can be optimized using policy gradient . Note that the demonstration data is fixed , and this ensemble can be trained once offline . We then interleave the supervised behavioral cloning updates and the policy gradient updates which minimize the variance of the ensemble . The full algorithm is shown in Algorithm 1 . We also found that dropout ( Srivastava et al. , 2014 ) , which has been proposed as an approximate form of ensembling , worked well ( see Appendix D ) . In practice , for the supervised loss we optimize the KL divergence between the actions predicted by the policy and the expert actions , which is an upper bound on the total variation distance due to Pinsker ’ s inequality . We also found it helpful to use a clipped uncertainty cost : CclipU ( s , a ) = { −1 if CU ( s , a ) ≤ q +1 else where the threshold q is a top quantile of the raw uncertainty costs computed over the demonstration data . The threshold q defines a normal range of uncertainty based on the demonstration data , and values above this range incur a positive cost ( or negative reward ) . The RL cost can be optimized using any policy gradient method . In our experiments we used advantage actor-critic ( A2C ) ( Mnih et al. , 2016 ) or PPO ( Schulman et al. , 2017 ) , which estimate the expected cost using rollouts from multiple parallel actors all sharing the same policy ( see Appendix C for details ) . We note that model-based RL methods could in principle be used as well if sample efficiency is a constraint .

Summary: The paper proposes an imitation learning algorithm that combines behavioral cloning with a regularizer that encourages the agent to visit states similar to the demonstrated states. The key idea is to use ensemble disagreement to approximate uncertainty, and use RL to train the imitation agent to visit states in which an ensemble of cloned imitation policies is least uncertain about which action the expert would take. Experiments on image-based Atari games show that the proposed method significantly outperforms BC and GAIL baselines in three games, and performs comparably or slightly better than the baselines in the remaining three games.