Provable Representation Learning for Imitation Learning via Bi-level Optimization

1 INTRODUCTION . Humans can often learn from experts quickly and with a few demonstrations and we would like our artificial agents to do the same . However , even for simple imitation learning tasks , the current state-of-the-art methods require thousand of demonstrations . Humans do not learn new skills from scratch . We can summarize learned skills , distill them and build a common ground , a.k.a , representation that is useful for learning future skills . Can we build an agent to do the same ? The current paper studies how to apply representation learning to imitation learning . Specifically , we want to build an agent that is able learn a representation from multiple experts ’ demonstrations , where the experts aim to solve different Markov decision processes ( MDPs ) that share the same state and action spaces but can differ in the transition and reward functions . The agent can use this representation to reduce the number of demonstrations required for a new imitation learning task . While several methods have been proposed ( Duan et al. , 2017 ; Finn et al. , 2017b ; James et al. , 2018 ) to build agents that can adapt quickly to new tasks , none of them , to our knowledge , give provable guarantees showing the benefit of using past experience . Furthermore , they do not focus on learning a representation . See Section 2 for more discussions . In this paper , we propose a framework to formulate this problem and analyze the statistical gains of representation learning . The main idea is to use bi-level optimization formulation where the “ outer ” optimization tries to learn the joint representation and the “ inner ” optimization encodes the imitation learning setup and tries to learn task-specific parameters . In particular , the inner optimization is flexible enough to allow the agent to interact with the environment . This framework allows us to do a rigorous analysis to show provable benefits of representation learning for imitation learning . With this framework at hand , we make the following concrete contributions : • We first instantiate our framework in the setting where the agent can observe experts ’ actions and tries to find a policy that matches the expert ’ s policy , a.k.a , behavior cloning . This setting can be viewed as a straightforward extension of multi-task representation learning for supervised learning ( Maurer et al. , 2016 ) . We show in this setting that with sufficient number of experts ( possibly optimizing for different reward functions ) , the agent can learn a representation that provably reduces the sample complexity for a new target imitation learning task . • Next , we consider a more challenging setting where the agent can not observe experts ’ actions but only their states , a.k.a. , the observation-alone setting . We set the inner optimization as a minmax problem inspired by Sun et al . ( 2019 ) . Notably , this min-max problem requires the agent to interact with the environment to collect samples . We again show that with sufficient number of experts , the agent can learn a representation that provably reduces the sample complexity for a target task where the agent can not observe actions from either source experts or the target expert . • We conduct experiments to verify our theoretical insights by learning a representation from multiple tasks using our framework and testing it using both behavior cloning and policy optimization . In these settings , we observe that by learning representations the agent can learn a good policy with fewer samples than needed to learn a policy from scratch . The key contribution is to connect existing literature on multi-task representation learning that deals with supervised learning ( Maurer et al. , 2016 ) to single task imitation learning methods with guarantees ( Syed & Schapire , 2010 ; Ross et al. , 2011 ; Sun et al. , 2019 ) . To our knowledge , this is the first work showing such guarantees for general losses that are not necessarily convex . 2 RELATED WORK . Representation learning has shown its great power in various domains . See Bengio et al . ( 2013 ) for a survey . Theoretically , Maurer et al . ( 2016 ) gave analysis showing representation can provably reduce the sample complexity in the multi-task supervised learning setting . Recently , Arora et al . ( 2019 ) analyzed the benefit of representation learning via contrastive learning . These papers all build representations for the agent / learner . We remark that researchers also try to build representations about the environment / physical world ( Wu et al. , 2017 ) . Imitation learning can help with sample efficiency of many problems ( Ross & Bagnell , 2010 ; Sun et al. , 2017 ; Daumé et al. , 2009 ; Chang et al. , 2015 ; Pan et al. , 2018 ) . Most existing work consider the setting where the learner can observe expert ’ s action . A general strategy is use supervised learning to learn a policy that maps the state to action that matches expert ’ s behaviors . The most straightforward one is behavior cloning ( Pomerleau , 1991 ) , which we also study in our paper . More advanced approaches have also been proposed ( Ross et al. , 2011 ; Ross & Bagnell , 2014 ; Sun et al. , 2018 ) . These approaches , including behavior cloning , often enjoy sound theoretical guarantees in the single task case . Our paper extends the theoretical guarantees of behavior cloning to the multitask representation learning setting . This paper also considers a more challenging setting , imitation learning from observation alone . Though some model-based methods have been proposed ( Torabi et al. , 2018 ; Edwards et al. , 2018 ) , these methods lack theoretical guarantees . Another line of work learns a policy that minimizes the difference between the state distributions induced by it and the expert policy , under a certain distributional metric ( Ho & Ermon , 2016 ) . Sun et al . ( 2019 ) gave a theoretical analysis to characterize the sample complexity of this approach and our method for this setting is inspired by their approach . A line of work uses meta-learning for imitation learning ( Duan et al. , 2017 ; Finn et al. , 2017b ; James et al. , 2018 ) . Our work is different from theirs as we want to explicitly learn a representation that is useful across all tasks whereas these work try to learn a meta-algorithm that can quickly adapt to a new task . For example , Finn et al . ( 2017b ) used a gradient based method for adaptation . Recently Raghu et al . ( 2019 ) argued that most of the power of MAML ( Finn et al. , 2017a ) like approaches comes from learning a shared representation . On the theoretical side of meta-learning and multi-task learning , Baxter ( 2000 ) performed the first theoretical analysis and gave sample complexity bounds using covering numbers . Bullins et al . ( 2019 ) provides an efficient algorithm that generalizes to new unseen tasks , but for linear representations . Another recent line of work analyzes gradient based meta-learning methods , similar to MAML ( Finn et al. , 2017a ) . Existing work on the sample complexity and regret of these methods ( Denevi et al. , 2019 ; Finn et al. , 2019 ; Khodak et al. , 2019 ) show guarantees for convex losses by leveraging tools from online convex optimization . In contrast , our analysis works for arbitrary function classes and the bounds depend on the gaussian averages of these classes . Recent work ( Rajeswaran et al. , 2019 ) uses a bi-level optimization framework for meta-learning and improves computation ( not statistical ) aspects of meta-learning through implicit differentiation . 3 PRELIMINARIES . Markov Decision Processes ( MDPs ) : LetM = ( S , A , P , C , ν ) be an MDP , where S is the state space , A is the finite action space with |A| = K , H ∈ Z+ is the planning horizon , P : S × A → 4 ( S ) is the transition function , C : S × A → R is the cost function and ν ∈ 4 ( S ) is the initial state distribution . We assume that cost is bounded by 1 , i.e . C ( s , a ) ≤ 1 , ∀s ∈ S , a ∈ A . This is a standard regularity condition used in many theoretical reinforcement learning work . A ( stochastic ) policy is defined as π = ( π1 , . . . , πH ) , where πh : S → 4 ( A ) prescribes a distribution over action for each state at level h ∈ [ H ] . For a stationary policy , we have π1 = · · · = πH = π . A policy π induces a random trajectory s1 , a1 , s2 , a2 , . . . , sH , aH where s1 ∼ ν , a1 ∼ π1 ( s ) , s2 ∼ Ps1 , a1 etc . Let νπh denote the distribution over S induced at level h by policy π . The value function V πh : S → R is defined as V πh ( sh ) = E [ H∑ i=h C ( si , ai ) | ai ∼ πi ( si ) , si+1 ∼ Psi , ai ] and the state-action function Qπh ( sh , ah ) is defined as Q π h ( sh , ah ) = Esh+1∼Psh , ah [ V π h ( sh+1 ) ] . The goal is to learn a policy π that minimizes the expected cost J ( π ) = Es1∼νV π1 ( s1 ) . We define the Bellman operator at level h for any policy π as Γπh : RS → RS , where for s ∈ S and g ∈ RS , ( Γπh g ) ( s ) : = Ea∼πh ( s ) , s′∼Ps , a [ g ( s ′ ) ] ( 1 ) Multi-task Imitation learning : We formally describe the problem we want to study . We assume there are multiple tasks ( MDPs ) sampled i.i.d . from a distribution η . A task µ ∼ η is an MDPMµ = ( S , A , H , Pµ , Cµ , νµ ) ; all tasks share everything except the cost function , initial state distribution and transition function . For simplicity of presentation , we will assume a common transition function P for all tasks ; proofs remain exactly the same even otherwise . For every task µ , π∗µ = ( π ∗ 1 , µ , . . . , π ∗ H , µ ) is an expert policy that the learner has access to in the form of trajectories induced by that policy . The trajectories may or may not contain expert ’ s actions . These correspond to two settings that we discuss in more detail in Section 5 and Section 6 . The distributions of states induced by this policy at different levels are denoted by { ν∗1 , µ , . . . , ν∗H , µ } and the average state dis- tribution as ν∗µ = 1 H H∑ h=1 ν∗h , µ . We define V ∗ h , µ to be the value function of π ∗ µ and Jµ to be the expected cost function for task µ . We will drop the subscript µ whenever the task at hand is clear from context . Of interest is also the special case where the expert policy π∗µ is stationary . Representation learning : In this work , we wish to learn policies from a function class of the form Π = F ◦ Φ , where Φ ⊆ { φ : S → Rd | ‖φ ( s ) ‖2 ≤ R } is a class of bounded norm representation functions mapping states to vectors and F ⊆ { f : Rd → ∆ ( A ) } is a class of functions mapping state representations to distribution over actions . We will be using linear functions , i.e . F = { x → softmax ( Wx ) |W ∈ RK×d , ‖W‖F ≤ 1 } . We denote a policy parametrized by φ ∈ Φ and f ∈ F by πφ , f , where πφ , f ( a|s ) = f ( φ ( s ) ) a . In some cases , we may also use the policy πφ , f ( a|s ) = I { a = arg max a′∈A f ( φ ( s ) ) a′ } 1 . Denote Πφ = { πφ , f : f ∈ F } to be the class of policies that use φ as the representation function . Given demonstrations from expert policies for T tasks sampled independently from η , we wish to first learn representation functions ( φ̂1 , . . . , φ̂H ) so that we can use a few demonstrations from an expert policy π∗ for new task µ ∼ η and learn a policy π = ( π1 , . . . , πH ) that uses the learned representations , i.e . πh ∈ Πφ̂h , such that has average cost of π is not too far away from π∗ . In the case of stationary policies , we need to learn a single φ by using tasks and learn π ∈ Πφ for a new task . The hope is that data from multiple tasks can be used to learn a complicated function φ ∈ Φ first , thus requiring only a few samples for a new task to learn a linear policy from the class Πφ . Gaussian complexity : As in Maurer et al . ( 2016 ) , we measure the complexity of a function class H ⊆ { h : X → Rd } on a set X = ( X1 , . . . , Xn ) ∈ Xn by using the following Gaussian average G ( H ( X ) ) = E sup h∈H d , n∑ i=1 j=1 γijhi ( Xj ) | Xj ( 2 ) where γij are independent standard normal variables . Bartlett & Mendelson ( 2003 ) also used Gaussian averages to show some generalization bounds . 1Break ties in any way

Summary: The paper tackles the representation learning problem where the aim is to learn a generic representation that is useful for a variety of downstream tasks. A two-level optimization framework is proposed: an inner optimization over the specific problem-at-hand, and an outer optimization over other similar problems. The problem is studied in two settings of the imitation learning framework with the additional aim of providing mathematical guarantees in terms of sample efficiency on new tasks. An extensive theoretical analysis is performed, and some preliminary empirical results are presented.