Mitigating Covariate Shift in Imitation Learning via Offline Data Without Great Coverage
Year: 2022

Abstract:
This paper studies offline Imitation Learning (IL) where an agent learns to imitate an expert demonstrator without additional online environment interactions. Instead, the learner is presented with a static offline dataset of state-action-next state transition triples from a potentially less proficient behavior policy. We introduce Model-based IL from Offline data (MILO): an algorithmic framework that utilizes the static dataset to solve the offline IL problem efficiently both in theory and in practice. In theory, even if the behavior policy is highly sub-optimal compared to the expert, we show that as long as the data from the behavior policy provides sufficient coverage on the expert state-action traces (and with no necessity for a global coverage over the entire state-action space), MILO can provably combat the covariate shift issue in IL. Complementing our theory results, we also demonstrate that a practical implementation of our approach mitigates covariate shift on benchmark MuJoCo continuous control tasks. We demonstrate that with behavior policies whose performances are less than half of that of the expert, MILO still successfully imitates with an extremely low number of expert state-action pairs while traditional offline IL methods such as behavior cloning (BC) fail completely. Source code is provided at https://github.com/jdchang1/milo.

heading: 1 Introduction
text: 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, cannot 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).

heading: 1.1 Related work
text: 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.

heading: 2 Setting
text: 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)] ] .

heading: 3 Algorithm
text: The core idea of MILO is to imitate the expert by optimizing an IPM distance between the agent and the expert with a penalty term for pessimism over the policy class. MILO consists of three steps:
1. Model learning: fit a model P̂ from the offline data Do to learn P , 2. Pessimistic penalty design: construct penalty function b(s, a) such that there is a high penalty
on state-action pairs that are not covered by the offline data distribution ρ.
3. Offline min-max model-based policy optimization: optimize Eq. (1)
Algorithm 1 provides the details of MILO. We explain each component in detail as follows.
Model learning and Penalty: Our framework assumes we can learn a calibrated model (P̂ , σ) from the dataset Do, in the sense that for any s, a, we have: ∥∥∥P̂ (·|s, a)− P (·|s, a)∥∥∥ 1 ≤ min{2, σ(s, a)}. Such model training is possible in many settings including classic discrete MDPs, linear models (KNR [32]), and non-parametric models such as GP. In practice, it is also common to train a model ensemble based on the idea of bootstrapping and then use the model-disagreement to approximate σ. With such a calibrated model, the penalty will simply be b(s, a) = O(Hσ(s, a)). We will formalize this model learning assumption in Section 4. We give several examples below.
For any discrete MDP, we use the empirical distribution, i.e., P̂ (s′|s, a) = N(s′, s, a)/(N(s, a)+λ), where N(s, a) is the number of (s, a) in Do, and N(s′, s, a) is the number of (s, a, s′) in Do, and λ ∈ R+. In this case, we can set σ(s, a) = Õ (√ |S|/N(s, a) ) . See example 1 for more details.
For continuous Kernelized Nolinear Regulator (KNR [32]) model where the ground truth transition P (s′|s, a) is defined as s′ = W ?φ(s, a) + , ∼ N (0,Σ), with φ being a (nonlinear) feature mapping, we can learn P̂ by classic Ridge regression on offline dataset Do. Here we can set
σ(s, a) = Õ ( β √ φ(s, a)>Σ−1no φ(s, a) ) for some β ∈ R+, where Σo is the data covariance matrix
Σno := ∑no i=1 φ(si, ai)φ(si, ai) > + λI. See example 2 for more details.
For non-parametric nonlinear model such as Gaussian Process (GP), under the assumption that P is in the form of s′ = g?(s, a) + , ∼ N (0,Σ) (here S ⊂ RdS ), we can simply represent P̂ using GP posteriors induced by Do, i.e., letting GP posterior be GP (ĝ, kno), we have P̂ (s′|s, a) being
represented as s′ = ĝ(s, a) + . Then, we can set σ(s, a) = Õ (βkno ((s, a), (s, a))) with some parameter β ∈ R+ (see example 3 for more details). GP is a powerful model and has been being widely used in robotics problems, see [36; 18; 8; 71; 25] for examples.
In practice, we can also use a model ensemble of neural networks with the maximum disagreement between models as σ. This has been widely used in practice (e.g., [47; 6; 48]). We leave the details to Section 5 where we instantiate a practical version of MILO, and the experiment section. As we can see from the examples mentioned above, in general, the penalty b(s, a) = O(Hσ(s, a)) is designed such that it has a high value in state-action space that is not covered well by the offline data Do, and has a low value in space that is covered by Do. Adding such a penalty automatically forces our policy to stay away from these regions where P̂ is not accurate. On the other hand, for regions where ρ has good coverage (thus P̂ is accurate), we force π to stay close to πe.
Pessimistic model-based min-max IL: Note Eq. 1 is purely computational, i.e., we do not need any real world samples. To solve such min-max objective, we can iteratively (1) perform the best response on the max player, i.e., compute the argmax discriminator f given the current π, and (2) perform incremental update on the min player, e.g., use policy gradient (PG) methods (e.g. TRPO) inside the learned model P̂ with cost function f(s, a) + b(s, a). We again leave the details to Section 5.

