TRAIL: N EAR -O PTIMAL I MITATION L EARNING WITH S UBOPTIMAL D ATA
A BSTRACT
The aim in imitation learning is to learn effective policies by utilizing near-optimal expert demonstrations.
However, high-quality demonstrations from human experts 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 ofﬂine datasets to facilitate provably improved downstream imitation learning?
In this work, we answer this question afﬁrmatively and present training objectives that use ofﬂine datasets to learn a factored transition model whose structure enables the extraction of a latent action space .
Our theoretical analysis shows that the learned latent action space can boost the sample-efﬁciency of downstream imitation learning, effectively reducing the need for large near-optimal expert datasets through the use of auxiliary non-expert data.
To learn the latent action space in practice, we propose
TRAIL (Transition-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-efﬁcient imitation learning.
We evaluate the practicality of our objective through experiments on a set of navigation and locomotion tasks.
Our results verify the beneﬁts suggested by our theory and show that TRAIL is able to improve baseline imitation learning by up to 4x in performance.
I NTRODUCTION
Imitation learning uses expert demonstration data to learn sequential decision making policies (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 ofﬂine data generated by suboptimal agents.
Since the ofﬂine data is suboptimal in performance, using it directly for imitation learning is infeasible.
While some prior works have proposed using suboptimal ofﬂine data for ofﬂine 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 ofﬂine datasets should contain useful information about the environment, if only we could distill that information into a useful form that can aid downstream imitation learning.
One approach to leveraging suboptimal ofﬂine datasets is to use the ofﬂine data to extract a lowerdimensional 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 performing 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 imitation 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., 2018a; Ajay et al., 2020; Pertsch et al., 2020;
Hakhamaneshi et al., 2021), such methods typically focus on the theoretical and practical beneﬁts of temporal abstraction by extracting temporally extended skills from data or experience.
That is, the main beneﬁt 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 on the question of action representation : instead of learning skills that provide for temporal abstraction, we aim to directly reparameterize the action space in a way that provides for more sample-efﬁcient downstream
𝒟 off {( s , a , s ′ )
} 𝒟 π
* {( s , a )}
T z ∘ ϕ ( s , a ) (1)
π α ( a | s , ϕ ( s , a ) ) (2) π Z
( ϕ ( s , a ) | s ) (3)
Pretraining
Downstream Imitation
Inference s z ∼ π Z ( s ) π α ( a | s , z ) a
Transition reparametrized actions imitation without the need to reduce control frequency. Unlike learning temporal abstractions, action reparamtrization does not have to rely on any hierarchical structures in the ofﬂine data, and can therefore utilize highly suboptimal datasets (e.g., with random actions).
Aiming for a provably-efﬁcient approach to utilizing highly suboptimal ofﬂine datasets, we use ﬁrst 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 ofﬂine dataset, and ﬁnally, (3) behavioral cloning (i.e., max-likelihood learning of latent actions) on an expert demonstration dataset.
The ﬁrst term in our bound immediately suggests a practical ofﬂine training objective based on a transition dynamics loss using an factored transition model.
We show that under speciﬁc factorizations (e.g., low-dimensional or linear), one can guarantee improved sample efﬁciency on 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-Reparametrized
Actions for Imitation Learning (TRAIL).
As shown in Figure 1, TRAIL consists of a pretraining stage (corresponding to the ﬁrst two terms in our bound) and a downstream imitation learning stage (corresponding to the last term in our bound).
During the pretraining stage, TRAIL uses an ofﬂine dataset to learn a factored transition model and a paired action decoder.
During the downstream imitation learning stage, TRAIL ﬁrst reparametrizes expert actions into the latent action space according 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 ﬂexibility and trains the EBM with a contrastive loss.
The EBM 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) a provably beneﬁcial objective for learning action representations without temporal abstraction and (ii) a practical algorithm for optimizing the proposed objective by learning an EBM or linear transition model.
An extensive evaluation on a set of navigation and locomotion tasks demonstrates the effectiveness of the proposed objective.
TRAIL’s empirical success compared to a variety of existing methods suggests that the beneﬁt of learning single-step action representations has been overlooked by previous temporal skill extraction methods.
Additionally, TRAIL signiﬁcantly improves behavioral cloning even when the ofﬂine dataset is unimodal or highly suboptimal (e.g., obtained from a random policy), whereas temporal skill extraction methods lead to degraded performance in these scenarios.
Lastly, we show that TRAIL, without using reward labels, can perform similarly or better than ofﬂine reinforcement learning (RL) with orders of magnitude less expert data, suggesting new ways for ofﬂine learning of squential decision making policies.
R ELATED W ORK
Learning action abstractions is a long standing topic in 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.
; Sharma et al.
; Gregor et al.
Warde-Farley et al.
Liu et al.
(2021) propose to learn a diverse set of skills by maximizing an information theoretic objective.
Online skill discovery is also commonly seen in a hierarchical framework that learns a continuous space (Vezhnevets et al., 2017;
Hausman et al., 2018; Nachum et al., 2018a; 2019) or a discrete set of lower-level policies (Bacon et al., 2017; Stolle & Precup, 2002; Peng et al., 2019), upon which higher-level policies are trained to solve speciﬁc tasks.
Different from these works, we focus on learning action representations ofﬂine from a ﬁxed suboptimal dataset to accelerate imitation learning.
Aside from online skill discovery, ofﬂine skill extraction focuses on learning temporally extended action abstractions from a ﬁxed ofﬂine dataset.
Methods for ofﬂine skill extraction generally involve maximum likelihood training of some latent variable models on the ofﬂine data, followed by downstream planning (Lynch et al., 2020), imitation learning (Kipf et al., 2019; Ajay et al., 2020;
Hakhamaneshi et al., 2021), ofﬂine RL (Ajay et al., 2020), or online RL (Fox et al., 2017; Krishnan et al., 2017; Shankar & Gupta, 2020; Shankar et al., 2019; Singh et al., 2020; Pertsch et al., 2020; 2021; Wang et al., 2021) in the induced latent action space.
Among these works, those that provide a theoretical analysis attribute the beneﬁt of skill extraction predominantly to increased temporal abstraction as opposed to the learned action space being any “easier” to learn from than the raw action space (Ajay et al., 2020; Nachum et al., 2018b).
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 ofﬂine 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 representation learning, there exists a disjoint set of work that focuses on accelerating RL with state representation 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 a latent state space from a learned dynamics model.
Analogous to our own derivations, these works attribute the beneﬁt of representation learning to a smaller latent state space reduced from a high-dimensional input state space (e.g., images).
https://www.overleaf.com/project/6137abc67863801e1d0aa
P RELIMINARIES
In this section, we introduce the problem statements for imitation learning and learning-based control, and deﬁne relevant notations.
Markov decision process.
Consider an MDP (Puterman, 1994)
M := (cid:104) S, A, R , T , µ, γ (cid:105) , consisting of a state space S , an action space A , a reward function R :
S ×
A → R , a transition function
T : S ×
A → ∆( S ) 1 , an initial state distribution µ ∈ ∆( S ) , and a discount factor γ ∈ [0 , 1)
A policy π : S → ∆( A ) interacts with the environment starting at an initial state s 0 ∼ µ .
An action a t ∼ π ( s t ) is sampled and applied to the environment at each step t ≥ 0 .
The environment produces a scalar reward R ( s t , a t ) and transitions into the next state s t +1 ∼ T ( s t , a t ) .
Note that we are speciﬁcally interested in the imitation learning setting, where the rewards produced by R are unobserved by the learner.
The state visitation distribution d π ( s ) induced by a policy π is deﬁned as d π ( s ) :
= (1 − γ ) (cid:80) ∞
t =0 γ t · Pr [ s t = s | π, M ] .
We relax the notation and use ( s, a ) ∼ d π to denote s ∼ d π , a ∼ π ( s ) .
Learning goal.
Imitation learning aims to recover an expert policy π ∗ with access to only a ﬁxed set of samples from the expert: D π ∗ =
{ ( s i , a i ) } ni =1 with s i ∼ d π ∗ and a i ∼ π ∗ ( s i ) .
One approach to imitation learning is to learn a policy π that minimizes some discrepancy between π and π ∗ .
In our analysis, we will use the total variation (TV) divergence in state visitation distributions,
[Equation]
as the way to measure the discrepancy between π and π ∗ .
Our bounds can be easily modiﬁed 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
[Equation]
which optimizes an upper bound of Diﬀ( π, π ∗ ) deﬁned above (Ross & Bagnell, 2010;
Nachum &
Yang, 2021):
[Equation]
BC with suboptimal ofﬂine data.
The standard BC objective (i.e., direct max-likelihood on
D π ∗ ) can struggle to attain good performance when the amount of expert demonstrations is limited (Ross et al., 2011; Tu et al., 2021).
We assume access to an additional suboptimal ofﬂine dataset D oﬀ =
{ ( s i , a i , s (cid:48) i ) } mi =1 , where the suboptimality is a result of (i) suboptimal action samples a i ∼ Unif A and (ii) lack of reward labels.
We use ( s, a, s (cid:48) ) ∼ d oﬀ as a shorthand for simulating ﬁnite sampling from D oﬀ via s i ∼ d oﬀ , a i ∼ Unif A , s (cid:48)
i ∼ T ( s i , a i ) , where d oﬀ is an unknown ofﬂine state distribution.
We assume d oﬀ sufﬁciently covers the expert distribution; i.e., d π ∗ ( s ) > 0 ⇒ d oﬀ ( s ) > 0 for all s ∈ S .
The uniform sampling of actions in D oﬀ is largely for mathematical convenience, and in theory can be replaced with any distribution uniformly bounded from below by η > 0 , and our derived bounds will be scaled by 1 |
A | η as a result.
This works focuses on how to utilize such a suboptimal D oﬀ to provably accelerate BC.
N EAR -O PTIMAL I MITATION L EARNING WITH R EPARAMETRIZED A CTIONS
In this section, we provide a provably-efﬁcient 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. P ERFORMANCE B OUND WITH R EPARAMETRIZED A CTIONS
Despite D oﬀ being highly suboptimal (e.g., with random actions), the large set of ( s, a, s (cid:48) ) tuples from D oﬀ 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 :=
T Z ◦ φ from the ofﬂine dataset D oﬀ , where φ : S ×
A → Z is an action representaiton function and
T Z : S × Z → ∆( S ) is a latent transition model.
Intuitively, good action representations should enable good imitation learning.
We formalize this intuition in the theorem below by establishing a bound on the quality of a learned policy based on (1) an ofﬂine pretraining objective for learning φ and T Z , (2) an ofﬂine decoding objective for learning an action decoder π α , and (3) a downstream imitation learning objective for learning a latent policy π Z with respect to latent actions determined by φ .
Theorem 1.
Consider an action representation function φ : S ×
A → Z , a factored transition model T Z : S × Z → ∆( S ) , an action decoder π α : S × Z → ∆( A ) , and a tabular latent policy π Z : S → ∆( Z ) .
Deﬁne the transition representation error as
[Equation]
the action decoding error as
[Equation]
and the latent behavioral cloning error as
[Equation]
Then the
TV divergence between the state visitation distributions of π [Equation] ◦ π Z :
S → ∆( A ) and π ∗ can be bounded as
Diﬀ( π α ◦ π Z , π ∗ ) ≤
Pretraining  
[Equation]
DownstreamImitation
[Equation]
where C 1 = γ | A | (1 − γ ) − [Equation] 1 + D χ [Equation]
d π ∗
[Equation]
d oﬀ [Equation] 12 ) , C 2 = γ [Equation]
1 − γ ) −
[Equation]
1 + D χ 2 ( d π ∗ (cid:107) d oﬀ ) 12 ) , C 3 = γ (1 − γ ) − 1 , π α ∗ is the optimal action decoder for a speciﬁc data distribution d oﬀ and a speciﬁc φ : [Equation] and π ∗ ,Z is the marginalization of π ∗ onto Z according to φ :
[Equation]
Theorem 1 essentially decomposes the imitation learning error into (1) a transition-based representation error J T , (2) an action decoding error J DE , and (3) a latent behavioral cloning error J BC ,φ .
Notice that only (3) requires expert data D π ∗ ; (1) and (2) are trained on the large ofﬂine data D oﬀ .
By choosing | Z | that is smaller than | A | , fewer demonstrations are needed to achieve small error in J BC ,φ compared to vanilla BC with J BC .
The Pearson χ 2 divergence term D χ 2 ( d π ∗ (cid:107) d oﬀ ) in Cand C 2 accounts for the difference in state visitation between the expert and ofﬂine data.
In the case where d π ∗ differs too much from d oﬀ , known as the distribution shift problem in ofﬂine RL (Levine et al., 2020), the errors from J T and J DE are ampliﬁed and the terms (1) and (2) in Theorem 1 dominate.
Otherwise, as J T → 0 and π α , φ → arg min J DE , optimizing π Z in the latent action space is guaranteed to optimize π in the original action space.
Sample Complexity To formalize the intuition that a smaller latent action space | Z | < | A | leads to more sample efﬁcient downstream behavioral cloning, we provide the following theorem in the tabular action setting.
First, assume access to an oracle latent action representation function φ orcl :=
OPT φ ( D oﬀ ) which yields pretraining errors (1) ( φ orcl ) and (2) ( φ orcl ) in Theorem 1.
For downstream behavioral cloning, we consider learning a tabular π Z on D π ∗ with n expert samples.
We can bound the expected difference between a latent policy π φ orcl ,Z with respect to φ orcl and π ∗ as follows.
Theorem 2.
Let φ orcl :=
OPT φ ( D oﬀ ) and π orcl,Z be the latent BC policy with respect to φ orcl .
We have,
[Equation]
where C 3 is the same as in Theorem 1.
We can contrast this bound to its form in the vanilla BC setting, for which | Z | = | A | and both (1) ( φ orcl ) and (2) ( φ orcl ) are zero.
We can expect an improvement in sample complexity from reparametrized actions when the errors in (1) and (2) are small and | Z | < | A | .
4. L INEAR T RANSITION M ODELS WITH D ETERMINISTIC L ATENT P OLICY
Theorem 1 has introduced the notion of a latent expert policy π ∗ ,Z , and minimizes the KL divergence between π ∗ ,Z and a tabular latent policy π Z .
However, it is not immediately clear, in the case of continuous latent actions, how to ensure that the latent policy π Z is expressive enough to capture any π ∗ ,Z .
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 ⊂ R d and a deterministic latent policy π θ ( s )
= θ s for some θ ∈ R d ×| S | .
While a deterministic θ in general cannot capture a stochastic π ∗ , we show that under a linear transition model T Z
[Equation]
φ ( s, a ) , there always exists a deterministic policy π θ : S → R d , such that θ s = π ∗ ,Z ( s ) , ∀ s ∈ S .
This means that our scheme for ofﬂine pretraining paired with downstream imitation learning can provably recover any expert policy π ∗ from a deterministic π θ , regardless of whether π ∗ is stochastic.
Theorem 3.
Let φ : S ×
A → Z for some Z ⊂ R d and suppose there exist w : S → R d such that
T Z ( s (cid:48) | s, φ ( s, a ))
=
w ( s (cid:48) ) (cid:62) φ ( s, a ) for all s, s (cid:48) ∈ S, a ∈ A .
Let π α : S × Z → ∆( A ) be an action decoder, π : S → ∆( A ) be any policy in M and π θ : S → R d be a deterministic latent policy for some θ ∈ R d ×| S | .
Then,
[Equation]
where C 4 = 14 | S |(cid:107) w (cid:107) ∞ , (1) and (2) corresponds to the ﬁrst and second terms in the bound in
Theorem 1.
By replacing term (3) in Theorem 1 that corresponds to behavioral cloning in the latent action space by term (4) in Theorem 3 that is a convex function unbounded in all directions, we are guaranteed that π θ is provably optimal regardless of the form of π ∗ and π ∗ ,Z .
Note that the downstream imitation learning objective implied by term (4) is simply the mean squared error between actions θ s chosen by π θ and reparameterized actions φ ( s, a ) appearing in the expert dataset.
4. TRAIL: R EPARAMETRIZED A CTIONS AND I MITATION L EARNING IN P RACTICE
In this section, we describe our learning framework, Transition-Reparametrized Actions for Imitation Learning (TRAIL).
TRAIL consists of two training stages: pretraining and downstream behavioral cloning.
During pretraining, TRAIL learns T Z and φ by minimizing J T ( T Z , φ ) =
E ( s,a ) ∼ d off [ D KL ( T ( s, a ) (cid:107)T Z ( s, φ ( s, a )))] .
Also during pretraining, TRAIL learns π α and φ by minimizing J DE ( π α , φ ) :
=
E ( s,a ) ∼ d off [ − log π α ( a | s, φ ( s, a ))] .
TRAIL parametrizes π α as a multivariate Gaussian distribution.
Depending on whether T Z is deﬁned according to Theorem 1 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 1, to ensure that the factored transition model T Z is ﬂexible to capture any complex (e.g., multi-modal) transitions in the ofﬂine dataset, we propose to use an energy-based model (EBM) to parametrize
T Z ( s (cid:48) | s, φ ( s, a )) ,
[Equation]
where ρ is a ﬁxed distribution over S and ψ : S → Z is a function of s (cid:48) .
In our implementation we set ρ to be the distribution of s (cid:48) in d oﬀ , which enables a practical learning objective for T Z by minimizing E ( s,a ) ∼ d off [ D KL ( T ( s, a ) (cid:107)T Z ( s, φ ( s, a )))] in Theorem 1 using a contrastive loss: E
[Equation]
During downstream behavioral cloning, TRAIL EBM learns a latent Gaussian policy π Z by minimizing J BC ,φ ( π Z )
=
E ( s,a ) ∼ ( d π ∗ ,π ∗ ) [ − log π Z ( φ ( s, a ) | s )] with φ ﬁxed.
During inference, TRAIL
EBM ﬁrst samples a latent action according to z ∼ π Z ( s ) , and decodes the latent action using a ∼ π α ( s, z ) to act in an environment.
Figure 1 describes this process pictorially.
TRAIL Linear for Theorem 3.
In the continuous latent action setting that corresponds to Theorem 3, we propose TRAIL linear, an approximation of TRAIL EBM, to enable learning linear transition models required by Theorem 3.
Speciﬁcally, we ﬁrst learn f, g that parameterize an energybased transition model T ( s (cid:48) | s, a ) ∝ ρ ( s (cid:48) )
exp {−|| f ( s, a ) − g ( s (cid:48) )
|| 2 / 2 } using the same contrastive loss as above (replacing φ and ψ by f and g ), and then apply random Fourier features (Rahimi et al., 2007) to recover ¯ φ ( s, a )
=
cos( Wf ( s, a ) + b ) , where W is a d × k matrix with entries sampled from a unit Gaussian and b a vector with entries sampled uniformly from [0 , 2 π ] .
W and b are implemented as an untrainable neural network layer on top of f .
This results in an approximate linear transition model,
[Equation]
During downstream behavioral cloning, TRAIL linear learns a deterministic policy π θ in the continuous latent action space determined by ¯ φ via minimizing (cid:13)(cid:13)
∂∂θ E s ∼ d π ∗ ,a ∼ π ∗
[Equation]
s ) [( θ s − ¯ φ ( s, a )) 2 ] (cid:13)(cid:13) 1 with ¯ φ ﬁxed.
During inference, TRAIL linear ﬁrst determines the latent action according to z
= π θ ( s ) , and decodes the latent action using a ∼
π α ( s, z ) to act in an environment.
[Figure] [Figure]
E XPERIMENTAL E VALUATION
We now evaluate TRAIL on a 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 ofﬂine data.
We evaluate the improvement attained by TRAIL over vanilla BC, and additionally compare TRAIL to previously proposed temporal skill extraction methods.
Since there is no existing benchmark for imitation learning with suboptimal ofﬂine data, we adapt existing datasets for ofﬂine RL, which contain suboptimal data, and augment them with a small amount of expert data for downstream imitation learning.
5. E VALUATING N AVIGATION WITHOUT T EMPORAL A BSTRACTION
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., 2021), and SPiRL (Pertsch et al., 2020), all of which use an ofﬂine 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, so we modify them to support downstream imitation learning as described in Appendix 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 et 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 state-of-the-art, especially in learning from suboptimal ofﬂine data based on D4RL.
To construct the suboptimal and expert datasets, we follow the protocol in Ajay et al.
(2020), which uses the full diverse or play D4RL AntMaze datasets as the suboptimal ofﬂine 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-corner navigation task, as they only contain data that navigates to random or ﬁxed locations different from task evaluation.
Implementation Details.
For TRAIL, we parameterize φ ( s, a ) and ψ ( s (cid:48) ) using separate feedforward 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 π α 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 as Gaussians following Ajay et al.
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 a search on the extend of temporal abstraction, and found t = 10 to work the best as reported in these papers’ maze experiments.
We also experimented with a version of vanilla
BC pretrained on the suboptimal data and ﬁne-tuned on expert data for fair comparison, which did not show a signiﬁcant difference from directly training vanilla BC on expert data.
Results.
Figure 3 shows the average performance of 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 temporal abstraction, we evaluate them both with the default temporal abstract, t = 10 , as well as without temporal abstraction, corresponding to t = 1 .
Note that TRAIL uses no temporal abstraction.
We
[Figure]
ﬁnd that on the simpler antmaze-medium task, TRAIL trained on a single-step transition model performs similarly to the set of temporal skill extraction methods with t = 10 .
However, these skill extraction methods experience a degradation in performance when temporal abstraction is removed ( t = 1 ).
This corroborates the existing theory in these works (Ajay et al., 2020), which attributes their beneﬁts predominantly to temporal abstraction rather than producing a latent action space that is “easier” to learn.
Meanwhile, TRAIL is able to excel without any temporal abstraction.
These differences become even more pronounced on the harder antmaze-large tasks.
We see that TRAIL maintains signiﬁcant improvements over vanilla BC, whereas temporal skill extraction fails to achieve good performance even with t = 10 .
These results suggest that TRAIL attains signiﬁcant improvement speciﬁcally 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 efﬁcient.
5. E VALUATING L OCOMOTION WITH H IGHLY S UBOPTIMAL O FFLINE D ATA
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 beneﬁt a broader set of tasks for which temporal abstraction may not be helpful, e.g., when the ofﬂine data is highly suboptimal (with near-random actions) or unimodal (collected by a single stationary policy).
In this section, we consider a Gym-MuJoCo task from D4RL using the same 8-DoF quadruped ant robot as the previously evaluated navigation task.
We ﬁrst learn action representations from the medium , medium-replay , 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 (exhibiting unimodal behavior), and the random dataset is collected by a randomly initialized policy and is hence highly suboptimal.
Implementation Details.
For this task, we additionally train a linear version of TRAIL by approximating the transition EBM using random Fourier features (Rahimi et al., 2007) and learn a deterministic latent policy following Theorem 3.
Speciﬁcally, we use separate feed-forward networks to parameterize f ( s, a ) and g ( s (cid:48) ) , and extract action representations using φ ( s, a )
=
cos( Wf ( s, a ) + b ) , where W, b are untrainable randomly initialized variables as described in Section 4.3.
Different from
TRAIL EBM which parametrizes π Z as a Gaussian, TRAIL linear parametrizes the deterministic π θ using a feed-forward neural network.
Results.
Our results are shown in Figure 4.
Both the EBM and linear versions of TRAIL consistently 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 a ﬂexible action decoder) resulted from unimodal ofﬂine datasets as discussed in (Ajay et al., 2020).
Surprisingly, TRAIL achieves a signiﬁcant performance boost even when latent actions are learned from the random dataset, suggesting the beneﬁt of learning action representations from transition models when the ofﬂine 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., 10 k), suggesting the beneﬁt of learning deterministic latent policies from Theoremwhen the environment is effectively approximated by a linear transition model.
[Figure]
5. E VALUATION ON D EEP M IND C ONTROL S UITE
Description.
Having witnessed the improvement TRAIL brings to behavioral cloning on AntMaze and MuJoCo Ant, we wonder how TRAIL perform on a wider spectrum of locomotion tasks with various degrees of freedom.
We consider 6 locomotion tasks from the DeepMind Control
Suite (Tassa et al., 2018) ranging from simple (e.g., 1-DoF cartople-swingup ) to complex (e.g., 21-DoF humanoid-run ) tasks.
Following the setup in Zolna et al.
(2020), we take 110 of the trajectories whose episodic reward is among the top 20% of the open source RL Unplugged datasets (Gulcehre et al., 2020) as expert demonstrations (see numbers of expert trajectories in Appendix C), and the bottom 80% of RL Unplugged as the suboptimal ofﬂine data.
For completeness, we additionally include comparison to Critic Regularized Regression (CRR)
(Wang et al., 2020), an ofﬂine RL method with competitive performance on these tasks.
CRR is trained on the full RL
Unplugged datasets (i.e., combined suboptimal and expert datasets) with reward labels.
Results.
Figure 5 shows the comparison results.
TRAIL outperforms temporal extraction methods on both low-dimensional (e.g., cartpole-swingup ) and high-dimensional ( humanoid-run ) tasks.
Additionally, TRAIL performs similarly to or better than CRR on 4 out of the 6 tasks despite not using any reward labels, and only slightly worse on humanoid-run and walker-walk .
To test the robustness of TRAIL when the ofﬂine data is highly suboptimal, we further reduce the size and quality of the ofﬂine 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-swim , walker-stand , and walker-walk due to this change in ofﬂine 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 ofﬂine data.
This set of results suggests a promising direction for ofﬂine 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 ofﬂine RL, this approach is applicable to settings where data quality cannot be easily expressed through a scalar reward.
C ONCLUSION
We have derived a near-optimal objective for learning a latent action space from suboptimal ofﬂine data that provably accelerates downstream imitation learning.
To learn this objective in practice, we propose transition-reparametrized actions for imitation learning (TRAIL), a two-stage framework that ﬁrst pretrains a factored transition model from ofﬂine 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 a random policy), providing a 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 speciﬁc to imitation learning, and insofar as the reparameterized action space simpliﬁes 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 learning-based control represents an exciting future direction, and we hope that our results provide initial evidence of such a potential.
A CKNOWLEDGMENTS
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.
Appendix A P ROOFS FOR F OUNDATIONAL L EMMAS
Lemma 4.
If π 1 and π 2 are two policies in M and d π 1 ( s ) and
d π 2 ( s ) are the state visitation distributions induced by policy π 1 and π 2 where d π ( s ) :
= (1 − γ ) (cid:80) t =
γ t · Pr [ s t = s | π, M ] .
Deﬁne Diﬀ( π 2 , π 1 ) =
D TV ( d π 2 (cid:107) d π 1 ) then
[Equation] [Equation]
where
[Equation]
is the TV-divergence between T ◦
π 1 ◦ d π 1 and T ◦ π 2 ◦ d π 1 .
Proof.
Following similar derivations in Achiam et al.
; Nachum et al.
(2018b), we express
D TV ( d π 2 (cid:107) d π 1 ) in linear operator notation:
[Equation]
where Π 1 , Π 2 are linear operators S
→ S ×
A such that Π
i ν ( s, a ) =
π i ( a | s ) ν ( s ) and 1 is an all ones row vector of size | S | .
Notice that d π 1 may be expressed in this notation as (1 − γ )
( I − γ T Π 1 ) − 1 µ .
We may re-write the above term as
[Equation]
Using matrix norm inequalities, we bound the above by
[Equation]
Since T Π 2 is a stochastic matrix, (cid:107) ( I − γ T Π 2 ) − 1 (cid:107) 1 , ∞ ≤ (cid:80) ∞ t =
γ t (cid:107)T Π 2 (cid:107) 1 , ∞ = (1 − γ ) − 1 .
Thus, we bound the above by γ
[Equation]
[Equation] and so we immediately achieve the desired bound in equation 6.
The divergence bound above relies on the true transition model T which is not available to us.
We now introduce an approximate transition model T to proxy Err d π 1 ( π 1 , π 2 , T ) .
Lemma 5.
For π 1 and π 2 two policies in M and any transition model T ( ·| s, a ) we have,
[Equation]
[Equation]
Proof.
[Equation]
[Equation]
[Equation]
15) ≤
[Equation]
and we arrive at the inequality as desired where the last step comes from D TV ( T ( s, a ) (cid:107)T
[Equation]
=
[Equation]
Now we introduce a representation function φ : S ×
A → Z and show how the error above may be reduced when T ( s, a )
=
T Z ( s, φ ( s, a )) :
Lemma 6.
Let φ : S ×
A → Z for some space Z and suppose there exists T Z : S × Z → ∆( S ) such that T ( s, a ) =
T Z ( s, φ ( s, a )) for all s ∈ S, a ∈ A .
Then for any policies π 1 , π 2 ,
[Equation]
where π k,Z ( z | s ) is the marginalization of π k onto Z :
[Equation]
for all z ∈ Z, k ∈ { 1 , 2 } .
Proof.
[Equation]
[Table]
[Equation]
and we arrive at the inequality as desired
[Equation]
Lemma 7.
Let d ∈ ∆( S, A ) be some state-action distribution, φ : S ×
A → Z
[Equation]
and π Z : S → ∆( Z ) .
Denote π α ∗ as the optimal action decoder for d, φ :
[Equation]
and π α ∗ ,Z as the marginalization of π α ∗ ◦ π Z onto Z :
[Equation]
Then we have [Equation] α
[Equation]
26) for all z ∈ Z and s ∈ S .
Proof.
[Equation]
and we have the desired equality
[Equation]
Lemma 8.
Let π Z : S → ∆( Z ) be a latent policy in Z and π α : S ×
Z
→ A be an action decoder, π α,Z be the marginalization of π α ◦ π Z onto Z :
[Equation]
Then for any s ∈ S we have
[Equation]
where π α ∗ is the optimal action decoder deﬁned in Lemma 7 (and this holds for any choice of d from
Lemma 7)
[Equation]
Proof.
[Equation] [Equation]
[Equation]
and we have the desired inequality
[Equation]
Lemma 9.
Let π 1 ,Z be the marginalization of π 1 onto Z as deﬁned in Lemma
[Equation]
and let π Z , π α , π α,Z be as deﬁned in Lemma
[Equation]
and let π α ∗ ,Z be as deﬁned in Lemma 7.
For any s ∈ S we have
[Equation]
1 ,Z ( s ) (cid:107)
[Equation]
Z ( s )) .
Proof.
The desired inequality is achieved by plugging the inequality from Lemma 8 into the following triangle inequality:
[Equation]
Our ﬁnal lemma will be used to translate on-policy bounds to off-policy.
Lemma 10.
For two distributions ρ 1 , ρ 2 ∈ ∆( S ) with ρ 1 ( s ) > 0 ⇒ ρ 2 ( s ) >
[Equation]
we have,
[Equation]
Proof.
The lemma is a straightforward consequence of Cauchy-Schwartz:
[Equation] [Equation]
Finally
[Equation]
to get the desired bound
[Equation]
we simply note that the concavity of the square-root function implies E ρ 2 [ h ( s )] ≤ E ρ
[ (cid:112) h ( s ) 2 ] ≤ (cid:112) E ρ
[ h ( s ) 2 ] .
B P ROOFS FOR M AJOR T HEOREMS B. P ROOF OF T HEOREM
Proof.
Let π 2 :=
π α ◦ π Z , we have π 2 ,Z ( z | s ) =
π α,Z ( z | s ) =
(cid:80) a ∈ A,φ ( s,a )
=
z ( π α ◦ π Z )( z | s ) .
By plugging the result of Lemma 9 into Lemma
[Equation]
we have
[Equation]
By plugging this result into Lemma [Equation] we have
[Equation]
By further plugging this result into Lemma 4 and let π 1 = π ∗
[Equation]
we have:
[Equation] [Equation]
Finally
[Equation]
by plugging in the off-policy results of Lemma 10 to the bound in Equation 54 and by applying Pinsker’s inequality D TV ( T ( s, a ) (cid:107)T Z ( s, φ ( s, a ))) 2 ≤ 12 D KL ( T
[Equation]
s, a ) (cid:107)T Z ( s, φ ( s, a ))) , we have
[Equation]
where C
[Equation]
and C
[Equation]
Since the max z ∈ Z is not tractable in practice, we approximate
[Equation]
using E ( s,a )
[Equation]
[ D
[Equation]
( π
[Equation]
( s, φ ( s, a )) (cid:107) π
α
( s, φ ( s, a )))] , which reduces to J DE ( π α , φ ) with additional constants.
We now arrive at the desired off-policy bound in Theorem 1.
B. P ROOF OF T HEOREM
Lemma 11.
Let ρ ∈ ∆( { 1 , . .
. , k } ) be a distribution with ﬁnite support.
Let ˆ ρ n denote the empirical estimate of ρ from n i.i.d.
samples X ∼ ρ .
Then,
[Equation]
Proof.
The ﬁrst inequality is Lemma 8 in Berend & Kontorovich (2012) while the second inequality is due to the concavity of the square root function.
Lemma 12.
Let D := { ( s i , a i ) } ni
=1 be i.i.d.
samples from a factored distribution x ( s, a ) :
=
ρ ( s ) π ( a | s ) for ρ ∈ ∆( S ) , π : S → ∆( A ) .
Let ˆ ρ be the empirical estimate of ρ in D and ˆ π be the empirical estimate of π in D .
Then,
[Equation]
Proof.
Let ˆ x be the empirical estimate of x in D .
We have,
[Equation]
Finally, the bound in the lemma is achieved by application of Lemma 11 to each of the TV divergences.
To prove Theorem 2, we ﬁrst rewrite Theorem 1 as
[Equation]
where (1) and (2) are the ﬁrst two terms in the bound of Theorem 1, and C
=
γ 1 − γ .
The result in Theorem 2 is then derived by setting φ = φ π orcl and π Z :=
π φ orcl ,Z and using the result of Lemma 12.
Note that the above sample analysis can be extended to the continuous latent action space characterized by Theorem 3 as follows.
Theorem 13.
Let φ orcl :=
OPT φ ( D oﬀ ) and π orcl,θ be the latent BC policy with respect to φ orcl .
Let d be the dimension of the continuous latent actions and (cid:107) φ (cid:107) ∞ be the l ∞ norm of φ orcl for any s, a .
We have
[Equation]
where (1), (2), and C 4 are the same as in Theorem 3.
Proof.
We use µ ∈ R d ×| S | to denote the optimal setting of θ which yields a zero l 1 -norm of
[Equation]
According to Theorem 3, we want to bound the l 1 -norm of ∂∂θ E s ∼ d π ,a ∼ π ∗ [( θ s − φ ( s, a )) 2 ] evaluated at the approximate solution ˆ µ ∈ R d ×| S | with respect to ﬁnite dataset D π ∗ ; i.e.,
[Equation]
with the convention that ˆ µ s =
if s does not appear in D π ∗ .
To this end, we have the following derivation, which uses E n to denote the expectation over realizations of ˆ µ due to n -size draws of the target dataset D π ∗ :
[Equation]
We now split up the inner expectation based on the number of times k that s appears in D π ∗ : E
[Equation]
where E k denotes the expectation over realizations of ˆ µ s over k -size draws of a ∼ π ∗ ( s ) .
By standard combinatorics
[Equation]
we know
[Equation]
Furthermore, for k = 0 , we have
[Equation]
while for k > 0 , since E k [ˆ µ s ]
= µ s , we have
[Equation]
Combining equations 74,
[Equation]
and 76 we have for any k ≥
[Equation]
and so by the binomial theorem, n
[Equation]
Plugging the above into equation 72 we deduce
[Equation]
and we have the convergence rate as desired
[Equation]
B. P ROOF OF T HEOREM
Proof.
The gradient term in Theorem 3 with respect to a speciﬁc column θ s of θ may be expressed as
[Equation]
and so,
[Equation]
Summing over s ∈ S , we have:
[Equation]
Thus
[Equation]
we have:
[Equation]
[Table]
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 E XPERIMENT D ETAILS C. A RCHITECTURE
We parametrize φ as a 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 we found TRAIL to be relatively robust to the latent dimension size as long as it is not too small (i.e., ≥ 64 ).
The latent skills in temporal skill extraction require a much smaller dimension size, e.g., 8 or 10 as reported by Ajay et al.
Pertsch et al.
We tried increasing the latent skill size for these work during evaluation, but found the reported value 8 to work the best.
We additionally experimented with different extend of skill extraction, but found the previously reported t = 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 β = 0 .
1 for the KL regularization term in the β VAE of OPAL (as reported).
We also use 0 .
1 as the weight for SPiRL and SkiLD’s KL divergence terms.
C. T RAINING AND E VALUATION
During pretraining, we use the Adam optimizer with learning rate 0 .
0003 for 200 k iterations with batch size 256 for all methods that require pretraining.
During downstream behavioral cloning, learned action representations are ﬁxed, but the action decoder is ﬁne-tuned on the expert data as suggested by Ajay et al.
Behavioral cloning for all methods including vanilla BC is trained with learning rate 0 .
0001 for 1 M iterations.
We experimented with learning rate decay of downstream BC by a factor of 3 at the 200 k boundary for all methods.
We found that when the expert sample size is small, decaying learning rate can prevent overﬁtting for all methods.
The reported results are with learning rate decay on AntMaze and without learning rate decay on other environments for all methods.
During the downstream behavioral cloning stage, we evaluate the latent policy combined with the action decoder every 10 k steps by executing π α ◦ π Z 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 ﬁgures.
C. M ODIFICATION TO S KI LD AND SP I RL
Since SkiLD (Pertsch et al., 2021) and SPiRL (Pertsch et al., 2020) are originally designed for RL as opposed to imitation learning, we replace the downstream RL algorithms of SkiLD and SPiRL by behavioral cloning with regularization (but keep skill extraction the same as the original methods).
Speciﬁcally, for SkILD, we apply a KL regularization term between the latent policy and the learned skill prior in the suboptimal ofﬂine 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 classiﬁer 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 SPiRL, we apply the KL divergence between latent policy and skill prior extracted from ofﬂine data (i.e., using the red term in Algorithm 1 of Pertsch et al.
as an additional term to latent behavioral cloning.
C. D ATASET D ETAILS
AntMaze.
For the expert data in AntMaze, we use the goal-reaching expert policies trained by Ajay et al.
(2020) (expert means that the agent is trained to navigate from the one corner of the maze to the opposite corner) to collect n = 10 trajectories.
For the suboptimal data in AntMaze, we use the full D4RL datasets antmaze-large-diverse-v0 , antmaze-medium-play-v0 , antmaze-medium-diverse-v0 , and antmaze-medium-play-v0 .
Ant.
For the expert data in Ant, we use a small set of expert trajectories selected by taking either the ﬁrst 10 k or 25 k transitions from ant-expert-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-medium-v0 , ant-medium-replay-v0 , and ant-random-v0 .
RL Unplugged.
For DeepMind Control Suite (Tassa et al., 2018) set of tasks, we use the RL
Unplugged (Gulcehre et al., 2020) dataset.
For the expert data, we take 110 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.
For the suboptimal data, we use the bottom 80% of the RL Unplugged dataset.
[Table]
[Figure]
D. A DDITIONAL BASELINES FOR RL U NPLUGGED
[Figure]
[Figure]
We conduct additional evaluation on an environment with tabular state and action spaces.
As shown in Figure 8, an agent is randomly placed into a four-room environment, and the task is to navigate to the target ‘T’.
The task reward is 1 at ‘T‘ and 0 elsewhere.
There are 12 discrete actions 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 1 or 2 grids (the action space is artiﬁcially blown up as suggested by the reviewer TRAIL to
[Figure] [Figure]
[Figure]
F V ISUALIZATION OF L ATENT A CTIONS