Mitigating Covariate Shift in Imitation Learning via Offline Data With Partial Coverage

1 Introduction . Covariate shift is a core issue in Imitation Learning ( IL ) . Traditional IL methods like behavior cloning ( BC ) [ 49 ] , while simple , suffer from covariate shift , learning a policy that can make arbitrary mistakes in parts of the state space not covered by the expert dataset . This leads to compounding errors in the agent ’ s performance [ 57 ] , hurting the generalization capabilities in practice . Prior works have presented several means to combat this phenomenon in IL . One line of thought utilizes an interactive expert , i.e . an expert that can be queried at an arbitrary state encountered during the training procedure . Interactive IL algorithms such as DAgger [ 59 ] , LOLS [ 15 ] , DART [ 40 ] , and AggreVaTe ( D ) [ 58 ; 66 ] utilize a reduction to no-regret online learning and demonstrate that under certain conditions , they can successfully learn a policy that imitates the expert . These interactive IL algorithms , however , can not provably avoid covariate shift if the expert is not recoverable . That is , ∗Equal contribution †Work done outside Amazon 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . Aπe ( s , a ) = Θ ( H ) where πe is the expert , Aπ is the usual ( dis ) advantage function,3 and H is the planning horizon [ 52 ; 4 , Chapter 15 ] . A second line of work that avoids covariate shift utilizes either a known transition dynamics model [ 2 ; 86 ] or uses real world interactions [ 27 ; 10 ; 65 ; 38 ; 56 ; 33 ] . Prior works have shown that with known transition dynamics or real world interactions , agents can provably avoid covariate shift in both tabular and general MDPs [ 4 ; 52 ] even without a recoverable expert . While these results offer strong theoretical guarantees and empirical performance , online interactions are often costly and prohibitive for real world applications where active trial-and-error exploration in the environment could be unsafe or impossible . A third perspective towards addressing this issue is to assume that the expert visits the entire state space [ 62 ] , where the expert effectively informs the learner what actions to take in every state . Unfortunately , such a full coverage expert distribution might be rare and holds only for special MDPs and expert policies ( for e.g . an expert that induces ergodicity in the MDP ) . In this work , we consider a new perspective towards handling the covariate shift issue in IL . In particular , we investigate a pure offline learning setting where the learner has access to neither the expert nor the environment for additional interactions . The learner , instead , has access to a small pre-collected dataset of state-action pairs sampled from the expert and a large batch offline dataset of state-action-next state transition triples sampled from a behavior policy that could be highly sub-optimal ( see Figure 1 where BC on the offline data results in a low-quality policy ) . Unlike prior works that require online interactions , our proposed method , MILO performs high fidelity imitation in an offline , data-driven manner . Moreover , different from interactive IL , we do not require the expert to be present during learning , significantly relieving the expert ’ s burden . Finally , in contrast to the prior work [ 62 ] that assumes the expert distribution covers the entire state-action space ( i.e. , maxπ maxs , a d π ( s , a ) /dπ e ( s , a ) < ∞ where dπ denotes the state-action distribution of policy π ) , we require the offline dataset to provide partial coverage , i.e. , it only needs to cover the expert ’ s state-actions ( i.e. , maxs , a dπ e ( s , a ) /ρ ( s , a ) < ∞ where ρ is the offline distribution of some behavior policy ) .4 In summary , we list our main contributions below : 1 . We propose Model based Imitation Learning from Offline data , MILO : a model-based framework that leverages offline batch data with only partial coverage ( see Section 4.1 for definition ) to overcome covariate shift in IL . 3in this work , we use cost instead of reward , thus we call Aπ the disadvantage function . 4In our analysis , we refine the density ratio dπ e ( s , a ) /ρ ( s , a ) via the concept of relative conditional number which allows us to extend it to large MDPs where the ratio is infinite but the relative condition number is finite . 2 . Our analysis is modular and covers common models such as discrete MDPs and linear models . Notably , our new result on non-parametric models ( e.g . Gaussian Processes ) with relative condition number is new even considering all existing results in offline RL ( see Remark 4,7,11 ) . 3 . The practical instantiation of our general framework leverages neural network model ensembles , and demonstrates its efficacy on benchmark MuJoCo continuous control problems . Specifically , even under low-quality behavior policies , our approach can successfully imitate using an extremely small number of expert samples while algorithms like BC completely fail ( Figure 1 ) . 1.1 Related work . Imitation Learning As summarized above , avoiding covariate shift in IL is an important topic . Another relevant line of research is IL algorithms that use offline or off-policy learning . ValueDICE [ 38 ] presents a principled way to leverage off-policy data for IL . In theory , the techniques from ValueDICE ( and more broadly , DICE [ 46 ; 84 ] ) require the data provided to the agent to have global coverage . Moreover in practice , ValueDICE uses online interaction and maintains an increasing replay buffer which may eventually provide global coverage . Instead , we aim to study offline IL without any online interactions and are interested in the setting where offline data does not have global coverage . Another line of work [ 29 ; 14 ] studies IL in an offline setting by only using the expert dataset . In contrast to these works , our goal is to study the use of an additional offline dataset collected from a behavior policy to mitigate covariate shift , as information theoretically any algorithm that relies solely on expert data will still suffer from covariate shift in the worst case [ 52 ] . Similar in setting , Cascaded Supervised IRL ( CSI ) [ 35 ] performs imitation offline with both an expert dataset and a static offline dataset by first fitting a reward function that is then used in Least Squares Policy Iteration . CSI , however , requires the expert data to have global coverage and does not mitigate covariate shift with partial coverage like MILO does . Finally , Variational Model-Based Adversarial IL ( V-MAIL ) [ 50 ] learns a dynamics model from a static offline dataset and performs offline imitation within the model . V-MAIL , however , studies zero-shot IL where the static offline data is samples from a variety of source tasks , and the expert dataset is samples from the transfer target task . In contrast , MILO investigates avoiding covariate shift in IL with an offline dataset collected by a potentially suboptimal policy from the same task . Offline RL In offline RL , algorithms such as FQI [ 22 ] have finite-sample error guarantees under the global coverage [ 45 ; 5 ] . Recently , many algorithms to tackle this problem have been proposed from both model-free [ 76 ; 68 ; 43 ; 26 ; 23 ; 39 ] and model-based perspectives [ 82 ; 34 ; 44 ] with some pessimism ideas . The idea of pessimism features in offline RL with an eye to penalize the learner from visiting unknown regions of the state-action space [ 54 ; 31 ; 79 ; 12 ] . We utilize pessimism within the IL context where , unlike RL , the learner does not have access to an underlying reward signal . Our work expands prior theoretical results by ( a ) formalizing the partial coverage condition using a notion of a relative condition number , and ( b ) offering distribution-dependent results when working with non-parametric models including gaussian processes . See Appendix C for a detailed literature review . 2 Setting . We consider an episodic finite-horizon Markov Decision Process ( MDP ) , M = { S , A , P , H , c , d0 } , where S is the state space , A is the action space , P : S × A → ∆ ( S ) is the MDP ’ s transition , H is the horizon , d0 is an initial distribution , and c : S × A → [ 0 , 1 ] is the cost function . A policy π : S → ∆ ( A ) maps from state to distribution over actions . We denote dπP ∈ ∆ ( S × A ) as the average state-action distribution of π under transition kernel P , that is , dπP = 1/H ∑H t=1 d π P , t , where dπP , t ∈ ∆ ( S × A ) is the distribution of ( s ( t ) , a ( t ) ) under π at t. Given a cost function f : S ×A 7→ [ 0 , 1 ] , V πP , f denotes the expected cumulative cost of π under the transition kernel P and cost function f . Following a standard IL setting , the ground truth cost function c is unknown . Instead , we have the demonstrations by the expert specified by πe : S → ∆ ( A ) ( potentially stochastic and not necessarily optimal ) . Concretely , we have an expert dataset in the form of i.i.d tuples De = { si , ai } nei=1 sampled from distribution d πe P . In our setting , we also have an offline static dataset consisting of i.i.d tuples Do = { si , ai , s′i } no i=1 s.t . ( s , a ) ∼ ρ ( s , a ) , s′ ∼ P ( s , a ) , where ρ ∈ ∆ ( S × A ) is an offline distribution resulting from some behavior policies . Note behavior policy could be a much worse policy than the expert πe . Our goal is Algorithm 1 Framework for model-based Imitation Learning with offline data ( MILO ) 1 : Require : IPM class F , model class P , policy class Π , datasets De = { s , a } , Do : = { s , a , s′ } 2 : Train Dynamics Model and Bonus : P̂ : S ×A → S and b : S ×A → R+ on offline data Do 3 : Pessimistic model-based min-max IL : with P̂ , b , De , obtain π̂IL by solving the following : π̂IL = argminπ∈Π maxf∈F [ E ( s , a ) ∼dπ P̂ [ f ( s , a ) + b ( s , a ) ] − E ( s , a ) ∼De [ f ( s , a ) ] ] ( 1 ) to only leverage ( De +Do ) to learn a policy π that performs as well as πe with regard to optimizing the ground truth cost c. More specifically , our goal is to utilize the offline static data Do to combat covariate shift and learn a policy that can significantly outperform traditional offline IL methods such as Behavior cloning ( BC ) , without any interaction with the real world or expert . Function classes We introduce function approximation . Since we do not know the true cost function c and transition kernel P , we introduce a cost function class F ⊂ S × A → [ 0 , 1 ] and a transition model class P : S × A → ∆ ( S ) . We also need a policy class Π . For the analysis , we assume realizability : Assumption 1. c ∈ F , P ∈ P , πe ∈ Π . We use Integral Probability Metric ( IPM ) as a distribution distance measure , i.e. , given two distributions ρ1 and ρ2 , IPM with F is defined as maxf∈F [ E ( s , a ) ∼ρ1 [ f ( s , a ) ] − E ( s , a ) ∼ρ2 [ f ( s , a ) ] ] .

Summary: This work considers an offline imitation learning setting where both expert and sub-optimal demonstrations are assumed to be known. While the covariate shift issue (the gap between training and test distributions) is a critical problem in offline imitation learning, e.g., behavioral cloning, authors address this issue by training a dynamics model through the sub-optimal demonstration (having state-action-successor state triples) and applying pessimistic penalty (incentivizing visiting the state-action pairs in the expert demonstration). With the theoretical derivations for various settings for dynamics modeling and the assumption that the demonstration does not have to cover the entire state-action space, authors show the theoretical guarantee of their proposed method. Also, the empirical results support their claim by showing that MILO (proposed algorithm) outperforms ValueDICE and BC.