heading: 3.1 Specialization to offline RL
text: In RL, the cost function c is given. The goal is to obtain π∗ = argmaxπ∈Π V π P,c. The pessimistic policy optimization procedure [81; 31] is π̂RL = argminπ∈Π E(s,a)∼dπ P̂
[c(s, a) + b(s, a)]. While this is not our main contribution, we will show a byproduct of our result is a novel non-parametric analysis for offline RL which does not assume ρ has global coverage (see Remarks 4,7,11).

heading: 4 Analysis
text: Our algorithm depends on the model P̂ estimated from the offline data. We provide a unified analysis assuming that P̂ is calibrated in that its confidence interval is provided. Specifically, we assume:
Assumption 2. With probability 1− δ, the estimate model P̂ satisfies the following: ‖P̂ (·|s, a)− P (·|s, a)‖1 ≤ min(σ(s, a), 2) ∀(s, a) ∈ S×A. We set the penalty as b(s, a) = H min(σ(s, a), 2).
We give the following three examples. For details, refer to appendix A. Example 1 (Discrete MDPs). Set uncertainty measure σ(s, a)= √ |S| log 2+log(2|S||A|/δ)
2{N(s,a)+λ} + λ N(s,a)+λ .
Example 2 (KNRs). In KNRs, the ground truth model is s′ = W ∗φ(s, a) + , ∼ N (0, ζ2I), where s ∈ RdS , a ∈ RdA , φ : S ×A 7→ Rd is some known state-action feature mapping. The estimator is
ĝ(·) = Ŵφ(·), Ŵ = argminW∈RdS×dA 1/no ∑ (s,a)∈Do [‖Wφ(s, a)− s ′‖22] + λ‖W‖2F ,
where ‖ · ‖F is a frobenius norm. We set the uncertainty measure σ(s, a):
σ(s, a) = (1/ζ)βno √ φ>(s, a)Σ−1no φ(s, a), Σno = ∑no i=1 φ(si, ai)φ >(si, ai) + λI
with βno = {2λ‖W ∗‖22 + 8ζ2[dS log(5) + log(1/δ) + Īno)]}1/2, where Īno = log(det(Σno/λI)). Example 3 (GPs). In GPs, the ground truth model is defined as s′ = g∗(s, a) + , ∼ N (0, ζ2I) where g? belongs to an RKHSHk with a kernel k(·, ·). Denote x := (s, a), we have GP posterior as ĝ(·) = S(Kno + ζ2I)−1k̄no(·), S = [s′1, · · · , s′no ] ∈ R
dS×no , k̄no(x) = [k(x1, x), · · · , k(xno , x)]>, {Kno}i,j = k(xi, xj) (1 ≤ i ≤ no, 1 ≤ j ≤ no), kno(x, x′) = k(x, x′)− k̄no(x)>(Kno + ζ2I)−1k̄no(x′),
with σ(·) = βnokno(·, ·)/ζ, βno = O((dS log 3(dSno/δ)Ino)1/2), Ino = log(det(I + ζ−2Kno)).
General results We show our general error bound results. For the proof, refer to appendix B. For analytical simplicity, we assume |F| is finite (but the bound only depends on ln(|F|)) 5.
5When |F| is infinite, we can show that the resulting error bound scales w.r.t its metric entropy.
Theorem 3 (Bound of MILO). Suppose assumptions 1,2. Then, with probability 1− 2δ,
V π̂ILP,c − V πe P,c ≤ Erro + Erre, Erro = 8H2E(s,a)∼dπeP [min(σ(s, a), 1)], Erre = 2H
√ log(2|F|/δ)
2ne .
We will show through a set of examples where E(s,a)∼dπeP [min(σ(s, a), 1)] shrinks to zero as no →∞ under the partial coverage, i.e., when ρ covers dπeP . Asymptotically, Erre will dominate the bound. Note that Erre has two components, a linear H and a term that corresponds to the statistical error related to expert samples and function class complexity. Comparing to BC, which has a rate O(H2 √ log(|Π|)/ne) [4, Chapter 14], we see that the horizon dependence is improved.
Before going to each analysis of Erro, we highlight two important points in our analysis. First, our bound requires only the partial coverage, i,e., it depends on πe-concentrability coefficient which measures the discrepancy between the offline data and expert data. This is the first work deriving the bound with πe-concentrability coefficient in IL with offline data. Second, our analysis covers nonparametric models. This is a significant contribution as previous pessimistic offline RL finite-sample error results have been limited to the finite-dimensional linear models or discrete MDPs [31; 54]. Remark 4 (Implications on offline RL). As in theorem 3, we have V π̂RLP,c − V π ∗ P,c = O(H2E(s,a)∼dπ∗P [σ(s, a)]) (appendix B). Note similar results have been obtained in [82; 34]. Since this term is Erro by just replacing πe with π∗, this offline RL result is a by-product of our analysis.

