Provable Representation Learning for Imitation Learning via Bi-level Optimization
Authors: Sanjeev Arora
Simon S. Du
Sham Kakade
Yuping Luo
Nikunj Saunshi
Venue:  (2020)
Keywords: imitation learning
representation learning
multitask learning
theory
behavioral cloning
imitation from observations alone
reinforcement learning

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 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 cases where the imitation setting being 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.

Under review as a conference paper at ICLR 2020
PROVABLE REPRESENTATION LEARNING FOR IMITA- TION LEARNING VIA BI-LEVEL OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
A common strategy in modern learning systems is learn 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 mul- tiple 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 learn task-specific parameters. We instantiate this framework for the imitation learning settings of behavior cloning and observation-alone Theo- retically, we provably show using our framework that representation learning can reduce the sample complexity of imitation learning in both settings. We also pro- vide proof-of-concepp experiments to verify our theoretical findings.
INTRODUCTION
Humans often learn from experts quickly and with few demonstrations and 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, represen- tation that is useful for learning future skills. Can we build an agent do the same?
The current paper studies how to apply representation learning to imitation learning, Specifically, we want build an agent that is able learn 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 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 representation. See Section for more diseussions.
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 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 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 viewed as a straightforward extension of multi-tast 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 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-alon setting. We set the inner optimization min- max 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

---

