TRAIL: Near-Optimal Imitation Learning with Suboptimal Data
Authors: Mengjiao Yang
Sergey Levine
Ofir Nachum
Venue: ICLR 2022 Poster (2022)
Keywords: Imitation Learning
Action Representations
Latent Dynamics Model
Offline Datasets

Abstract: In imitation learning, one aims to learn task-solving policies using access to near-optimal expert trajectories collected from the task environment. However, high-quality trajectories -- e.g., from human experts -- can be expensive to obtain in practical settings. On the contrary, it is often much easier to obtain large amounts of suboptimal trajectories which can nevertheless provide insight into the structure of the environment, showing what \emph{could} be done in the environment even if not what \emph{should} be done. Is it possible to formalize these conceptual benefits and devise algorithms to use offline datasets to yield \emph{provable} improvements to the sample-efficiency of imitation learning? In this work, we answer this question affirmatively and present training objectives which use an offline dataset to learn an approximate \emph{factored} dynamics model whose structure enables the extraction of a \emph{latent action space}. Our theoretical analysis shows that the learned latent action space can boost the sample-efficiency of downstream imitation learning, effectively reducing the need for large near-optimal expert datasets through the use of auxiliary non-expert data. We evaluate the practicality of our objective through experiments on a set of navigation and locomotion tasks. Our results verify the benefits suggested by our theory and show that our algorithms is able to recover near-optimal policies with fewer expert trajectories.

Published as a conference paper at ICLR 2022
TRAIL: NEAR-OPTIMAL IMITATION LEARNING WITH SUBOPTIMAL DATA
Mengjiao Yang
Sergey Levine
Ofir Nachum Google Brain
UC Berkeley, Google Brain sherryy@google.. com
UC Berkeley, Google Brain
ABSTRACT
The aim in imitation learning is to learn effective policies by utilizing near-optimal expert demonstrations. However, high-quality demonstrations from human ex- perts can be expensive to obtain in large number. On the other hand, it is often much easier to obtain large quantities of suboptimal or task-agnostic trajectories, which are not useful for direct imitation, but can nevertheless provide insight into the dynamical structure of the environment, showing what could be done in the environment even if not what should be done. We ask the question, is it possible to utilize such suboptimal offline datasets to facilitate provably im- proved downstream imitation learning? In this work, affirmatively and present training objectives that use offline datasets to learn a factored transition model whose structure enables the extraction of latent ac- tion space. Our theoretical analysis shows that the learned latent action space can boost the sample-effccinccy of downstream imitation learning, effectively re- ducing the need for large near-optimal expert datasets through the use of aux- iliary non-expert data. To learn the latent action space in practice, we propose TRAIL (Transition-Reey) Reparametrized Actions for Imitation Learning). an algorithm that learns an energy-based transition model contrastively, and uses the transition model to reparametrize the action space for sample-efficient imitation learning. We evaluate the practicality of our objective through experiments on a set of nav- igation and locomotion tasks. Our results verify the benefits suggested by our theory and show that TRAIL is able to improve baseline imitation learning by up to 4x in performance.
answer this question
INTRODUCTION
Imitation learning uses expert demonstration data to learn sequential decision making poli- cies (Schaal, 1999). Such demonstrations, often produced by human experts, can be costly to obtain in large number. On the other hand, practical application domains, such as recommendation (Afsar et al., 2021) and dialogue (Jiang et al., 2021) systems, provide large quantities of offline data gen- erated by suboptimal agents. Since the offline data is suboptimal in performance, using it directly for imitation learning is infeasible. While some prior works have proposed using suboptimal offline data for offline reinforcement learning (RL) (Kumar et al., 2019; Wu et al., 2019; Levine et al., 2020), this would require reward information, which may be unavailable or infeasible to compute from suboptimal data (Abbeel & Ng, 2004). Nevertheless, conceptually, suboptimal offline datasets should contain useful information about the environment, only we could distill that information into a useful form that can aid downstream imitation learning.
One approach to leveraging suboptimal offline datasets is to use the offline data to extract lower- dimensional latent action space, and then perform imitation learning on an expert dataset using this latent action space. If the latent action space is learned properly, one may hope that perform- ing imitation learning in the latent space can reduce the need for large quantities of expert data. While a number of prior works have studied similar approaches in the context of hierarchical im- itation and RL setting (Parr & Russell, 1998; Dietterich et al., 1998; Sutton et al., 1999; Kulkarni et al., 2016; Vezhnevets et al., 2017; Nachum et al., 2018ar Ajay et al., 2020; Perisch et al., 2020; Hakhamaneshi et al., 2021), such methods typically focus on the theoretical and practical benefits of temporal abstraction by extracting temporally extended skills from data or experience. That is, the main benefit of these approaches is that the latent action space operates at a lower temporal frequency than the original environment action space. We instead focus directly the question of action representation: instead of learning skills that provide for temporal abstraction, we aim to di- rectly reparameterize the action space in a way that provides for more sample-efficient downstream

---

