Title: PROVABLE REPRESENTATION LEARNING FOR IMITA-TION LEARNING VIA BI-LEVEL OPTIMIZATION Anonymous

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.

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 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.

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; Daumé 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): Let M = (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 → (S) is the transition function, C : S × A → R is the cost function and ν ∈ (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 → (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 s 1 , a 1 , s 2 , a 2 , . . . , s H , a H where s 1 ∼ ν, a 1 ∼ π 1 (s), s 2 ∼ P s1,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 (s h ) = E H i=h C(s i , a i ) | a i ∼ π i (s i ), s i+1 ∼ P si,ai

and the state-action function

Q π h (s h , a h ) is defined as Q π h (s h , a h ) = E s h+1 ∼Ps h ,a h [V π h (s h+1 )].

The goal is to learn a policy π that minimizes the expected cost J(π) = E s1∼ν V π 1 (s 1 ). We define the Bellman operator at level h for any policy π as Γ π h : R S → R S , where for s ∈ S and

EQUATION

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 MDP M µ = (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 → R d | φ(s) 2 ≤ R} is a class of bounded norm representation functions mapping states to vectors and F ⊆ {f : R d → ∆(A)} is a class of functions mapping state representations to distribution over actions. We will be using linear functions, i.e.

F = {x → softmax(W x) | W ∈ R K×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 → R d } on a set X = (X 1 , . . . , X n ) ∈ X n by using the following Gaussian average

EQUATION

where γ ij are independent standard normal variables. Bartlett & Mendelson (2003) also used Gaussian averages to show some generalization bounds.

4. BI-LEVEL OPTIMIZATION FRAMEWORK:

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

EQUATION

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.

5. REPRESENTATION LEARNING FOR BEHAVIORAL CLONING:

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 τ = (s 1 , a 1 , . . . , s H , a H ) 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 (π) = E s∼ν * µ ,a∼π * µ (s) I{π(s) = 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

EQUATION

where is any surrogate to 0-1 loss I{a = 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) = {(s t j , a t j )} n j=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 {s t j }. 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 (π(s t j ), a t j ) = 1 T T t=1 min π∈Π φ ˆ x (t) (π) (5)

where x is loss on samples x = {(s j , a j )} n j=1 defined as

x (π) = 1 n n j=1 (π(s j ), a j ).

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 µ (π * µ ) ≤ H 2 (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 ≥ c 1 max H 4 R 2 log(|Φ|) 2 , H 4 ln(4/δ) 2

, and number of samples (expert trajectories) per task

satisfies n ≥ c 2 H 4 R 2 K 2 for small constants c 1 , c 2 , then with probability 1 -δ, E µ∼η E x∼µ n J µ (π φ,x ) -E µ∼η J µ (π * µ ) ≤ H 2 γ + 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 η. 

= O max H 4 R 2 log(|Φ|) 2 , H 4 R 2 K

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 on H 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.

5.1. PROOF SKETCH:

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(Φ({s t j })) 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 µ (π * µ ) ≤ H 2

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.

6. REPRESENTATION LEARNING FOR OBSERVATION-ALONE SETTING:

Now we consider the setting where we cannot observe experts' actions but only their states. As in 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 ∼ P s,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 = {X 1 , . . . , X H } where X h is the dataset for level h. Specifically,

EQUATION

H } where x

(i) h = {(s i j , a i j , si j , si j )} n j=1 ∼ (µ (i) ) n . Additionally we denote S h = {s i j } T,n i=1,j=1

to be all the s-states in X h , Sh and Sh 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 φ∈Φ Lh (φ), where

Lh (φ) = 1 T T i=1 min π∈Π φ max g∈G 1 n n j=1 [Kπ(a i j |s i j )g(s i j ) -g(s i j )] = 1 T T i=1 min π∈Π φ ˆ x (i) h (π)

where for dataset x = {(s j , a j , sj , sj )} n j=1 , ˆ x h (π) := max

g∈G 1 n n j=1

[Kπ(a j |s j )g(s j ) -g(s j )]. Note that because of the max g∈G , ˆ x h is no longer an unbiased estimator of µ h when x ∼ µ n h . However we can still show generalization bounds.

Evaluating representations φ1 , . . . , φH : Learned representations are tested on a new task µ ∼ η as follows: get samples x = (x 1 , . . . , x H )6 for all levels using trajectories from π * µ , where x h ∼ µ n h . For each level h, learn π φh ,x h = arg min π∈Π φ ˆ x h h (π) and consider the policy π φ,x = (π φ1,x1 , . . . , π φH ,x H ). 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

EQUATION

We make the following two assumptions for the subsequent theorem. 

∈ Φ such that π * h,µ ∈ Π φ * h for every h ∈ [H], µ ∈ support(η).

Now we present our main theorem for the observation-alone setting. 5 If G contains all bounded functions, then it reduces to minimizing TV between ν π h and ν * h .

Theorem 6.1. Let φh ∈ arg min φ∈Φ Lh (φ). Under Assumptions 6.1,6.2, with probability 1 -δ over sampling of X = (X 1 , . . . , X H ), we have

E µ∼η E x J(π φ,x ) -E µ∼η J(π * µ ) ≤ H h=1 (2H -2h + 1) gen,h + O(H 2 ) φ be where φ be = E µ∼η E x [ π φ,x

be ] is the average inherent Bellman error and

gen,h = O KG(Φ(S h )) 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 ≥ c 1 max H 4 R 2 K 2 log(|Φ|) 2 , H 4 ln(H/δ) 2

, and number of samples (trajectories) per task satisfies

n ≥ c 2 max H 4 K 2 log(|G|) 2 , H 4 R 2 K 3 2

for small constants c 1 , c 2 , then with probability 1 -δ,

E µ∼η E x J(π φ,x ) -E µ∼η J(π * µ ) ≤ O(H 2 ) φ 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. R 2 log(|Φ|) max{log(|G|), R 2 K}. 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 & Szepesvá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.

7. EXPERIMENTS:

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 = {(s t j , a t j )} n j=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:

EQUATION

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.

8. CONCLUSION:

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.

A PROOFS FOR BEHAVIORAL CLONING:

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. that A = [K]. Remember that : (A)×A → R is defined as (v, a) = -log(v a ) for some v ∈ R K and v a 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 → W x | W ∈ R K×d , W F ≤ 1} (9) (v, a) = -log(softmax(v) a ), v ∈ R K , 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 min f ∈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

EQUATION

where Q = sup

y∈R dn \{0} 1 y E sup f ∈F n,K i=1,j=1 γ ij f (y i ) 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∈R dn \{0} 1 y E sup f ∈F n,K i=1,j=1 γ ij f (y i ) j ≤ √ K Proof. Q := sup y∈R dn \{0} 1 y E sup f ∈F n,K i=1,j=1 γ ij f (y i ) j = sup y∈R dn \{0} 1 y E sup W F ≤1 n,K i=1,j=1 γ ij W j , y i = sup y∈R dn \{0} 1 y E sup W F ≤1 K j=1 W j , n i=1 γ ij y i = sup y∈R dn \{0} 1 y E K j=1 n i=1 γ ij y i 2 ≤ sup y∈R dn \{0} 1 y K j=1 E n i=1 γ ij y i 2 = sup y∈R dn \{0} 1 y K j=1 n i=1 y i 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 = π * µ (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 µ (π * ) ≤ H 2 µ (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.

B PROOFS FOR OBSERVATION-ALONE:

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

Lh (φ h ) = E µ∼η E x∼µ n h µ h (φ, f φ h x )

We first show a guarantee on the performance of representations ( φ1 , . . . , φH ) as measured by the functions L1 , . . . , LH . Theorem B.1. With probability at least 1 -δ in the draw of X = (X (1) , . . . ,

X (H) ), ∀h ∈ [H] Lh ( φh ) ≤ min φ∈Φ L h (φ) + c gen,h (Φ) + c gen,h (F, G) + c ln(H/δ) T where gen,h (Φ) = KG(Φ(S h )) 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 Lh to the expected cost on the tasks. Theorem B.2. Consider representations (φ 1 , . . . , φ H ) with Lh (φ h ) ≤ h . Let x = (x 1 , . . . , x H ) be samples at different levels for a newly sampled task µ ∼ η such that x h ∼ µ n h . Let π φ,x = (π φ1,x1 , . . . , π φ H ,x H ) 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(H 2 ) φ 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 φ∈Φ L h (φ) = 0 for every h ∈ 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 x h (φ, f, g) =

1 n n j=1 [Kπ φ,f (a j |s j )g(s j ) -g(s j )]. Let mx (φ) = min f ∈F max g∈G ˆ x h (φ, f, g) = ˆ x h (φ, f φ x , ĝφ x ), mµ,x (φ) = max g∈G µ h (φ, f φ x , g), m µ (φ) = min f ∈F max g∈G µ h (φ, f, g). Note that L h (φ) = E µ∼η m(φ), Lh (φ) = E µ∼η E x∼µ n mµ,x (φ). Define the distri- bution ρ h where x ∼ ρ h is the same as µ ∼ η and then x ∼ µ n h . 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 φ ∈ Φ, Lh (φ) -E x∼ρ h mx (φ) ≤ gen,h (F, G) Lemma B.5. With probability 1 -δ, for every φ ∈ Φ, E x∼ρ h mx (φ) - 1 T i mx (i) (φ) ≤ gen,h (Φ) + O   log( 1 δ ) T  

We prove these lemmas later. First we prove Theorem B.1 using them. If φ * h = arg min φ∈Φ L h (φ), then ) by Hoeffding's inequality and fifth is bounded by the following argument: let f φ , g φ = arg min 

Lh ( φh ) -L h (φ * h ) = Lh ( φh ) -E x∼ρ h mx (φ) + E x∼ρ h mx (φ) - 1 T i mx (i) ( φh ) + 1 T i mx (i) ( φh ) - 1 T i mx (i) (φ * h ) + 1 T i mx (i) (φ * h ) -E x∼ρ h mx (φ * h ) + E µ∼η [ E x∼µ n mx (φ * h ) -m µ (φ * h )] ≤ 2 gen,h (F, G) + gen,h (Φ) + O   log( 1 δ ) T  

f ∈F arg max g∈G µ (φ, f, g) E x∼µ n mx (φ * 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 ) +

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 Γ * h V * h+1 (s) -E s∼ν π h Γ * h V * h+1 (s)] ≤ h (x h ) + 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 (x h ) + 2∆ h + 2 π be . We now bound ∆ h

∆ h = max g∈G E s∼ν π h-1 E a∼π h-1 s ∼Ps,a g(s ) -E s∼ν * h g(s) = max g∈G E s∼ν π h-1 E a∼π h-1 s ∼Ps,a g(s ) -E s∼ν * h-1 E a∼π h-1 s ∼Ps,a g(s) + max g∈G E s∼ν * h-1 E a∼π h-1 s ∼Ps,a g(s ) -E s∼ν * h g(s) = max g∈G E s∼ν π h-1 Γ π h-1 g(s ) -E s∼ν * h-1 Γ π h-1 g(s) + h-1 (x h-1 ) ≤ ∆ h-1 + 2 π be + h-1 (x h-1 ) Thus ∆ h ≤ 2(h -1) π be + 1:h-1 (x 1:h-1

) and so ∆h ≤ 1:h (x 1:h ) + 1:h-1 (x 1:h-1 ) + (4h -2) π be . This implies that

J(π φ,x ) -J(π * ) = H h=1 ∆h ≤ H h=1 (2H -2h + 1) h (x h ) + O(H 2 ) π φ,x

be Taking expectation wrt µ ∼ η and x ∼ µ n completes the proof.

B.3 PROOFS OF LEMMAS:

Proof of Lemma B.3. Again we define F 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 ˆ Maurer et al. (2016) , we get that

x h (φ, f , g) = ˆ x h (φ, softmax(f ), g). Notice that (•, α, β, a) is 2K- lipschitz, (v, •, β, a) is K-lipschitz and (v, α, •, a) is 1-lipschitz, Using Theorem 8(i) from

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π Ex G( (F (φ(s h )), G(s h ), G(s h ), a)) n ≤ 2 √ 2πKG(F (φ(s h )), a) n + E µ∼η E x∼µ n √ 2πKG(G(s h )) n + √ 2πG(G(s h )) n ≤ c RK √ 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.

Proof of Lemma B.4.

Lh (φ) -E x∼ρ h mx (φ) = E µ∼η E x∼µ n mµ,x (φ) -E µ∼η E x∼µ n mx (φ) = 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) -ˆ x h (φ, 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 Lh , 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(X s -X t ) 2 ≤ E(Y s -Y t ) 2 , ∀s, t ∈ S Then E sup s∈S X s ≤ E sup s∈S Y s

Using Theorem 8(ii) from Maurer et al. (2016) , we get that

EQUATION

where S = {( m(φ) x1 , . . . , m(φ) x T ) : φ ∈ Φ}. We bound the Gaussian average of S using Slepian's lemma. Define two Gaussian processes indexed by Φ as

X φ = i γ i m(φ) x (i) and Y φ = 2K √ n i γ ijk φ(s i j ) k For x = {(s j , a j , sj , sj )}, 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 | ˆ x h (φ, f, g) -ˆ x h (φ , f, g)|) 2 =   sup f ∈F ,g∈G 1 n j [Kπ φ,f (a j |s j )g(s j ) -Kπ φ ,f (a j |s j )g(s j )]   2 = K 2   sup f ∈F ,g∈G 1 n j (f (φ(s j )) aj -f (φ (s j )) aj )g(s j )   2 ≤ K 2 n sup f ∈F j f (φ(s j )) aj -f (φ (s j )) aj 2 ≤ 4K 2 n j |φ(s j ) -φ (s j )| 2 = 4K 2 n j,k (φ(s j ) k -φ (s j ) 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 (φ(s i j ) k -φ (s i j ) k ) 2 = E(Y φ -Y φ ) 2

Thus by Slepian's lemma, we get

G(S) = E sup φ∈Φ X φ ≤ E sup φ∈Φ Y φ = 2K √ n G(Φ({s i j }))

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) = ˆ x h (φ, f, g) -ˆ x h (φ , f , g ) ≤ ˆ x h (φ, f , g ) -ˆ x h (φ , f , g ) ≤ ˆ x h (φ, f , g ) -ˆ x h (φ , f , g ) ≤ sup f ∈F ,g∈G | ˆ x h (φ, f, g) -ˆ x h (φ , 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 | ˆ x h (φ, f, g)- ˆ x h (φ , f, g)|. Proof of Lemma B.6. Let ḡ = arg max g∈G E s∼ν π h Γ π h g(s) -E s∼ν * h Γ π h g(s) and g = arg min g∈G |g - Γ π h ḡ| (ν π h +ν * h )/2 . max g∈G E s∼ν π h Γ π h g(s) -E s∼ν * h Γ π h g(s) = E s∼ν π h Γ π h ḡ(s) -E s∼ν * h Γ π h ḡ(s) ≤ | E s∼ν π h g (s) -E s∼ν * h g (s)| + | E s∼ν π h [g (s) -Γ π h ḡ(s)]| + | E s∼ν * h [g (s) -Γ π h ḡ(s)]| ≤ max g∈G | E s∼ν π h g(s) -E s∼ν * h g(s)| + 2 E s∼(ν π h +ν * h )/2 [|g (s) -Γ π h ḡ(s)]|] ≤ ∆ h + 2 π be C DATA SET COLLECTION DETAILS C.1 DATASET FROM TRAJECTORIES

Given n expert trajectories for a task µ, for each trajectory τ = (s 1 , a 1 , . . . , s H , a H ) we can sample an h ∼ U([H]) and select the pair (s h , a h ) from that trajectory7 . This gives us n i.i.d. pairs {(s j , a j )} n j=1 for the task µ. We collect this for T tasks and get datasets x (1) , . . . , x (T ) .

C.2 DATASET FROM TRAJECTORIES AND INTERACTION:

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 from A and then get a state s from P s,a by resetting the environment to s and playing action a. We collect this for T tasks and get datasets X (i) = {x 

D EXPERIMENT DETAILS:

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. Experiment Setup We first describe the construction of the NoisyCombinationLock environment. The state space is R 40 . Each state s is in the form of [s real , s noise ], while s real ∈ R 20 is either a onehot vector or a zero vector, and s noise ∈ R 20 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 s real = e i for some i and the action is a * i , then s real = e i+1 and we'll get reward 1. Otherwise s real will be all zero and the reward is 0. s noise will always be sampled from the Gaussian distribution. Note that once s real 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 s real is always e 1 .

The representation has dimension of 10. We limit the function φ to be a linear mapping from R 40 to R 10 . 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 R 8 . The original action space in Swimmer is R 2 , 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. -log(π φ,ft (s t j ) a t j ).

Optimization:

(13)

Then we minimize L(φ, f 1 , . . . , f T ) and obtain the optimal φ.

:

Break ties in any way

We can easily extend the theory to non-stationary policies

We abuse notation and use π * µ (s) instead of arg max a∈A π * µ (s)a

These statements are qualitative since we are comparing upper bounds.

Note that we do not need the datasets x h at different levels to be independent of each other

In practice one can use all pairs from all trajectories, even though the samples are not strictly i.i.d.

Figures and Tables:

Figure 1: Experiments for verifying theory. Left: validation loss on DirectedSwimmer. Right: average return on NoisyCombinationLock

[H]. Thus from Theorem B.1, we get that Lh ( φ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. B.1 PROOF OF THEOREM B.1

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

} for every i ∈ [T ], where each dataset x (i) h a set of n tuples obtained level h. Rearranging, we can construct the datasets X h = {x

Figure 3: The total rewards by different algorithms in DirectedSwimmer.

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(φ, f 1 , . . . , f T )

The first term H 2 γ 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 H 4 R 2 log(|Φ|)

Table content: <table><tr><td>2</td><td colspan="2">tasks sampled from η, we can use them to</td></tr><tr><td colspan="2">learn a representation such that for a new task we only need n = O H 4 R 2 K 2</td><td>samples (expert</td></tr><tr><td colspan="3">demonstrations) to learn a linear policy with good performance. In contrast, without access to tasks,</td></tr><tr><td>we would need n</td><td/></tr></table>

Sun et al. (2019), we also solve a problem at each level; consider a level h ∈ [H]. = (π 1 , . . . , π H ) that matches the distributions ν π h with ν * h for every h, w.r.t. a discriminator class G 5 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

Table content: <table><tr><td>µ h (π) = max g∈G [ E s∼ν  *  h,µ</td><td>E a∼π(s) s∼Ps,a</td><td colspan="2">g(s) -</td><td>E h+1,µ s∼ν  *</td><td colspan="2">g(s)]</td><td>(6)</td></tr><tr><td>= max g∈G [ E s∼ν  *  h,µ</td><td colspan="2">E a∼U (A) s∼Ps,a</td><td colspan="3">Kπ(a|s)g(s) -</td><td>E h+1,µ s∼ν  *</td><td>g(s)]</td></tr><tr><td>Choice of µ</td><td/><td/><td/><td/><td/></tr></table>

Consider a task µ. For simplicity of notation, we use π h instead π φ h ,x h , π instead of π φ,x . Let ν π h and ν * h be the state distributions at level h induced by π φ,x and π * µ respectively. Let be the loss of policy π h at level h. By definition, h = E

Table content: <table><tr><td/><td/><td/><td/><td/><td>gen,h (F, G)</td></tr><tr><td colspan="3">where the second inequality uses Lemma B.3.</td><td/><td/></tr><tr><td>B.2 PROOF OF THEOREM B.2</td><td/><td/><td/><td/></tr><tr><td>h (x h ) = max g∈G</td><td>E s∼ν  *  h</td><td>[ E a∼π h s ∼Ps,a</td><td colspan="3">s ∼Ps,a g(s ) -E h a∼π  *</td><td>g(s )]</td></tr><tr><td/><td/><td/><td>µ∼η</td><td>E x∼µ n h</td><td>h (x). Using Lemma C.1 from Sun</td></tr><tr><td>et al. (2019), we have</td><td/><td/><td/><td/></tr></table>

Peer review (iclr2020):
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. Pros: 1. The paper presents a theoretical investigation into multi-task/meta learning for RL via learning representations. While, the main theoretical contributions are perhaps marginal with regard to prior work, the problem setting (RL) seems novel and the two theorems in this context are interesting. 2. The paper is well-written and appears to be very rigorous. I did not check for the correctness of all technical parts. There are several *abuse of notations* in the main text, which sometimes impact the otherwise smooth read of the technical parts. 3. A good and concise review of RL concepts is provided. Cons: 1. The paper lacks a good literature review to place this work in the right context. For example, while the paper refers to the works of Maurer 2016 and Sun et al. 2019 at several places, it is not formally and clearly mentioned anywhere what are the similarities to these prior works and what are the new contributions. For example, Maurer 2016 proposes a multi-task learning setup using representation learning, and most of Theorem 4.1 in this paper is taken from the results in that paper. While, the current paper uses bilevel optimization setting in an RL context, it is not clear to me if this (bilevel + RL) setting has any significant bearing against the theoretical results furnished by Maurer 2016. For example, the theorems in this paper (as far as I see) show that the bounds are scaled by a constant defined by H^2, the trajectory length? If there is something beyond this, then the paper needs to explicitly point it out. The same comment goes with the results against Sun et al. 2019 in Theorem 5.1. 2. Given that the main goal of this paper is to connect sample complexity with representation learning, it is important the paper provide a theorem stating this precisely. Theorems 4.1 and 5.1 provide a general bound, and the sample complexity is being described in the explanations of this theorem, which is very informal. Also, against what is claimed in the abstract, it appears that representation learning helps reduce sample complexity only when the number of tasks are larger, which perhaps needs to be explicitly mentioned. Also, note that there is a bearing of the bound on the trajectory length H (Theorem 4.1, and 5.1). Shouldn*t this factor be also accounted for when explaining the sample complexity? 3. There has been several recent works on model-agnostic meta-learning (that also uses bilevel optimization and implicit gradients), however, older works on meta-learning (only for imitation learning) have been cited. The paper should include more recent works in this area and contrast against their theoretical findings. Apart from these, below are some minor comments that could help improve the reading of this paper: a. Theorem 3.1 is not really a theorem, since it is very informal. Also, fix the Theorem numbers. b. Bullet 1. after Theorem 3.1, ell^x(pi) concentrates to ell^x(pu*). Also, the mention about sample complexity here should be backed with some reference/citation. c. Assumption 4.2: The notation pi_mu(s)_{pi*(mu(s)) is unclear, shouldn*t the subscript contain an argmax over the actions for pi*? d. Theorem 4.1 and 5.1, what does it mean by *probability 1-delta over the choice of the dataset X* ? Also, mu^n seems undefined. e. The first two terms in Theorem 4.1 are claimed standard, provide the citations? f. Theorem 4.1, perhaps use some other notation for c, which is defined as the cost/reward in the RL setting. Overall, the paper has some interesting theoretical results, and is mathematically rigorous, however lacks a clear distinction from prior and more recent works in this area.