TRAIL: Near-Optimal Imitation Learning with Suboptimal Data

1 INTRODUCTION . 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 offline data generated 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 , if 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 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 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 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-efficient downstream 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 offline data , and can therefore utilize highly suboptimal datasets ( e.g. , with random actions ) . Aiming for a 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-likelihood learning of latent actions ) on an expert demonstration dataset . The first term in our bound immediately suggests a practical offline training objective based on a transition dynamics loss using an factored transition model . We show that under specific factorizations ( e.g. , low-dimensional or linear ) , one can guarantee improved sample efficiency 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 first 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 offline dataset to learn a factored transition model and a paired action decoder . During the downstream imitation learning stage , TRAIL first 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 flexibility 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 beneficial 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 benefit of learning single-step action representations has been overlooked by previous temporal skill extraction methods . Additionally , TRAIL significantly improves behavioral cloning even when the offline 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 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 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 . ( 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 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 specific tasks . Different from these works , we focus on learning action representations offline from a 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 involve maximum likelihood training of some latent variable models on 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 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 benefit 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 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 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 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 . 3 PRELIMINARIES . In this section , we introduce the problem statements for imitation learning and learning-based control , and define relevant notations . Markov decision process . Consider an MDP ( Puterman , 1994 ) M : = 〈S , A , R , T , µ , γ〉 , consisting of a state space S , an action spaceA , a reward functionR : 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 s0 ∼ µ . An action at ∼ π ( st ) is sampled and applied to the environment at each step t ≥ 0 . The environment produces a scalar reward R ( st , at ) and transitions into the next state st+1 ∼ T ( st , 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 dπ ( s ) induced by a policy π is defined as dπ ( s ) : = ( 1− γ ) ∑∞ t=0 γ t · Pr [ st = 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 fixed set of samples from the expert : Dπ∗ = { ( si , ai ) } ni=1 with si ∼ dπ∗ and ai ∼ π∗ ( si ) . 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 , Diff ( π , π∗ ) = DTV ( d π‖dπ∗ ) , as the way to measure the discrepancy between π and π∗ . 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 ( π ) : = E ( s , a ) ∼ ( dπ∗ , π∗ ) [ − log π ( a|s ) ] , which optimizes an upper bound of Diff ( π , π∗ ) defined above ( Ross & Bagnell , 2010 ; Nachum & Yang , 2021 ) : Diff ( π , π∗ ) ≤ γ 1− γ √ 1 2 Edπ∗ [ DKL ( π∗ ( s ) ‖π ( s ) ) ] = γ 1− γ √ const ( π∗ ) + 1 2 JBC ( π ) . 1∆ ( X ) denotes the simplex over a set X . BC with suboptimal offline 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 offline dataset Doff = { ( si , ai , s′i ) } mi=1 , where the suboptimality is a result of ( i ) suboptimal action samples ai ∼ UnifA and ( ii ) lack of reward labels . We use ( s , a , s′ ) ∼ doff as a shorthand for simulating finite sampling from Doff via si ∼ doff , ai ∼ UnifA , s′i ∼ T ( si , ai ) , where doff is an unknown offline state distribution . We assume doff sufficiently covers the expert distribution ; i.e. , dπ∗ ( s ) > 0 ⇒ doff ( s ) > 0 for all s ∈ S. The uniform sampling of actions in Doff 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 Doff to provably accelerate BC .

Summary: 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.  The 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.  The 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.