Under review as a conference paper at ICLR 2020
experts, the agent can learn representation that provably reduces the sample complexity for target task where the agent cannot observe actions from either source experts or the target expert. We conduct experiments to verify theoretical insights by learning representation from multi- ple tasks using our framework and testing it using both behavior cloning and policy optimization. In these settings, observe that by learning representations the agent can learn 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 guar- antees (Syed & Schapire 2010] et al. [011] Sun al., 2019). To our knowledge, this 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. (0113) for survey. Theoretically, |auuer 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 al., 2017 Daumé et al.) 2009); Chang al.] 2015) Pan et al.) 2018). Most existing work con- sider 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 (Pomerlcau) 1991). which we also study in our paper. More advanced approaches have also been proposed (Ross et al. 2011 Ross & Bagnell 2014, [un 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 multi- task representation learning setting.
This paper also considers a more challenging setting, imitation learning from observation alone. Though some model- methods have been proposed (Torabi 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 certain distri- butional 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- weta-lcarning for imitation learning DDuan et al. 2017 Finn et al. 2017b James] al. 2018). Our work is different from theirs as we want to explictlly learn a representation that is useful across all tasks whereas these work try to learn a meta-algorithm that can quickly adapt to new task. For example. Finn et al. (20176) 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. [Buliins et al. 2019) provides an efficient algorithm that generalizes to new unseen tasks, but for linear repre- sentations. Another recent line of work analyzes gradient based meta-lcarning methods, similar to MAML (Finn et al.] (017a). Existing work on the sample complexity and regret of these meth- ods (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
Markoy Decision Processes (MDPs): Let M- = (S.A, be an MDP, where S is the state space, A is the finite action space with |A| K, H Z4 is the planning horizon, P: A
2

---

Under review as a conference paper at ICLR 2020
A (S) is the transition function, : A + is the cost function and v E 4(S) is the initial state distribution. We assume that cost is bounded by 1, i.. a( Σ 1s E S. E A. This is a standard regularity condition used in many theoretical reinforcement learning work. A (stochastic) policy is defined as π = (T1,..... TH), where Th : S + (A) preseribes a distribution over action have स = t induces a random trajectory $1, $2, .2..... SH, aH where 81 ~ 2, ~ (1(s), 82 ~ Pa 2, etc. Let vi denote the distribution over S induced level h by policy æ. The value function Vi S + R is defined as
for each state at level h E [ For stationary policy,
и n. A policy
(sh) - E Σ C(s a:) a (3(3)), ~ 8141 Par,a
and the state-action function QT (sh, ap) is defined as Q7 (sh: an) Esnty VV (sh- 1)] The goal is to learn a policy π that minimizes the expected cost J ((TT) = Esioe Vi (81). We define the Bellman operator at level h for any policy π as 1 RS + RS, where for 8 E and E RR. (T))) (O) [g(s))
Multi-task Imitation learning: We formally describe the problem we want to study. We as- sume there are multiple tasks (MDPs) sampled i.d. from a distribution 7. A task μl ~ 17 is an MDP M₁ = (S, A. H, Pp: Uy); all tasks share everything except the cost function, initial state distribution and transition function. For simplicity of presentation, we will assume a common tran- sition function P for all tasks; proofs remain exactly the same even otherwise. For every task j4, T = (Ti....... IH 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 (Vi, V' and the average state dis-
tribution as V H
1 Σ
"h We define Vi to be the value function of T and Ju to be the h=1
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 Tr. is stationary.
Representation learning: In this work, we wish to learn policies from a function class the form I F Q where & C fo S + Rd [|o(s)]|2 I R} is class of bounded norm representation functions mapping states to vectors and F C { Rd + A(A)) is class of functions mapping state representations to distribution over actions. We will be using linear functions, i.e. F {
softmax(War W E R xx |[||F I 1}. We denote a policy parametrized by 6 E 4 and f F f(o(s)) )a- In some cases, we may also use the policy π (a)
by 0 where T T
llce = arg max ff (o(s)) a'EA
Denote 11 = {
F} to be the class of policies that use Φ as
the representation function.
Given demonstrations from expert policies for T tasks sampled independently from 17, we wish to first learn representation functions (fi...., ,H)) so that we can use a few demonstrations from an expert policy xt for new task μ 17 and learn policy representations, i.e. Fa E I such that has average cost of π is not too far away from n², In the case of stationary policies, we need to leam single by using tasks and learn E II for new task. The hope is that data from multiple tasks can be used to learn a complicated function ø e & first, thus requiring only a few samples for new task learn linear policy from the class TT
(T...... TH) that uses the learned
Gaussian complexity: As in Maurer et al. H. C {h: X + Rd) on set X (X1..... Xn) )
2016).
we measure the complexity of a function class X by using the following Gaussian average
d,n
G(H(X)) = E
sup NCH
Ry hyyh((,)) X3
(2)
where Yij are independent standard normal variables. Bartlett & Mendelson 2003) also used Gaus- sian averages to show some generalization bounds.
Break ties in any way
3

---

Under review as a conference paper at ICLR 2020
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 statistical guarantees. Our main idea is to phrase learning represen- tations for imitation learning as the following bi-level optimization
min L(Ø) == E min (TT) 6E
p
Here d is the inner loss function that penalizes ı being different from TH for the task μ.. In general, can use any loss & that is used for single task imitation learning, e.g. for the behavioral cloning setting (cf. Section5), 8 is a classification like loss that penalizes the mismatch between predictions by x² and π,, while for the observation-alone setting (ef. Section 66) 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 7) in practice we only have access to few samples for 7 tasks; let ,.... x(T) from 7). We thus learn the representation Ø by minimizing empirical version L of Equation
x(T) denote samples from tasks (1)), , (((2)) sampled i.i.d.
T
T
Σ ( T
1
(8) (TT (110)))
min TEIL
T
where 1X is the empirical loss on samples and π. arg mell ex (It) corresponds to task specific policy that uses a fixed representation ø. Our goal then is to show that for a new task 1), the policy x の learned by using samples x from the task has low expected cost Jp. i.c.
Informal Theorem 4.1. With high probability over the sampling of train task data and with suffi- cient number of tasks and samples per task,
E E (own
E (x( is small 10m
At high level, in order to prove such a theorem for particular choice of 6e, we would need to prove the following three properties about f and (x:
1. (*x)) concentrates to E (T) simultaneously for all E TI (for fixed Φ), with sample complex-
ity depending on some complexity measure of 11 rather than being polynomial in [S];
2. small value of (*() implies a small value for Jµ(z) J (T());
3. if s and ø' induce "similar" representations then mintgem (* and min (T)) are close.
The first property ensures that learning policy for a single task by fixing the representation is sample efficient, thus making representation learning useful problem to solve. The second property ensures that matching the behavior of the expert as measured by the loss f ensures low average cost ie., fd is meaningful for the average cost; any standard imitation learning loss will satisfy this. The third property is specific to representation learning and requires fd to use representations in a smooth way. This ensures that the empirical loss for T tasks is a good estimate for the average loss tasks sampled from 7)- We prove these three properties for the cases where 8 is the either behavioral cloning loss observation-alone loss, with natural choices for the empirical loss fx 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. (201T). Later in Section d 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 the bi-level optimization framework. In the single task setting, the goal of behavioral cloning (Syed & Schapire 2010 Ross et al. 2011)
4

---

Under review as a conference paper at ICLR 2020
is to expert trajectories of the form 7 (51, Ul....,,811 0.) to learn a stationary policy tries to mimic the decisions of the expert policy on the states visited by the expert. For a task 14, this reduces to supervised classification problem that minimizes surrogate to the following loss E (T) = Beweggannz
P that
awaz (=) [<(()) # a). We abuse notation and denote this distribution over (8, a)
for task μu as (u; so (8, a) is the same
Ujes (s). Prior work (Syed & Schapire 2010
Ross al. 2011) have shown that small value of 6o (I)) implies small difference J(x) Thus for our setting, we choose 8 to be of the following form
(T)
(T)()))))=
E ((T((),,)
(4)
S (a)
(s,a)rup
where é is any surrogate to 0-1 loss Ifa * arg max that is Lipschitz in @(s). In this work a'EA
consider the logistic loss ((T(()),) a) (s)()())a))
Learning é from samples: Given expert trajectories for T tasks μ()) .-... ..... ( we construct
dataset X [xx2) ,..., x()), where x()) = {(5), at) E= (n(0)" is the dataset for task L.
Details of the dataset construction provided in Section C.1 Let S denote the set of states {s)
Instantiating our framework, we learn good representation by solving é = arg min (()), where
( = T
Σ min
Σ
(m(()) at)
1 Σ
min sEll
(T)
(5)
wells n
T
t=1
t=1
where 6x is loss on samples x = (ss) az) a,)}y=1 defined as (T) = + ト ((T())))gg)
Evaluating representation oz A learned representation & is tested on a new task μ ~ v) as follows: draw samples x ~ 1t" using trajectories from Tī and solve -., arg min (Tr). Does -. have
=ET
expected cost J (T)x not much larger than (TT,))) The following theorem answers this question. We make the following two assumptions to prove the theorem.
Assumption 5.1. The expert policy nú is deterministic for every E support(n).
Assumption 5.2 (Policy realizability). There is representation 6² E & such that for every μ E support(n), Tp E II such that <p($));((>)
B > 1 7, Vs E S for some ry < 1/2.
The first assumption holds if TH is aiming to maximize some cost function. The second assumption for representation learning to make sense: we need assume the existence of a common repre- sentation o* that can approximate all expert policies and y measures this expressiveness of 4. Now we present our first main result.
Theorem 5.1. Let & arg min L(o). Under Assumptions 5.1 5.2 PE
with probability 8 over the
sampling of dataset X, we have
E u
E x
Jm(TT)
(2 Egen)
pomm
where Egen
C
G()(S) TV=
RVK + c²
In(4/8) T
for some small constants d.
To gain intuition for what the above bound means, we give 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 's lemma.
Corollary 5.1. In the same setting as Theorem 5.1 suppose t finite. If number of tasks satis- fies T 2 C1 max HH' R² log((4)
H* Im(1/8) {. and number of samples (expert trajectories) per task R²K for small constants C1, then with probability 1 - 6,
satisfies n 2 02
E E powty
Ju (T( x
E. pow
Jn(TT) S H2y +e
We can casily extend the theory non-stationary policies 3 We abuse notation and use (s) instead of arg max n (8) aeA
5

---

Under review as a conference paper at ICLR 2020
Discussion: The above bound says that as long as we have enough tasks learn representation from d and sufficient samples per task to learn a linear policy, the learned policy will have small average cost a task from 77 The first tenm H²y is small if the representation class @ is expressive enough to approximate the expert policies (see Assumption 5.2). 'he results says that we have access data from T = H it (og(1())
tasks sampled from 77, we can use them to (H")")oe((1))
learn a representation such that for a new task we only need n = O =* samples (expert demonstrations) to learn a linear policy with good performance. In contrast, without access to tasks, +1) H*R 1) samples from the task to learn a good policy π IT 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 efficien Note that the dependence sample complexity on H comes from the error propagation when going from ! to Jni this is also observed in single task imitation
we would need π = O max H log(
2²
learning (Ross et al. 2011 Sun et al. 2019).
We give 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 9. that learned on samples μº is evaluated the distribution μ and the average loss incurred by representation & across tasks is L(ő) E jamrry
E (*(xx X~1
On the other hand, if the learner had complete access to the distribution 7 and distributions μ for ev- ery task, then the loss minimizer would be o* arg minse L()), where L(Ø) = E min l (T).
paw f xello
Using results from Maurer al. (0116), prove the following about a Lemma 5.2. With probability over the choice of X, & E arg min (Ø) GE
satisfies
((Φ) min L(Ø) + G(D(Es))) GE
RVK Vn
+ +c" In(1/5)
+
Tyn
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 Jn of the policies induced by é for tasks μ This can obtained by using the connection between the surrogate 0-1 loss 6 and the cost Jj. 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. (201T)
Lemma 5.3. Given a representation é with L(d)= S Let x ~ p. be samples for new task u 1). Let xx be the policy learned by behavioral cloning the samples, then under Assumption 5.1
E poo X
J
E Ju(TT) H2 pmm
This suggests that making L small is good enough. A simple implication of Assumption 5.2 that minge L(d) S L(Ø*) < 7, along with the above two lemmas completes the proof.
6 REPRESENTATION LEARNING FOR OBSERVATION-ALON SETTING
Now we consider the setting where we cannot observe experts' actions but only their states. As in Sun al. (2019), also solve problem each level; consider level h E [M].
Choice of e:: Let TI, = i 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 nH. The goal in imitation learning with observations alone (Sun et al. 2019 is
*These statements are qualitative since we are comparing upper bounds.

---

Under review as a conference paper at ICLR 2020
to learn a policy ( (Ti......: .H. that matches the distributions Vi with vi for every h, w.r.t discriminator class O that contains the true value functions Vi,, ..... Vi and is approximately closed under the Bellman operator of πτ. Instead, in this work we learn π that matches the distributions Thh 1/ and 1 for every h w.r.t to a class G S {g S - R. Ig S 1) that contains the value functions and has stronger Bellman operator closure property. For every task μ., 2 is defined as
(T) = max E
E gEg sn anxx(e) Row Pa.a
((8)
E h4t.
g(5)
max E
E KT(als)g($)
E S 1+
9(3)
DEG
S his
A)
Br.P..
where we rewrite t by importance sampling in the second equation; this will be useful to get an empirical estimate. While our definition of f differs slightly from the one used in Sun et al. (201)), using similar techniques, we will show that small values for (Tx) for every h E [ will ensure that the policy π (mi,..., .,,...... HH) will have expected cost (x((T) close to (TT). We abuse notation, and for a task y we denote μ (J1.....,1a where Hú is the distribution of (s,0,8,8) used in f thus (s, 3) ~ μ4h is equivalent to ~ V/ a ~ U()A), 8 ~ Ps.a, 3 ~ 1/+11k
Learning on from samples: We assume, 1) access to 2n expert trajectories for T independent train tasks, 2) ability to reset the environment at any state and sample from the transition P(-|s, a) for any a E A. The second condition is satisfied in many problems equipped with simulators. Using the sampled trajectories for the T tasks (u(1) , ,...,,.................. m,() and doing some interaction with environment, we get the following dataset X ( } where XA the dataset for level H where x(") X = [(s). Say,,,,, 8) 80)}= ~ (u())), Additionally we denote S, = {sf (= to be all the s-states X Sn and Sn are similarly defined. Details about how this dataset is constructed from expert trajectories and interactions with environment is provided in Section C.2
h. Specifically, X,, = (x/ : x (2))
We learn the representation on arg min (Ø), where
dE@
T
T
(Ø)
Σ
min max- ell gEg n. j=1
1 [KK(())))))) s(3)) (KTT(a)s)))(33) g(3,)]- Note /
+
Σ
min (x Elllo
c (I))
T
where for dataset x= {1s aj 8y 8,),ا (T) = max gEğ
that because of the max,ec. 2 is no longer unbiased estimator of ty when x ph However we can still show generalization bounds.
Evaluating representations Ś2....,áu: .h: Learned representations are tested on a new task j c v as follows: get samples
(x1, (xi,..., X (6 for all levels using trajectories from The where
X4 pha For each level h, learn n arg min (x (T)) and consider the policy πo (Tr1.x11 ...... THn .H)). ). Before presenting the guarantee for - we introduce notion of Bellman error that will show up in our results. For a policy π = (TTI. ( ..) and an expert policy ा
(Tf ......... Tit), we define the inherent Bellman error
==
max max min AE(H) gEg g'EG sm +2g)
E
[lg - (s)]
We make the following two assumptions for the subsequent theorem. These are standard assump- tions in theoretical reinforcement learning literature.
Assumption 6.1 (Value function realizability). Vi,, E G for every h E [H], μ E su)port(n).
Assumption 6.2 (Policy realizability). There are representations oi...., Oū E & such that Th E 4 for every h E [H] μ support(n)-
Now we present our main theorem for the observation-alone setting.
'If & contains all bounded functions, then it reduces to minimizing TV between vi and vj,-
NNoe that we do need the datasets x,, different levels he independent of each other
7

---

Under review as a conference paper at ICLR 2020
Theorem 6.1. Let Φi arg min (o). Under Assumptions 6.7 6.2 6.2 gE
with probability 1 - 8 over
sampling of = (X,,..., X11), we have -X1,,..........
H1 2h + ggen + )
E E J xx)
E
J(T)))
255222 (2H
janny
Hmm
h=1
where be
E Elc
u is the average inherent Bellman error and
pewny
Egen,h = (G()(())) TVn
KG(G(SN)
+ G(((()))
RKVK K
In H//))
DEE X
T
We again give a PAC-style guarantee for the special case where the class of representation functions T and value function class G are finite. It follows from the above theorem and Massart's 's lemma.
Corollary 6.1. In the setting of Theorem 6.1 suppose 4, 9 are finite. If number of tasks satisfies
T 2 01 max
11 n² n² x² log((())
11 (n(1166)
and number of samples (trajectories) per task satisfies
H** K2 19)((9)
H for small constants Cr Cz. then with probability 1-
C2 max
E EJ
E J(T S ((H))) à +e.
Discussion: As in the previous section, the number of samples required for new task after learn- ing 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² log(( >> maxxlog((1)), R² R) In the above bounds, bb is Bell- man error term. This type of error terms occur commonly 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 optimal policy. Le et al. (2019) used similar notion that they call inherent Bellman evaluation error.
The proof of Theorem 6. follows a similar outline to that 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
EXPERIMENTS
In this section we present experimental results on the DirectedSwimmer environment (modified from the Swimmer environment from OpenAl gym (Brockman et al. 2016)) with Todorov et al. (2012) simulator and a MoisyCombinationn.co two aims: 1) verify the benefit of representation learning predicted by our theory, 2) test the power of representations learned via our framework by using the representation and doing policy optimization instead of imitation learning. In periments we learn representations using Equation
environment designed by ourself. These experiments have
broader context: we learn a policy for new task
5 Experiment details are deferred to Section D
Our method: Given access to dataset X
{(s), a))) of state acteon pairs each for T tasks, For any task learn linear policy ı from the class 1
learn å according Equation
T
f
arg min 6E
min FTFF
log(1 (s))at)
t=1
Baseline: For task we learn a policy from the class II without learning representation first.
Verification of theory: In Figure I 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 few experts, can significantly reduce the sample complexity. On the right, report the average
8

---

Under review as a conference paper at ICLR 2020
DirectedSwimmer
NoisyCombinatinlLock
baseline
20
expert
Validation loss
20
2 experts
2000000 returr
experts
baseline
experts
20
raperts
64 experts
5-
150
175
+samples task
pairs
Figure Experiments for verifying theory. Left: validation loss on DireetedSwimmer.
Right:
average return on NoisyCombinntionl..ck
PPO on DirectedSwimmer
PPO on NoisyCombination..ck
3,0
awerage ntturn
20
2.11
expart
32
20
expert
beseline
expart
3K
10C
K5K
26K
2bK
30K
20K 40K
60K
D0K
100K
steps
steps
Figure 2: Experiments policy Optimization with representation trained by imitation learning Left: average return on the DirectedSwimmer. Right: average return on the NoisyCombinationL..co
reward of the trained policies on the environment. Here number of experts is small (4 or 16), the baseline method can beat policies trained using represen- tation 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 out- perform the baseline method. This behavior is expected as when the number of experts is small, we may learn 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 still far behind the policy trained using representation learning with 64 experts.
different phenomenon: when the
Policy optimization with representations trained by imitation learning: We next test the utility of representations learned via our framework for RL. After training representation, we use simpli- fied proximal policy optimization method that learns linear policy over the learned representation. Results are reported in Figure 2] For DirectedSwimmer and NoisyCombiintionnLock, 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 represen- tation 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 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.
observe
CONCLUSION
The current paper proposes a bi-level optimization framework formulate and analyze representa- tion learning for imitation learning using multiple demonstrators. Theoretical guarantees are pro- vided 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 rep- resentation learned via imitation learning is also useful for policy optimization in the reinforcement learning setting. We believe it is interesting theoretical question explain this phenomenon. Ad- ditionally, extending this bi-level optimization framework to incorporate methods beyond imitation learning is interesting future direction.

