Year: 2019

Abstract:
A common strategy in modern learning systems is to learn a representation which is useful for many tasks, a.k.a. representation learning. We study this strategy in the imitation learning setting for Markov decision processes (MDPs) where multiple experts’ trajectories are available. We formulate representation learning as a bi-level optimization problem 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. We instantiate this framework for the imitation learning settings of behavior cloning and observation-alone. Theoretically, we provably show using our framework that representation learning can reduce the sample complexity of imitation learning in both settings. We also provide proof-of-concept experiments to verify our theoretical findings.

heading: 1 INTRODUCTION
text: 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 cannot 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 cannot 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES
text: 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

heading: 4 BI-LEVEL OPTIMIZATION FRAMEWORK
text: In this section we introduce our framework and give a high-level description of the conditions under which this framework gives us statistical guarantees. Our main idea is to phrase learning representations for imitation learning as the following bi-level optimization
min φ∈Φ L(φ) := E µ∼η min π∈Πφ
`µ(π) (3)
Here `µ is the inner loss function that penalizes π being different from π∗µ for the task µ. In general, one can use any loss `µ that is used for single task imitation learning, e.g. for the behavioral cloning setting (cf. Section 5), `µ is a classification like loss that penalizes the mismatch between predictions by π∗ and π, while for the observation-alone setting (cf. Section 6) it is some measure of distance between the state visitation distributions induced by π and π∗. The outer loss function is over the representation φ. The use of bi-level optimization framework naturally enforces policies in the inner optimization to share the same representation.
While Equation 3 is formulated in terms of the distribution η, in practice we only have access to few samples for T tasks; let x(1), . . . ,x(T ) denote samples from tasks µ(1), . . . , µ(T ) sampled i.i.d. from η. We thus learn the representation φ̂ by minimizing empirical version L̂ of Equation 3.
L̂(φ) = 1
T T∑ i=1 min π∈Πφ `x (i) (π) = 1 T T∑ i=1 `x (i) (πφ,x (i) )
where `x is the empirical loss on samples x and πφ,x = arg minπ∈Πφ `x(π) corresponds to a task specific policy that uses a fixed representation φ. Our goal then is to show that for a new task µ ∼ η, the policy πφ̂,x learned by using samples x from the task µ has low expected cost Jµ, i.e., Informal Theorem 4.1. With high probability over the sampling of train task data and with sufficient number of tasks and samples per task,
E µ∼η E x Jµ(π φ̂,x)− E µ∼η Jµ(π ∗ µ) is small
At a high level, in order to prove such a theorem for a particular choice of `µ, we would need to prove the following three properties about `µ and `x:
1. `x(π) concentrates to `µ(π) simultaneously for all π ∈ Πφ (for a fixed φ), with sample complexity depending on some complexity measure of Πφ rather than being polynomial in |S|; 2. a small value of `µ(π) implies a small value for Jµ(π)− Jµ(π∗µ); 3. if φ and φ′ induce “similar” representations then minπ∈Πφ `µ(π) and minπ∈Πφ′ ` µ(π) are close.
The first property ensures that learning a policy for a single task by fixing the representation is sample efficient, thus making representation learning a useful problem to solve. The second property ensures that matching the behavior of the expert as measured by the loss `µ ensures low average cost i.e., `µ is meaningful for the average cost; any standard imitation learning loss will satisfy this. The third property is specific to representation learning and requires `µ to use representations in a smooth way. This ensures that the empirical loss for T tasks is a good estimate for the average loss on tasks sampled from η. We prove these three properties for the cases where `µ is the either behavioral cloning loss or observation-alone loss, with natural choices for the empirical loss `x. However the general proof recipe can be used for potentially many other settings and loss functions.
In the next section, we will describe representation learning for behavioral cloning as an instantiation of the above framework and describe the various components of the framework. Furthermore we will describe the results and give a proof sketch to show how the aforementioned properties help us show our final guarantees. The guarantees for this setting follow almost directly from results in Maurer et al. (2016) and Ross et al. (2011). Later in Section 6 we describe the same for the observations alone setting which is more non-trivial.

heading: 5 REPRESENTATION LEARNING FOR BEHAVIORAL CLONING
text: Choice of `µ: We first specify the inner loss function in the bi-level optimization framework. In the single task setting, the goal of behavioral cloning (Syed & Schapire, 2010; Ross et al., 2011)
is to use expert trajectories of the form τ = (s1, a1, . . . , sH , aH) to learn a stationary policy2 that tries to mimic the decisions of the expert policy on the states visited by the expert. For a task µ, this reduces to a supervised classification problem that minimizes a surrogate to the following loss `µ0−1(π) = Es∼ν∗µ,a∼π∗µ(s)I{π(s) 6= a}. We abuse notation and denote this distribution over (s, a) for task µ as µ; so (s, a) ∼ µ is the same as s ∼ ν∗µ, a ∼ π∗µ(s). Prior work (Syed & Schapire, 2010; Ross et al., 2011) have shown that a small value of `µ0−1(π) implies a small difference J(π)−J(π∗). Thus for our setting, we choose `µ to be of the following form
`µ(π) = E s∼ν∗µ,a∼π∗µ(s) `(π(s), a) = E (s,a)∼µ `(π(s), a) (4)
where ` is any surrogate to 0-1 loss I{a 6= arg max a′∈A π(s)a′} that is Lipschitz in φ(s). In this work we consider the logistic loss `(π(s), a) = − log(π(s)a).
Learning φ from samples: Given expert trajectories for T tasks µ(1), . . . , µ(T ) we construct a dataset X = {x(1), . . . ,x(T )}, where x(t) = {(stj , atj)}nj=1 ∼ (µ(t))n is the dataset for task t. Details of the dataset construction are provided in Section C.1. Let S denote the set of states {stj}. Instantiating our framework, we learn a good representation by solving φ̂ = arg min
φ∈Φ L̂(φ), where
L̂(φ) := 1
T T∑ t=1 min π∈Πφ 1 n n∑ j=1 `(π(stj), a t j) = 1 T T∑ t=1 min π∈Πφ ˆ̀x(t)(π) (5)
where `x is loss on samples x = {(sj , aj)}nj=1 defined as `x(π) = 1n ∑n j=1 `(π(sj), aj).
Evaluating representation φ̂: A learned representation φ̂ is tested on a new task µ ∼ η as follows: draw samples x ∼ µn using trajectories from π∗µ and solve πφ̂,x = arg min
π∈Πφ̂ ˆ̀x(π). Does πφ̂,x have
expected cost Jµ(πφ̂,x) not much larger than Jµ(π∗µ)? The following theorem answers this question. We make the following two assumptions to prove the theorem. Assumption 5.1. The expert policy π∗µ is deterministic for every µ ∈ support(η). Assumption 5.2 (Policy realizability). There is a representation φ∗ ∈ Φ such that for every µ ∈ support(η), πµ ∈ Πφ ∗ such that πµ(s)π∗µ(s) 3 ≥ 1− γ,∀s ∈ S for some γ < 1/2.
The first assumption holds if π∗µ is aiming to maximize some cost function. The second assumption is for representation learning to make sense: we need to assume the existence of a common representation φ∗ that can approximate all expert policies and γ measures this expressiveness of Φ. Now we present our first main result.
Theorem 5.1. Let φ̂ ∈ arg min φ∈Φ L̂(φ). Under Assumptions 5.1,5.2, with probability 1 − δ over the sampling of dataset X, we have
E µ∼η E x∼µn
Jµ(π φ̂,x)− E
µ∼η Jµ(π
∗ µ) ≤ H2(2γ + gen)
where gen = c G(Φ(S)) T √ n + c′R √ K√ n
+ c′′ √
ln(4/δ) T , for some small constants c, c ′, c′′.
To gain intuition for what the above bound means, we give a PAC-style guarantee for the special case where the class of representation functions Φ is finite. This follows directly from the above theorem and the use of Massart’s lemma. Corollary 5.1. In the same setting as Theorem 5.1, suppose Φ is finite. If number of tasks satisfies T ≥ c1 max { H4R2 log(|Φ|) 2 , H4 ln(4/δ) 2 } , and number of samples (expert trajectories) per task
satisfies n ≥ c2H 4R2K 2 for small constants c1, c2, then with probability 1− δ,
E µ∼η E x∼µn
Jµ(π φ̂,x)− E
µ∼η Jµ(π
∗ µ) ≤ H2γ +
2We can easily extend the theory to non-stationary policies 3We abuse notation and use π∗µ(s) instead of argmax
a∈A π∗µ(s)a
Discussion: The above bound says that as long as we have enough tasks to learn a representation from Φ and sufficient samples per task to learn a linear policy, the learned policy will have small average cost on a new task from η. The first term H2γ is small if the representation class Φ is expressive enough to approximate the expert policies (see Assumption 5.2). The results says that if we have access to data from T = O ( H4R2 log(|Φ|)
2
) tasks sampled from η, we can use them to
learn a representation such that for a new task we only need n = O ( H4R2K 2 ) samples (expert
demonstrations) to learn a linear policy with good performance. In contrast, without access to tasks, we would need n = O ( max { H4R2 log(|Φ|) 2 , H4R2K 2 }) samples from the task to learn a good
policy π ∈ Π from scratch. Thus if the complexity of the representation function class Φ is much more than number of actions (log(|Φ|) K in this case), then multi-task representation learning might be much more sample efficient4. Note that the dependence of sample complexity onH comes from the error propagation when going from `µ to Jµ; this is also observed in single task imitation learning (Ross et al., 2011; Sun et al., 2019).
We give a proof sketch for Theorem 5.1 below, while the full proof is deferred to Appendix A.

heading: 5.1 PROOF SKETCH
text: The proof has two main steps. In the first step we bound the error due to use of samples. The policy πφ,x that is learned on samples x ∼ µn is evaluated on the distribution µ and the average loss incurred by representation φ across tasks is L̄(φ) = E
µ∼η E x∼µn `µ(πφ,x).
On the other hand, if the learner had complete access to the distribution η and distributions µ for every task, then the loss minimizer would be φ∗ = arg minφ∈Φ L(φ), where L(φ) := E
µ∼η min π∈Πφ
`µ(π).
Using results from Maurer et al. (2016), we can prove the following about φ̂
Lemma 5.2. With probability 1− δ over the choice of X, φ̂ ∈ arg min φ∈Φ L̂(φ) satisfies
L̄(φ̂) ≤ min φ∈Φ
L(φ) + c G(Φ({stj}))
T √ n
+ c′ R √ K√ n
+ c′′ √ ln(1/δ)
T
The proof of this lemma is provided in the appendix for completeness.
The second step of the proof is connecting the loss L̄(φ) and the average cost Jµ of the policies induced by φ for tasks µ ∼ η. This can obtained by using the connection between the surrogate 0-1 loss `µ and the cost Jµ that has been established in prior work (Ross et al., 2011; Syed & Schapire, 2010). The following lemma uses the result for deterministic expert policies from Ross et al. (2011). Lemma 5.3. Given a representation φ with L̄(φ) ≤ . Let x ∼ µn be samples for a new task µ ∼ η. Let πφ,x be the policy learned by behavioral cloning on the samples, then under Assumption 5.1
E µ∼η E x∼µn
Jµ(π φ,x)− E
µ∼η Jµ(π
∗ µ) ≤ H2
This suggests that making L̄ small is good enough. A simple implication of Assumption 5.2 that minφ∈Φ L(φ) ≤ L(φ∗) ≤ γ, along with the above two lemmas completes the proof.