Published as a conference paper at ICLR 2022
Transition reparametrized actions
(I) T. a)
(2) Fg (als, (()
me (bls, a))s)
a z Is. 2)
as)) goff
{(s, a) D π
mz(3))
Inference S
Pretraining
Downstream Imitation
The TRAIL framework. Pretraining learns a factored transition model 'Tz Φ and an action decoder To on Doff_ Downstream imitation learns a latent policy T2 on D with expert actions reparametrized by ø. During inference, Tz and Ta are combined to sample an action.
Figure 1:
imitation without the need to reduce control frequency. Unlike learning temporal abstractions, ac- tion reparamtrization does not have to rely on any hierarchical structures in the offline data, and can therefore utilize highly suboptimal datasets (c.g., with random actions).
Aiming for provably-efficient approach to utilizing highly suboptimal offline datasets, we use first principles to derive an upper bound on the quality of an imitation learned policy involving three terms corresponding to (1) action representation and (2) action decoder learning on a suboptimal offline dataset, and finally, (3) behavioral cloning (i.e., max-likelihodd learning of latent actions) on an expert demonstration dataset. The first term in our bound immediately suggests a practical offline training objective based on transition dynamics loss using an factored transition model. We show that under specific factorizations (c.g., now-dimensional or linear), one can guarantee improved sample efficiency the expert dataset. Crucially, our mathematical results avoid the potential shortcomings of temporal skill extraction, as our bound is guaranteed to hold even when there is no temporal abstraction in the latent action space.
We translate these mathematical results into an algorithm that we call Transition Transition-Reparameenni Reparametrized Actions for Imitation Learning (TRAIL). As shown in Figure 1, TRAIL consists of pretraining stage (corresponding the first terms in our bound) and a downstream imitation learning stage (corresponding to the last term our bound). During the pretraining stage, TRAIL uses an offline dataset to learn factored transition model and a paired action decoder. During the downstream imitation learning stage, TRAIL first reparametrizes expert actions into the latent action space ac- cording to the learned transition model, and then learns a latent policy via behavioral cloning in the latent action space. During inference, TRAIL uses the imitation learned latent policy and action decoder in conjunction to act in the environment. In practice, TRAIL parametrizes the transition model as an energy-based model (EBM) for flexibility and trains the EBM with a contrastive loss. The FBM enables the low-dimensional factored transition model referenced by our theory, and we also show that one can recover the linear transition model in our theory by approximating the EBM with random Fourier features (Rahimi et al., 2007).
To summarize, our contributions include (i) provably beneficial objective for learning action rep- resentations without temporal abstraction and (ii) practical algorithm for optimizing the proposed objective by learning an EBM or linear transition model. An extensive evaluation on a set of nav- igation and locomotion tasks demonstrates the effectiveness of the proposed objective. TRAIL's empirical success compared to a variety of existing methods suggests that the benefit of lcarning single-step action representations has been overlooked by previous temporal skill extraction meth- ods. Additionally, TRAIL significantly improves behavioral cloning even when the offline dataset is unimodal or highly suboptimal (e.g.. obtained from random policy), whereas temporal skill extrac- tion methods lead to degraded performance in these scenarios. Lastly, we show that TRAIL, without using reward labels, can perform similarly or better than offline reinforcement learning (RL) with orders of magnitude less expert data, suggesting new ways for offline learning of squential decision making policies.
2 RELATED WORK
Learning action abstractions a long standing topic the hierarchical RL literature (Parr & Russell, 1998; Dietterich et al., 1998; Sutton et al., 1999; Kulkarni et al., 2016; Nachum et al., 2018a). A large body of work focusing on online skill discovery have been proposed as a means to improve exploration and sample complexity in online RL. For instance, Eysenbach et al. (2018); Sharma et al. (2019); Gregor et al. (2016); Warde-Farley et al. (2018); Liu et al. (2021) propose to learn a diverse set of skills by maximizing information theoretic objective. Online skill discovery is also
commonly seen in hierarchical framework that learns a continuous space (Vezhnevets et al., 2017; Hausman et al., 2018 8; Nachum et al., 2018a; 2019) or discrete set of lower-level policies (Bacon
2

---

Published as a conference paper at ICLR 2022
et al,, 2017; Stolle & Precup, 2002; Peng et al., 2019), upon which higher-level policies are trained to solve specific tasks. Different from these works, we focus on learning action representations offline from fixed suboptimal dataset to accelerate imitation learning.
Aside from online skill discovery, offline skill extraction focuses on learning temporally extended action abstractions from a fixed offline dataset. Methods for offline skill extraction generally in- volve maximum likelihood training of some latent variable models the offline data, followed by downstream planning (Lynch et al., 2020), imitation learning (Kipf et al., 2019; Ajay et al., 2020; Hakhamaneshi et al., 2021), offline RL (Ajay et al., 2020; Zhou al., 2020), online RL (Fox et al., 2017; 7; Krishnan et al., 2017; Shankar & Gupta, 2020; Shankar et al., 2019; Singh et al., 2020; Pertsch al., 2020; 2021; Wang et al., 2021) in the induced latent action space. Among these works, those that provide a theoretical analysis attribute the benefit of skill extraction predominantly to increased temporal abstraction as opposed to the learned action space being any "casier" to learn from than the raw action space (Ajay et al., 2020; Nachum et al., 201 Unlike these methods, our analysis focuses on the advantage of a lower-dimensional reparametrized action space agnostic to temporal abstraction. Our method also applies to offline data that is highly suboptimal (e.g., contains random actions) and potentially unimodal (e.g., without diverse skills to be extracted), which have been considered challenging by previous work (Ajay et al., 2020).
While we focus on reducing the complexity of the action space through the lens of action representa- tion learning, there exists disjoint set of work that focuses on accelerating RL with state represen- tation learning (Singh et al., 1995; Ren & Krogh, 2002; Castro & Precup, 2010; Gelada et al., 2019; Zhang et al., 2020; Arora et al., 2020; Nachum & Yang. 2021), some of which have proposed to extract latent state space from learned dynamics model. Analogous to our own derivations, these works attribute the benefit of representation learning to a smaller latent state space reduced from a high-dimensional input state space (e.g., images). Lastly, there exist model-based approaches that utilizes offline data to learn model dynamics which in tern accelerates imitation (Chang et al., 2021; Rafailov et al., 2021). These work differ from our focus of using the offline data to learn latent action space.
PRELIMINARIES
In this section, we introduce the problem statements for imitation learning and learning- based con- trol, and define relevant notations.
Markoy decision process. Consider an MDP (Puterman, 1994) M == (S. A, μ. 7). con- sisting of state space S, an action space A, a reward function R: S A - IR, transition function T S x A- + A an initial state distribution fa E A A(S), and a discount factor E [0, 1) A policy π : + (AA) interacts with the environment starting at an initial state so μ. An action a, ~ (ss) is sampled and applied to the environment at each step & 2 0. The environment pro- duces scalar reward R(s,, a,) and transitions into the next state St+1 ~ T(s,, at). Note that we are specifically interested in the imitation learning setting, where the rewards produced by R are unobserved by the learner. The state visitation distribution (s) induced by a policy π is defined as d' ds((s) (1 -7) -i 7th Pr = s/r.M]. We relax the notation and use (s, a) ~ d to denote 8 d= T (s).
Learning goal. Imitation learning aims to recover an expert policy n. with access to only fixed set of samples from the expert: = {(s,, a))}²= with s; dā and a; ~ (s;).. One approach to imitation learning is to learn policy π that minimizes some discrepancy between π and T. In
our analysis, we will use the total variation (TV) divergence in state visitation distributions,
Diff( TT) = Drv (d" I(d"),,
as the way to measure the discrepancy between æ and T., Our bounds can be easily modified to apply to other divergence measures such as the Kullback-Leibler (KL) divergence or difference in expected future returns. Behavioral cloning (BC) (Pomerleau, 1989) solves the imitation learning problem by learning æ from D=. via a maximum likelihood objective
JBC (TT := (ss.a) ...) log π (a]ss)]
which optimizes an upper bound of Diff (IT, n.) defined above (Ross & Bagnell, 2010; Nachum & Yang, 2021): Diff T.) < 1- 2
7 (A())) denotes the simplex over set
7
Ea.
const( (T.) + JBC(TT.
1-7V
2
3

---

Published as a conference paper at ICLR 2022
BC with suboptimal offline data. The standard BC objective (i.e., direct max-likelihood on DF=) can struggle to attain good performance when the amount of expert demonstrations is lim- ited (Ross et al., 2011; Tu et al., 2021). We assume access dataset Doff = ((8; aj x{) ). where the suboptimality a result of (i) suboptimal action sam- ples a, ~ UnifA and (ii) lack of reward labels. We use (s, a, s') ~ duff as shorthand for sim- ulating finite sampling from Doff via 8; doff 0; Unif S/. as),, where doff is an unknown offline state distribution. We assume doff sufficiently covers the expert distribution; i.e., d (s) > 0 = duff (s) > 0 for all s E S. The uniform sampling of actions Doff is largely for mathematical convenience, and in theory can be replaced with 1 distribution uniformly bounded from below by 1) > O, and our derived bounds will be scaled by i as a result. This works focuses on how to utilize such suboptimal Doff to provably accelerate BC.
additional suboptimal offline
4 NEAR-OPIE OPAR-OL IMITATION LEARNING WITH REPARAMETRIZED ACTIONS In this section, we provide provably-efficient objective for learning action representations from suboptimal data. Our initial derivations (Theorem 1) apply to general policies and latent action spaces, while our subsequent result (Theorem 3) provides improved bounds for specialized settings with continuous latent action spaces. Finally, we present our practical method TRAIL for action representation learning and downstream imitation learning.
4.1 PERFORMANCE BOUND WITH REPARAMETRIZED ACTIONS
Despite Duff being highly suboptimal (e.g., with random actions), the large set of (s, a, s') tuples from Duff reveals the transition dynamics of the environment, which a latent action space should support. Under this motivation, we propose to learn a factored transition model T Tz from the offline dataset Doff, where Φ S A + Z is an action representaiton function and
Tz : S Z A(S) latent transition model. Intuitively, good action representations should enable good imitation learning.
We formalize this intuition in the theorem below by establishing bound on the quality of learned policy based on (1) an offline pretraining objective for learning é and Tz, (2) an offline decoding objective for learning an action decoder Ta and (3) downstream imitation learning objective for learning latent policy a with respect to latent actions determined by Φ.
Theorem 1. Consider an action representation function o : S x A + Z. factored transition (S),, an action decoder Ta 7 - ((A), and a tabular latent policy m S (Z)). Define the transition representation error as
model Tz
Jr(T2, Φ) = (e,,))rdott [TKL(T(s, a) | ()s,))))),
the action decoding error (TR x) == ((na))aorr log na (als, (())
and the latent behavioral cloning error as
(T)) = (d log ( ass)
Then tions
the TV of
divergence
between the state visitation be bounded as
distribu-
Ta o a2
A(A) and n
can
Diff(Ta T.).
C V
E Eis.a) [DKL(T a) T2 ((s,
=TI(Tz,
Pretraining
+C2 V2
E Egngoot mma (Ta(Fa. (8, 2), zEL Fs const(doft .)) + JDDE(TA.
(2)
Downstream Imitation
+Ca Ex-ato [DKI. (T.,, (8))
(3)
= const(TT.; 2) + JBC,0
where C1 7|A[(1 -1 (xa (ldofft C2 7(1 -1 (da )
C3 7(1 - -7) -1, na- is the optimal action decoder for specific data distribution goll and
4

---

Published as a conference paper at ICLR 2022
specific o:
dott (s.a) 1[2 a()
Ta- Tal(a] En'es doft (s, a') - 1[z = o(s,
and 2 is the marginalization of T, onto Z according to Φ:
N. No ==
Σ T. (als).
(EA,a=6(s.a)
Theorem essentially decomposes the imitation learning error into (1) transition-based represen- lation eπror Jr, (2) an action decoding eπror JDE. and (3) latent behavioral cloning error JBC, Notice that only (3) requires expert data D= (1) and (2) are trained on the large offline data Doff. By choosing |2| that is smaller than |A], fewer demonstrations are needed to achieve small error in JBC. compared to vanilla BC with JBC- The Pearson X2 divergence term "d= and C2 accounts for the difference in state visitation between the expert and offline data. In the case where d= differs too much from doff known the distribution shift problem in offline RI. (Levine et al., 2020), the errors from Jr and JDE are amplified and the terms (1) and (2) in Theorem 1 dom- inate. Otherwise, as JT - and Ta é + arg min JDE, optimizing z in the latent action space is guaranteed to optimize in the original action space.
in C1
Sample Complexity To formalize the intuition that smaller latent action space |Z| |A| leads to more sample efficient downstream behavioral cloning, we provide the following theorem in the tabular action setting. First, assume access Doret == OPT a (Dofl) which yields pretraining errors (1)(grrd) and (2)or in Theorem 1. For downstream behavioral cloning, we consider learning a tabular TT on D with n expert samples. We can bound the expected difference between a latent policy Tore with respect to Øorot and T. as follows.
oracle latent action representation function
Theorem 2. Let dord OPT We have,
df (Doff) and Ford be the latent BC policy with respect to Pordl
| n
Ep.. DDff a T.)I & (1)(Gora) + (2)(oorl) + C3
where C3 is the same as in Theorem 1.
We can contrast this bound to its form in the vanilla BC setting, for which |Z| -AA| and both (1)(orr and (2)(o))))) are zero. We can expect an improvement in sample complexity from reparametrized actions when the errors (1) and (2) are small and |Z| < |/|.
4.2 LINEAR TRANSITION MODELS WITH DETERMINISTIC LATENT POLICY
Theorem has introduced the notion of a latent expert policy т, 2. and minimizes the KL divergence between т.,Z and tabular latent policy TT. However, it is not immediately clear, in the case of continuous latent actions, how to ensure that the latent policy πे expressive enough to capture any T. In this section, we provide guarantees for recovering stochastic expert policies with continuous latent action space under a linear transition model.
Consider a continuous latent space Z C Id and a deterministic latent policy Te Te(s) = 62 for some 0 d8| While a deterministic 0 in general cannot capture a stochastic R., we show that under a linear transition model (z' (ssss, Ø(s, a)) w(s')T o(s, a).. there always exists deterministic policy xe + Rd, such that 0, T Vs E S. This means that our scheme for offline pretraining paired with downstream imitation learning can provably recover any expert policy from a deterministic Ta, regardless whether æ. is stochastic.
for some Z Rd and suppose there exist w S Rd such that Tz(s'(s, ((s, a)) w(s')T a(s, a) for all s, E S,a E A. Let Ta : x Z + A(A) be an action decoder, π. + (AA any policy M and me - RRd he a deterministic latent policy for some Rdx(s] Then,
Theorem 3. Let Φ S A
Diff(na o no, n.) I (1)(T2,
Φ) + (2) ( Φ)
Downstream
+Ca
θ E Bgnd"= arm,
0(s, a))
(4)
Imitation
00
where C4 = ||S|||/]|o. (I) and (2) corresponds to the first and second terms in the bound in Theorem 1.
5

---

Published as a conference paper at ICLR 2022
By replacing term (3) in Theorem that corresponds to behavioral cloning in the latent action space by term (4) in Theorem 3 that a convex function unbounded in all directions, we are guaranteed that Ty is provably optimal regardless of the form of T, and .. Note that the downstream imitation learning objective implied by term (4) is simply the mean squared error between actions Og chosen by wo and reparameterized actions (s8, a) appearing in the expert dataset.
4.3 TRAIL: REPARAMETRIZED ACTIONS AND IMITATION LEARNING IN PRACTICE
In this section, we describe our learning framework, Transition- maaniition-Repaaametriiee Actions for Im- itation Learning (TRAIL). TRAIL consists of two training stages: pretraining and downstream behavioral cloning. During pretraining. TRAIL learns Tz and Φ by minimizing Jr(Tz.0) Φ) (a.a))=dtf (T,,(((s.))) (2, 0(s, a)))] Also during pretraining, TRAIL learns To and ø by minimizing JDE TT. 6) := (dtt) log To (a|s, ((s, a))]. TRAIL parametrizes Ta as mul- tivariate Gaussian distribution. Depending on whether Tz is defined according to Theorem or Theorem 3, we have either TRAIL EBM or TRAIL linear.
TRAIL EBM for Theorem 1. In the tabular action setting that corresponds to Theorem I, to ensure that the factored transition model Tz is flexible to capture any complex (e.g., multi-modal) transitions in the offline dataset, we propose to use an energy-based model (EBM) to parametrize (z(5'ss, a)),
(s's's, ((s, p(s") -expppp(((,,))
where p is fixed distribution over S and S -> Z is function of s'. In our implementation we set to be the distribution of s' in doff which enables a practical learning objective for Tz by minimizing (ssaa)a [DKL T) (s. Ø(s, a))] in Theorem using contrastive loss:
B (- log Tz ((5))) (d 2 E [l] a) (s))
+ log Ex lexpl - (0, (8)))))))
During downstream behavioral cloning, TRAIL EBM learns latent Gaussian policy a by mini- mizing JBC. (T2) = [( log 2 (0(s, a}}s)] with ø fixed. During inference, TRAIL EBM first samples a latent action according to z ~ (s),,, and decodes the latent action using Ta (s, 2) to act in environment. Figure 1 describes this process pictorially.
TRAIL Linear for Theorem 3. In the continuous latent action setting that corresponds to Theo- rem 3, we propose TRAIL linear, an approximation of TRAIL EBM, to enable learning linear tran- sition models required by Theorem 3. Specifically,,
first learn f, that parameterize an energy- a) - g(s')l [[))([(/2)) using the same contrastive loss above (replacing é and 1i by and g), and then apply random Fourier features (Rahimi al., 2007) to recover o(s, a) = cos (W f(, a) + b). where W a d x & matrix with entries sampled from unit Gaussian and 6 vector with entries sampled uniformly from [0, 2T... W and are im- plemented as an untrainable neural network layer on top of This results in an approximate linear transition model,
based transition model (ss'ss, a) x p(s')
T(s'ls, a) x p(s') expf (f. -aaoo))gsss)))) g(s))11//2) x T(s') a).
During downstream behavioral cloning, TRAIL linear learns a deterministic policy me in the continu- latent action space determined by ø via minimizing fixed. During inference, TRAIL linear first determines the latent action according to = = (x)),, and decodes the latent action using a ~ Ta (s, z) to act in environment.
B.-dd ,awm) (s) (0, o(s, a)) with
antmaze
antmaae
cartpole swingup cheetah
fish swim
walker
walker
Figure 2: Tasks for our empirical evaluation. We include the challenging AntMaze navigation tasks from D4RL (Fu et al., 2020) and low (I-DoF) to high (21-DoF) dimensional locomotaion tasks from DeepMind Control Suite (Tassa et al., 2018).

