On Computation and Generalization of Generative Adversarial Imitation Learning

1 INTRODUCTION . As various robots ( Tail et al. , 2018 ) , self-driving cars ( Kuefler et al. , 2017 ) , unmanned aerial vehicles ( Pfeiffer et al. , 2018 ) and other intelligent agents are applied to complex and unstructured environments , programming their behaviors/policy has become increasingly challenging . These intelligent agents need to accommodate a huge number of tasks with unique environmental demands . To address these challenges , many reinforcement learning ( RL ) methods have been proposed for learning sequential decision-making policies ( Sutton et al. , 1998 ; Kaelbling et al. , 1996 ; Mnih et al. , 2015 ) . These RL methods , however , heavily rely on human expert domain knowledge to design proper reward functions . For complex tasks , which are often difficult to describe formally , these RL methods become impractical . The Imitation Learning ( IL , Argall et al . ( 2009 ) ; Abbeel & Ng ( 2004 ) ) approach is a powerful and practical alternative to RL . Rather than having a human expert handcrafting a reward function for learning the desired policy , the imitation learning approach only requires the human expert to demonstrate the desired policy , and then the intelligent agent ( a.k.a . learner ) learns to match the demonstration . Most of existing imitation learning methods fall in the following two categories : • Behavioral Cloning ( BC , Pomerleau ( 1991 ) ) . BC treats the IL problem as supervised learning . Specifically , it learns a policy by fitting a regression model over expert demonstrations , which directly maps states to actions . Unfortunately , BC has a fundamental drawback . Recall that in supervised learning , the distribution of the training data is decoupled from the learned model , whereas in imitation learning , the agent ’ s policy affects what state is queried next . The mismatch between training and testing distributions , also known as covariate shift ( Ross & Bagnell , 2010 ; Ross et al. , 2011 ) , yields significant compounding errors . Therefore , BC often suffers from poor generalization . • Inverse Reinforcement Learning ( IRL , Russell ( 1998 ) ; Ng et al . ( 2000 ) ; Finn et al . ( 2016 ) ; Levine & Koltun ( 2012 ) ) . IRL treats the IL problem as bi-level optimization . Specifically , it finds a reward function , under which the expert policy is uniquely optimal . Though IRL does not have the error compounding issue , its computation is very inefficient . Many existing IRL methods need to solve a sequence of computationally expensive reinforcement learning problems , due to their bi-level optimization nature . Therefore , they often fail to scale to large and high dimensional environments . More recently , Ho & Ermon ( 2016 ) propose a Generative Adversarial Imitation Learning ( GAIL ) method , which obtains significant performance gains over existing IL methods in imitating complex expert policies in large and high-dimensional environments . GAIL generalizes IRL by formulating the IL problem as minimax optimization , which can be solved by alternating gradient-type algorithms in a more scalable and efficient manner . Specifically , we consider an infinite horizon Markov Decision Process ( MDP ) , where S denotes the state space , A denotes the action space , P denotes the Markov transition kernel , r∗ denotes the reward function , and p0 denotes the distribution of the initial state . We assume that the Markov transition kernel P is fixed and there is an unknown expert policy π∗ : S → P ( A ) , where P ( A ) denotes the set of distributions over the action space . As can be seen , { st } T−1t=0 essentially forms a Markov chain with the transition kernel induced by π∗ as Pπ ∗ ( s , s′ ) = ∑ a∈A π ∗ ( a | s ) ·P ( s′ | s , a ) . Given n demonstration trajectories from π∗ denoted by { s ( i ) t , a ( i ) t } T−1t=0 , where i = 1 , ... , n , s0 ∼ p0 , at ∼ π∗ ( · | st ) , and st+1 ∼ P ( · | st , at ) , GAIL aims to learn π∗ by solving the following minimax optimization problem , minπ maxr∈R [ Eπr ( s , a ) − Eπ∗nr ( s , a ) ] , ( 1 ) where Eπ [ r ( s , a ) ] = limT→∞ E [ 1T ∑T−1 t=0 r ( st , at ) |π ] denotes the average reward under the policy π when the reward function is r , and Eπ∗n [ r ( s , a ) ] = 1 nT ∑n i=1 ∑T−1 t=0 [ r ( s ( i ) t , a ( i ) t ) ] denotes the empirical average reward over the demonstration trajectories . As shown in ( 1 ) , GAIL aims to find a policy , which attains an average reward similar to that of the expert policy with respect to any reward belonging to the function classR . For large and high-dimensional imitation learning problems , we often encounter infinitely many states . To ease computation , we need to consider function approximations . Specifically , suppose that for every s ∈ S and a ∈ A , there are feature vectors ψs ∈ RdS and ψa ∈ RdA associated with a and s , respectively . Then we can approximate the policy and reward as π ( ·|s ) = π̃ω ( ψs ) and r ( s , a ) = r̃θ ( ψs , ψa ) , where π̃ and r̃ belong to certain function classes ( e.g . reproducing kernel Hilbert space or deep neural networks , Ormoneit & Sen ( 2002 ) ; LeCun et al . ( 2015 ) ) associated with parameters ω and θ , respectively . Accordingly , we can optimize ( 1 ) with respect to the parameters ω and θ by scalable alternating gradient-type algorithms . Although GAIL has achieved significant empirical progresses , its theoretical properties are still largely unknown . There are three major difficulties when analyzing GAIL : 1 ) . There exists temporal dependency in the demonstration trajectories/data due to their sequential nature ( Howard , 1960 ; Puterman , 2014 ; Abounadi et al. , 2001 ) ; 2 ) . GAIL is formulated as a minimax optimization problem . Most of existing learning theories , however , focus on empirical risk minimization problems , and therefore are not readily applicable ( Vapnik , 2013 ; Mohri et al. , 2018 ; Anthony & Bartlett , 2009 ) ; 3 ) . The minimax optimization problem in ( 1 ) does not have a convex-concave structure , and therefore existing theories in convex optimization literature can not be applied for analyzing the alternating stochastic gradient-type algorithms ( Willem , 1997 ; Ben-Tal & Nemirovski , 1998 ; Murray & Overton , 1980 ; Chambolle & Pock , 2011 ; Chen et al. , 2014 ) . Some recent results suggest to use stage-wise stochastic gradient-type algorithms ( Rafique et al. , 2018 ; Dai et al. , 2017 ) . More specifically , at every iteration , they need to solve the inner maximization problem up to a high precision , and then apply stochastic gradient update to the outer minimization problem . Such algorithms , however , are rarely used by practitioners , as they are inefficient in practice ( due to the computationally intensive inner maximization ) . To bridge such a gap between practice and theory , we establish the generalization properties of GAIL and the convergence properties of the alternating mini-batch stochastic gradient algorithm for solving ( 1 ) . Specifically , our contributions can be summarized as follows : •We formally define the generalization of GAIL under the “ so-called ” R-reward distance , and then show that the generalization of GAIL can be guaranteed under reward distance as long as the class of the reward functions is properly controlled ; •We provide sufficient conditions , under which an alternating mini-batch stochastic gradient algorithm can efficiently solve the minimax optimization in ( 1 ) , and attains sublinear convergence to a stationary solution . To the best of our knowledge , these are the first results on statistical and computational theories of imitation learning with reward/policy function approximations . Our work is related to Syed et al . ( 2008 ) ; Cai et al . ( 2019 ) . Syed et al . ( 2008 ) study the generalization and computational properties of apprenticeship learning . Since they assume that the state space of the underlying Markov decision process is finite , they do not consider any reward/policy function approximations ; Cai et al . ( 2019 ) study the computational properties of imitation learning under a simple control setting . Their assumption on linear policy and quadratic reward is very restrictive , and does not hold for many real applications . Notation . Given a vector x = ( x1 , ... , xd ) > ∈ Rd , we define ‖x‖22 = ∑d j=1 x 2 j . Given a function f : Rd 7→ R , we denote its ` ∞ norm as ‖f‖∞ = maxx |f ( x ) | . 2 GENERALIZATION OF GAIL . To analyze the generalization properties of GAIL , we first assume that we can access an infinite number of the expert ’ s demonstration trajectories ( underlying population ) , and that the reward function is chosen optimally within some large class of functions . This allows us to remove the maximum operation from ( 1 ) , which leads to an interpretation of how and in what sense the resulting policy is close to the true expert policy . Before we proceed , we first introduce some preliminaries . Definition 1 ( Stationary Distribution ) . Note that any policy π induces a Markov chain on S × A . The transition kernel is given by Pπ ( s ′ , a′ | s , a ) = π ( a′ | s′ ) · P ( s′ | s , a ) , ∀ ( s , a ) , ( s′ , a′ ) ∈ S ×A . When such a Markov chain is aperiodic and recurrent , we denote its stationary distribution as ρπ . Note that a policy π is uniquely determined by its stationary distribution ρπ in the sense that π ( a | s ) = ρπ ( s , a ) / ∑ a∈A ρπ ( s , a ) . Then we can write the expected average reward of r ( s , a ) under the policy π as Eπ [ r ( s , a ) ] = limT→∞ E [ 1 T ∑T−1 t=0 r ( st , at ) ∣∣π ] = Eρπ [ r ( s , a ) ] = ∑ ( s , a ) ∈S×A ρπ ( s , a ) · r ( s , a ) . We further define theR-distance between two policies π and π′ as follows . Definition 2 . Let R denote a class of symmetric reward functions from S × A to R , i.e. , if r ∈ R , then −r ∈ R. Given two policy π′ and π , theR-distance for GAIL is defined as dR ( π , π ′ ) = supr∈R [ Eπr ( s , a ) − Eπ′r ( s , a ) ] . The R-distance over policies for Markov decision processes is essentially an Integral Probability Metric ( IPM ) over stationary distributions ( Müller , 1997 ) . For different choices ofR , we have variousR-distances . For example , we can chooseR as the class of all 1-Lipschitz continuous functions , which yields that dR ( π , π′ ) is the Wasserstein distance between ρπ and ρπ′ ( Vallender , 1974 ) . For computational convenience , GAIL and its variants usually choose R as a class of functions from some reproducing kernel Hilbert space , or a class of neural network functions . Definition 3 . Given n demonstration trajectories from time 0 to T − 1 obtained by an expert policy π∗ denoted by ( s ( i ) t , a ( i ) t ) T−1 t=0 , where i = 1 , ... , n , a policy π̂ learned by GAIL generalizes under the R-distance dR ( · , · ) with generalization error , if with high probability , we have |dR ( π∗n , π̂ ) − dR ( π∗ , π̂ ) | ≤ , where dR ( π∗n , π̂ ) is the empiricalR-distance between π∗ and π̂ defined as dR ( π ∗ n , π̂ ) = supr∈R [ Eπ∗nr ( s , a ) − Eπ̂r ( s , a ) ] with Eπ∗n [ r ( s , a ) ] = 1nT ∑n i=1 ∑T−1 t=0 [ r ( s ( i ) t , a ( i ) t ) ] . The generalization of GAIL implies that theR-distance between the expert policy π∗ and the learned policy π̂ is close to the empirical R-distance between them . Our analysis aims to prove the former distance to be small , whereas the latter one is what we attempts to minimize in practice . We then introduce the assumptions on the underlying Markov decision process and expert policy . Assumption 1 . Under the expert policy π∗ , ( st , at ) T−1t=0 forms a stationary and exponentially βmixing Markov chain , i.e. , β ( k ) = supn EB∈σn0 supA∈σ∞n+k |P ( A|B ) − P ( A ) | ≤ β0 exp ( −β1k α ) , where β0 , β1 , α are positive constants , and σ j i is the σ-algebra generated by ( st , at ) j t=i for i ≤ j . Moreover , for every s ∈ S and a ∈ A , there are feature vectors ψs ∈ RdS and ψa ∈ RdA associated with a and s , respectively , and ψs and ψa are uniformly bounded , where ‖ψs‖2 ≤ 1 and ‖ψa‖2 ≤ 1 , ∀s ∈ S and ∀a ∈ A . Assumption 1 requires the underlying MDP to be ergodic ( Levin & Peres , 2017 ) , which is a commonly studied assumption in exiting reinforcement learning literature on maximizing the expected average reward ( Strehl & Littman , 2005 ; Li et al. , 2011 ; Brafman & Tennenholtz , 2002 ; Kearns & Singh , 2002 ) . The feature vectors associated with a and s allow us to apply function approximations to parameterize the reward and policy functions . Accordingly , we write the reward function as r ( s , a ) = r̃ ( ψs , ψa ) , which is assumed to be bounded . Assumption 2 . The reward function class is uniformly bounded , i.e. , ‖r‖∞ ≤ Br for any r ∈ R. Now we proceed with our main result on generalization properties of GAIL . We useN ( R , , ‖ · ‖∞ ) to denote the covering number of the function classR under the ` ∞ distance ‖ · ‖∞ . Theorem 1 ( Main Result ) . Suppose Assumptions 1-2 hold , and the policy learned by GAIL satisfies dR ( π ∗ n , π̂ ) − infπ dR ( π∗n , π ) < , where the infimum is taken over all possible learned policies . Then with probability at least 1 − δ over the joint distribution of { ( a ( i ) t , s ( i ) t ) T−1t=0 } ni=1 , we have dR ( π ∗ , π̂ ) − inf π dR ( π ∗ , π ) ≤ O ( Br√ nT/ζ √ logN ( R , √ ζ nT , ‖ · ‖∞ ) +Br √ log ( 1/δ ) nT/ζ ) + , where ζ = ( β−11 log β0T δ ) 1 α . Theorem 1 implies that the policy π̂ learned by GAIL generalizes as long as the complexity of the function class R is well controlled . To the best of our knowledge , this is the first result on the generalization of imitation learning with function approximations . As the proof of Theorem 1 is involved , we only present a sketch due to space limit . More details are provided in Appendix A.1 . Proof Sketch . Our analysis relies on characterizing the concentration property of the empirical average reward under the expert policy . For notational simplicity , we define φ = Eπ∗r ( s , a ) − 1nT ∑n i=1 ∑T−1 t=0 r ( s ( i ) t , a ( i ) t ) . The key challenge comes from the fact that ( s ( i ) t , a ( i ) t ) ’ s are dependent . To handle such a dependency , we adopt the independent block technique from Yu ( 1994 ) . Specifically , we partition every trajectory into disjoint blocks ( where the block size is of the order O ( ( log ( T ) + log ( 1/δ ) ) 1/α ) , and construct two separable trajectories : One contains all blocks with odd indices ( denoted by Bodd ) , and the other contains all those with even indices ( denoted by Beven ) . We define φ1 = Eπ∗r ( s , a ) − 2nT ∑n i=1 ∑ ( s ( i ) t , a ( i ) t ) ∈Bodd r ( s ( i ) t , a ( i ) t ) , and analogously for φ2 with ( s ( i ) t , a ( i ) t ) ∈ Beven . Then we have P ( supr∈R φ ≥ ε ) ≤ P ( supr∈R φ12 + supr∈R φ2 2 ≥ ε ) ≤ P ( supr∈R φ1 ≥ ε ) + P ( supr∈R φ2 ≥ ε ) . We consider a block-wise independent counterpart of φ1 denoted by φ̃1 , where each block is sampled independently from the same Markov chain as φ1 , i.e. , φ̃1 has independent blocks of samples from the same exponentially β-mixing Markov chain . Accordingly , we denote φ̃1 as φ̃1 = Eπ∗r ( s , a ) − 2nT ∑n i=1 ∑ ( s ( i ) t , a ( i ) t ) ∈B̃odd r ( s ( i ) t , a ( i ) t ) , where B̃odd denotes i.i.d . blocks of samples . Now we bound the difference between φ1 and φ̃1 by P ( sup r∈R φ1 − E [ sup r∈R φ̃1 ] ≥ ε− E [ sup r∈R φ̃1 ] ) ≤ P ( sup r∈R φ̃1 − E [ sup r∈R φ̃1 ] ≥ ε− E [ sup r∈R φ̃1 ] ) + CβT/ ( log ( T ) + log ( 1/δ ) ) 1/α , where C is a constant , and β is the mixing coefficient , and P ( supr∈R φ̃1 − E [ supr∈R φ̃1 ] ≥ ε − E [ supr∈R φ̃1 ] ) can be bounded using the empirical process technique for independent random variables . The details of the above inequality can be found in Corollary 3 in Appendix A.1 , where the proof technique is adapted from Lemma 1 in Mohri & Rostamizadeh ( 2009 ) . Let φ̃2 be defined analogously as φ̃1 . With a similar argument further applied to φ2 and φ̃2 , we obtain P ( sup r∈R φ ≥ ε ) ≤ 2P ( sup r∈R φ̃1 − E [ sup r∈R φ̃1 ] ≥ ε− E [ sup r∈R φ̃1 ] ) + 2CβT/ ( log ( T ) + log ( 1/δ ) ) 1/α . The rest of our analysis follows the PAC-learning framework using Rademacher complexity and is omitted ( Mohri et al. , 2018 ) . We complete the proof sketch . Example 1 : Reproducing Kernel Reward Function . One popular option to parameterize the reward by functions is the reproducing kernel Hilbert space ( RKHS , Kim & Park ( 2018 ) ; Li et al . ( 2018 ) ) . There have been several implementations of RKHS , and we consider the feature mapping approach . Specifically , we consider g : RdS × RdA → Rq , and the reward can be written as r ( s , a ) = r̃θ ( ψs , ψa ) = θ > g ( ψs , ψa ) , where θ ∈ Rq . We require g to be Lipschitz continuous with respect to ( ψa , ψs ) . Assumption 3 . The feature mapping g satisfies g ( 0 , 0 ) = 0 , and there exists a constant ρg such that for any ψa , ψ′a , ψs and ψ ′ s , we have ‖g ( ψs , ψa ) − g ( ψ′s , ψ′a ) ‖22 ≤ ρg √ ‖ψs − ψ′s‖22 + ‖ψa − ψ′a‖22 . Assumption 3 is mild and satisfied by popular feature mappings , e.g. , random Fourier feature mapping1 ( Rahimi & Recht , 2008 ; Bach , 2017 ) . The next corollary presents the generalization bound of GAIL using feature mapping . Corollary 1 . Suppose ‖θ‖2 ≤ Bθ . For large enough n and T , with probability at least 1 − δ over the joint distribution of { ( a ( i ) t , s ( i ) t ) T−1t=0 } ni=1 , we have dR ( π ∗ , π̂ ) − inf π dR ( π ∗ , π ) ≤ O ( ρgBθ√ nT/ζ √ q log ( ρgBθ √ nT/ζ ) + ρgBθ √ log ( 1/δ ) nT/ζ ) + . Corollary 1 indicates that with respect to a class of properly normalized reproducing kernel reward functions , GAIL generalizes in terms of theR-distance . Example 2 : Neural Network Reward Function . Another popular option to parameterize the reward function is to use neural networks . Specifically , let σ ( v ) = [ max { v1 , 0 } , ... , max { vd , 0 } ] > denote the ReLU activation for v ∈ Rd . We consider a D-layer feedforward neural network with ReLU activation as follows , r ( s , a ) = r̃W ( ψs , ψa ) = W > Dσ ( WD−1σ ( ... σ ( W1 [ ψ > a , ψ > s ] > ) ) ) , whereW = { Wk | Wk ∈ Rdk−1×dk , k = 1 , ... , D − 1 , WD ∈ RdD−1 } and d0 = dA + dS . The next corollary presents the generalization bound of GAIL using neural networks . 1More precisely , Assumption 3 actually holds with overwhelming probability over the distribution of the random mapping . Corollary 2 . Suppose ‖Wi‖2 ≤ 1 , where i = 1 , ... , D. For large enough n and T , with probability at least 1− δ over the joint distribution of { ( a ( i ) t , s ( i ) t ) T−1t=0 } ni=1 , we have dR ( π ∗ , π̂ ) − inf π dR ( π ∗ , π ) ≤ O ( 1√ nT/ζ √ d2D log ( D √ dnT/ζ ) + √ log ( 1/δ ) nT/ζ ) + . Corollary 2 indicates that with respect to a class of properly normalized neural network reward functions , GAIL generalizes in terms of theR-distance . Remark 1 ( The Tradeoff between Generalization and Representation of GAIL ) . As can be seen from Definition 2 , the R-distances are essentially differentiating two policies . For the Wassersteintype distance , i.e. , R contains all 1-Lipschitz continuous functions , if dR ( π , π′ ) is small , it is safe to conclude that two policies π and π′ are nearly the same almost everywhere . However , when we choose R to be the reproducing kernel Hilbert space or the class of neural networks with relatively small complexity , dR ( π , π′ ) can be small even if π and π′ are not very close . Therefore , we need to choose a sufficiently diverse class of reward functions to ensure that we recover the expert policy . As Theorem 1 suggests , however , that we need to control the complexity of the function class R to guarantee the generalization . This implies that when parameterizing the reward function , we need to carefully choose the function class to attain the optimal tradeoff between generalization and representation of GAIL .

Summary: This paper presents theoretical analysis of the generalization properties of GAIL, as well as the local convergence of the traditional minibatch SGD applied to its min-max optimization problem (without assuming convex-concave structure of the game).  Specifically, the authors first prove a generalization bound for GAIL that characterizes how well optimizing the empirical GAIL objective minimizes the true population version (in terms of an "R-distance" they use the characterize the expected distributions induced by the underlying policies for a given reward class).  Second, they prove a convergence result for minibatch SGD applied to the min-max game, showing that the method will converge to a stationary point regardless of any convex-concave structure.