heading: 6 REPRESENTATION LEARNING FOR OBSERVATION-ALONE SETTING
text: Now we consider the setting where we cannot observe experts’ actions but only their states. As in Sun et al. (2019), we also solve a problem at each level; consider a level h ∈ [H].
Choice of `µh: Let π ∗ µ = {π∗1,µ, . . . , π∗H,µ} be the sequence of expert policies (possibly stochastic) at different levels for the task µ. Let ν∗h,µ be the distribution induced on the states at level h by the expert policy π∗µ. The goal in imitation learning with observations alone (Sun et al., 2019) is
4These statements are qualitative since we are comparing upper bounds.
to learn a policy π = (π1, . . . , πH) that matches the distributions νπh with ν ∗ h for every h, w.r.t. a discriminator class G5 that contains the true value functions V ∗1 , . . . , V ∗H and is approximately closed under the Bellman operator of π∗. Instead, in this work we learn π that matches the distributions πh · ν∗h and ν∗h+1 for every h w.r.t. to a class G ⊆ {g : S → R, |g|∞ ≤ 1} that contains the value functions and has a stronger Bellman operator closure property. For every task µ, `µh is defined as
`µh(π) = maxg∈G [ E s∼ν∗h,µ E a∼π(s) s̃∼Ps,a g(s̃)− E s̄∼ν∗h+1,µ g(s̄)] (6)
= max g∈G [ E s∼ν∗h,µ E a∼U(A) s̃∼Ps,a Kπ(a|s)g(s̃)− E s̄∼ν∗h+1,µ g(s̄)]
where we rewrite `µh by importance sampling in the second equation; this will be useful to get an empirical estimate. While our definition of `µh differs slightly from the one used in Sun et al. (2019), using similar techniques, we will show that small values for `µh(πh) for every h ∈ [H] will ensure that the policy π = (π1, . . . , πH) will have expected cost Jµ(π) close to Jµ(π∗µ). We abuse notation, and for a task µ we denote µ = (µ1, . . . , µH) where µh is the distribution of (s, a, s̃, s̄) used in `µh; thus (s, a, s̃, s̄) ∼ µh is equivalent to s ∼ ν∗h,µ, a ∼ U(A), s̃ ∼ Ps,a, s̄ ∼ ν∗h+1,µ.
Learning φh from samples: We assume, 1) access to 2n expert trajectories for T independent train tasks, 2) ability to reset the environment at any state s and sample from the transition P (·|s, a) for any a ∈ A. The second condition is satisfied in many problems equipped with simulators. Using the sampled trajectories for the T tasks {µ(1), . . . , µ(T )} and doing some interaction with environment, we get the following dataset X = {X1, . . . ,XH} where Xh is the dataset for level h. Specifically, Xh = {x(1)h , . . . ,x (T ) H } where x (i) h = {(sij , aij , s̃ij , s̄ij)}nj=1 ∼ (µ(i))n. Additionally we denote Sh = {sij} T,n i=1,j=1 to be all the s-states in Xh, S̃h and S̄h are similarly defined. Details about how this dataset is constructed from expert trajectories and interactions with environment is provided in Section C.2. We learn the representation φ̂h = arg min
φ∈Φ L̂h(φ), where
L̂h(φ) = 1
T T∑ i=1 min π∈Πφ max g∈G 1 n n∑ j=1 [Kπ(aij |sij)g(s̃ij)− g(s̄ij)] = 1 T T∑ i=1 min π∈Πφ ˆ̀x(i) h (π)
where for dataset x = {(sj , aj , s̃j , s̄j)}nj=1, ˆ̀xh(π) := max g∈G 1 n n∑ j=1 [Kπ(aj |sj)g(s̃j) − g(s̄j)]. Note that because of the maxg∈G , ˆ̀xh is no longer an unbiased estimator of ` µ h when x ∼ µnh . However we can still show generalization bounds.
Evaluating representations φ̂1, . . . , φ̂H : Learned representations are tested on a new task µ ∼ η as follows: get samples x = (x1, . . . ,xH)6 for all levels using trajectories from π∗µ, where xh ∼ µnh . For each level h, learn πφ̂h,xh = arg minπ∈Πφ̂ ˆ̀ xh h (π) and consider the policy π φ̂,x = (πφ̂1,x1 , . . . , πφ̂H ,xH ). Before presenting the guarantee for πφ̂,x, we introduce a notion of Bellman error that will show up in our results. For a policy π = (π1, . . . , πH) and an expert policy π∗ = (π∗1 , . . . , π ∗ H), we define the inherent Bellman error
πbe := max h∈[H] max g∈G min g′∈G E s∼(ν∗h+ν π h )/2 [|g′(s)− (Γπh g)(s)|] (7)
We make the following two assumptions for the subsequent theorem. These are standard assumptions in theoretical reinforcement learning literature. Assumption 6.1 (Value function realizability). V ∗h,µ ∈ G for every h ∈ [H], µ ∈ support(η). Assumption 6.2 (Policy realizability). There are representations φ∗1, . . . , φ∗H ∈ Φ such that π∗h,µ ∈ Πφ ∗ h for every h ∈ [H], µ ∈ support(η).
Now we present our main theorem for the observation-alone setting. 5If G contains all bounded functions, then it reduces to minimizing TV between νπh and ν∗h. 6Note that we do not need the datasets xh at different levels to be independent of each other
Theorem 6.1. Let φ̂h ∈ arg min φ∈Φ L̂h(φ). Under Assumptions 6.1,6.2, with probability 1 − δ over sampling of X = (X1, . . . ,XH), we have
E µ∼η E x J(πφ̂,x)− E µ∼η J(π∗µ) ≤ H∑ h=1 (2H − 2h+ 1) gen,h +O(H2) φ̂be
where φ̂be = E µ∼η E x [ π φ̂,x be ] is the average inherent Bellman error and
gen,h = O
( KG(Φ(Sh))
T √ n
+ E µ∼η E x∼µn
[ KG(G(s̃h))
n + G(G(s̄h)) n
] + RK √ K√
n +
√ ln(H/δ)
T
)
We again give a PAC-style guarantee for the special case where the class of representation functions Φ and value function class G are finite. It follows from the above theorem and Massart’s lemma. Corollary 6.1. In the setting of Theorem 6.1, suppose Φ,G are finite. If number of tasks satisfies T ≥ c1 max { H4R2K2 log(|Φ|) 2 , H4 ln(H/δ) 2 } , and number of samples (trajectories) per task satisfies
n ≥ c2 max { H4K2 log(|G|) 2 , H4R2K3 2 } for small constants c1, c2, then with probability 1− δ,
E µ∼η E x J(πφ̂,x)− E µ∼η J(π∗µ) ≤ O(H2) φ̂ be + .
Discussion: As in the previous section, the number of samples required for a new task after learning a representation is independent of the class Φ but depends only on the value function class G and number of actions. Thus representation learning is very useful when the class Φ is much more complicated than G, i.e. R2 log(|Φ|) max{log(|G|), R2K}. In the above bounds, φ̂be is a Bellman error term. This type of error terms occur commonly in the analysis of policy iteration type algorithms (Munos, 2005; Munos & Szepesvári, 2008). We remark that unlike in Sun et al. (2019), our Bellman error is based on the Bellman operator of the learned policy rather than the optimal policy. Le et al. (2019) used a similar notion that they call inherent Bellman evaluation error.
The proof of Theorem 6.1 follows a similar outline to that of behavioral cloning. However we cannot use the results from Maurer et al. (2016) directly since we are solving a min-max game for each task. We provide the proof in Appendix B.