heading: 4.1 Analysis: Discrete MDPs
text: We start from discrete MDP as a warm up. Denote Cπe = max(s,a) d πe P (s, a)/ρ(s, a). Theorem 5. Suppose λ = Ω(1) and the partial coverage Cπe <∞. With probability 1− δ,
Erro ≤ c1H2 (√
Cπe |S|2|A| no + C πe |S||A| no ) · log(|S||A|c2/δ),
where c1, c2 are universal constants.
The error does not depend on supπ∈Π C π or C̄ = supπ∈Π max(s,a) d π P (s, a)/d πe P (s, a). We only require the partial coverage Cπe <∞, which is much weaker than supπ∈Π Cπ <∞ (ρ has global coverage) and C̄ <∞ (dπeP has global coverage [62]). When Cπe is small and no is large enough, Erre = O ( H √ |S||A|/ne ) dominates Erro in theorem 3. Then, the error is linear in horizon H .

heading: 4.2 Analysis: KNRs and GPs for Continuous MDPs
text: Now we move to continuous state-action MDPs. In continuous MDPs, assuming the boundedness of density ratio Cπe is still a strong assumption. As we dive into the KNR and the nonparametric GP model, we will replace the density ratio with a more refined concept relative condition number.
KNRs Let Σρ = E(s,a)∼ρ[φ(s, a)φ(s, a)>] and Σπe = E(s,a)∼dπeP [φ(s, a)φ(s, a) >]. We define the relative condition number as Cπe = supx∈Rd ( x>Σπex x>Σρx ) . Even when density ratio is infinite, this number could still be finite as it concerns subspaces on φ(s, a) rather than the whole S ×A.
To further gain its intuition, we can consider discrete MDPs and the feature mapping φ(s, a) ∈ R|S||A| which is a one-hot encoding vector that has zero everywhere except one at the entry corresponding to the pair (s, a). In this case, the relative condition number is reduced to max(s,a) d πe P (s, a)/ρ(s, a). Theorem 6 (Error for KNRs). Suppose sups,a ‖φ(s, a)‖ ≤ 1, λ = Ω(1), ζ2 = Ω(1), ‖W ∗‖2 = Ω(1) and the partial coverage Cπe <∞. With probability 1− δ,
Erro ≤ c1H2 ( rank2(Σρ) + rank(Σρ) log( c2 δ ) )√ dSCπe no · log1/2(1 + no), (2)
where c1 and c2 are some universal constants. Theorem 6 suggests Erro is Õ(H2rank[Σρ]2 √ d|S|Cπe/no). In other words, when Cπe , rank[Σρ] are small and the offline sample size no is large enough, Erre dominates Erro in theorem 3. Again, in this case, Erre = O ( H √ ln(F)/ne ) , and we see that it grows linearly w.r.t horizon H .
Our result is distribution dependent and captures the possible low-rankness of the offline data, i.e., rank[Σρ] depends on ρ and could be much smaller than the ambient dimension of feature φ(s, a). The quantity Cπe corresponds to the discrepancy measured between the batch data and expert data. This is much smaller than the worst-case concentrability coefficient: C̃ = supπ∈Π C
π . Remark 7. In RL, a similar quantity has been analyzed in [31], which studies the error bound of linear FQI with pessimism. Comparing to our result only requiring partial coverage, [31, Corollary 4.5] assumes the global coverage, i.e., Σρ is full-rank, which is stronger than C̃ <∞.
GPs Now we specialize our main theorem to non-parametric GP models. For simplicity, following [63], we assume S ×A is a compact space. We also suppose the following. Recall x := (s, a). Assumption 8. k(x, x) ≤ 1,∀x ∈ S ×A. k(·, ·) is a continuous and positive semidefinite kernel.
Under the assumption 8, we can use Mercer’s theorem [73], which shows that there exists a set of pairs of eigenvalues and eigenfunctions {µi, ψi}∞i=1, where ∫ ρ(x)ψi(x)ψi(x)dx = 1 for all i
and ∫ ρ(x)ψi(x)ψj(x)dx = 0 for i 6= j. Eigenfunctions and eigenvalues essentially defines an infinite-dimensional feature mapping φ(x) := [ √ µ1ψ1(x), . . . , √ µ∞ψ∞(x)]
>. Here, k(x, x) = φ(x)>φ(x), and any function f ∈ Hk can be represented as f(·) = α>φ(·). Note that the eigenvalues and eigenfunctions are defined w.r.t the offline data ρ, thus our result here is still distribution dependent rather than a worst case analysis which often appears in online RL/IL settings [63; 32; 78; 17].
Assume eigenvalues {µ1, . . . , µ∞} are in non-increasing order, we define the effective dimension, Definition 9 (Effective dimension). d∗ = min{j ∈ N : j ≥ B(j + 1)no/ζ2}, B(j) = ∑∞ k=j µk.
The effective dimensions d∗ is widely used and calculated for many kernels [85; 7; 72; 28]. In finitedimensional linear kernels {x 7→ a>φ(x); a ∈ Rd} (k(x, x) = φ>(x)φ(x)), we have d∗ ≤ rank[Σρ]. Thus, d∗ is considered to be a natural extension of rank[Σρ] to infinite-dimensional models.
Theorem 10 (Error for GPs). Let Σπe = Ex∼dπeP [φ(x)φ(x) >],Σρ = Ex∼ρ[φ(x)φ(x)>]. Suppose assumption 8, ζ2 = Ω(1) and the partial coverage Cπe = sup‖x‖2≤1(x >Σπex/x
>Σρx) <∞. With probability 1− δ,
Erro ≤ c1H2 ( (d∗)2 + d∗ log(c2/δ) )√ dSCπe no · √ log3(c2dSno/δ) log(1 + no), (3)
where c1, c2 are universal constants. The theorem suggests that Erro is Õ(H2d∗2 √ dSCπe/no). Thus, when Cπe , d∗ are not so large and no is large enough, Erre asymptotically dominates Erro in theorem 3 (again Erre is linear in H).
While we defer the detailed proof of the above theorem to Appendix C.3, we highlight some techniques we used here. The analysis is reduced to how to bound the information gain Ino and Ex∼dπeP [kno(x, x)]. In both cases, we analyze them into two steps: transforming them into the variational representation and then bounding them via the uniform low with localization (Lemma 39). Remark 11 (Implication to Offline RL). As related literature, in model-free offline RL, [70; 21] obtain the finite-sample error bounds using nonparametric models. Though their bounds can be characterized by the effective dimension, their bounds assume full coverage, i.e., max(s,a) 1/ρ(s, a) <∞.

