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: The authors present MILO an offline imitation learning algorithm that uses a learned model of the environment to generate trajectories from arbitrary policies. This allows to reuse similar algorithms than in the online case such as GAIL or Wasserstein GAIL.  In addition, the authors use a bonus (pessimism) to penalize data that is generated in places with high model uncertainty. This can be achieved using an ensemble. They present experimental results as well as an ablation study on OpenAI Gym.