heading: 7 EXPERIMENTS
text: In this section we present experimental results on the DirectedSwimmer environment (modified from the Swimmer environment from OpenAI gym (Brockman et al., 2016)) with Todorov et al. (2012) simulator and a NoisyCombinationLock environment designed by ourself. These experiments have two aims: 1) verify the benefit of representation learning predicted by our theory, 2) test the power of representations learned via our framework in a broader context: we learn a policy for a new task by using the representation and doing policy optimization instead of imitation learning. In our experiments we learn representations using Equation 5. Experiment details are deferred to Section D.
Our method: Given access to a dataset X = {(stj , atj)}nj=1 of n state-action pairs each for T tasks, we learn a φ̂ according to Equation 8. For any new task we learn a linear policy π from the class Πφ̂.
φ̂ = arg min φ∈Φ min f1,...,fT∈F
1
T T∑ t=1 1 n n∑ j=1 − log(πφ,ft(stj)atj ) (8)
Baseline: For a task we learn a policy π from the class Π without learning a representation first.
Verification of theory: In Figure 1 we verify our theoretical findings. On the left, we test on the DirectedSwimmer environment and report the logistic loss on the validation, which measures how close the trained policy is to the target expert policy. We find that learning representations, even with a few experts, can significantly reduce the sample complexity. On the right, we report the average
reward of the trained policies on the environment. Here we see a different phenomenon: when the number of experts is small (4 or 16), the baseline method can beat policies trained using representation learning, though the baseline method requires more samples to do so. When the number of experts is large (64), we see the policy trained using representation learning can significantly outperform the baseline method. This behavior is expected as when the number of experts is small, we may learn a sub-optimal representation and because we fix this representation for training the policy, more samples for the test task cannot make this policy better, whereas more samples always make the baseline method better. Nevertheless, when the number of experts is large, we can significantly reduce the sample complexity. With 60 samples, the base line method is still far behind the policy trained using representation learning with 64 experts.
Policy optimization with representations trained by imitation learning: We next test the utility of representations learned via our framework for RL. After training a representation, we use a simplified proximal policy optimization method that learns a linear policy over the learned representation. Results are reported in Figure 2. For DirectedSwimmer and NoisyCombinationLock, we observe a common pattern. When the number of experts to learn the representation is small, the baseline method enjoys better performance than the policies trained using representation learning. As the number of experts to learn the representation increases, we see the policy trained using representation learning can initially outperform baseline, sometime significantly. However, unsurprisingly, the baseline method performs very well with a large number of samples, since it is allowed to learn a representation from scratch. This experiment suggests that representations trained via imitation learning can be useful beyond imitation learning, especially when the target task has few samples.