heading: 5 Practical Algorithm
text: In this section we instantiate a practical version of MILO using neural networks for the model class P and policy class Π. We use the Maximum Mean Discrepancy (MMD) with a Radial Basis Function kernel as our discriminator class F . Note using MMD as our discrepancy measure allows us to compute the exact maximum discriminator argmaxf∈F in closed form (and is detailed in appendix). We use a KL-based trust-region formulation for incremental policy update inside the learned model P̂ . Based on Eq. (1), we first formalize the following constrained optimization framework:
minπ∈Π maxf∈F ( E(s,a)∼dπ
P̂ [f(s, a) + b(s, a)]− E(s,a)∼De [f(s, a)]
) s.t.E(s,a)∼De [` (a, s, π)] ≤ δ,
where ` : A×S ×Π 7→ R is a loss function (e.g., negative log-likelihood or any supervised learning loss one would use in BC). Essentially, since we have De available, we use it together with any
Algorithm 2 A practical instantiation of MILO 1: Require: expert dataset De = {s, a}, offline dataset Do := {s, a, s′} 2: Train an ensemble of neural network models {ĝ1, . . . , ĝn} where each Pi starts with different
random initialization; 3: Set bonus b(s, a) = maxi,j ‖gi(s, a)− gj(s, a)‖2 and initialize πθ0 . 4: for t = 0→ T − 1 do 5: wt = arg max‖w‖2≤1 w > ( E(s,a)∼dπ
P̂ [φ(s, a)]− E(s,a)∼De [φ(s, a)]
) , ft(s, a) := w>t φ(s, a)
6: θt+1 = θt − ηF−1θt
( E
(s,a)∼d πθt P̂
[ ∇ lnπθt(a|s)A
πθt P̂ ,ft+b
(s, a) ] + λE(s,a)∼De [∇`(a, s, πθt)] )
7: end for
supervised learning loss to constrain the policy hypothesis space. Note for a deterministic expert πe, the expert policy is always a feasible solution. Thus adding this constraint reduces the complexity of the policy class but does not eliminate the expert policy, and our analysis in Section 4 still applies.
In our practical instantiation, we replace the hard constraint instead by a Lagrange multiplier, i.e. we use the behavior cloning objective as a regularization term when solving the min-max problem:
minπ∈Π maxf∈F [ E(s,a)∼dπ
P̂ (f(s, a) + b(s, a))− E(s,a)∼De [f(s, a)]
] + λ · E(s,a)∼De [` (a, s, π)].
Given policy πθt (θt denotes the parameters), we update the discriminator ft. Then, with a fixed ft, in order to update policy π we use NPG as in line 6 in Algorithm 2, where Aπθ
P̂ ,f+b is the disadvantage
function of πθ and Fθt := E(s,a)∼dπθt P̂ [∇ lnπθt(a|s)∇ lnπθt(a|s)>] is the fisher information matrix.

heading: 6 Experiments
text: We aim to answer the following questions with our experiments: (1) How does MILO perform relative to other offline IL methods, (2) What is the impact of pessimism on MILO’s performance? (3) How does the behavior policy’s coverage impact MILO’s performance? (4) How does MILO’s result vary when we increase the number of samples drawn from the expert policy?
We evaluate MILO on five environments from OpenAI Gym [11] simulated with MuJoCo [67]: Hopper-v2, Walker2d-v2, HalfCheetah-v2, Ant-v2, and Humanoid-v2. We compare MILO against the following baselines: (1) ValueDICE [38], a state-of-the-art off-policy IL method modified for the offline IL setting; (2) BC on the expert dataset; and (3) BC on both the offline and expert dataset.
For the expert dataset, we first train expert policies and then randomly sample (s, a)-pairs from a pool of 100 expert trajectories collected from these expert policies. We randomly sample to create very small expert (s, a)-pair datasets where BC struggles to learn. Note that BC is effective at imitating the expert for MuJoCo tasks even with a single trajectory; prior works [27; 37; 38] have used similar sub-sampling strategies to create expert datasets to make it harder for BC to learn. The offline datasets are
collected building on prior Offline RL works [76; 34]; each dataset contains 1 million samples from the environment. We first train behavior policies with mean performances often less than half of the expert performance (Table 1, column 2). All results are averaged over five random seeds. See appendix for details on hyperparameters, environments, and dataset composition.

heading: 6.1 Evaluation on MuJoCo Continuous Control Tasks
text: Figure 2 presents results comparing MILO against benchmarks. MILO is able to achieve close to expert level performance on three out of the five environments and outperforms both BC and ValueDICE on all five environments. Both MILO and ValueDICE were warmstarted with one epoch of BC on the offline dataset. We significantly outperform BC’s performance when trained on the expert dataset, suggesting MILO indeed mitigates covariate shift through the use of a static offline dataset of (s, a)-pairs. BC on both the offline and expert dataset does improve the performance, but this still cannot successfully imitate the expert since BC has no way of differentiating random/sub-optimal trajectories from the expert samples. ValueDICE, on the other hand, does explicitly aim to imitate the expert samples; however, in theory, it would require either the offline data (i.e. the replay buffer) or the expert samples to have full coverage over the state-action space. Since our offline dataset is mainly collected from a sub-optimal behavior policy and our expert samples are from a high quality expert, neither our offline nor our expert dataset is likely to have full coverage globally; thus potentially hurting the performance of algorithms like ValueDICE. Note that MILO is still able to perform reasonably well across environments even with these offline and expert datasets.

heading: 6.2 Ablation
text: Impact of Pessimism Figure 3 (Left 2) presents MILO’s performance on two representative environments with and without pessimism (i.e., setting penalty to be zero) added to the imitation objective. Pessimism stabilizes and improves the final performance for MILO. In general, MILO consistently outperforms benchmarks and/or achieves expert level performance for a given set of hyperparameters. See the Appendix for evaluation on other environments.
Behavior with more expert samples We investigate whether MILO is able to achieve expert performance with more expert samples in the two environments (walker and humanoid) that it did not
solve with very small expert datasets in Figure 2. Figure 3 (Right 2) shows that with one trajectory worth of expert samples, MILO is able to achieve expert performance on walker and humanoid.
Impact of Coverage As our analysis suggests, MILO’s performance degrades as the offline data’s coverage over the expert’s state-action space decreases. We use the behavior policy’s value as a surrogate for lower coverage, i.e. a lower value suggests lower coverage. We generate two additional offline datasets for each environment by lowering the performance of the behavior policy. The three datasets are: (1) the original offline datasets used in Table 1 (≈ 25% for Hopper-v2 and≈ 50% for others); (2) ones
that have roughly half the performance of (1) (12% for Hopper-v2 and ≈ 25% for others); and (3) ones collected from a random behavior policy (Random). Table 2 shows that MILO performs reasonably on three environments even with a lower coverage dataset (second column) and achieves more than 20% of the expert performance on three environment even with the Random dataset.

heading: 7 Conclusion
text: MILO investigates how to mitigate covariate shift in IL using an offline dataset of environment interactions that has partial coverage of the expert’s state-action space. We show the effectiveness of MILO both in theory and in practice. In future works, we hope to scale to image-based control.
We want to highlight the potential negative societal/ethical impacts our work. An IL algorithm is only as good as the expert that it is imitating, not only in terms of performance but also with regards to the negative biases and intentions that the demonstrator has. When designing real-world experiments/applications for MILO we believe the users should do their due diligence on removing any negative bias or malicious intent in the demonstrations that they provide.

Review summaries:
This work proposes a new method for offline imitation learning. The setup is slightly different than usual as it supposes that, additionally to the expert dataset we want to imitate, we have access to a large dataset of non-expert interactions. The method is the following:   1) learn a model of the transition kernel from the data.  2) Along with the model, learn an “uncertainty penalty” (the converse of an exploration bonus) that will encourage the agent to stay on the expert support.  3) Train *offline* a min-max model-based policy-optimization with the uncertainty penalty.   The contributions also include theoretical analysis on the proposed method in simplified cases.

This paper studies imitation learning without access to the environment. The agent is provided with two datasets: one containing expert policy state-action pairs and one containing state-action-next state triplets collected from the environment. The authors propose the algorithm MILO that leverages pessimism to mitigate covariate shift and provide performance upper bound in discrete MDPs, kernelized nonlinear regulators, and Gaussian processes. The authors show improved performance achieved by MILO compared to other baselines such as behavioral cloning in MuJoCo control tasks. 

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.  

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.