---

Under review as a conference paper at ICLR 2020
REFERENCES
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Peter L.. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. J. Mach. Learn. Res., 2003.
Jonathan Baxter. A model inductive bias learning. J. Artif. Int. Res., 2000.
Y. Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new per- spectives. IEEE transactions on pattern analysis and machine intelligence, 08 2013.
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
Brian Bullins, Elad Hazan, Adam Kalai, and Roi Livni. Generalize across tasks: Efficient algo- rithms for linear representation learning. In Proceedings of the 30th International Conference Algorithmic Learning Theory, 2019.
Kai-Wei Chang, Akshay Krishnamurthy, Alekh Agarwal, Hal Daumé, III, and John Langford. Learning to search better than your teacher. In Proceedings of the 32nd International Confer- on International Conference on Machine Learning Volume 37, ICML'15. JMLR.org, 2015.
Hal Daumé, lii, John Langford, and Daniel Marcu. Search-based structured prediction. Mach. Learn., 2009.
Giulia Denevi, Carlo Ciliberto, Riccardo Grazzi, and Massimiliano Pontil. Learning-to-ler stochastic gradient descent with biased regularization. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford,
John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. Openai baselines. https: ines. 2017.
//github. www.bb/ccoooppaaaaaaa//
Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. In Advances in Neural Information Processing Systems 30. 2017.
Ashley D. Edwards, Himanshu Sahni, Yannick Schroccker, and Charles Lee Isbell. Imitating latent policies from observation. arXiv preprint arXiv:1805.07911 2018.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. Proceedings of the 34th International Conference on Machine Learning, 2017a.
Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imita- tion learning via meta-learning. 09 2017b.
Chelsea Finn, Aravind Rajeswaran, Sham Kakade, and Sergey Levine. Online meta-learning. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS, 2016.
Stephen James, Michael Bloesch, and Andrew Davison. Task-embedded control networks for few- shot imitation learning. 10 2018.
Mikhail Khodak, Maria-Florina Balcan, and Ameet Talwalkar. Adaptive gradient-based meta- learning methods. arXiv preprint arxiv:1906.02717.22 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
10