heading: 8 CONCLUSION
text: The current paper proposes a bi-level optimization framework to formulate and analyze representation learning for imitation learning using multiple demonstrators. Theoretical guarantees are provided to justify the statistical benefit of representation learning. Some preliminary experiments verify the effectiveness of the proposed framework. In particular, in experiments, we find the representation learned via imitation learning is also useful for policy optimization in the reinforcement learning setting. We believe it is an interesting theoretical question to explain this phenomenon. Additionally, extending this bi-level optimization framework to incorporate methods beyond imitation learning is an interesting future direction.

heading: A PROOFS FOR BEHAVIORAL CLONING
text: We prove Theorem 5.1 in this section by proving Lemma 5.2,5.3. In this section, we abuse notation and define `µ(φ, f) := `µ(πφ,f ), where `µ is defined in Equation 4. Let f̂φx = arg min
f∈F `x(φ, f) be
the optimal task specific parameter for task µ by fixing representation φ. Thus by our definitions in Section 5, we get πφ,x = πφ,f̂ φ x . We assume w.l.o.g. thatA = [K]. Remember that ` : 4(A)×A → R is defined as `(v, a) = − log(va) for some v ∈ RK and va is the coordinate corresponding to action a ∈ A = [K]. We define a new function class and loss function that will be useful for our proofs
F ′ = {x→Wx |W ∈ RK×d, ‖W‖F ≤ 1} (9)
`′(v, a) = − log(softmax(v)a),v ∈ RK , a ∈ A (10)
We basically offloaded the burden of computing softmax from the class F to the loss `. We can convert any function f ′ ∈ F ′ to one in F by transforming it to softmax(f ′). We now proceed to proving the lemmas
Proof of Lemma 5.2. We can then rewrite the various loss functions from Section 5 as follows
L̂(φ) = 1
T T∑ i=1 min f∈F ′ 1 n n∑ j=1 `′(f(φ(s)), a)
L(φ) = E µ∼η min f∈F ′ E (s,a)∼µ
`′(f(φ(s)), a)
L̄(φ) = E µ∼η E x∼µn E (s,a)∼µ
`′(f̂ ′ φ
x(φ(s)), a)
where f̂ ′ φ
µ ∈ arg minf ′∈F ′ `x(φ,softmax(f ′)). It is easy to show that both `′(·, a) `′(f ′(·), ·) are 2-lipschitz in their arguments for every a ∈ A and f ′ ∈ F ′. Using a slightly modified version of Theorem 2(i) from Maurer et al. (2016), we get that for φ̂ ∈ arg minφ∈Φ L̂(φ), with probability at least 1− δ over the choice of X
L̄(φ̂)−min φ∈Φ
L(φ) ≤ 2 √ 2πG(Φ(S))
T √ n
+ √
2πQ′ sup φ∈Φ √√√√ Eµ∼η,(s,a)∼µ ‖φ(s)‖2 n + √ 8 log(4/δ) T
L̄(φ̂)−min φ∈Φ L(φ) ≤ cG(Φ(S)) T √ n + c′ Q′R√ n
+ c′′ √ log(4/δ)
T (11)
where Q′ = sup y∈Rdn\{0} 1 ‖y‖E sup f∈F ′ n,K∑ i=1,j=1 γijf ′(yi)j . First we discuss why we need a modified version of their theorem. Our setting differs from the setting for Theorem 2 from Maurer et al. (2016) in the following ways
• F ′ is a class of vector valued function in our case, whereas in Maurer et al. (2016) it is assumed to contain scalar valued. The only place in the proof of the theorem where this shows up is in the definition of Q′, which we have updated accordingly. • Maurer et al. (2016) assumes that `′(·, a) is 1-lipschitz for every a ∈ A and that f ′(·) is L lipschitz for every f ′ ∈ F ′. However the only properties that are used in the proof of Theorem 16 are that `′(·, a) is 1-lipschitz and that `′(f ′(·), a) is L-lipschitz for every a ∈ A, which is exactly the property that we have. Hence their proof follows through for our setting as well.
Lemma A.1. Q′ := sup y∈Rdn\{0} 1 ‖y‖E sup f∈F ′ n,K∑ i=1,j=1 γijf ′(yi)j ≤ √ K

heading: Proof.
text: Q′ := sup y∈Rdn\{0}
1
‖y‖ E sup f∈F ′ n,K∑ i=1,j=1 γijf ′(yi)j
= sup y∈Rdn\{0}
1
‖y‖ E sup ‖W‖F≤1 n,K∑ i=1,j=1 γij〈Wj , yi〉
= sup y∈Rdn\{0}
1
‖y‖ E sup ‖W‖F≤1 K∑ j=1 〈Wj , n∑ i=1 γijyi〉
= sup y∈Rdn\{0}
1
‖y‖ E √√√√ K∑ j=1 ∥∥∥∥∥ n∑ i=1 γijyi ∥∥∥∥∥ 2
≤ sup y∈Rdn\{0}
1
‖y‖ √√√√ K∑ j=1 E ∥∥∥∥∥ n∑ i=1 γijyi ∥∥∥∥∥ 2
= sup y∈Rdn\{0}
1
‖y‖ √√√√ K∑ j=1 n∑ i=1 ‖yi‖2 = 1 ‖y‖ √ K‖y‖2 = √ K
where we use Jensen’s inequality and linearity of expectation for the first inequality and properties of standard normal gaussian variables for the equality after that.
Plugging in Lemma A.1 into Equation 11 completes the proof.
We now proceed to prove the next lemma.
Proof of Lemma 5.3. Suppose L̄(φ) = E µ∼η E x∼µn `µ(πφ,x) ≤ . Consider a task µ ∼ η and samples x ∼ µn and let µ(x) = `µ(πφ,x) so that L̄(φ) = E µ∼η E x∼µn µ(x). Since π∗µ is deterministic, we get
E s∼ν∗µ E a∼πφ,x I{a 6= π∗µ(s)} = E s∼ν∗µ [1− πφ,x(s)π∗µ(s)]
≤ E s∼ν∗µ [− log(1− (1− πφ,x(s)π∗µ(s)))]
= E s∼ν∗µ [− log(πφ,x(s)π∗µ(s))] = µ(x)
where we use the fact that x ≤ − log(1 − x) for x < 1. for the first inequality. Thus by using Theorem 2.1 from Ross et al. (2011), we get that Jµ(πφ,x)−Jµ(π∗) ≤ H2 µ(x). Taking expectation w.r.t. µ ∼ η and x ∼ µn completes the proof.
Proof of Theorem 5.1. By using Assumption 5.2, we first get that
L(φ∗) = E µ∼η min π∈Πφ∗ E s∼ν∗µ − log(π(s)π∗µ(s))
≤ E µ∼η E s∼ν∗µ − log(πµ(s)π∗µ(s))
≤ E µ∼η E s∼ν∗µ − log(1− γ) ≤ 2γ
where in the last step we used − log(1 − x) ≤ 2x for x < 1/2. Hence from Lemma 5.2 we get L̄(φ̂) ≤ 2γ + gen,h, which combining with Lemma 5.3 gives the desired result.

heading: B PROOFS FOR OBSERVATION-ALONE
text: Before proving Theorem 6.1, we introduce the following loss functions, as we did in the proof sketch for the behavioral cloning setting. We again abuse notation and define `µ(φ, f) := `µ(πφ,f ), where `µ is defined in Equation 6. Let f̂φx = arg min
f∈F `x(φ, f) be the optimal task specific parameter for
task µ by fixing representation φ. As before, we define the following
L̄h(φh) = E µ∼η E x∼µnh
`µh(φ, f̂ φh x )
We first show a guarantee on the performance of representations (φ̂1, . . . , φ̂H) as measured by the functions L̄1, . . . , L̄H . Theorem B.1. With probability at least 1− δ in the draw of X = (X(1), . . . ,X(H)), ∀h ∈ [H]
L̄h(φ̂h) ≤ min φ∈Φ Lh(φ) + c gen,h(Φ) + c ′ gen,h(F ,G) + c′′
√ ln(H/δ)
T
where gen,h(Φ) = KG(Φ(Sh))
T √ n
and gen,h(F ,G) = E µ∼η E x∼µn
[ KG(G(s̃h))
n + G(G(s̄h)) n
] + RK √ K√
n
We then connect the losses L̄h to the expected cost on the tasks. Theorem B.2. Consider representations (φ1, . . . , φH) with L̄h(φh) ≤ h. Let x = (x1, . . . ,xH) be samples at different levels for a newly sampled task µ ∼ η such that xh ∼ µnh . Let πφ,x = (πφ1,x1 , . . . , πφH ,xH ) be policies learned using the samples, then under Assumption 6.1,
E µ∼η E x J(πφ,x)− E µ∼η J(π∗µ) ≤ H∑ h=1 (2H − 2h+ 1) h +O(H2) φbe
where φbe = E µ∼η E x [ π φ,x be ] is the average inherent Bellman error.
It is easy to show that under Assumption 6.2, minφ∈Φ Lh(φ) = 0 for every h ∈ [H]. Thus from Theorem B.1, we get that L̄h(φ̂h) ≤ gen,h, where gen,h = gen,h(Φ)+ gen,h(F ,G)+c′′
√ ln(H/δ)
T .
Invoking Theorem B.2 on the representations {φ̂h} completes the proof.

