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