---

Published as a conference paper at ICLR 2022
Export D
esgert: traje
expert: traje
eepert: traje
espert: traje
Suboptimal 1"
ontmaze- lazge-diverse
antmaze- antzaze-laare-phe large-play
TRAIL (EBM)
SkilD (t=10
SkilD [=1)
SEJJU (tm10)
SPRR [ts1]
OPAL (t=10)
OPAL (t-1)
Buseline BC
20
40
20 40
60
20 00 60
80 20
40
60
80
Figure 3: Average success rate (%) over seeds of TRAIL EBM (Theorem 1) and temporal skill extraction methods SkiLD (Pertsch et al., 2021), SPiRL (Pertsch et al., 2020), and OPAL (Ajay et al., 2020) pretrained on suboptimal Doff' Baseline BC corresponds to direct behavioral cloning of expert D without latent actions.
5 EXPERIMENTAL EVALUATION
We now evaluate TRAIL on set of navigation and locomotion tasks (Figure 2). Our evaluation is designed to study how well TRAIL can improve imitation learning with limited expert data by leveraging available suboptimal offline data. We evaluate the improvement attained by TRAIL over vanilla BC, and additionally compare TRAIL to previously proposed temporal skill extraction meth- ods. Since there is no existing benehmark for imitation learning with suboptimal offline data, we adapt existing datasets for offline RL, which contain suboptimal data, and augment them with small amount of expert data for downstream imitation learning.
5.1 EVALUATING NAVIGATION WITHOUT TEMPORAL ABSTRACTION
Description and Baselines. We start our evaluation on the AntMaze task from D4RL, (Fu et al., 2020), which has been used as a testbed by recent works on temporal skill extraction for few-shot imitation (Ajay et al., 2020) and RL (Ajay et al., 2020; Pertsch et al., 2020; 2021). We compare TRAIL to OPAL (Ajay et al., 2020), SkilD (Pertsch et al., 202 1), and SPiRL (Pertsch et al., 2020), all of which use offline dataset to extract temporally extended (length t 10) skills to form a latent action space for downstream learning. SkiLD and SPiRL are originally designed only for downstream RL,
modify them to support downstream imitation learning described in Ap- pendix C. While a number of other works have also proposed to learn primitives for hierarchical imitation (Kipf et al., 2019; Hakhamaneshi et al., 2021) and RL (Fox et al., 2017; Krishnan al., 2017: Shankar et al., 2019; Shankar & Gupta, 2020; Singh et al., 2020), we chose OPAL, SkiLD, and SPiRL for comparison because they are the most recent works in this area with reported results that suggest these methods are ofte-oof-hhe-att especially in learning from suboptimal offline data based on D4RL. To construct the suboptimal and expert datasets, follow the protocol in Ajay et al. (2020), which uses the full diverse play D4RL AntMaze datasets the suboptimal offline data, while using a set of n = 10 expert trajectories (navigating from one corner of the maze to the opposite corner) as the expert data. The diverse and play datasets are suboptimal in the corner-to-comee corner navigation task, as they only contain data that navigates to random or fixed locations different from task evaluation.
Implementation Details. For TRAIL, we parameterize a(s, a) and (ss') using separate feed- forward neural networks (see details in Appendix C) and train the transition EBM via the contrastive objective described in Section 4.3. We parametrize both the action decoder Ta and the latent иz using multivariate Gaussian distributions with neural-network approximated mean and variance. For the temporal skill extraction methods, we implement the trajectory encoder using a bidirectional RNN and parametrize skill prior, latent policy, and action decoder Gaussians following Ajay et al. (2020). We adapt SPiRL and SkiLD for imitation learning by including the KL Divergence term between the latent policy and the skill prior during downstream behavioral cloning (see details in Appendix C). We do search on the extend of temporal abstraction, and found f. = 10 to work the best as reported in these papers' maze experiments. We also experimented with version of vanilla BC pretrained on the suboptimal data and fine-tuned on expert data for fair comparison, which did show significant difference from directly training vanilla BC on expert data.
Results. Figure shows the average performance TRAIL in terms of task success rate (out of 100%) compared to the prior methods. Since all of the prior methods are proposed in terms of tempo- ral abstraction, we evaluate them both with the default temporal abstract, = 10, as well as without temporal abstraction, corresponding
Note that TRAIL
temporal abstraction. We

---

Published as a conference paper at ICLR 2022
Regurt
unt-expoet 10%
dat-expezs
anc-expert 10%
ast-oxpezt 25%
ant-cagees
ant-expeet
ant-raddimm
ant-redium
ant-swedoun-rpplay
ant-rancon
TRAIL Onwer
Sup (m0)
5700,
GPS (1m)
Figure 4: Average rewards (over 4 seeds) of TRAIL EBM (Theorem 1), TRAIL linear (The- orem 3), and baseline methods when using variety of unimodal (ant-medium) um), low-quality (ant-medis um-uummrpppp weppaay.. and random (ant-randon) offline datasets Doff paired with smaller expert dataset D (either 10k or 25k expert transitions).
find that on the simpler antmaze performs similarly to the set of temporal skill extraction methods with = 10. However, these skill extraction methods experience a degradation in performance when temporal abstraction is removed
um task, TRAIL trained on a single-step transition model
= 1). This corroborates the existing theory in these works (Ajay et al., 2020), which attributes their benefits predominantly to temporal abstraction rather than producing latent action space that "easier" to learn. Meanwhile, TRAIL able to excel without any temporal abstraction.
These differences become even more pronounced the harder antmaze- age tasks. We that TRAIL maintains significant improvements over vanilla BC, whereas temporal skill extraction fails achieve good performance even with 10. These results suggest that TRAIL attains significant improvement specifically from utilizing the suboptimal data for learning suitable action representations, rather than simply from providing temporal abstraction. Of course, this does not mean that temporal abstraction is never helpful. Rather, our results serve as evidence that suboptimal data can be useful for imitation learning not just by providing temporally extended skills, but by actually reformulating the action space to make imitation learning easier and more efficient.
5.2 EVALUATING LOCOMOTION WITH HIGHLY SUROPTIMAL OFFLINE DATA
Description. The performance of TRAIL trained on a single-step transition model in the previous section suggests that learning single-step latent action representations can benefit broader set of tasks for which temporal abstraction may not be helpful, e.g., when the offline data is highly sub- optimal (with near-random actions) or unimodal (collected by a single stationary policy). In this section, we consider Gym-MuJoCo task from D4RI. using the same 8-DoF quadruped ant robot as the previously evaluated navigation task. We first learn action representations from the med! um, medi um uenrr, or random datasets, and imitate from 1% or 2.5% of the expert datasets from D4RL. The medium dataset represents data collected from a mediocre stationary policy (ex- hibiting unimodal behavior), and the random dataset is collected by randomly initialized policy and is hence highly suboptimal.
Implementation Details. For this task. we additionally train linear version of TRAIL by approx- imating the transition EBM using random Fourier features (Rahimi et al., 2007) and learn determin- istic latent policy following Theorem 3. Specifically, we use separate forwar networks to pa- rameterize [s, a) and g(s'), and extract action representations using Φ(s, a) = cos(W a) + b), where W. b are untrainable randomly initialized variables described in Section 4.3. Different from TRAIL EBM which parametrizes Z as a Gaussian, TRAIL linear parametrizes the deterministic m using a feew-fd neural network.
Results. Our results are shown in Figure 4. Both the EBM and linear versions of TRAIL con- sistently improve over baseline BC, whereas temporal skill extraction methods generally lead to worse performance regardless of the extent of abstraction, likely due to the degenerate effect (i.e., latent skills being ignored by flexible action decoder) resulted from unimodal offline datasets as discussed in (Ajay et al., 2020). Surprisingly, TRAIL achieves a significant performance boost even when latent actions are learned from the random dataset, suggesting the benefit of learning action representations from transition models when the offline data is highly suboptimal. Additionally, the linear variant of TRAIL performs slightly better than the EBM variant when the expert sample size is small (i.e., 10k), suggesting the benefit of learning deterministic latent policies from Theorem when the environment is effectively approximated by linear transition model.
8

---

Published as a conference paper at ICLR 2022
(onerge)
TRAIL (linear)
Baseline
OPAL
CRR
D D"Cartpole-pvingoo -224
chestab-run -208
tiab-mxin -208
we-ler-atand -208
wal.er-voll
bumanoid-rus -203
1 oarmpol compplle-eetngee enteipg
Cheetab-run BOK
Fish-akts BOX
s00 ker--tand B0N
Sumanold-ra B0
1003
203
691
530
200
490
230
204
10002 39.6
1200
200k
100
52A
Figure 5: Average task rewards (over seeds) of TRAIL EBM (Theorem 1), TRAIL linear (Theo- rem 3), and OPAL (other temporal methods are included in Appendix D) pretrained on the bottom 80% of the RL Unplugged datasets followed by behavioral cloning in the latent action space on To of the top 20% of the RL Unplugged datasets following the setup in Zolna et al. (2020). Baseline BC achieves low rewards due to the small expert sample size. Dotted lines denote the performance of CRR (Wang et al., 2020), an offline RL method trained the full RL Unplugged datasets with reward labels.
to
5.3 EVALUATION ON DEEPMIND CONTROL SUITE
Description. Having witnessed the improvement TRAIL brings to behavioral cloning on AntMaze and MuJoCo Ant, we wonder how TRAIL perform on wider spectrum of locomotion tasks with various degrees of freedom. We consider locomotion tasks from the Deep.ini Control Suite (Tassa et al., 2018) ranging from simple (e.g., 1-DoF (e.g.. 21-DoF (amanoid-run) tasks. Following the setup in Zolna al. (2020), we take TO the trajectories whose episodic reward is among the top 20% of the open source RL Unplugged datasets (Gulcehre al., 2020) as expert demonstrations (see numbers of expert trajectories in Ap- pendix C), and the bottom 80% of RL Unplugged as the suboptimal offline data. For completeness, we additionally include comparison to Critic Regularized Regression (CRR) (Wang et al., 2020), an offline RL method with competitive performance on these tasks. CRR is trained the full RL. Unplugged datasets (i.e., combined suboptimal and expert datasets) with reward labels.
-witopppsswinup)) to complex I of
Results. Figure shows the comparison results. TRAIL outperforms temporal extraction methods on both -owdiimesioonal (e.g., cartpole- partpole-swingup) and high- nigh-dimensional (humanoid d run) tasks. Additionally, TRAIL performs similarly or better than CRR 4 out of the tasks despite not using any reward labels, and only slightly worse on humanoid run and walker-war test the robustness of TRAIL when the offline data is highly suboptimal, we further reduce the size and quality of the offline data to the bottom 5% of the original RL Unplugged datasets. As shown in Figure 6 in Appndix D, the performance of temporal skill extraction declines in fish-sw i.m, wall 1 -er--staaeee nd, and wa 1 ker-wall k due to this change in offline data quality, whereas TRAIL. maintains the same performance as when the bottom 80% data was used, suggesting that TRAIL is more robust to low-quality offline data.
To
This set of results suggests a promising direction for offline learning of sequential decision making policies, namely to learn latent actions from abundant low-quality data and behavioral cloning in the latent action space on scarce high-quality data. Notably, compared to offline RL, this approach is applicable to settings where data quality cannot be easily expressed through a scalar reward.
CONCLUSION
We have derived near-optimal objective for learning a latent action space from suboptimal offline data that provably accelerates downstream imitation learning. To learn this objective in practice, we propose cransttin--eparrmeeriiee actions for imitation learning (TRAIL), a two-stage framework that first pretrains factored transition model from offline data, and then uses the transition model to reparametrize the action space prior to behavioral cloning. Our empirical results suggest that TRAIL can improve imitation learning drastically, even when pretrained on highly suboptimal data (e.g., data from random policy), providing new approach to imitation learning through a combination of pretraining on task-agnostic or suboptimal data and behavioral cloning on limited expert datasets. That said, our approach to action representation learning is not necessarily specific to imitation learning, and insofar as the reparameterized action space simplifies downstream control problems, it could also be combined with reinforcement learning in future work. More broadly, studying how learned action reparameterization can accelerate various facets of saaring-bbsee control represents an exciting future direction, and we hope that our results provide initial evidence of such a potential.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
We thank Dale Schuurmans and Bo Dai for valuable discussions. We thank Justin Fu, Anurag Ajay, and Konrad Zolna for assistance in setting up evaluation tasks.
REFERENCES
Pieter Abbeel and Andrew Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference Machine learning, pp. 2004.
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International Conference on Machine Learning, pp. 22-31. PMLR, 2017.
M Mehdi Afsar, Trafford Crump, and Behrouz Far. Reinforcement learning based recommender systems: A survey. arXiv preprint arXiv:2101.0622 2021.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline prim- itive discovery for accelerating offline reinforcement learning. arXiv preprint 0et:22000..61... 2020.
Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable representation learning for imitation learning via bi-level optimization. In International Conference on Machine Learning. pp. 367-376. PMLR, 2020.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Daniel Berend and Arych Kontorovich. On the convergence of the empirical distribution. arXiv preprint arXiv: 1rXiv::205.6711, 2012.
Pablo Castro and Doina Precup. Using bisimulation for policy transfer in mdps. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, 2010.
Jonathan Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 34, 2021.
Thomas G Dietterich al. The maxq method for hierarchical reinforcement learning. ICML, volume 98, pp. 118 126. Citeseer, 1998.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint D0XXvv/882.0000000 2018.
Roy Fox, Sanjay Krishnan, Ion Stoica, and Ken Goldberg. Multi-level discovery of deep options. arXiv preprint arXiv:1703.0000
2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp. 2170-2179, PMLR, 2019.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint 161 1.05577, 2016
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine, Sergio Gomez Colmenarejo, Kon- rad Zolna, Rishabh Agarwal, Josh Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI un- plugged: Benchmarks for offline reinforcement learning. arXiv e-prints, pp. arXiv-2006, 2020.
Kourosh Hakhamaneshi, Ruihan Zhao, Albert Zhan, Pieter Abbeel, and Michael I Laskin. Hierarchi- cal few-shot imitation with skill transition models. arXiv preprint arXiv:2 00.00098...11. 2021.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Leam- ing Representations, 2018
10

---

Published as a conference paper at ICLR 2022
Haoming Jiang, Bo Dai, Mengjiao Yang, Tuo Zhao, and Wei Wei. Towards automatic evaluation of dialog systems: A model-free off off-po evaluation approach. arXiv preprint arXiv:2102.10242, 2021.
Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Gonzaez- Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. Compile: Compositional imitation learning and execution. In International Conference on Machine Learning, pp. 3418-3428. PMLR, 2019.
Sanjay Krishnan, Roy Fox, Ion Stoica, and Ken Goldberg. Ddco: Discovery of deep continuous options for robot learning from demonstrations. In Conference on robot learning, pp. 418-437. PMLR, 2017.
Tejas Kulkarni, Karthik Narasimhan, Ardavan Saeedi, and Josh Tenenbaum. Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems, 29:3675-3683, 2016.
Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. arXiv preprint arXiv: 00949, 2019.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives open problems. arXiv preprint arXiv:2005.01643, 2020.
Jinxin Liu, Donglin Wang, Qiangxing Tian, and Zhengyu Chen. Learn goal-conditioned policy with intrinsic motivation for deep reinforcement learning. arXiv preprini arXiv:2104.050433 2021.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, pp. 1113- 1132. PMLR, 2020.
Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. arXiv preprint arXiv:2105. 2021.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforce- ment learning. arXiv preprint arXiv: 000.. 2018a.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning. arXiv preprint arXiv:/810.01257, 2018b.
Ofir Nachum, Michael Ahn, Hugo Ponte, Shixiang Gu, and Vikash Kumar. Multi-agent manipula- tion via locomotion using hierarchical sim2real. arXiv preprint 1xXiv:9908.05224 2019.
Ronald Parr and Stuart Russell, Reinforcement learning with hierarchies machines. Advances in neural information processing systems, pp. 1043-1049, 1998.
Xue Bin Peng. Michael Chang, Grace Zhang, Pieter Abbeel, and Sergey Levine. Mep: Learn- ing composable hierarchical control with multiplicative compositional policies. arXiv preprint arXiv:1905.0980 099008 2019.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. arXiv preprint arXiv:2010. 1944, 2020.
Karl Pertsch, Youngwoon Lee, Yue Wu, and Joseph J Lim. Guided reinforcement learning with learned skills. In Self-Supervisionn for Reinforcement Learning Workshop-ICLR 2021, 2021.
Dean Pomerleau. Alvinn: An autonomous land vehicle in neural network. Technical report, CARNEGIE-MELLON UNIV PITTSBURGH PA ARTIFICIAL INTELLIGENCE AND PSY- CHOLOGY .... 1989.
Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John
Wiley Sons, Inc., 1994.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imitation learning using variational models. Advances in Newral Information Processing Systems, 34, 2021.

---

Published as a conference paper at ICLR 2022
Ali Rahimi, Benjamin Recht, al. Random features for large-scale kernel machines. In NIPS, volume 3. pp. 5. Citeseer, 2007.
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2017.
Zhiyuan Ren and Bruce Krogh. State aggregation in markov decision processes. In Proceedings of the 41st IEEE Conference Decision and Control, 2002., volume 4, pp. 3819-3824. IEEE, 2002.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. Workshop and Conference Proceedings, 2010.
JMLR
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction no-regret online learning. In Proceedings of the fourteenth international con- ference on artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Stefan Schaal. Is imitation learning the route humanoid robots? Trends in cognitive sciences, (00:233-442, 1999.
Tanmay Shankar and Abhinav Gupta, Learning robot skills with temporal variational inference. In International Conference Machine Learning, pp. 8624-8633 PMLR, 2020.
Tanmay Shankar, Shubham Tulsiani, Lerrel Pinto, and Abhinav Gupta. Discovering motor programs by recomposing demonstrations. In International Conference on Learning Representations, 2019.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint arXiv: 1977v0990..01657, 2019.
Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. arXiv preprint arXiv:2011.1002 2020.
Satinder P Singh, Tommi Jaakkola, and Michael Jordan. Reinforcement learning with soft state aggregation. Advances in neural information processing systems, pp. 361-368, 1995.
Martin Stolle and Doina Precup. Learning options in reinforcement learning. In International Symposium on abstraction, reformulation, and approximation, pp. 212-223. Springer, 2002.
Richard Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181- 211, 1999.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmalcki, Josh Merel, Andrew Lefraane, al. Deepmind control suite. arXiv preprint arXiv:1801.00690 2018.
Stephen Tu, Alexander Robey, and Nikolai Matni. Closing the closed-loop distribution shift in safe imitation learning. arXiv preprint arXiv:2102.09161 2021.
Alexander Sasha Vezhnevets, Simon Osindero. Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. In International Conference Machine Learning, pp. 1540-3549, PMLR, 2017.
Xiaofei Wang, Kimin Lee, Kourosh Hakhamaneshi, Pieter Abbeel, and Michael Laskin. Skill pref- erences: Learning to extract and execute robotic skills from human feedback. arXiv preprint M.XV:22008.05382. 2021.
Ziyu Wang. Alexander Novikov, Konrad Zolna, Jost Tobias Springenberg. Scott Reed, Bobak Shahriari. Noah Siegel, Josh Merel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. arXiv preprint arXiv:2006.1 250.:2000../1344 1544, 2020.
David Warde-Farley, Tom Van de Wiele, Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through non- con-paramttric discriminative rewards. arXiv preprint arXiv: 1359, 2018
12

---

Published as a conference paper at ICLR 2022
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: 111111991..1331,,2222 2019.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.
Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforce- ment learning. arXiv preprint arXiv:2011.07213, 2020.
Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcchre, Ziyu Wang, Yusuf Ay- tar, Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience. arXiv preprint arXiv:201 I. 13885, 2020.
13

---

Published as a conference paper at ICLR 2022
Appendix
A PROOFS FOR FOUNDATIONAL LEMMAS
Lemma 4. If γH and i2 are two policies in M and (s) and (s)(s) are the state visitation
distributions induced by policy Tı and T2 where d= == (1- -1-7) Ei= Pr = M M). Define Diff (T2, T1) DrDDDrd (dda)lla)) then
Difs(
1-7 T-
7 Erram Erran (TI, T2, T2, 7),
(6)
where
1
B (s)) Eria (TI, n2. T) (h (7 -((')))))) az)ll
2 E'ES
the TV-divergence between T o π1 d and T o T2 o
Proof. Following similar derivations in Achiam et al. (2017); Nachum et al. (2018b), we express
(d d in linear operator notation:
1
Diff T2, T (rv (I 2 -101--NU--TU))------1000000
plp (8)
where II, II, are linear operators S + S x A such that I1, a) = (a(a))) and is an all ones
row vector of size (S). Notice that d may be expressed in this notation as (1 - (
We may re-write the above term as
1 () 7TI1
((T -II (TII) - - -I
=Y 1|(1 77112) -1112) TII, )
(9)
Using matrix norm inequalities, we bound the above by
? 2 11(T - 77112)
1 - d
(10)
Since Th2 stochastic matrix, I1(T - 7TT112) 111,00 S Ei 7 (1 we bound the above by
Thus,
2(1 7) ( -TI, 1-7
7 Erraz (2,,
12, T),
(11)
and so we immediately achieve the desired bound equation
The divergence bound above relies on the true transition model which not available to us. We now introduce approximate transition model T to proxy ErIÆ=I (T1, T2, T).
Lemma 5. For πi and n2 two policies M and any transition model (-|s,a) a) we have, Erra (2, Proof.
n2; T) < (Drraa/
Err (T., n2, T). (12)
Erig (TI, 12, T)
((ssss a (1 (T(s'ss, T(s'ss,,
Σ
(13)
2
FES
(14)
2
S'ES
DEA
Σ
Σ
E (TTslss.)))77///))))) m((d)) +
- (55 (4)))
2
.'ES
aEA
(15)
1
Exadr - -() (((s)) Errar (T1, T2, T) (16)
2
E'ES
aeA
1
Eaner (T(('s T (m)(ss =2(a|s))]]
[ + Errgs T2, (17)
2
JES
|A
a))] ErI.5 (TI, T),
(18)
14

---

Published as a conference paper at ICLR 2022
and we arrive at the inequality desired where the last step comes from Div(T(s.a)(((ss.a) Eves (s TTs'ss,
Now introduce representation function o x A Z and show how the error above may be reduced when T(s, a) = T2(s, ((,, a)):
Lemma 6. Let & 55x A for some space and suppose there exists Tz S + A(S) such that T(s, a) = (2(s, for all E S, A. Then for any policies 71; 22.
Erræ-1 (TT, n2, T) S Band" (DIv((T1,2112.22.
(19)
where (z|s) the marginalization of T&k onto Z:
(2|0) ==
E
me(als)
(20)
a aEA.z=o(s.a)
for all 2 E k E (1,
Proof.
E Egnd (e -T(s'ss, a2)]
(21)
2
s'ES
Σ
T2(6'|s, ((s))s,,
Σ
Ta(s'ss, ((5 (()))(((() (s)
2
S'ES
aES,
sES,eEA
Σ
Σ
Tz(3'{8,2)
Σ (a(assdd
Σ
Tz(5'(s,z)
Σ
(2(a(s)d'
2
S'ES
eES
aEA.
SES
ae
((x.0)=2
((.n)==
1
Σ
Σ
T2(5"1s,2)m,,22111/)
E
2
S'ES
sES,2EZ
sEB,cEZ
1
Σ
Exnarr
Σ
(5 2)) -2,,(((8))
(22)
2
S'ES
LEEZ
1
2 Exaut
Σ
T2(s'1ss 2) (z)())))) n2, (2|s))
(23)
LEEZ d'ES
1 E
Σ -2 LEEZ
(24)
2
Eggddr [Drv (T1,2 and we arrive at the inequality desired.
(25)
Lemma 7. Let d E A) be some state-action distribution, & S x A Z, and mz S - A(Z). Denote Ta as the optimal action decoder for d, 6z d(s, a) 1/2 = Φ(s,
Ta (ala, 2) Eves d(s, a') d(s,a')
1[2 Φ(s,
and Fet z as the marginalization of ah o onz onto
Ta" 1 [2|s)) ==
Σ
(Ta (2)ss
Σ
Σ
Tar (als,,)Tz(((()))
(EA,)=0(3,a)
aE FEZ
FEZ
Then we have
Taa 2 (2|s) (z(s)))
(26)
for all E Z and E S.
15

---

Published as a conference paper at ICLR 2022
Proof.
Ta ((2|))=
Σ
Σ
Fa (a|s, (2)s)
(27)
oεA,==((e,)) FEZ
d(s,a) 1| o(s
Σ
Σ
Lce 1[ o(s (2(())
(28)
aEA,2=((3,0) BEZ
d(s,a) 1/2 11==111=1:0 ((8.0) a)
Σ
CdEAd(s,a) 1/2 o(s a') (2|8)
(29)
(EAis=((E.0)
d(s,a) 1|== ((..a) a)
(2(ss)
Σ
(30)
Ewedd(s,d) 1/2 - o(s
aEA.==o(s.a)
= >z(2s))
(31)
and we have the desired equality.
Lemma 8. Let S + ((Z) be latent policy in Z and Ta S x Z - A be an action decoder, Ta,2 be the marginalization of Ta o Tz onto Z:
No. (z)s))
= Σ
(Ta (als)
Σ Σ
())((s)
at aEA,,==(s.a)
aEA, ((5,a) BEZ
Then for any E S we have
Drv (Ty(TT(s))[Tag (s)) S max Drv (s, (Ta 2),
(32)
=EZ
where aa- is the optimal action decoder defined in Lemma 7 (and this holds for any choice of d from Lemma 7).
Proof.
(33)
[ - (a..(22(s))
(34)
ZEZ
Σ
TT (z(s))
Σ
Σ
Ta (als, (8)
(35)
2
ZEZ
aEA.==0(spp SEZ
Σ
TZ (2(s))
Σ
Σ
(Ta(als; Ta (als, za (als, 2)) x2 (2)s)
(36)
ZEZ
aEA.:=ó(s.a) EE7
Σ
Σ Σ (Fa(a)s, na- (als, E)) T (21s)
(by Lemma 7)
(37)
sEZ
(=s,,==(s.0) BEZ
E Ez)zz(s)
Σ
Σ
(Ta To (a|s, 2)1
(38)
2
=EZ A,EA, ((s.n)
Esnrz(()
Σ [als, Ta (als,
(39)
LaEd
=b (s) [DV (Ta 2) Ta (s,
(40)
< max DrV (T (s, 2) Ta- (s. =EZ
(41)
(42)
and we have the desired incquality.
Lemma 9. Let Tı, be the marginalization of γi onto as defined in Lemma 6, and uz, Na.
Ta,Z be as defined in Lemma 8, and let Ta- 2 be defined in Lemma 7. For any E S we have Drv (TI, [s)] (s)) max DTV(T (s, a)a (s) + DTv(TI.Z()))))))))))
(43)
16

---

Published as a conference paper at ICLR 2022
Proof. The desired inequality is achieved by plugging the inequality from Lemma 8 into the follow- ing triangle inequality:
(T) & (Iv(T2(5)]=m..2))) + Drv(m))(()))T2(())
(44)
Our final lemma will be used to translate on-policy bounds to off-policy.
Lemma 10. For two distributions Pl, P2 E ASS) with pi (1() p2(s) 0, we have,
Ep, [b()]] (1 (xe(oullez))) (Ep
(45)
Proof. The lemma is straightforward consequence of Stuchy-Schwartz
Ep, [h(s)] Ep, [h()) + (Ep, [h(s)] Bp, [h())])
(46)
= =p [x()] + Σ (i({) p2(0)
P2(8) p2(5)=n(s)
(47)
8ES
(pi(s) P2(s))) 1
- EEpa [x()] +
(48)
Pz(()
BES
BES
Ep, [h(s] + xa (p )
Epg (h()))
(49)
Finally, to get the desired bound, simply note that the concavity of the square-root function implies Ep [h(s)] Ppp (v()))) S (ss))
B PROOFS FOR MAJOR THEOREMS
B.I PROOF OF THEOREM 1
Proof. Let 72 == Ta o mz, we have (z)2(= = no,z (x|s) EnEA, (z)(s))) By
plugging the result of Lemma 9 into Lemma 6, we have
Errgy (TI, 7) < E E
max #EZ
Drv =T) (
(50)
By plugging this result into Lemma 5, have
Erray (TI, T2, T) S (s.)) ) /D
(51)
+ Eandri
max Drv (Ta (Ta- (s. ZE
(52)
+ Egndrt [Div(T1.2
[2()]
(53)
By further plugging this result into Lemma 4 and let म = T., we have:
7|A|
Dif T2, ) <
((s.a))((<1, D (1.
Eg [max Drv () 2))
-
#EZ
7
Ex DTV (z)
(54)
1-
Finally, by plugging in the off-pooicy results of Lemma 10 to the bound in Equation 54 and by ap- (ss, a))),
plying Pinsker's inequality Drv(T(s, (z)), a()) have
Diff T2, T.) C1
V
E
(e.j (2 ((.,
(T(
+ C2
Eyngutt [max (TTa (s, 2))1 =))
=E2
ss cof ②)) + JDD
+ C3
Eswar. (III.(TT) (8)1
(55)
const( Φ) + (BC.
17

---

Published as a conference paper at ICLR 2022
where C1 aA[(1 (1 + Dys ) C2 7(1 71 +
(d= ) and C3
7(1 7)- Since the max not tractable Egndor [max,ez DKL (8, (s, using
in practice, we approximate B(a,a)aaur DKI(FFe (s, d(s, [ (s, o(s, a)) a])]], which reduces to (DE 6) with additional constants. We now arrive at the desired off policy bound in Theorem 1.
B.2 PROOF OF THEOREM? 2
Lemma 11. Let E A(f)...,, kk})) be distribution with finile support. Let Pn. denote the empirical estimate of p f.rrmm i.idd samples X p. Then,
k 2 1=1
1 2
k
Ex[PTv(pllen) S
p())
<
(56)
Proof. The first inequality is Lemma 8 in Berend & Kontorovich (2012) while the second inequality is due to the concavity of the square root function.
Let D == {s,,,,)) ). be i.i.d samples from a factored distribution x(s, a) == Lemma 12.
((s))(a)s)) for p A(S) : S (A)) Let be the empirical estimate of p in DD and a be the empirical estimate of π in D. Then,
SIIA
Ep E E (Dr(p s
(57)
Proof. Let & be the empirical estimate of in D. We have, E E (1)(
1
Σ p(s)
|T(als)
[[ass)]
(58)
8,0
()
a(s,a)
Σ p(s)
(59)
p(s)
p(s)
8,8
(()
a(s.a)
((s.a)
z(s, a)
$
Σ
p(s)
Σ
p(s)
2
p(s)
P(s)
p(s)
p(s)
8,8
8,a
(60)
((,,a) a)
a(s,a)
Σ
p(s)
+ Drv(cll()
(61)
2
P(s)
p(s)
8,0
Σ
p(s)
1
Σ a(ssa) + Drv(=lle)
(62)
2
p(8)
p(s)
1
Σ p(s)
1
1
((s)
(av(
(63)
2
((s)
((s)
Drv (pllp) + Drv(xll))
(64)
Finally, the bound in the lemma is achieved by application of Lemma 11 to each of the TV diver- gences.
To prove Theorem 2, we first rewrite Theorem as Diff π.)= < (L)(Ø) + (2)()) + C3- Es-ar. (T. where (1) and (2) are the first two terms in the bound of Theorem 1, and C3 12.. The result in Theorem result of Lemma 12.
(65)
then derived by setting á ΦiTere and az = Tooret and using the
Note that the above sample analysis can be extended to the continuous latent action space character- ized by Theorem 3 as follows.
Theorem 13. Let doret == OPT o (Doff) and Torel be the latent BC policy with respect to Doret- Let be the dimension of the continuous latent actions and [1oo be the Los norm of Qore for any s, We have
2
Ep.. (Dff 0, R.)] S (1)())rd) (2)(dorat) C4 dll0
n+11'
where (1), (2), and Cs are the same as in Theorem 3.
18

---

Published as a conference paper at ICLR 2022
Proof. We use μ E Rdx[s] to denote the optimal setting of 0 which yields zero 4 1o-mrrm of 0 E Benddd [[0g Φ(s, i.e.. 80
,, B.=Bowr. (a) a)].
(66)
According to Theorem 3, we want to bound the I no-orrm of 8 Rendd ama. [10g - ((8, a)}²] evaluated at the approximate solution ji E [Rdx [S] with respect to finite dataset D i.e.,
20
i. Ba- (Ø(s.a). a)),
(67)
with the convention that P. = if s does not appear D To this end, we have the following derivation, which uses EE,, to denote the expectation over realizations of ja due to n-size draws of the target dataset =
0 00 o=p
|.] En [Badd" [lles ENNn, [o(s, a)ll
En
E Band" arn, [00 o(s, a))
(68)
En [Bsuat [||is M4s [1]]
(69)
= Egnd" [E, [I]e, Palla]]-
(70)
We now split up the inner expectation based on the number of times ls that appears in D
E Egwdd
[E,,
[1les Mss /=]|1]]]
Egand
Pr(count(s) k]] E& Ille, Malla]
(71)
Es
Pr[count(s) = k] E& ll malll]]
(72)
(73)
where Eg denotes the expectation over realizations of Aúo over k-size draws of a ~ Tx(s).
By
standard combinatorics, we know
Pr|count(s) k] (x)
dt (s))
(74)
k
Furthermore, for = 0, we have
ER [l - p=lli]² = || 17 S dL 11b11
(75)
while for l > 0, since E& [R=] = Ms, we have
dl 2d 1d
E& [ll Ms 111 s d. E& [llĀ_ [=1 = d Vark [p,]] S
(76)
+1
Combining equations 74, 75, and 76 we have for any &e 2 0
d3 (s) Pr|count(s) = k]] E& [ll uollal2 v
2d 11
11
de (s) k++ (1 - d= (s))
k + 1
k
2d ,
(n + 1 kk+1)
dtt (5) (s) dth ds)))
n + 1
and so by the binomial theorem,
2d 11/1
(5)(s) PP-PPPo(/ots] = k:]
Ex [lli, - mssl1)² S
(78)
n +
Plugging the above into equation 72 we deduce
2|S1
Ex-ar Em, Alle - Mellall s dl 0o
(79)
n 1
and we have the convergence rate as desired.
B.3 PROOF OF THEOREM? 3
Proof. The gradient term in Theorem with respect to a specific column 0, of A may be expressed
as
i
Bandy an (n) ((5,a)) a
00,
2 (d (s)o(s, a)] + 2d (s)0. (() (d (s)o(s, a)] + 2E,_,, (s) 22),
(80)
19

---

Published as a conference paper at ICLR 2022
and so,
To(s)T θ
Bands (w)r)) - o(3,a))))
00,
2n (d' (s) ] + 2R 1dt [dd (s)
(81)
Summing over 8 E S, we have:
Σ w(s') Rind o0
00,
#ES
aa 2B gn a (5),2=0, (2((11)))
(82)
Thus, we have:
1
Erra- (T, nT. T)
LE Egdd (a0 -0 + Tz(s' s,
2
x'E8
1
w(s')
E Bgnddd 053 a))
4
00,
e'ES
EES
1
4 ||||||||>>
E Egnd (o) 008 a))
(83)
30
Then by combining Lemmas 4, 5, 10, and apply Equation 83 (as opposed to Lemma 6 as in the tabular case), we arrive at the desired bound in Theorem 3.
C EXPERIMENT DETAILS C.I ARCHITECTURE
We parametrize two-hidden layer fully connected neural network with 256 units per layer. A Swish (Ramachandran et al., 2017) activation function is applied to the output of each hidden layer. We use embedding size 64 for AntMaze and 256 for Ant and all DeepMind Control Suite (DMC) tasks after sweeping values of 64, 256, and 512, though the latent dimension size as long as is not too small (i.c., 2 64). The latent skills temporal skill extraction require much smaller dimension size, e.g.. 10 reported by Ajay al. (2020); Pertsch et al. (2021). We tried increasing the latent skill size for these work during evaluation, but found the reported value 8 work the best. We additionally experimented with different extend of skill extraction, but found the previously reported = 10 to also work the best. We implement the trajectory encoder in OPAL, SkiLD, and SPiRL using a bidirectional LSTM with hidden dimension 256. We use 3 = 0.1 for the KL regularization term in the 3 VAE of OPAL (as reported). We also use 0.1 as the weight for SPiRL and SkiLD's KL divergence terms.
found TRAIL to be relatively robust to
C.2 TRAINING AND EVALUATION
During pretraining, we use the Adam optimizer with learning rate 0.0003 for 200k iterations with batch size 256 for all methods that require pretraining. During downstream behavioral cloning, learned action representations are fixed, but the action decoder is fine-tuned on the expert data as suggested by Ajay al. (2020). Behavioral cloning for all methods including vanilla BC is trained with learning rate 0.0001 for 1M iterations. We experimented with learning rate decay of down- stream BC by factor of 3 at the 200k boundary for all methods. We found that when the expert sample size is small, decaying learning rate can prevent overfitting for all methods. The reported results are with learning rate decay on AntMaze and without learning rate decay on other envi- ronments for all methods. During the downstream behavioral cloning stage, we evaluate the latent policy combined with the action decoder every 10k steps by executing Ta Rz in the environment for 10 episodes and compute the average total return. Each method is run with 4 seeds where each seed corresponds to one set of action representations and downstream imitation learning result on that set of representations. We report the mean and standard error for all methods in the bar and line figures.
C.3 MODIFICATION TO SKILD AND SPIRL
Since SkiLD (Pertsch et al., 2021) and SPiRL (Pertsch et al., 2020) are originally designed for RL opposed imitation learning, replace the downstream RL algorithms of SkiLD and SPiRL by behavioral cloning with regularization (but keep skill extraction the same as the original methods). Specifically, for SkILD, apply KI. regularization term between the latent policy and the learned
20

---

Published as a conference paper at ICLR 2022
skill prior in the suboptimal offline dataset during pretraining. and another KL regularization term between the latent policy and a learn "skill posterior" on the expert data as done in the original paper during downstream behavioral cloning. We do not need to train the binary classifier that SkiLD trains to decide which regularizer to apply because we know which set of actions are expert versus suboptimal in the imitation learning setting. For SPiRI., we apply the KI. divergence between latent policy and skill prior extracted from offline data (i.e., using the red term in Algorithm of Pertsch al. (2020)) as an additional term to latent behavioral cloning.
C.4 DATASET DETAILS
AntMaze. For the expert data in AntMaze, we use the goal-reaching expert policies trained by Ajay et (2020) (expert means that the agent is trained to navigate from the one corner of the maze to the opposite corner) to collect 10 trajectories. For the suboptimal data in AntMaze, use the full D4RL datasets antmaze- lwwgllaaaeeeeeeeeeeeeeeee diverse-
mede-m-ddduppppp v0,
and wtw.maaze-meiin---l-...
Ant. For the expert data in Ant, we use small of expert trajectories selected by taking either the first 10k or 25k transitions from ant ant-expert-vve -v0 in D4RL, corresponding to about 10 and 25 expert trajectories, respectively. For the suboptimal data in Ant, we use the full D4RL datasets ant-medi um-v0, ant-medi rm-replyy-00, and aat-rradden ant-random-vO.
RL Unplugged. For DeepMind Control Suite (Tassa al., 2018) set of tasks, Unplugged (Gulcehre et al., 2020) dataset. For the expert data, we take TО of the trajectories whose episodic reward is among the top 20% of the open source RL Unplugged datasets following the setup in Zolna et al. (2020). For the suboptimal data, we use the bottom 80% of the RL Unplugged dataset. Table records the total number of trajectories available in RL Unplugged for each task (80% of which are used as suboptimal data), and the number of expert trajectories used in our evaluation.
the RL
Task
# Total
# D
cartpole-swingup
40
2
cheetah-run
300
3
fish-swim
200
1
humanoid-run
3000
53
walker-stand
200
4
walker-walk
200
6
Table I: Total number of trajectories from RL Unplugged (Gulcehre et al., 2020) locomotion tasks used to train CRR (Wang et al., 2020) and the number of expert trajectories used to train TRAIL. The bottom 80% of # Total is used to learn action representations by TRAIL.
21

---

Published as a conference paper at ICLR 2022
D ADDITIONAL EMPIRICAL RESTULS
D.I ADDITIONAL BASELINES FOR RL UNPLUGGED
TRAIL (energy)
TRAIL (incar)
Bascline BC
OPAL (1=10)
SPURL (1=10)
Sail (1110)
CRUR
1466
1200
600
www.tpppoooowgggg..... 5%
dro-cheetab--un
dm- fill-swwim 59%
dm- wwo-hhmmnncideeee
1020
220
3000.
Figure 6: Average task rewards (over 4 seeds) of TRAIL EBM (Theorem 1), TRAIL linear (The- orem 3), and OPAL, SkiLD, SPiRL trained on the bottom 80% (top) and bottom 5% (bottom) of the RL Unplugged datasets followed by behavioral cloning in the latent action space. Baseline BC achieves low rewards due to the small expert sample size. Dotted lines denote the performance of CRR (Wang et al., 2020) trained on the full dataset with reward labels.
D.2 FRANKAKITCHEN RESULTS
Expert D
kitchen--omplete
kitchen-complete
SuboptimalD
kitchen-mixed
kitchen-partial
TRAIL (EBM)
TRAIL (linear)
SkilD (t~10)
SkilD (t=1)
SPIRL (t=1)
OPAL (t=10)
Baseline BC
0
25
50
75
100
25
50
75
100
Figure 7: Average rewards (over 4 seeds) of TRAIL EBM (Theorem 1), TRAIL linear (Theorem 3), and bascline methods pretrained on kitchen-mix xed and kit kitchhn-part ial from D4RL to imitate kit than SKiLD and OPAL. with temporal abstraction over 10 steps.
TRAIL linear without temporal abstraction performs slightly better
22

---

Published as a conference paper at ICLR 2022
D.3 DISCRETE MAZE RESULTS
Vanilla BC
TRAIL
1,
0.45
U'X
0.40
0.35
TT
0.30
target demos (N 20)
Figure 8: Average task rewards (over seeds) of TRAIL EBM (Theorem 1) and vanilla BC (right) in discrete four-room maze environment (left) where an agent is randomly placed in the maze and tries to reach the target "T" TRAIL learns discrete latent action space of size from the discrete original action space of size 12 on 500 uniform random trajectories of length 20 shows clear benefit over vanilla BC on expert data.
We conduct additional evaluation in Figure 8, an agent is randomly placed into four-room environment, and the task nav- igate to the target "T" The task reward is at "T" and 0 elsewhere. There are 12 discrete ac- tions corresponding to rotating clockwise by 90, 180, 270, 360 degrees, rotating counterclockwise by 90, 180, 270, 360 degrees, moving forward by 1 or 2 grids, and moving backward by or 2 grids (the action space is artificially blown up as suggested by the reviewer). TRAIL is pretrained on 500 trajectories of length 20 with uniform action selection. The expert demonstration always navigates to the target 'T' from any random starting location. TRAIL's latent action dimension is set to 4. We that TRAIL with smaller latent action space offers benefits over vanilla BC.
environment with tabular state and action spaces. As shown
23

---

Published as a conference paper at ICLR 2022
E ABLATION STUDY
Export 1
expert 10 kzaje
expezt 10 trate
expezt 10 tzols
expezt 10 trajs
Suboptimal Der
antzaze-acdium sedeeevvrrdddddddeeeeeee
antzaze-acdiup www.-adddddddddeeeeeeeeeee
TRAIL tano-False
TRAIL tume-True
TRAIL emb
TRAIL embod-512
Finelune BC
Daselino
Figure 9: Ablation study on action decoder finetuning, latent dimension size, and pretraining base- line BC on suboptimal data in the AntMaze environment. TRAIL with default embedding dimension 64 and finetuning the action decoder corresponds to the second row. Other dimension size (256 and 512) lead to worse performance. Finetuning the action decoder on the expert data some small benefits. Pretraining BC on suboptimal data before finetuning expert does not lead to significantly better performance.
Repars D"
ant-expart 15k
ant-expart 25k art-madion
And ant-aopart 18k
ant-eepprrt 25%
ant-expart 10k
ant-exxert 25k
ant-nedfon
Cheert unbod- 200
TEAL (1500)
Plotaan 30
Figure 10: Ablation study latent dimension size in the Ant environment. TRAIL is generally robust to the choices of the latent action dimension (64, 256, 512) for the Ant task.
Repart Dx
AGC snt-oxpact 15k
ant-expert 25k
ant-oxpars 18k
ant-eepars 25%
ant-expert 10k
ant-expezz 25k
Subh "p
TRAIL IEIS teszethiae
SKALL danur taaandske
TAN (9900 soeettee
TRAIL Seaat Thee
Figure 11: Ablation study finetuning the action decoder in the Ant environment. Finetuning the action decoder leads to slight benefit.
24

---

Published as a conference paper at ICLR 2022
F VISUALIZATION OF LATENT ACTIONS
PCA Raw Action
TSNE Raw Action
PCA Latent Action
TSNE Latent Action
middes
runden fendaxe
1.5
30
20
20
0.5
10
0-
0.0
-10
-0.5
-10
-20
-20
-30-
reddm
madezs-mpho
-1.5
-00
-2
2 -20 20
50
Figure 12: PCA and t-SNE visualizations of the random, medium-replay medium, and expert D4RL, Ant datasets. Without action representation learning (left), the distinction between expert and suboptimal actions is not obvious. The latent actions of TRAIL (right), on the other hand, results in the expert latent actions being more visually separable from suboptimal actions.
25


Reviews:
[{"review_id": "wAgzrlXuDBa", "reviewer": "Reviewer_xchs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "# Strengths:\n\n(1) The method is tested in several challenging environments. In all of them, TRAIL performs better than or is on par with the baselines.  \n\n(2) The paper is well written, and the theory appears to be sound\n\n# Weaknesses: \n\n(1) The bound of the imitation learning error derived in theorem 1 and 2 relies on the fact that $|Z| < |A|$. In all of the experiments, the action embedding ($> 64$) is bigger than the original action space ($=8$ for the ant). Could the authors provide an intuition of why such big action embeddings are required? This is somewhat contradictory to the theory.\n\n(2) The comparison to the offline reinforcement learning setting is interesting. \nHowever, the big promise of offline RL is to extract an optimal policy of the MDP induced by the offline dataset. That means, in principle, offline RL should be able to recover an optimal policy just from random data (given a sufficient state and action space coverage in the offline dataset) or a multitude of behaviors.\nI would like to see a small paragraph in section 5.3 discussing this setting and how this affects BC. \n\n(3) In C.2, the authors write that the action decoder was fine-tuned on the expert data while keeping the action representations fixed. How does that affect the final performance?\n\n# Needs clarification and general comments\n\n(1) The authors claim that “[a] good action representation should enable good imitation learning.” \nCould the authors give an intuition (independent of the empirical results) on why a good BC action representation should emerge from a good transition-based representation?  \n\n(2) In theorem 1 $\\pi_{Z}$ is defined as $\\pi_{Z} \\colon S \\times Z \\to \\Delta(A)$. If I am not mistaken, that should be $\\pi_{Z} \\colon S \\to Z$\n\n(3) I think in the sample complexity paragraph after (1) and (2) $J_{T}$ and $J_{DE}$ are missing, respectively.\n\n(4) In equation 5 $\\psi(s^{\\prime})$ appears for the first time. However, I am failing the find the definition of $\\psi(s)$\n\n(5) The BC baseline is only trained with the expert data. Is there a way of incorporating the suboptimal data in the training of the vanilla BC baseline? Maybe by first training with the suboptimal data and later fine-tuning on the expert data.\n\n(6) I would like to see a plot showing the effects of the different design choices (size of the latent embedding, fine-tuning the action encoder) on the performance.\n\n(7) To empirically prove the theoretical claims, the authors could run another experiment in which they artificially blow up the action dimensionality and check if they can learn a much smaller embedding.\n", "summary_of_the_paper": "The paper proposes a method to accelerate behavioral cloning (BC) (especially in the low data regime) by utilizing a (much larger) auxiliary dataset of suboptimal behaviors. \nThe authors claim that learning a good latent action representation (in this case, by learning a transition-based action representation from the auxiliary dataset) should make the imitation learning problem easier. \nThe training is divided into two phases: (1) learning a factorized transition model in conjunction with an action embedding and action decoder, (2) learning a latent policy via the BC objective hoping that the transition-based representation accelerates the BC learning.", "summary_of_the_review": "At the moment, I tend towards a weak reject because I don’t find the connection between theory and empirical evaluation too convincing. The reduction in the BC error shown in theorem 1 and 2 are based on the assumption that $|Z| < |A|$. However, if I am not mistaken, that is not the case in any of the experiments where $|Z|$ is usually much bigger than $|A|$. \nOn the other side, the empirical results show a clear improvement over the baselines and vanilla BC, indicating that the learned latent embedding is helpful. \nI am willing to change my score if the authors strengthen the connection between the theory and the empirical results. For instance, they could explain why the transition-based representation might help BC, although the latent embedding is bigger than the original action space.\n\n### Decision due to rebuttal\n\nThe authors addressed all of my concerns sufficiently well. Therefore I am happy to raise my score from a 5 to a 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Vf_rXKuXKG", "reviewer": "Reviewer_M6PX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "## Strengths\n\n- The paper proposes a theoretically justified method to learn from expert and sub-optimal data. The method is quite flexible in the sense that it requires a minimal assumption on the sub-optimality of the data (i.e., only that sub-optimal data should have an overlapped support with the expert data). Though, the sample complexity result relies on the assumption of the optimal action encoder which may not be realizable in practice.\n\n- The experiments are done thoroughly to evaluate the efficiently and robustness of TRAIL. A sufficient number of baselines is compared against TRAIL in a sufficient number of benchmark tasks. Still, the results could be made stronger by considering more practical tasks such as robot arm manipulation.  \n\n## Weakness \n\n- The main issue is on the clarity of the paper. In particular, I do not understand $\\theta_s$ and $\\theta$ is section 4.2. Is $\\theta_s$ the $s$-th column of the matrix $\\theta \\in \\mathbb{R}^{d \\times |S|}$? If so, is $\\theta : S \\mapsto \\mathbb{R}^d$ a typo or you assume it is a linear map? I also do not fully understand about the linear transition model in section 4.3 and the cosine non-linearity. Why is the cosine applied here when linearity is the goal? \n\n## questions\n- How are the number and dimensionality of latent action chosen for TRAIL EBM and TRAIL linear in each task? Are they tuned differently for each task or the same value is used across tasks?\n", "summary_of_the_paper": "The paper considers an imitation learning (IL) problem with both expert and suboptimal demonstrations. The paper claims that sub-optimal demonstrations can be used to learn latent action abstractions which can improve the efficiency of down-stream IL. To solve this problem, the paper proposes TRAIL, which pre-trains an action encoder-decoder and a latent transition model using sub-optimal data, and performs behavioral cloning (BC) with expert data to learn a policy in the latent action space. The paper derives error bounds showing that the pre-training step and the down-stream BC step contribute to solving an IL problem in the original space from the view of divergence minimization. The paper also derives a bound showing that an optimal action abstraction may improve the sample complexity of BC. Experiments  indicate that TRAIL effectively learns from a limited number of expert demonstrations and is robust to the sub-optimality of sub-optimal demonstrations. \n\nMain contributions: An effective method to learn an expert policy from few expert data and a large set of sub-optimal data. The method is theoretically justified and empirically well-supported. \n", "summary_of_the_review": "I find the problem setting and proposed method very interesting. The theoretical analyses are valuable, and the experiments are convincing. Though, there is an issue regarding the clarity which make it difficult to fully understand the paper. I rate the paper as above the borderline. \n\n** Update after authors' rebuttal **\nThe rebuttal and the revision address my concerns on the paper. I think more positive about the paper and have increased the score. The new result (Figure 7) in the appendix is a good addition and I suggest including it in the main paper as well. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TR4u1Byy8ij", "reviewer": "Reviewer_nwvX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "### Strength ###\n- This paper considers a practical imitation learning setting in which there’s a very small amount of high-quality demonstrations that are expensive to obtain and there’s a large amount of (task-independent) suboptimal data that is cheaper to get. \n- The proposed algorithm is simple and theoretically-grounded. \n- This method of learning action reparameterization is applicable beyond imitation learning. Potential applications include offline reinforcement learning and transfer learning. \n\n### Suggestions and Questions ###\n- In the analysis, the paper considers the discrete state (and discrete action) setting. Is it possible to extend the analysis to continuous spaces? \n  - In Theorem 3, when the latent action space is continuous, how does \\pi_{alpha^*} in the second error term adapt from the first equation in page 5 which requires the latent action space to be discrete? \n  - Sample complexity is studied under the tabular case in page 5. What’s the sample complexity (e.g. Theorem 2) when latent action space is continuous? \n- In the TRAIL EBM for Theorem 1 paragraph in Section 4.3, you wrote “We allow \\phi to be updated while optimizing J_{DE}”. Would this increase the first error term (Eq. (1))? Have you tried to optimize \\phi, T, and \\pi_\\alpha simultaneously?\n- I am excited about the experimental results. They look very promising. \n  - However, It’s not clear to me whether learning action reparameterization is superior to learning state encoding and learning transition models without factorization. For example, a naive baseline may be interesting to consider is learning a transition model from the large amount of suboptimal data and then applying GAIL or AIRL to match the state distribution of the small amount of demonstrations using the learned transition model. \n  - Some form of visualizing the latent action space learned may be interesting. \n\n### Other Comments ###\n- \\pi_k in the second equation in page 5, \\pi_5 has not been introduced \n- In Theorem 1. Is \\pi_Z S -> \\Delta(Z)?\n- In Section 4.3 TRAIL EBM for Theorem 1, “we set \\rho to be the distribution of s’ in d^{off}”, does that mean you sample both the samples in the contrastive loss from the suboptimal data?\n\n### Post-rebuttal comments ###\nThank the authors very much for their response in a short amount of time. All of my questions have been clarified. As before, I am leaning towards accepting the paper. \n", "summary_of_the_paper": "The paper proposes an imitation learning algorithm, TRAIL, that can benefit from a large amount of suboptimal demonstrations besides a small amount of high-quality demonstrations. This is achieved through learning a factored transition model with action reparameterization from the suboptimal or even random demonstrations before doing behavior cloning. The authors analyze the error bound of the algorithm and show improved sample complexity with action reparameterization. Moreover, TRAIL is verified on a set of navigation and locomotion imitation learning tasks and it outperforms baselines based on temporal action abstraction in terms of task success rate. ", "summary_of_the_review": "Because of the strengths discussed above, I am leaning towards accepting the paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper investigates what we can learn from _suboptimal_ demonstrations for imitation learning. It suggests that we can learn about the structure of the environment by finding a factored dynamics model including a latent action space. It demonstrates both theoretically and empirically that this information can reduce sample requirements for downstream IL.

The reviewers praised the simplicity of the method (including its minimal assumptions), the theoretical analysis, and the breadth of the experimental validation. The authors were helpful during the discussion period, and addressed any questions or concerns the reviewers raised.

Overall, this is an interesting idea and a well-executed paper.
Decision: Accept (Poster)