heading: B.1 PROOF OF THEOREM B.1
text: Before proving the theorem, we discuss important lemmas. In yet another abuse of notation, we define `µh(φ, f, g) = E(s,a,s̃,s̄)∼µh [Kπφ,f (a|s)g(s̃) − g(s̄)] and `xh(φ, f, g) = 1 n n∑ j=1 [Kπφ,f (aj |sj)g(s̃j)− g(s̄j)].
Let m̂x(φ) = min f∈F max g∈G ˆ̀x h(φ, f, g) = ˆ̀x h(φ, f̂ φ x , ĝ φ x), m̄µ,x(φ) = max g∈G `µh(φ, f̂ φ x , g), mµ(φ) =
min f∈F max g∈G `µh(φ, f, g). Note that Lh(φ) = E µ∼η m(φ), L̄h(φ) = E µ∼η E x∼µn m̄µ,x(φ). Define the distribution ρh where x ∼ ρh is the same as µ ∼ η and then x ∼ µnh . Lemma B.3. For every φ ∈ Φ and h ∈ [H],
E µ∼η E x∼µn sup f∈F sup g∈G
[ ˆ̀x h(φ, f, g)− ` µ h(φ, f, g) ] ≤ gen,h(F ,G)
Lemma B.4. With probability 1− δ, for every φ ∈ Φ, L̄h(φ)− E
x∼ρh m̂x(φ) ≤ gen,h(F ,G)
Lemma B.5. With probability 1− δ, for every φ ∈ Φ,
E x∼ρh
m̂x(φ)− 1
T ∑ i m̂x(i)(φ) ≤ gen,h(Φ) +O √ log( 1δ ) T 
We prove these lemmas later. First we prove Theorem B.1 using them. If φ∗h = arg min φ∈Φ Lh(φ), then
L̄h(φ̂h)− Lh(φ∗h) = ( L̄h(φ̂h)− E
x∼ρh m̂x(φ) ) + ( E
x∼ρh m̂x(φ)−
1
T ∑ i m̂x(i)(φ̂h)
)
+
( 1
T ∑ i m̂x(i)(φ̂h)− 1 T ∑ i m̂x(i)(φ ∗ h)
)
+
( 1
T ∑ i m̂x(i)(φ ∗ h)− E x∼ρh m̂x(φ ∗ h) ) + E µ∼η [ E x∼µn m̂x(φ ∗ h)−mµ(φ∗h)]
≤ 2 gen,h(F ,G) + gen,h(Φ) +O √ log( 1δ ) T  where for the first part we use Lemma B.4, second part we use Lemma B.5, third part is upper
bounded by 0 by optimality of φ̂h, fourth is upper bounded by O( √ log( 1δ )
T ) by Hoeffding’s inequality and fifth is bounded by the following argument: let fφ, gφ = arg min
f∈F arg max g∈G `µ(φ, f, g)
E x∼µn
m̂x(φ ∗ h) = E
x∼µn min f∈F max g∈G
ˆ̀x h(φ ∗ h, f, g)
≤ E x∼µn max g∈G ˆ̀x h(φ ∗ h, f φ∗h , g) = E x∼µn ˆ̀x h(φ ∗ h, f φ∗h , g̃)
≤ `µh(φ ∗ h, f φ∗h , g̃) + gen,h(F ,G) ≤ `µh(φ ∗ h, f φ∗h , gφ ∗ h) + gen,h(F ,G) = mµ(φ∗h) + gen,h(F ,G)
where the second inequality uses Lemma B.3.

heading: B.2 PROOF OF THEOREM B.2
text: Consider a task µ. For simplicity of notation, we use πh instead πφh,xh , π instead of πφ,x. Let νπh and ν∗h be the state distributions at level h induced by π φ,x and π∗µ respectively. Let
h(xh) = max g∈G E s∼ν∗h [ E a∼πh s′∼Ps,a g(s′)− E a∼π∗h s′∼Ps,a g(s′)]
be the loss of policy πh at level h. By definition, h = E µ∼η E x∼µnh h(x). Using Lemma C.1 from Sun et al. (2019), we have J(πφ,x)− J(π∗µ) = H∑ h=1 ∆̄h = H∑ h=1 E s∼νπh [ E a∼πh(·|s),s′∼Ps,a V ∗h+1(s ′)− E a∼π∗h(·|s),s ′∼Ps,a V ∗h+1(s ′)
] Observe that
∆̄h = E s∼νπh [ E a∼πh(·|s),s′∼Ps,a
V ∗h+1(s ′)− E
a∼π∗h(·|s),s ′∼Ps,a
V ∗h+1(s ′)]
≤ max g∈G E s∼ν∗h [ E a∼πh(·|s),s′∼Ps,a g(s′)− E a∼π∗h(·|s),s ′∼Ps,a g(s′)]+
max g∈G [ E s∼νπh Γπh g(s)− E s∼ν∗h Γπh g(s)] + [ E s∼ν∗h Γ∗hV ∗ h+1(s)− E s∼νπh Γ∗hV ∗ h+1(s)]
≤ h(xh) + max g∈G [ E s∼νπh Γπh g(s)− E s∼ν∗h Γπh g(s)] + max g∈G [ E s∼νπh g(s)− E s∼ν∗h g(s)]
Lemma B.6. Defining ∆h = max g∈G | E s∼νπh g(s)− E s∼ν∗h g(s)|, we have
max g∈G [ E s∼νπh Γπh g(s)− E s∼ν∗h Γπh g(s)] ≤ ∆h + 2 πbe
Using the above lemma, we get ∆̄h ≤ h(xh) + 2∆h + 2 πbe. We now bound ∆h
∆h = max g∈G ∣∣∣∣∣∣ Es∼νπh−1 Ea∼πh−1 s′∼Ps,a g(s′)− E s∼ν∗h g(s) ∣∣∣∣∣∣ = max
g∈G ∣∣∣∣∣∣ Es∼νπh−1 Ea∼πh−1 s′∼Ps,a g(s′)− E s∼ν∗h−1 E a∼πh−1 s′∼Ps,a g(s) ∣∣∣∣∣∣+ maxg∈G ∣∣∣∣∣∣ Es∼ν∗h−1 Ea∼πh−1 s′∼Ps,a g(s′)− E s∼ν∗h g(s) ∣∣∣∣∣∣ = max
g∈G ∣∣∣∣∣ Es∼νπh−1 Γπh−1g(s′)− Es∼ν∗h−1 Γπh−1g(s) ∣∣∣∣∣+ h−1(xh−1)
≤ ∆h−1 + 2 πbe + h−1(xh−1)
Thus ∆h ≤ 2(h− 1) πbe + 1:h−1(x1:h−1) and so ∆̄h ≤ 1:h(x1:h) + 1:h−1(x1:h−1) + (4h− 2) πbe. This implies that
J(πφ,x)− J(π∗) = H∑ h=1 ∆̄h ≤ H∑ h=1 (2H − 2h+ 1) h(xh) +O(H2) π φ,x be
Taking expectation wrt µ ∼ η and x ∼ µn completes the proof.

heading: B.3 PROOFS OF LEMMAS
text: Proof of Lemma B.3. Again we defineF ′ as in Equation 9. Let `(v, α, β, a) = Ksoftmax(v)aα− β, and let `′µh (φ, f ′, g) = `′µh (φ,softmax(f ′), g) = E
(s,a,s̃,s̄)∼µh `(f ′(φ(s)), g(s̃), g(s̄), a) for f ′ ∈
F ′ and similarly define ˆ̀′ x
h(φ, f ′, g) = ˆ̀xh(φ,softmax(f ′), g). Notice that `(·, α, β, a) is 2Klipschitz, `(v, ·, β, a) is K-lipschitz and `(v, α, ·, a) is 1-lipschitz, Using Theorem 8(i) from Maurer et al. (2016), we get that
E µ∼η E x∼µn sup f∈F sup g∈G
[ ˆ̀x h(φ, f, g)− ` µ h(φ, f, g) ] = E µ∼η E x∼µn sup f ′∈F ′ sup g∈G [ ˆ̀′x h(φ, f ′, g)− `′µh (φ, f ′, g)
] ≤ √
2π ExG(`(F ′(φ(sh)),G(s̄h),G(s̃h),a)) n
≤ 2 √
2πKG(F ′(φ(sh)),a) n + E µ∼η E x∼µn
[√ 2πKG(G(s̄h))
n +
√ 2πG(G(s̃h))
n
]
≤ cRK √ K√
n + c′ E µ∼η E x∼µn
[ KG(G(s̄h))
n + G(G(s̃h)) n
] ≤ gen,h(F ,G)
where we used lipschitzness and Slepian’s lemma for second inequality and a similar computation to Lemma A.1 for the third.