---

Under review as a conference paper at ICLR 2020
Hoang Le, Cameron Voloshin, and Yisong Yue. Batch policy learning under constraints. In Pro- ceedings of the 36th International Conference on Machine Learning, pp. 5703-3712, 2019.
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representation learning. The Journal of Machine Learning Research, 17(1):2853-2844, 2016.
Rémi Munos. Error bounds for approximate value iteration. Proceedings of the 20th National Conference on Artificial Intelligence Volume 2, AAAI005, AAAI Press, 2005.
Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. J. Mach. Learn. Res., 2008.
Yunpeng Pan, Ching-An Cheng, Kamil Saigol, Keuntaek Lee, Xinyan Yan, Evangelos Theodorou, and Byron Boots. Agile autonomous driving using end-to-end deep imitation learning. In Pro- ceedings of Robotics: Science and Systems, 2018.
D. A. Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation, 3, 1991.
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? towards understanding the effectiveness of maml. arXiv preprint arXi:: 09197, 2019.
Aravind Rajeswaran, Chelsea Finn, Sham Kakade, and Sergey Levine. Meta-learning with implicit gradients. arXiv preprint 1906.02777 2019.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 661-668, 2010.
Stéphane Ross and Andrew Bagnell. Reinforcement and imitation learning interactive no- regret learning. arXiv preprint arXiv:1406.5977 2014.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international con- ference on artificial intelligence and statistics, 2011.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 177.007777.06347, 2017.
Ju Sun, Qing Qu, and John Wright. Complete dictionary recovery over the sphere I: Overview and the geometric picture. IEEE Transactions on Information Theory, 63(2):853-884, 2017.
Wen Sun, J. Andrew Bagnell, and Byron Boots. Truncated horizon policy search: Combining reinforcement learning and imitation learning. arXiv preprint arXiv: /240, 2018.
Wen Sun, Anirudh Vemula, Byron Boots, and Andrew Bagnell. Provably efficient imitation learn- ing from observation alone. arXiv preprint arXiv: 19t199955.1944 10948, 2019.
Umar Syed and Robert E Schapire. A reduction from apprenticeship learning to classification. In Advances in Neural Information Processing Systems 23, pp. 2253-2261. 2010.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based con- trol. pp. 5036-5033. IEEE, 2012. URL bttp:///blpp.un-trree.r. http://cerr.eeedcccoot/oorooo
ros Froo2012.hir h #odorovvT12
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In IJCAI, 2018.
Jiajun Wu, Erika Lu, Pushmeet Kohli, Bill Freeman, and Joshua Tenenbaum. Learning to see physics via visual de-animation. In NIPS, 2017.
11