heading: Proof of Lemma B.4.
text: L̄h(φ)− E x∼ρh m̂x(φ) = E µ∼η E x∼µn m̄µ,x(φ)− E µ∼η E x∼µn m̂x(φ)
= E µ∼η E x∼µn max g∈G
`µh(φ, f̂ φ x , g)−max g∈G ˆ̀x h(φ, f̂ φ x , g)
≤ E µ∼η E x∼µn max g∈G [`µh(φ, f̂ φ x , g)− ˆ̀xh(φ, f̂φx , g)]
≤ E µ∼η E x∼µn max f∈F max g∈G [`µh(φ, f, g)− ˆ̀ x h(φ, f, g)]
≤ gen,h(F ,G)
where we use the definition of L̄h, obviousness for the first inequality and Lemma B.3 for the last.
Proof of Lemma B.5. We wil be using Slepian’s lemma
Lemma B.7 (Slepian’s lemma). Let {X}s∈S and {Y }s∈S be zero mean Gaussian processes such that
E(Xs −Xt)2 ≤ E(Ys − Yt)2,∀s, t ∈ S
Then
E sup s∈S Xs ≤ E sup s∈S Ys
Using Theorem 8(ii) from Maurer et al. (2016), we get that
sup φ∈Φ
[ E
x∼ρh m̂x(φ)−
1
T ∑ i m̂x(i)(φ)
] ≤ √ 2π
T G(S) +
√ 9 ln(2/δ)
2T (12)
where S = {(m̂(φ)x1 , . . . , m̂(φ)xT ) : φ ∈ Φ}. We bound the Gaussian average of S using Slepian’s lemma. Define two Gaussian processes indexed by Φ as
Xφ = ∑ i γim̂(φ)x(i) and Yφ = 2K√ n ∑ i γijkφ(s i j)k
For x = {(sj , aj , s̃j , s̄j)}, consider 2 representations φ and φ′,
(m̂(φ)x − m̂(φ′)x)2 = (min f∈F max g∈G ˆ̀x h(φ, f, g)−min f∈F max g∈G ˆ̀x h(φ ′, f, g))2
≤ ( sup f∈F,g∈G |ˆ̀xh(φ, f, g)− ˆ̀xh(φ′, f, g)|)2
=  sup f∈F,g∈G ∣∣∣∣∣∣ 1n ∑ j [Kπφ,f (aj |sj)g(s̃j)−Kπφ ′,f (aj |sj)g(s̃j)] ∣∣∣∣∣∣ 2
= K2  sup f∈F,g∈G ∣∣∣∣∣∣ 1n ∑ j (f(φ(sj))aj − f(φ′(sj))aj )g(s̃j) ∣∣∣∣∣∣ 2
≤ K 2
n sup f∈F ∑ j ( f(φ(sj))aj − f(φ′(sj))aj )2 ≤ 4K 2
n ∑ j |φ(sj)− φ′(sj)|2 = 4K2 n ∑ j,k (φ(sj)k − φ′(sj)k)2
where we prove the first inequality later, second inequality comes from g being upper bounded by 1 and by Cauchy-Schwartz inequality, third inequality comes from the 2-lipschitzness of f .
E(Xφ −Xφ′) = ∑ i (m̂(φ)x(i) − m̂(φ′)x(i))2
≤ 4K 2
n ∑ i,j,k (φ(sij)k − φ′(sij)k)2 = E(Yφ − Yφ′)2
Thus by Slepian’s lemma, we get
G(S) = E sup φ∈Φ Xφ ≤ E sup φ∈Φ Yφ = 2K√ n G(Φ({sij}))
Plugging this into Equation 12 completes the proof. To prove the first inequality above, notice that
min f∈F max g∈G
ˆ̀x h(φ, f, g)−min
f∈F max g∈G
ˆ̀x h(φ ′, f, g) = ˆ̀xh(φ, f, g)− ˆ̀xh(φ′, f ′, g′)
≤ ˆ̀xh(φ, f ′, g′′)− ˆ̀xh(φ′, f ′, g′) ≤ ˆ̀xh(φ, f ′, g′′)− ˆ̀xh(φ′, f ′, g′′) ≤ sup f∈F,g∈G |ˆ̀xh(φ, f, g)− ˆ̀xh(φ′, f, g)|
By symmetry, we also get that min f∈F max g∈G ˆ̀x h(φ, f, g)−min f∈F max g∈G ˆ̀x h(φ ′, f, g) ≤ sup f∈F,g∈G |ˆ̀xh(φ, f, g)− ˆ̀x h(φ ′, f, g)|.
Proof of Lemma B.6. Let ḡ = arg max g∈G
( E
s∼νπh Γπh g(s)− E s∼ν∗h Γπh g(s)
) and g′ = arg min
g∈G |g −
Γπh ḡ|(νπh+ν∗h)/2.
max g∈G
( E
s∼νπh Γπh g(s)− E s∼ν∗h Γπh g(s) ) = E s∼νπh [ Γπh ḡ(s)− E s∼ν∗h Γπh ḡ(s) ] ≤ | E
s∼νπh g′(s)− E s∼ν∗h g′(s)|+ | E s∼νπh [g′(s)− Γπh ḡ(s)]|+ | E s∼ν∗h [g′(s)− Γπh ḡ(s)]|
≤ max g∈G | E s∼νπh g(s)− E s∼ν∗h g(s)|+ 2 E s∼(νπh+ν ∗ h)/2 [|g′(s)− Γπh ḡ(s)]|] ≤ ∆h + 2 πbe

heading: C DATA SET COLLECTION DETAILS

heading: C.1 DATASET FROM TRAJECTORIES
text: Given n expert trajectories for a task µ, for each trajectory τ = (s1, a1, . . . , sH , aH) we can sample an h ∼ U([H]) and select the pair (sh, ah) from that trajectory7. This gives us n i.i.d. pairs {(sj , aj)}nj=1 for the task µ. We collect this for T tasks and get datasets x(1), . . . ,x(T ).

heading: C.2 DATASET FROM TRAJECTORIES AND INTERACTION
text: Given 2n expert trajectories for a task µ, we use first n trajectories to get independent samples from the distributions ν∗1,µ, . . . , ν ∗ H,µ respectively for the s̄ states in the dataset. Using the next n trajectories, we get samples from ν∗0,µ, . . . , ν ∗ H−1,µ for the s states in the dataset, and for each such state we uniformly sample an action a fromA and then get a state s̃ from Ps,a by resetting the environment to s and playing action a. We collect this for T tasks and get datasets X(i) = {x(i)1 , . . . ,x (i) H } for every i ∈ [T ], where each dataset x(i)h a set of n tuples obtained level h. Rearranging, we can construct the datasets Xh = {x(1)h , . . . ,x (T ) h }.

heading: D EXPERIMENT DETAILS
text: For the policy optimization experiments, we use 4 random seeds to evaluate our algorithm. We show the results for 1 test environment as the results for other test environments are also showing the algorithm works but the magnitude of reward might be different, so we do not average the numbers over different test environments.
7In practice one can use all pairs from all trajectories, even though the samples are not strictly i.i.d.
Experiment Setup We first describe the construction of the NoisyCombinationLock environment. The state space is R40. Each state s is in the form of [sreal, snoise], while sreal ∈ R20 is either a onehot vector or a zero vector, and snoise ∈ R20 is sampled from N (0, I). The action space is discrete and has size 2. For each MDP, we have a sequence of actions a∗ ∈ [2]20. This is the sequence of optimal actions. We use different a∗ to define different environments. The transition model is that: If sreal = ei for some i and the action is a∗i , then s′real = ei+1 and we’ll get reward 1. Otherwise s′real will be all zero and the reward is 0. snoise will always be sampled from the Gaussian distribution. Note that once sreal is all zero, it will not change and the reward will always be 0. The maximum horiozn is set to 20 and therefore, the optimal policy has return 20. The initial sreal is always e1.
The representation has dimension of 10. We limit the function φ to be a linear mapping from R40 to R10. Although the dimension of representation is smaller than the number of states, there still exists a linear mapping from states to representation such that we can find a linear optimal policy. For each expert, we collect 200 state-action pairs to train the representation φ. The trajectories are generated by the optimal policy.
When training the policy using an RL algorithm, to reduce the impact of initialization, the last full connected layer is initialized to 0. We use the PPO (Schulman et al., 2017) algorithm to train our policy with code from Dhariwal et al. (2017).
DirectedSwimmer A DirectedSwimmer environment is the same as Swimmer in OpenAI Gym (Brockman et al., 2016), except the following: the reward function is parametrized by a direction d with ‖d‖ = 1, and is defined as the traveled distance along the direction d. For each task, we sample a random direction. The state space is still R8. The original action space in Swimmer is R2, and we discretize the action space, such that each entry can be only one of {−1,−0.5, 0, 0.5, 1}. We also reduce the maximum horizon from 1000 to 100. We trained the experts for 1 million steps by PPO to make sure it converges.
The function φ we use has two fully connected layers and two ReLU layers. The number of hidden units is 100, so is the dimension of representation. We also include the total rewards that each algorithm can get in Figure 3. Note even though the baseline has a high validation loss, its performance can be quite good. This does not indicate a failure of representation learning, but it shows that lower logistic loss does not always imply higher reward.
Optimization All optimization, including training φ, π and behavior cloning baseline, is done by Adam (Kingma & Ba, 2014) with learning rate 0.001 until it converges. To solve equation 8, we build a joint loss over φ and all f ’s in each task,
L(φ, f1, . . . , fT ) = 1
nT T∑ t=1 n∑ j=1 − log(πφ,ft(stj)atj ). (13)
Then we minimize L(φ, f1, . . . , fT ) and obtain the optimal φ.

Review summaries:
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. 

This paper theoretically explores reinforcement/imitation learning via representation learning. The key theoretical question being investigated is the relationship between representation learning in a multi-task/meta learning setup and its dependence to the sample/task complexity. The paper sets up the problem in bilevel optimization framework, where the inner optimization learns/optimizes task specific losses, while the outer optimization learns the representation used in the inner level tasks. The main takeaway from the two theorems (which are the core contributions of this paper) are that when the number of tasks is higher than the number of samples, representation learning can reduce the sample complexity. The paper explores two scenarios in imitation learning, namely behavioral cloning and when only the states of the experts are available (and not their actions). Some experiments are provided to empirically validate the theory.