---

Under review as a conference paper at ICLR 2020
A PROOFS FOR BEHAVIORAL CLONING
We prove Theorem 5.1 in this section by proving Lemma 5.3 In this section, we abuse notation 8x (Ø, J) be FEF
and define ed (t, = t (T+:), where fe is defined in Equation Let f arg min
the optimal task specific parameter for task μ by fixing representation Thus by definitions in Section 5 we get T - T -- We assume w.l.o.g. that = [K]. Remember that e ((A) x + R is deñned E(v, a) log (a) for some RK and va is the coordinate corresponding to action a E A = [K]. We define a new function class and loss function that will be useful for our proofs
F F= fx Wx W E kxd || 1
(v, log( (a)offmax(())),veEl
A
(10)
We basically offloaded the burden of computing softmau from the class F to the loss f. We can convert any function f' F to one in F by transforming it softmax(f").
We now proceed to proving the lemmas
Proof of Lemma
5.2
We can then rewrite the various loss functions from Section as follows
n
[(6) T
min (((()())).a) j=1
((Ø)
min (((()(s)),a)
pom FEF (8.0)~p
I(Φ)
E E Hom X (k,0)mp
((''((()))
where fi arg FF (0, softmax((")) It is easy to show that both (',,, (f((()), -)
2- 2-lipschitz in their arguments for every A and f E F. Using slightly modified version of Theorem 2(i) from Maurer al. (0016), get that for ò arg min L()), with probability at least 1 - over the choice of X
E p deb
|
2 G
8 8lo((4/8)
I(δ) min L(Ø) S
VTTQ
sup
Tvn
n
T
L()) < G($(S)) PEB
Q'R Vn
log(4/8)
I(Φ)
min
+
+
d"
(11)
T/n
T
a.K
where Q'
sup yERd" \fop
|ل E sup Σ feJ" i=رj=1
(y):)- 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) 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. (0116) assumes that (' a) is 1-lipschitz for every a E and that f-(- is L lipschitz for every I E F'. However the only properties that are used in the proof of Theorem 16 are that (((.a) 1-lipschitz and that e' (')) a) L-lipschitz for every a A, which is exactly the property that we have. Hence their proof follows through for our setting as well.
n E sup
Lemma A.I. Q'
sup VERdM/(0)
M
Σ
mal'( S VK
12

---

Under review as a conference paper at ICLR 2020
Proof.
R.K
0=
sup
E sup feF
Tig (V:)i
1yll
yER
[00
im 1,
n.K
sup
E
sup
Ves WW Vi)
sER
{o)
W
11
xm 1,
K
15
sup
E
sup [W
W
Ys)
1
{0) wER
11
f=1
1
K
n
sup
(0) |lvll
E
Yig
yi
yER
sup K
||v||
Vi
yER (0)
1
sup ((0
|lvil
1
122
M
K VK
where we use Jensen's ' inequality and lincarity 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 I(Ø) E E
('(x0x)) S Consider a task μ ~ 1 and samples
x
11th and let (x) (xxxxx so that I(d)
E (x(x). Since TÃ is deterministic, we get pr xm
E
E E
Ha + (+))
E
|| - n xx (8) m<(())
smalt amapo
swar
log(1 (1 T (s))(())()))))
E s
|- log( x 2x (s)= = (x)
where we the fact that < <- log(1 x) for 1. for the first inequality. Thus by using (m*)
Theorem 2.1 from Ross et al. 2201 we get that (=( - S H (x). Taking expectation w.r.t. μ 1) and x n² completes the proof.
Proof of Theorem 5.1 By using Assumption 5.2 we first get that
L(d")
min 1og(T())=;(0)) rell" evUn
u
< E
E.
log (8)
wmm swat
< E
E - log(1 - 7) S 27
pmn sw
where in the last step used log((1 x) < 2x for 1/2. Hence from Lemma 5.2 we get T(()) < 2y + fgen,h, which combining with Lemma 5.3 gives the desired result.
13

---

Under review as a conference paper at ICLR 2020
B PROOFS FOR OBSERVATION-ALONEE
Before proving Theorem 6.I we introduce the following loss functions, as we did in the proof sketch
for the behavioral cloning setting. We again abuse notation and define 8" {"(Ø, f) = e' (r)r01)), where go defined in Equation 6 Let f arg min ex (ó. f) be the optimal task specific parameter for
JEF
task by fixing representation Φ. As before, we define the following
(n(nn)= = E E
(6, fx
X
We first show guarantee the performance of representations (Ō1,...,áH) as measured by the functions LH. Theorem B.1. With probability at least 1 - 8 in the draw of X (X X vh [M]]
In(H/88
(dn)))= I min (o) + CEgen, ( + c'e Eynn, ΦEO
(F,5)+
T
where Egen, (T) KG()(Ss)) and Egen, (F = E E faw xmup
KG(G(B)))
+ G(G((n))
+ RKvK vT
Tvn
We then connect the losses La to the expected cost on the tasks.
Theorem B.2. Consider representations (dI..... .H) with (n(dr) S Ch. Let XI,...,XH! be samples at different levels for newly sampled task μ such that Xh. ~ pb. Let πI (T xH)*H) be policies learned using the samples, then under Assumption 6.1
H
E (T=
EE J(T poron
E (21 2h 1)E + )(H
h=1
④ where te
E Elege Howny
1 is the averuge inherent Bellman error.
is easy to show that under Assumption 6.2 minget (Ø) = for every h E [H]. Thus from
Theorem B.I get that (Øx) S Eyen,h, where Egen,h Eyen,h i ( ( (F, +c" [ln(H/6)
Invoking Theorem B.2 the representations {Φr) completes the proof.
B.1 PROOF OF THEOREM B.1
Before proving the theorem, we discuss important lemmas. In yet another abuse of no- tation, we define f.9) ((,,,,35)) P (als)g(5) g(3) and (o,
K (a,ls)) ((5))].
Let mx( min max f (0. f.g) (0) fo 02), max
max (0, fC g). m2(
JEF gEG
DEG
min max ( f.g). Note that Ln(Ø)= = E m(d). (n( E panxaaa
m)x( Define the distri-
JEF DEG
pm-
bution Ph where x Ph the same as fr 1 and then x ~ phf.
Lemma B.3. For every & E & and l E [H].
E sup sup x FEF gEG
(2. f.g) - f.. S Spen,n
Lemma B.4. With probability 1 8, for every b
Ln(d)
E
mx S Eg (F,))
x~pa
Lemma B.5, With probability 1- for every & s
log(t) T
E m
Σ ()) S Egen.h (D)
(1) +
Xeupk
T
14

---

Under review as a conference paper at ICLR 2020
We prove these lemmas later. First prove Theorem B.1 using them. If of arg min (o) then DE
In Ln
rix
xmppi
mx(())
(0)
T Σ mixtes (Ōm)
xmpp
x (on)
1
Σ mixco (6%)
T
1 Σ (1) (0%)
mx(
x-pp
E E my (0%) m
pan Xeyee
())()) T
S 2c gen
( (F,9) + 'gen h (T) +
where for the first part we use Lemma bounded by 0 by optimality of d fourth is upper bounded by O(y log(h) ity and fifth is bounded by the following argument: let 1 arg min arg max e" ( S.9)
B.4 second part we use Lemma
B.5 third part is upper
by Hoeffding's inequal-
JEF g6g
E Xmujan
mx
E min max 07 (6,, x~up" fer geg
< E max 2a ( 5 (hh
E
(hh ,kh
x-je gEG
xmpp"
S el (0h, f g) Egen, (F,9)
$ 4% (8h, f のh + Egen,h (F, = mmp + Cgen,h
where the second inequality uses Lemma B.3
B.2 PROOF OF THEOREM B
Consider task μ. For simplicity of notation, we use Th instead x4, π instead of π* Let Vi and v be the state distributions at level h induced by π. and π respectively. Let
e/ (xx) = max E
EIE E amta s'mPa.
g(
E g(s')] a~ PP
9E9 swu
be the loss of policy Ti at level h. By definition, es
E en (x). Using Lemma C.I from Sun x
al. (011)), have
A
Vitz(
V=+)('")
Observe that
A
E. swap a
E
Viti(s')-
ONTT(GE)),' Viti(''))
max E. gEg swyh am)) (1)
9(8)
e-TT(- ~P.. g(s')}+
max E I gT)((s)
E [Tg(s)] +[ E V
N (s) E I V
V (s)
DEG sweg
awaag
<
C) (xh) max E. (gg(s) [Tg(s)]
max
g(s) E g())
gEst swang
Swug
gEG swaet
s-0t
15

---

Under review as a conference paper at ICLR 2020
Lemma B.6. Defining
An max
E
g(s)
g(s)l, we have
gEg
swen
max gEG s
E TT
[Kg(s)] S
An 26
Using the above lemma, we get án S (xx)
2DA + 2e We now bound An
A max .EG
E smup
9(6')
E g(s)
a 8 Pz..
snuny
max gEG
E
E 0--8n-1 PP...
g(5')
E
E
g(s)
+ max gEG
E SNoL
E ammth-1
g(s')
E g(s) smuk
Rownh1
s'
o' Pai
max
E
F 11-19(5)=
E
h 19(s)
Eh (x
gEg
Smpp 6-1
< A + 2chr + Ch ch- (X8-1)
Thus A. < 2(h 1) + 11:h-1 and so A. S (xI(())))) -una + (4h - 2) This implies that
J(T - J(x')
Σ
A S
(H 2h + O((H)))
h=1
h=1
Taking expectation wrt j ~ 1) and x ~ yC" completes the proof.
B.3 PROOFS OF LEMMAS
Proof of Lemma B.3
Again we define F' in Equation Let Let - 2/ (0. softmaxff))))
a, (5.0. (()) for fl
3, and let 6% (o.
F' and similarly define (6. f. (6, softsaxmmx Notice that E( a, 3, 2K- lipschitz, l( 3, a)) is K K-lipschitz and l(v, a, a, ;, a) is 1-lipschitz, Using Theorem 8(i) from Maurer al. (2016), we get that
E sup sup (0, f.9) (6. f.9)
poon x fEF gEG
( Ch .)) - (o, f' .9)
E E
sup sup PEF gEG
p x
2 Fx G(((F)(((s)), (Sa)) ((S)). a))
n
22 a)
VTKKG(((Ss)
VTTG(G(SA)
E
E
n
x
n
n
< RK K
KG()(S)))
G(G(S)))
S Egen,h (F,9)
POW X
where we used lipschitzness and Slepian's lemma for second inequality and a similar computation to Lemma A. for the third.
Proof of Lemma B.4
(o)
E x
(x)
E (O)
E mx(
p x
p x
E max ( 1 (g9 - max ix f2.g)
x DEG
gEG
<
E
E
maxi (o, fe. g) f2.
paim x
gEg
16

---

Under review as a conference paper at ICLR 2020
< E max max Heony xrop" JEF gE9
( (o. - (...f..))
S Egen, (F.9)
where we use the definition of Ln. 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), that
Let {X)ses and sES be zero mean Gaussian processes such
B(X, - x1) -x)² s B(Y, Yo),, Vs
Then
sup X, I E sup Y SES
sES
Using Theorem 8(ii) from Maurer et al. (006)), we get that
2 S T
9 (n(2/8
sup
E mx (Ø)
T Σ (
G(S)+
(12)
x-pph
2T
where S (o) m( E }). We bound the Gaussian average of S using Slepian's lemma. Define two Gaussian processes indexed by + as
2K vn Σ InSAO(s))
X
Σ 7mll)) x(o and Y,
For x as. 8,)}, consider representations and d',
(m(()x - =x)'))x²² = (min max (1, -min max , ,f,)))² fEF gEG
fer gEG ( (6.)
< sup fEF.gEB
K ( Kx ((gl(s)))()))
sup SeF.geG
E(f(O(si)aa - f(o'(s))) ))
2
k2
sup SEF.EE
K2 sup
Σ (f(6(s)))o,- - -s(o's))))
<
feF
4K
4K²
Σ (d(s)) 8 22,,))1
[d((;;)& ='()))))
j.b
where we prove the first inequality later, second inequality comes from g being upper bounded by and by Stuchy-Schwarth inequality, third inequality comes from the 2-lipschitzness of J.
E(X, Xx) = Σ (mi())22 -m((())))) )2
< 4K2
E(o(s) x²(s))a))= E(Yo - Yz
Thus by Slepian's lemma, we get
2K
G(S) Eup X < Esup Yd G(0((s))))
6
6c4
Vn
17

---

Under review as a conference paper at ICLR 2020
Plugging this into Equation 12 completes the proof. To prove the first inequality above, notice that
min max (o. f. min max (d', ff JEF
- - f'
JEF 9E9
f. 9²") ('
- f f' - (.d'""
s sup
(hl....0 1,00 f.9))
SEF.gE
By symmetry, we also get that min max (6, f.9) 9)0---))0)) min max JEF SEg
(d', f.9) S
sup JEF.EEG
f.9)
JeF SEG
Proof of Lemma B.6 Let g = arg max
E
(g((s)
E hg((s) and g' = arg min g gEG
gEG
envr
I [o +2///2-
max gEg
1Tg(s)
E
170()))
[T](()) TKg(s)
5
$ E swug
g
E
g((s
+|
E (g'(s)
[x)(s)] +1 E [g'( (
< max gEg
E g(5)
E [(8)
+2
E
- (5)))l
a +
S Dn + 20
C DATA SET COLLECTION DETAILS
C.I DATASET FROM TRAJECTORIES
Given n expert trajectories for a task for each trajectory 7= ($1 01 ....... 8H. a.) we can sample an U([H)) and select the pair (sh, as) from that trajectory This gives us i.dd. pairs (s3, (5),,;)) a;)) for the task μ. We collect this for tasks and get datasets x(1) ,..., (()) (2).
C.2 DATASET FROM TRAJECTORIES AND INTERACTION
Given 2n expert trajectories for task μ4, we use first trajectories to get independent samples from the distributions vi,,.-- HH,,, respectively for the states in the dataset. Using the next n trajecto- ries, we get samples from vi.,..... V- for the s states the dataset, and for each such state we uniformly sample an action a from A and then get state & from P..a by resetting the environment to
s and playing action a. We collect this for tasks and get datasets (
{xf ..... x for every set of n tuples obtained level h. Rearranging, (1) (I)
(4)
i (T), where each dataset x, (i)
construct
the datasets X
(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.
In practice one all pairs from all trajectories, even though the samples strictly i.i.d.
18

---

Under review as a conference paper at ICLR 2020
DirectedSwimmer
30
20
baseline
expert
experts
4 experts
125
200
# samples for test task
Figure 3: The total rewards by different algorithms in DirectedSwimmer.
Experiment Setup We first describe the construction of the NoisyCombination..com environment. The state space is R:40 Each state s is in the form of [Sreal, Smuise), while Sreal E TR is either a one- hot vector or zero vector, and &noise Io is sampled from (0, I). The action space is discrete and has size For each MDP, we have a sequence of actions E [2]20 This is the sequence of optimal actions. We use different a" to define different environments. The transition model is that: 8real = e; for some and the action is a,, then 5'zal = 6 and we'll get reward 1. Otherwise 8'eal will be all zero and the reward is O. noose will always be sampled from the Gaussian distribution. Note that once Freal is all zero, it will not change and the reward will always be The maximum horiozn is set 20 and therefore, the optimal policy has return 20. The initial Seaal is always e1.
The representation has dimension of 10. We limit the function & to be a linear mapping from 1R to R.00 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 train the representation ø. The trajectories are generated by the optimal policy.
When training the policy using an RI. algorithm, to reduce the impact of initialization, the last full connected layer is initialized to O. We use the PPO (Schulman et al. 2017) algorithm to train our policy with code from Dhariwal al. (0177).
DirectedSwimmer A DirectedSwimmer environment is the same as Swimmer in OpenAl Gym (Brockman et al., 2016), except the following: the reward function is parametrized by a direction sample
with | = 1, and is defined as the traveled distance along the direction d. For each task, a random direction. The state space is still R The original action space in Swimmer is IR,, and we discretize the action space, such that each entry be only one of {-1, -0.5, O, 0.5, 1, 1}. We also reduce the maximum horizon from 1000 to 100. We trained the experts for million steps by PPO make sure converges.
The function d we use has two fully connected layers and two ReLU layers. The number of hidden units is 100, is the dimension of representation. We also include the total rewards that each algo- rithm can get in Figure Note even though the baseline has a high validation loss, its performance can be quite good. This does not indicate 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 & Bal 2014) with learning rate 0.001 until it converges. To solve equation g build a joint loss over o and all f's in each task,
T
Llo, fi,...,,ry
log(T o. (s)))s)))))
(13)
nT
Then we minimize L(ø, fi,... fr) and obtain the optimal ø.
19


Reviews:
[{"review_id": "rJx5HLM-oB", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overview:\n\nThe 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. \n\nDecision:\n\nIn its current form, the paper should be rejected because (1) the empirical analysis is incomplete – the baseline isn't very appropriate, the results are not conclusive, details are scattered or not included, (2) the literature survey does not connect the proposed approach with existing approaches, and does not convince the reader why all the existing approaches have not been compared against empirically, (3) the paper is generally unpolished and needs more work before being considered for acceptance.\n\nDetails:\n\nThe paper makes both theoretical and empirical claims. I did not have the time to thoroughly verify the theoretical claims and took them at face value. I consider the theoretical guarantees associated with the proposed approach a welcome and valuable contribution to this field that has recently been relying primarily on limited empirical work to assess any method. \n\nThe empirical results presented in the paper do not sufficiently support the claims of sample efficiency. One of the main issues with the empirical analysis is the choice of the baseline, which learns a policy from scratch. This does not help make conclusions about the sample efficiency of the proposed method on new tasks. A better baseline would be one that learns some representation from the T previous tasks, which would help infer if the proposed method to learn representations is actually more sample efficient on new tasks or not. There is also no comparison with existing approaches that are mentioned in the Related Work section. If those aren’t appropriate baselines for this problem, a small explanation of the reasons why would help readers understand why they haven’t been compared against. Additionally, an analysis of statistical significance of the results is missing and would significantly help in gauging the efficacy of the proposed approach. \n\nThe paper notes that these are some preliminary experiments. The completion of the empirical analysis would definitely make a stronger case for this paper to be accepted.\n\nMinor comments to improve the paper:\n\n- Error bars in the plot, specification of number of runs, and other such experimental details would be very helpful in interpreting the results. \n- It would help a reader if the paper was more self-contained, e.g., if terms like supp(\\eta), \\bar{s}, \\tilde{s} are defined more clearly.\n- It would also help to say what the proofs intuitively mean, e.g., for a new task drawn from this particular distribution of tasks, the agent would achieve close-to-X performance within Y samples – something along those lines.\n- There are some typos, e.g., 'possibility'->'possibly' on page 1, missing $H$ in specification of MDPs on page 2, 'exiting'->'exciting' on page 8, some latex symbols in Appendix D, etc.\n- The bibliography has a lot of issues – some references are incorrectly parsed (e.g., Yan Duan, Marcin Andrychowicz, Bradly Stadie, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. 03 2017), others are inconsistent (e.g., \"In NIPS\" and \"In Advances in Neural…”; the arXiv ones).\n\n   ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "r1lKVtETYH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nPros: \n1. 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.\n2. 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. \n3. A good and concise review of RL concepts is provided.\n\nCons:\n1. 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.\n\n2. 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? \n\n3.  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. \n\nApart from these, below are some minor comments that could help improve the reading of this paper:\na. Theorem 3.1 is not really a theorem, since it is very informal. Also, fix the Theorem numbers. \n\nb. 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.\n\nc. 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*? \n\nd. 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.\n\ne. The first two terms in Theorem 4.1 are claimed standard, provide the citations?\n\nf. Theorem 4.1, perhaps use some other notation for c, which is defined as the cost/reward in the RL setting.\n\nOverall, 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. \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJxHdQmqFS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "(I bid on the paper thinking that bi-level optimisation would play a major role in the paper. Unfortunately, it does not, so my expertise in bi-level optimisation is not much use, I am afraid.)\n\nThe authors study applying policies learned on one task to another task, while considering practical finite-sample limitations. They call this the \"representatition learning for imitation learning\". Unfortunately:\n\nThe make extensive assumptions, summarised on page 3, but not formalised as Assumptions 1, 2, 3 anywhere, as far as I can tell. They assume: \n-- concentration of the loss, \"which guarantees within-task sample efficiency\" (but does not seem easy to support in practice? actually, one may observe samples from a stochastic process, rather than iid samples with any concentration what-so-ever?), \n-- that the \"loss\" they use with is somehow close to the optimum of the expected value function J (for which I again see no justification, empirical or otherwise).\n\nThen they reuse results of Maurer et al (2016) and extend them to a case where the actions are not observable. The results are plausible, given the assumptions. Given the assumptions, they also do not seem to be particularly relevant to the practice of RL?\n\nThe empirical results involve only benchmarks of the authors own coinage, and hence are hard to evaluate. It seems plausible, again, however, that the approach may work in some cases. \n\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a methodology for learning a representation given multiple demonstrations, by optimizing the representation as well as the learned policy parameters. The paper includes some theoretical results showing that this is a sensible thing to do, and an empirical evaluation.

Post-discussion, the reviewers (and me!) agreed that this is an interesting approach that has a lot of promise. But there was still concern about he empirical evaluation and the writing. Hence I am recommending rejection.
Decision: Reject