Minimax Optimal Online Imitation Learning via Replay Estimation
Authors: Gokul Swamy
Nived Rajaraman
Matt Peng
Sanjiban Choudhury
Drew Bagnell
Steven Wu
Jiantao Jiao
Kannan Ramchandran
Venue: NeurIPS 2022 Accept (2022)
Keywords: imitation learning

Abstract: Online imitation learning is the problem of how best to mimic expert demonstrations, given access to the environment or an accurate simulator. Prior work has shown that in the \textit{infinite} sample regime, exact moment matching achieves value equivalence to the expert policy. However, in the \textit{finite} sample regime, even if one has no optimization error, empirical variance can lead to a performance gap that scales with $H^2 / N_{\text{exp}}$ for behavioral cloning and $H / N_{\text{exp}}$ for online moment matching, where $H$ is the horizon and $N_{\text{exp}}$ is the size of the expert dataset. We introduce the technique of ``replay estimation'' to reduce this empirical variance: by repeatedly executing cached expert actions in a stochastic simulator, we compute a smoother expert visitation distribution estimate to match. In the presence of general function approximation, we prove a meta theorem reducing the performance gap of our approach to the \textit{parameter estimation error} for offline classification (i.e. learning the expert policy). In the tabular setting or with linear function approximation, our meta theorem shows that the performance gap incurred by our approach achieves the optimal $\widetilde{O} \left( \min( H^{3/2} / N_{\text{exp}}, H / \sqrt{N_{\text{exp}}} \right)$ dependency, under significantly weaker assumptions compared to prior work. We implement multiple instantiations of our approach on several continuous control tasks and find that we are able to significantly improve policy performance across a variety of dataset sizes.

Minimax Optimal Online Imitation Learning via Replay Estimation
Gokul Swamy" Carnegie Mellon University gswamy@cmu edu
Nived Rajaraman" UC Berkeley
Matthew Peng UC Berkeley
nived. m..ved.rajaaraaanbberkkk.. edu
Sanjiban Choudhury Comell University
J. Andrew Bagnell
Aurora Innovation and Carnegie Mellon University
Zhiwei Steven Wu Carnegie Mellon University
Jiantao Jiao UC Berkeley
Kannan Ramchandran
UC Berkeley
Abstract
Online imitation learning is the problem of how best mimic expert demonstra- tions, given access to the environment shown that in the infinite sample regime, exact moment matching achieves value equivalence to the expert policy. However, in the finite sample regime, even if one has no optimization error, empirical variance can lead to performance gap that scales with H² /Nexp for behavioral cloning and H/ / /Nexp for online moment matching, where H is the horizon and Noxp is the size of the expert dataset. We introduce the technique replay estimation to reduce this empirical variance: by repeatedly executing cached expert actions in a stochastic simulator, we compute smoother expert visitation distribution estimate to match. In the presence para- metric function approximation, we prove meta theorem reducing the performance gap of our approach to the parameter estimation error for offline classification (i.e. learning the expert policy). In the tabular setting with linear function approximation, our meta theorem shows that the performance gap incurred by our approach achieves the optimal Ō (min( 3// ///2/Nexm Nexps H/ /Nexp) dependency, under significantly weaker assumptions compared prior work. We implement multiple instantiations of our approach on several continuous control tasks and find that we are able to significantly improve policy performance across variety of dataset sizes.
accurate simulator. Prior work has
Introduction
In online imitation learning (IL). one is given access to (a) fixed set of expert demonstrations and (b) an environment simulator perform rollouts in. Many online IL approaches fall under the umbrella of solving moment matching problem between learerr and expert trajectory distributions
Ziebart et al. 2008 Ho and Ermon 2016).
min sup [=[ - (f
(1)
well feF
where [.] denotes the expectation
random trajectory {(51, (1)1,a)),- a,),,.,........ ... (SH an)) generated by rolling out म. Swamy et al. [021 show that for an appropriate choice of F, approximate solutions
Eqqual contribution. Correspondence gswamy@cmu.
nived. .ooddrajaaa aaaaarrre@bkkeeeeeccooo edu. edu.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

sampling
moment matching
S
replay estimation
new data
Figure 1: Top: Attempting to exactly match a sanite-sample approximation of expert moments can cause learner to reproduce chance occurrences (e.g. the relatively unlikely flight through the trees). This can lead to policies that perform poorly at test time (c.g. because the learner flies through the trees relatively often). Bottom: Replay estimation reduces the empirical variance in expert demonstrations by repeatedly executing observed expert actions in a stochastic simulator. By generating new trajectories (e.g. 81 + 81 > 82 on the right) that are consistent with expert actions, one can augment the original demonstration set and compute expert moments more accurately.
of (T) have a performance gap linear in the horizon, the gold standard for sequential problems. However, key assumption in their work is that expert moments (f,, a)]) can be estimated arbitrarily well from the available demonstrations. The resulting bounds are therefore purely a function of optimization error. Moving the finite-sample regime introduces an additional concern: the statistical error that stems from the randomness in the data-generating process. When a small set of demonstrations are used in adversarial optimization procedure like (I). the learner may choose to take incorrect actions order to match noisy moments estimated from the dataset, leading policies that perform poorly test-time. Ideally, one would solve this problem by querying the expert for more demonstrations, as in the work of Ross et al. [201I). However, when we are unable to do so, we still have grapple with the question of "how can we smooth out a noisy empirical estimate of expert moments?"
Our answer to this question is the technique of replay estimation. In its most basic form, replay estimation consists of repeatedly executing observed expert actions within a stochastic simulator, terminating rollouts whenever one ventures out of the support of the expert demonstrations. Effec- tively, this approach stitches together parts of different trajectories to generate a smoothed estimate of expert moments. By using the simulator where we know the expert's actions, we can generate more diverse training data that is nevertheless consistent with the expert demonstrations. We argue once conceptually simple, practically feasible, and minimax optimal in
that this technique several settings. Formally, we prove that in the worst case, behavioral cloning has a performance gap x H²2 /Nxpp. online empirical moment matching, x /Nexp= and our approach of replay estimation & min /13/// Nexp, H/ / /Nexp) in the tabular setting well with linear function ap- proximation. Our key insight is that we can use a combination of simulated and empirical rollouts to optimally estimate expert moments. We can then plug this improved estimate into a variety of moment-matche algorithms, for strong test-time policy performance. More explicitly, our work makes the following three contributions:
1. We extend replay estimation (RE) Rajaraman et al. 2020 beyond the tabular and determin- istic setting by introducing the notion of a soft membership oracle and prefix weights.
2. We show how to instantiate the membership oracles for with parametric function ap- proximation and prove a meta-theorem relating the imitation gap of RE to the parameter estimation error for offline classification on the dataset (Theorem 3 Instantiating our main result in the case of linear function approximation, we show how to achieve the best known imitation gap of Ô (H3/2 ///4////////////////// /Nexp) under significantly weaker assumptions compared to prior work Rajaraman et al. 2021
3. We give multiple practical options for constructing performant membership oracles. We then use these approximate oracles to significantly improve the performance of online II. on several continuous control tasks across a variety of dataset sizes. We also investigate the differences between our proposed oracles.

---

We sketch the benefits of replay estimation before providing theoretical and empirical evidence to support our claims.
2 The Replay Estimator
We begin with a tabular vignette to illustrate our key insight in greater detail. We compare two algorithms: offline behavioral cloning (BC) PPoerleau [1989) and online moment matching (MM)
Swamy et al. [021]. Throughout, we focus on learning policies from finite samples.
Suboptimality of Empirical Moment Matching. Consider the MDP in Fig, 2 where the expert always takes the green action. Doing so puts them in 81 or 52 with equal probability. Given that the expert is deterministic and there are few states, BC could easily recover the expert's policy by learning to simply output the observed green action on both states, even when there are very few demonstrations.
Now, what would happen if we tried to match moments of the expert's state-action visitation distribution for this problem? It is rather unlikely that we see exactly equal probabilities for both states in the observed data. If by chance we see $2 more than we see $1, the learner might realize that the only way to match the observed state dis- tribution (a prerequisite for matching the observed state- action distribution) is to occasionally take the red action at $2. In general, this could cause the learner to spend an unnecessary amount of time in $2 which may be un- desirable (e.g. 52 corresponds to the tree-filled area in Fig. (top)). The core issue we hope to illustrate in this example is that by treating the empirical estimate of the expert's 's behavior as perfectly accurate, distribution matching can force the learner to take incorrect actions to minimize training error, leading test-time performance
Unif({s1, 823)
expert empirical replay
P(s)
81
82
S
$/
An MDP where the expert always takes the green action that puts them in the uniform distribution over $1 and 82- Because of full expert support, BC will learn to always take this action both states. However, if the empirical state distribution is more tilted towards 82, will take the incorrect red action.
Figure 2:
degradation. As we will discuss in Sec. 3. this can lead to slow statistical rates x H/ ////VN /Nexp.
Suboptimality of Behavioral Cloning. Because it does not account for the covariate shift that results from policy action choices, behavioral cloning can lead to a quadratic compounding of errors and poor test time performance Ross et al. [2011]. Consider, for example, the MDP in Fig, B
Let us assume that the expert always takes the green action, dropping them in state in the top row with uniform probability. In small demonstration set, we might not see expert actions at some states in the top row. At all such states, will have no idea of what to do. In contrast, MM will take the green action as doing so might send the learner back to state with positive demonstration support. Thus for this problem, MM will recover the optimal policy while BC will not. As we will discuss in Sec. 3] this leads errors x H² /Nexp in the worst case.
Unif({ (11 I sn
91
82
STi
expert empirical replay
P(s)
S
5
5
8
Figure 3: The expert always takes the green action, which places it in a uniform distribution over .1,..., .1,....,,........ At states where we have demonstrations (e.g. 81, 82), both BC and MM will take the same, correct action. However, at states where we have no demonstrations (e.g. 83), MM will correctly take the green action get back states with demonstration support, while BC might not.
Replay Estimation. The previous two examples show us that there exist simple MDPs for which BC or MM will not recover the expert's policy. This begs the question: is it possible to do better than both worlds and recover the optimal policy on both problems with single algorithm? turns out it is indeed possible to do so, via the technique of replay estimation. In its simplest form,
3

---

replay estimation involves playing a cached expert action whenever possible and rt-starting the rollout if one ventures out of the support of the demonstrations. Then, one appends these rollouts to the demonstration set, treating them as additional training data while biased, they are consistent with observed expert behavior. Intuitively, repeated simulation has a smoothing effect on the training data as doing so marginalizes out the statistical error that comes from the stochasticity of the dynamics. We can see this point more explicitly by considering the above two MDP examples: in Fig. 2 repeatedly playing the green action and appending these rollouts to the expert dataset would bring us much closer to a uniform distribution over 81 and 82. Similarly, in Fig. [ replay estimation would bring us
toward uniform distribution over the states (s1, Sn,) } in the expert demonstrations.
We could then plug this improved distribution estimate into the MM procedure (T). Notice how doing so would cause MM to be highly likely to recover the optimal policy on both MDPs. For example, in Fig- 2 replay estimation would make the learner much less likely to play the red action in 82, turns out this fact is sufficient to establish statistical optimality in the tabular setting and with linear function approximation, with an error rate x min( 1///////////// /Nexp//expppppppp H1 /Nxp) In short, replay estimation is practical technique for reducing some of the finite- finite-sample variance in expert demonstrations that enables MM to perform optimally in the finite sample regime. We now provide some intuition on how to generalize this approach to beyond the tabular setting.
Leaving the Tabular Setting. The prior work of Rajaraman et al. 2020) considers the tabular setting; this characteristic makes it easy answer the question of "on what states do we know the expert's 's action?" To enable us to answer this question more generally, we introduce the notion of a membership oracle M S + {0, 13. Explicitly, M(s)
for states where we know the expert's
action well (e.g. states where we have lots of similar demonstrations) and M(s) = 0 otherwise. Alternatively, .Ms)) on states where BC, which attempts directly output expert actions, is accurate and 0 otherwise.
We can then compute expert moments by splitting on the output of the membership oracle:
E [x=[[(s,a)] = F&E [((s, 1)] + F [f (()( =)) 0]
(2)
(4x)
Note that the indicators in (i) and (ii) are complements of each other, rendering the above sum a valid estimate of the expert moment. As we know the expert action well wherever M(s) 1, simulated rollouts of the BC policy approximates (i) well; on the other hand we resort to naive empirical estimate to approximate (ii), as we do not know enough about the expert's ' action at these states to accurately generate additional demonstrations via BC rollouts. In general, we relax A to a soft membership oracle Zadeh (1965). in order to handle uncertainty in how well we know the expert's s action at a given state. We proceed by first analyzing the statistical properties of applying MM to this bipartite estimator before discussing practical constructions of performant membership oracles.
3 Theoretical Analysis
The proofs of all results in this section are deferred to appendix A We begin by introducing some notation.
labular Sample Complexity
Notation. Let A(X) denote the probability simplex over set X and let S, respectively denote greater than, lesser than and equality up constants. We study the IL problem in the episodic MDP setting with state space S, action space A and horizon H. We assume that the transition, reward function and policies can be non-stationary. The MDP transition is denoted / {p, PI. ... P1- where is the initial state distribu- tion and P. 8 x - A(S). while the reward function is denoted
min| 185
Hy
[(
SHH Number of Trajs., Nexp
fri.... .TH where r : S x A + j0, 1]. In the online imitation learning setting, the learner has access to a finite dataset D of Nexp trajectories (i.e. the sequences of states visited and actions played) generated by rolling out expert policy E. Importantly, the learner does not observe rewards during rollouts. The fundamen- tal goal of the learner is to learn a policy सt such that the
Figure 4: In the tabular setting, RE in- herits the superior low-data performance of moment-matching approaches and is able to perform better than both MM and BC with enough data.

---

imitation gap, - is small. Here, (T) denotes the expected value of the policy т, [Ex - rx( ar) Ep[:] denotes the empirical expectation computed using trajectories from the dataset D.
Behavior cloning. A standard approach for imitation learning is behavioral cloning (BC) Pomerlea [989], which trains classifier from expert states to expert actions. More formally, for the 0-1 loss (or a continuous proxy) BC minimizes the empirical classification error,
arg min D Tlll
Σ "1 H
1
((To(s)))
It is known from Ross et al. imitation gap for BC is always S [S] HT The fast 1/Nexp rate comes from the fact that on
2011 and Rajaraman et al. [2020] that in the tabular setting, the expected
states where the learner observes the deterministic expert's 's actions, BC simply replays them. Thus, the performance gap is proportional to the mass of unseen states, which decays 1/Nuxp. The H² dependence comes from the fact that single mistake can take the learner out of the distribution of the expert, causing to make mistakes for the rest of the horizon. Importantly, this bound tight- - i.e. there exists an MDP instance on which BC incurs this imitation gap (See Appendix A.3)
Theorem (Theorem 6.1 of Rajaraman al. (0200). -(=)) "[]] min H, /
). Then, there exists tabular MDP instance
such that the BC incurs, [J n -
Moment matching, Many standard algorithms in the IL literature fall under the empirical moment matching framework (e.g. HH and Ermon [2016]. MaxEnt IRL Ziebart et al. [008 ); see Table 3 of SSwamy et al. [2021] for more examples. Reward moment matching corresponds to finding a policy which best matches the state-action visitation measure of x in the sense of minimizing an
Integral Probability Metric (IPM) Müller [1997]. In the finite sample setting, the empirical moment
matching learner Th! attempts to best match the empirical state-visitation measure. Namely,
Σ/ f((s), az)
E fels,, as)
.N π
E arg min sup E,, fEl
-ED
(4)
JEF
H
H
First we show an upper bound on the imitation gap incurred by empirical moment matching.
Theorem 3.1. Consider the empirical moment matching learner π H8 (eq. (4), ), instantiated with an ap-
propriate discriminator class F. The imitation gap satisfies E [J(m ") - [==")) S //S//Noox
The proof of this result can be found in Appendix A.2 It turns out that this guarantee is essentially tight for empirical moment matching, answering an open question from Rajaraman et al. 2202)).
Theorem 3.2. H 2 4, there is tabular IL instance with 2 states and actions on which with constant probability, the empirical moment matching learner (eq. 44 incurs, J ( (TT Z H/ HIV Newp-
The proof of this result is deferred to appendix A.4 The proof of this lower bound exploits the fact that the data generation process in the dataset is inherently random. Consider slight modification of the MDP instance shown fig. 2 where the reward function is 0 for = 1. For Z 2, the transition function is absorbing at both states; the reward function equals 1 at the state 81 for any action and is 0 everywhere else. Then, the expert state distribution at time 2 and every time thereon in uniform across the two states, {1/2, 1/2]. However, in the dataset D, the learner sees noisy realization of this distribution in the dataset of the form {1/2 8, 1/2 + 5} for [8] =1/ this noise, the empirical moment matching learner may be encouraged to deviate from the expert's 's observed behavior and pick the red action at 82 as this results in better match to the empirical state visitation measures at every point in the rest of the episode- state-action visitation measure. The learnrr is willing to pick action different from what the expert played in order to better match the inherently noisy empirical state-action visitation distribution.
Because of
prerequisite to matching the empirical
Remark 1. Theorems I and 3.2 are separate lower bound IL instances against the performance of BC and empirical moment matching. On the uniform mixture of the two MDPs (i.e. deciding the underlying MDP based on the outcome of a fair coin), with constant probability, both J "
[ Z S)H22 /Nexp and {x (TT Z H/ /H/NNess On this mixture instance, training both and empirical moment matching and choosing the better of the two is also statistically suboptimal.
5

---

3.1 Replay Estimation
A natural question at this point is whether there is an algorithm that is better than both worlds, i.c.
algorithm which can outperform the worst case imitation gap of empirical moment matching and BC The answer to this question in the tabular setting was recently provided by Rajaraman et al. (2020) who propose Mimic-MD, achieving better performance than both BC and MM This improvement is possible because BC does not use any dynamics information and MM does not leverage the knowledge of where expert actions are known. However, is unclear how to extend this approach beyond the tabular setting the algorithm relies large measure of states being visited in the demonstrations.
Algorithm 1:
Replay Estimation (RE)
Input:
Expert demonstrations D. policy class 11, moment class F Di- Fr. simulator SIM, ALG which returns a membership oracle given a dataset; Partition the dataset D into D2 and D2;
2
Using ALG, train the membership oracle M on D1;
3
Train BC using behavior cloning on D1;
Roll out 8C in SIM Nreplay times to construct a new dataset, Dreplay:
5 Define prefix weights P(s1...)= = II M(sr, t'):
Define,
F(S) = IF E Deeplny
1 H
H E at) (P()I...)) + ED
H Em -1-
P(SL...))
7
Output: x solution to the moment mattenig problem:
arg min sup E =enl JEF
1 H
Σ Se(se, ai)
- B())
(5)
To handle this challenge, we introduce the notion of soft membership oracle [adeh 1965)] .M S x [H] + [0, 1] which captures the learner's 's inherent uncertainty in the expert's 's actions at a state at each point in an episode. The soft membership oracle assigns high weight to state if BC is likely to closely agree with the expert policy and gives a lower weight to states where BC is likely to be inaccurate. By this definition, if the membership oracle is consistently large at all the states visited in a trajectory, we can be confident that trajectory generated by BC is as though it was rollout from the expert policy. Formally, for any function g and time 1,... H, we have the decomposition,
Ex=lg(s a:)] E [g(s,, at (S1 Exe [g(s,, at (1 -P(ss...)))
(4)
(4)
where P(s)...) is defined as the prefix weight TE.=IM([s f'). We need to use prefix weights instead of the single-sample weights sketched in the previous section to account for the probability of BC getting to the current state in the same manner the expert would have. Because of the high accuracy of BC on segments with high prefix weights, in eq. (6). (t) can be approximated by replacing the expectation over n by that over 8C i.e. replay estimation. On the other hand, since the prefix weight is low on the remaining trajectories in (ii), we know that BC is inaccurate, so we resort to using a simple empirical estimate estimate this term.
While we leave the particular choice of the soft membership oracle flexible, intuitively, states at which BC closely agrees with the expert policy should be given high weight while where those where BC is inaccurate should be weighted lower. In Section 4 we discuss several practical approaches to designing such a soft membership oracle. We first prove a generic policy performance guarantee for the outputs of our algorithm as a function of the choice of M.
Theorem 2. Consider the policy SE returned by Algorithum / Assume that F E Il and the ground truth reward function Yt E Fi, which is assumed to be symmetric (St E F, > -Si E F.) and bounded (For all fo E Fu. lIfelloo S 1). Choose Dil, D2l @(Nexp) and suppose Nreplay oo.

---

With probability 2 1 - 36.
log (Fmax H
)(x²) - Si + L2 +
Nexp
where F Fmax - m [F,l, and,
(V (19())))))
L1 A H2 E
(8)
H
log (Fmax H/5) E E M(s. 0]
C2 2 11322
NoxP
We discuss a proof of this result Appendix B and include bounds when N Nreplay is finite.
Remark 2. Note that Theorem 2 can be extended infinite function families using the standard technique of replacing (Fl by the log-covering number of Fe in say, the L2 norm, for Newv" For ease of exposition here, we stick to the case where Ft is finite.
The term Li measures how accurate BC is on states from expert trajectories where M(s () is large. Intuitively, we set M(s. 1) = on states where BC is accurate and ((ts. () 0 elsewhere, we would expect this term to be small. Le can be thought of a measure of BC is coverage: it tells us how much of the expert's visitation distribution we believe BC to inaccurate on. BC has good coverage (i.e. M(s is small on expert trajectories), we expect this term to be small.
Prima facie, one might think that because Ci resembles the imitation gap of BC and L2 resembles that of MM RE can only perform well the best of BC (xx H² /Nexp) and MM (xx H/ /Nexp) on a given instance. However, with a careful choice of M, one can achieve "better than both worlds" statistical rates. In particular, since RE generalization of Mimic-Me of RRaaraman et al.) [2020) in the tabular setting, an appropriately initialized version of RE achieves the optimal imitation gap of min
H
log
[S|
and strictly improves over both BC and MM.
Nesp
exp
We now show how to extend this result and instantiate the membership oracle for parametric function approximation and provide a statistical guarantee under a particular margin assumption.
3.2 Parametric Function Approximation: Reduction to Offline Classification
While BC can be thought of as reduction of IL offline classification, the algorithm does not take into account the knowledge of the transition of the MDP. This is reflected in the quadratic dependency in the horizon, еrror compounding Ross al. [2011. In this section, we study a novel reduction II. with parametric function approximation to parameter estimation in offline classification, which we define formally. We provide a provable guarantee, assuming the learner has access a classification oracle and the underlying function class admits a Lipschitz parameterization.
Definition (IL with function- approximation). In this setting, for each t E [H), there is parameter class O C [B, the wnit La ball in dimensions, and an associated function class {Se, Or E ),). For each & [H] there exists an unknown 0f O such that Vs S.
T E (8) = arg max Joz (g,(8,0) a).
aEA
Definition 2 (Lipschitz parameterizationn. function class G {go 0 E ) where ga(+) X R is said to satisfy L-Lipschitz parameterization if. (19a(.) Go (-)llo S LI - 0'1/2- In other words, for each x X, go(x) is an Lilipsschtzzuuncttioo
in the L2 norm.
Assumption 1. For each the class (So, Oc E O.) is L-Lipschitz in its parameterization, 0, O
To deal with parametric function approximation, we assume that the learner has access to an offline classification oracle which, given a dataset of classification examples, approximately returns the underlying ground truth parameter. More formally,
7

---

Assumption 2 (Offline classification oracle). We assume that the learner has access to a multi- class classification oracle, which given n. examples of the form, (s', ad) where o² Lid. DD and a² arg max uEA J (s', a), returns a 0 E θ such that, with probability 2 - 6, 110 0 ll2 S E
We assume that this classification oracle is used by RE to train the BC policy in Line 3 of Algorithm
A careful reader might note that Assumption 2 asks for a slightly stronger requirement than just finding a classifier with small generalization error (which need not be close to the ground truth 8*). The latter problem was studied in Danicly et al. (2013) who show that the Natarajan dimension, up to log- log-factors in the number of classes (i.c. number of actions) captures the gencralization error of the best learner, which scales as (1/nn) given n classification examples. Under some mild regularity (coverage) assumptions on the input distribution D, we show that the generalization guarantee carries over to learning the parameter 8, for example in linear classification (Appendix (B.4). In particular, show that Eo,n,8 S d+Hog(1/5) when so = (0,-): The membership oracle we study is defined below,
+1
3a E A such that, Val E A, Sapo (8, a) Sape (8, d') 2 2T.E,, 5/xp.5/HH (10)
M(s
0 otherwise.
Intuitively, M assigns a state as +1 if BC classifies it with a significant margin as some action.
Finally, we introduce the main assumption on the IL instances we study. We assume that the classification problems solved by BC each E [M] satisfy a margin condition.
Assumption (Weak margin condition). The weak margin condition assumes that for each t. there is no classifier 0 O2 such that for large mass of states, fo(st, T. (s:) - max (st) fo a), i.e. the "margin" from the nearest classification boundary, is small. Formally, the weak-margin condition with parameter μ states that, for any E Oe and 1j
S 1/p,
Prą (f((s::T(s))
max a
fo(st,a) a)> PM
(1I)
The weak margin condition only assumes that there is at least an exponentially small (in n) mass of states with margin at least 1J. Smaller y indicates larger mass away from any decision boundary.
suffices to assume that eq. (11 is only true for 0 as the classifier in Assumption for our guarantees (Theorem By to hold.
Under these three assumptions and using the membership oracle defined above, we can provide a strong guarantee for RE:
Theorem 3.
For IL with parametric function appreximation, under Assumptions D instatiating RE ensures that with probability 2 1 48,
appropriately
LL log (Fiax H/8) Et
Ne..Nocp.6/11 6/HH
log Fmaxx//0
[( - J( S H312 $
(12)
Nexp
H
Nexp
In Appendix B.4 we explicitly instantiate these guarantees for the special case of where the expert follows a linear classifier acting on a known set of feature representations of the state-action pairs. The prior work of Rajaraman et al. [2021] shows an imitation gap of Ō (H3/2 /Nexp) for RE in the linear expert setting, under the restrictive assumptions of binary action space, and a certain uniform distribution assumption on the feature distribution. In contrast, the guarantee of Theorem 3 can be used to achieve the same guarantee, for general action spaces and under a significant weakening of the uniform feature distribution constraint to multi-class analog of the strong distribution assumption Audibert and Tsybakov| [2007). Under this assumption, we show that the best achievable parameter recovery eπтor, Nop.No and the best achievable 0-1 generalization error (L [mmx..e fg(s, a) & max,e A fe- (s, a)]. in the notation of Assumption 2) for classification, Eglass
NexpNo match up to problem dependent constants.
The best known statistical guarantees on imitation gap for BC (Ros et al. (011] and MM are,
E/n
oclass Cr. Neep. 5/H
J( - J(x < G)) A H²
and
(13)
H
log (Fmax H/18)
B( -J( < Gap( A H.
(14)
Nexp
8

---

The guarantee in Theorem 3 for RE therefore grows as // whenever we are able to establish that &e,, Nexp...) = Colass
Ne..N up to problem dependent constants. Our guarantees for RE in the presence of parametric function approximation thus give reasons to expect RE to improve the performance of MM. We now turn our attention to validating this in practice.
4 Practical Algorithm
When considering RE (Alg. DT. two main questions arise: (i) How does construct a membership oracle in practice, especially when action spaces may be continuous?, and (ii) How does one solve the moment matching problem (eq. (5 We provide potential answers to both of these questions.
Walller21BBulletEinwo Nepp
Noiry Willeer2DBllleein Neep
Rallout Prefix Wiight
3500
1500
PaD Pas. Plc Pop
0:94
0.92
0.92
1000
0.94
0.996
0.06
1000
500
0.92
0.96
0.92
0:98
0.99
1000000 200000 20000 4000000 550000
200000
800000
600000
Number of Trn. Steps
Number of Erv. Steps
Figure 5: Left: All variants of RE are able to nearly match expert performance while MM struggles to make any progress. Center: We add i.i.d. noise to the environment to make the control problem more challenging. RE is still able to match expert performance, unlike MM Right: We compute correlations between the idealized prefix weights of MExx and the other oracles and see Muax correlate most.
Nikey RollerDDBulleEer
Nitsy HoppeeBllleessv
WulleeDD By
HoppetBullstEm -0 /2y m
2300
1100
250
200-
30
Figure 6: We see RE with Mux improve the performance of MM on the Noisy Walker2DBullerFan and HopperBulletEaw tasks. We RE (and MM out-perform BC on the initial-state-pertuie pertture Walker2DBulletEaw and HopperBulletEny tasks.
Membership Oracle. M represents how uncertain BC about the expert's action. Ideally, we would be able to query the expert policy for an action and weight by the distance between the BC and expert actions, e.g. for sigmoid function a and constants μ, 3,
n7 BC (s) x (s)
MExP (s, () = |4
(15)
3
However, in the non-interactive setting, we can only approximate this quantity. The first approxima- tion we propose using inspired by Random Network Distillation (RND) [Burda et al.) 2018). used by Wang et al. [019 to estimate the support of the expert policy. We instead propose using RND as measure of how uncertain BC is about expert actions. That is,
11
n
BC
BC (8)
MEND (8, =
(16)
where BC is network trained to imitate 8C expert states. To train 80 evaluate HC expert states and plug this new dataset into their standard IL pipeline.
We can also utilize other uncertainty measures like the disagreement of an ensemble, which has shown success on various simulated sequential decision making tasks (Pathak et al., 2019]. Past work by Brantley et al. 2009 utilizes the variance of set of BC learners as regularizer top of standard

---

BC error. We instead propose using it to weight the two halves of our estimator. Instead, we define,
Var(m BC(A) (s)) (s))
MvAn (s,
(17)
B
where { x Bo(k)) } are BC policies trained with different initializations, which produces sufficient diversity when using deep networks as function approximators. Lastly, we can also use the maximum difference in the ensemble Kidambi et al. [202T] as disagreement measure:
max, T. (s() BC
". (
MHxx (s,
(18)
B
For computing prefix weights, we use the average of distances up till the current timestep instead of the sum of distances that would follow directly from the idealized algorithm presented in Sec. 3 This modification serves to improve the numerical stability of our method.
Moment Matching. We implement approximate Nash equilibrium computation of (5) by running a
no-regret learner against
counterpart (Swamy ct al. [021]] Our approach is related to the GAlL. algorithm of Ho and Ermon (2066) which we improve in ways: (i) we use a general Integral Probability Metric Müller original paper which improves the representation power of the the discriminator, (ii) we add in
1997 instead of the
Divergence used in the
gradient penalties to the discriminator, which improves convergence rates Gulrajani et al. (2017). (iii we solve the entropy-regulariie policy optimizer, for improved sample efficiency, and (iv) we use optimistic mirror descent instead of gradient descent as our optimization algorithm for both players, giving us faster convergence to Nash equilibria, both in theory SSrggaanis al. [2015) and in practice Daskalakis et al. [2017]. Together, these changes lead to an implementation which significantly out out-performs the original, giving us a strong baseline to compare against. We include an ablation to confirm this fact in Appendix C We emphasize that the RE technique can be used to improve any online moment matching algorithm and that the above description is merely the approach we chose for this paper.
forward problem via Soft-Actor Critic Haarnoja et al. [0118] as the
5 Experimental Results
We now quantify the empirical benefits of RE on several continuous control tasks from the the PyBullet suite Coumans and Bai (2016-2019). All the task consider have long horizons (H 1000) and we use relatively few demonstrations. (Nexp < 20). We set Nreplay as 100 BC rollouts (Line 4 of Algorithm D. We test all four membership oracles from the previous section (MEXP as an idcalized target, MEDD MyAR. and MMax as practical solutions). In Fig. (left). we see that with only twelve trajectories. is able reliably match expert performance for all oracles considered, while MM is not. The environment considered in this experiment is nearly deterministic, indicating that RE can help even when the environment is not stochastic. We hypothesize that the randomness in the initial state is sufficient for replay estimation to generate significant improved estimate of the state-action visitation measure. This improvement is especially interesting considering both of the examples we study in Sec. 3 were heavily stochastic. We see similar result in Fig. (center) where we add in i.idd. noise to the environment dynamics at each timestep. This makes the problem significantly more challenging than the standard version of the Walker task. RE is still able to match expert performance, with Mpax working notably well. The correlation plot in Fig. 3 (right) shows us MMAx appears to be best correlated with the idealized prefix weights, MEXP under the state distribution induced by BC. Because of its superior performance, we use MMAX for the rest of our experiments. In the left half of Fig. 6 see RE improve the performance of MM In the right half, BC in responding to an extremely tiny amount of noise added to the initial velocity of the agent (similar to the experiments of Reddy al. (2019)-
RE out-perform
defer more details to Appendix C). These results indicate that RE can out-perform MM and BC, agreeing with our theory. We release our code at
https:// tigbtth..... com/ www.cooooooceeecccc
est
2 After the initial publication of this paper, we continued to tune our baseline and fix bugs in our method's implementation. We then updated the result plots for both of the Noisy environments. Compared to their original performances, MM scores higher and RE scores lower. However, the RE still significantly out-performs MM.
10

---

Acknowledgements
ZSW is supported in part by the NSF FAI Award #1939606, Google Faculty Research Award, J.P. Morgan Faculty Award, a Facebook Research Award, an Okawa Foundation Research Grant, and Mozilla Research Grant. GS is supported by GPU award from NVIDIA. NR and JJ were partially supported by NSF Grants 11S-1901252, and CCF-1909499. KR is supported by ARO fund 051242-002.
References
Jean-Y Yves Audibert and Alexandre B. Tsybakov. Fast learning rates for plug-in classifiers. The Annals of Statistics, 35(2):608 - 633, 2007. doi: 10.2114009055500000000001 URI,tttps
//doi. 10.g//10.00 000.24440000555660000001117
Kianté Brantley, Wen Sun, and Mikael Henaff. Disagreemeent-regulaaiizee imitation learning. In International Conference on Learning Representations, 2019.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 12894, 2018.
Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. http://pybbuleet. org 2016-2019.
Amit Daniely and Shai Shalev-Shwartz. Optimal learners for multiclass problems. CoRR, abs/1405.2444 2420, 2014.
wrw./oog/////ae..e 2420
Amit Daniely, Sivan Sabato, Shai the ERM principle. CoRR, abs/ htt//300828999 2013. URL ltttp://axxi, org/ abs//308/1112 2893
and Shai Shalev- Multiclass learnability and
Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng. Training gans with optimism. arXiv preprint arXiv:1711.0011 2017.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein gans, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor, 2018.
Yanjun Han, Jiantao Jiao, and Tsachy Weissman. Minimax estimation of discrete distributions under & loss, 2015.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning, 2016.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims, Morel Model- based offline reinforcement learning, 2021.
Enno Mammen and Alexandre B. Tsybakov. Smooth discrimination analysis. The Anmals of Statistics, 1786):188-1 1999. ISSN 00905364. URL http::ttww.wwwjette wttp://jssoooooossttaibbos.e wtt.://rrt///////050000
David McAllester and Robert E. Schapire. On the convergence rate of good-turing estimators. In Proceedings of the Thirleenth Annual Conference Computational Learning Theory, COLT '00, page 1-6, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc. ISBN 155860703X.
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429 443, 1997.
Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagreement, 2019.
Dean A Pomerleau, Alvinn: An autonomous land vehicle in neural network. Technical report, CARNEGIE-MELLON UNIV PITTSBURGH PA ARTIFICIAL INTELLIGENCE AND PSY. CHOLOGY 1989.
Antonin Raffin. RI baselines3 zoo. bttts://github. www/DLR-RR///// Btt.a/RRr/-r-eelelleee-- basel Saaselines3-aan zoo
2020.

---

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines: Reliable reinforcement learning imp.emmttoiipllll Journal of Machine Learning Research, 22(268):1-8 2021. URI.
www.e/orrpppppp/eeeeeeeeee 1364. html
Nived Rajaraman, Lin F Yang, Jiantao Jiao, and Kannan Ramachandran. Toward the fundamental limits of imitation learning. arXiv preprint arXiv:2009.05990 2020.
Nived Rajaraman, Yanjun Han, Lin Yang. Jingbo Liu, Jiantao Jiao, and Kannan Ramchandran, On the value of interaction and function approximation in imitation learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL attps://openreviiet http://eww..eeeeoooooooa//aaaaaa
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards, 2019.
Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 66 -- 668, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL WRRLhttpps//pprccediinn mlr. wttppeeee/rrrsslaaaa.aeaaaa html
Stephane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning, 2011.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707.06347 2017. URL hRtp::///arx aat////////// 063477
Gokul Swamy, Sanjiban Choudhury, J. Andrew Bagnell, and Zhiwei Steven Wu. Of moments and matching: A game-theoretic framework for closing the imitation gap. 2021. URL https
7/arxiv .rgg http:/abs21100...0 03236
Vasilis Syrgkanis, Alekh Agarwal, Haipeng Luo, and Robert E Schapire. Fast convergence of regularized learning in games. arXiv preprint arXiv: 1507.00407 2015.
Ruohan Wang, Carlo Ciliberto, Pierluigi Vito Amadori, and Yiannis Demiris. Random expert distillation: Imitation learning via expert policy support estimation. In International Conference on Machine Learning, pages 6536-6544, PMLR, 2019.
L.A. Zadeh. Fuzzy sets. Information and Control, 8(3):338-353, 1965. ISSN 0019-9958. doi: http:////....1100.o0.o0/0/99/aa0 sctooc///rttaoo///////0999588119
URL https://ww//w http://wwwwcieenneerrt....ooo
Brian D) Ziebart, Andrew I. Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438. Chicago, IL, USA, 2008.
12


Reviews:
[{"review_id": "r15l4eBMWcF", "reviewer": "Reviewer_hKQj", "summary": "This paper studies the problem of online imitation where the environment interaction is allowed. This paper focuses on the moment matching method. The vanilla method directly estimates the expert's state-action distribution and therefore suffers a sample barrier issue when the sample size is sufficiently large. To address this issue, this paper extends the idea presented in Rajaraman et al. (2020) with new algorithm designs and analyses. Specifically, the proposed algorithm in this paper can work with function approximation, whereas the MIMIC-MD algorithm in Rajaraman et al. (2020) is only applied to the tabular setting. In terms of the theoretical analysis, this paper shows an imitation gap $\\widetilde{O}(H^{3/2} d^{5/4}/N_{\\exp})$ in the linear function approximation setting, under weaker assumptions compared with that in Rajaraman et al. (2021a). Finally, experiments on tasks from the PyBullet suite show the efficacy of the proposed algorithm.\n\n=====\n\nAfter the rebuttal, most of my concerns are well addressed. I recommend accepting this paper for two reasons. First, it makes a significant step in improving the worst-case bound of MM-style algorithms. Second, it investigates an interesting experiment setup. The main criticism of this paper includes insufficient related work (refer to the Limitations section) and weak experiment results (refer to my last reply on the discussion page). Based on the assumption that the authors can fix all these issues before the final version, I recommend accepting this paper and increase my score from 5 to 7. \n\n ", "questions": "\n- Can the authors help explain why MM has different algorithm behaviors in Figure 2 and Figure 3?\n\n This paper claims that MM might fail on the MDP shown in Figure 2, but later on, it says that it can take the green action (expert action) on the MDP shown in Figure 3. However, the theoretical guarantee in Theorem 3.1 for MM cannot tell me why MM is sometimes good and sometimes bad. I realize algorithm properties of MM are much studied in the prior work [1]. Are there new insights in this paper? \n\n- Why is the performance of MM bad on Pybullet Env? Is the superior performance of RE really benefited from a better estimation?\n\nI observe that the performance of MM is bad in Figure 5, and I am surprised by this result because I believe that MM should perform well. To support my thought, I conducted an experiment by running the DAC algorithm without any tunning (the experiment configuration follows the code given by the authors). DAC is a particular case of MM when the function class allows a KL-divergence to measure the state-action distribution matching. Then, I find that with 10 expert trajectories, DAC can achieve the performance of 1500 ($\\pm 500$) on HopperBulletEnv and 2000 ($\\pm 200$) on Walker2dBulletEnv. I believe further tuning can boost the performance. \n\nThe main concern is when can RE be better than MM? MM-style algorithms are known to perform well on many tasks, even though it may suffer a sample barrier issue in the worst case. The authors are recommended for the discussion in Section 3.3 of [1], in which it is shown that RE has no clear advantage over MM in instances similar to locomotion control. ", "limitations": "I appreciate and like the theoretical results developed in this paper. However, I still have some concerns. \n\n**Important references are missing**\n\nI believe the current references are somehow biased. I would suggest that the authors add a section of related work, if the authors agree. For example, similar instances in Figure 2 and Figure 3 and associated upper bound and lower bound for MM have been studied in [1]. Furthermore, the idea of RE (i.e., learning a BC policy to rollout) is also presented in [1] (though [1] studied the tabular setting). I think the connection with [1] should be clarified. It would be great to give credits for prior studies of MM algorithms (in other papers, it is called adversarial imitation learning algorithms) [2, 3, 4]. Otherwise, readers cannot give a full picture of MM algorithms. \n\n\n\n**Gaps Between Theory and Practice**\n\n- It is not clear for what kind of instances RE is better.  Note that we cannot claim that an algorithm is better than another if the upper bound is better.  This is because upper bounds could be loose. Even for instances in Figure 2 and Figure 3 of this paper, I do not know whether RE can be better than BC or MM, since there are no experiments or theoretical guarantees. Compared with the lower bound instances for BC and MM, I know RE can show an advantage in the large sample regime. But I do not know whether these hard cases occur in practice (or at least, whether these hard instances are similar to tasks in the benchmark considered in this paper). \n- The improvement of RE is in the large sample regime but the experiment is in a small sample regime. From Figure 4, it seems that MM is better in a large sample regime ($N_{\\exp} > |S|H$). However, the improvement magnitude should be minor in this regime, as the imitation gap of MM is already smaller than $\\mathcal{O}(\\sqrt{H})$. To better see this, let us consider the locomotion control tasks, in which H=1000. Then, the imitation gap of MM is smaller than 31, meaning that the improvement of a better algorithm can not be larger than 31. However, in experiments, it is shown that RE can be significantly better than MM in a small sample regime. \n- MM should have good performance. This paper shows that MM has bad performance on Walekr2dBulletEnv. Though I do not find a good reference from the literature, my experiments show that MM can perform well; refer to my comments in the Questions section.\n\nI would suggest that this paper can be rephrased to emphasize the theoretical contributions, if the authors agree. The motivation in section 2 in terms of ''better than both worlds'' is not very sound to me (the authors are encouraged to explain more about this motivation). Instead, I believe that a study of minimax optimality is sufficient to motivate this paper. \n\n**Minor Issues**\n\n- A right bracket is missing in Line 15 of the abstract.\n\n- The MM family is too large and it is not clear which algorithm is studied in the main paper. In Appendix, it seems that TV distance is considered. I believe this point should be clarified in the main paper.  It is recommended to give more explanations about MM (though I prefer the name of state-action distribution matching). Otherwise, it is hard to follow the discussion in Section 2.\n\n- The sample regime in theoretical guarantees should be clarified. For example, the bound in Theorem 3.1 should be $\\min \\\\{ H, H\\sqrt{|\\mathcal{S}|/N_{\\exp}} \\\\}$. Otherwise, one may be confused that if $N_{\\exp}= 1$, the theoretical guarantee in Theorem 3.1 is meaningless. In other words, the bound in Theorem 3.1 is only meaningful in a large sample regime where $N_{\\exp} > |\\mathcal{S}|$.\n\n- In experiments, this paper claims that gradient penalty can improve the convergence rate. But, I do not realize this point. Instead, I know that the gradient penalty can help reduce the generalization error; see the discussion about the discriminator in [3-4].\n\n\nI would like to recommend accepting this paper if the above concerns are well addressed. Basically, some misleading claims should be clarified. ", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- This paper introduces several new algorithm designs (soft membership and prefix weight). These components make the idea in Rajaraman et al. (2020)  more generable.\n- This paper provides new and better theoretical guarantees for the proposed algorithm. These theoretical results are essential for studying the minimax optimality on a statistical side.\n\n**Weakness**\n\n- Some relevant references are missing. The discussion in Section 2 has been presented in the prior work; see [1]. The idea of learning a BC policy and rolling out it to collect more trajectories to improve the estimation is also considered in [1]. The theoretical analysis of MM (Theorem 3.1 and Theorem 3.2 in this paper) has been conducted in [1]. Prior studies [2-4] about MM (and variants) are not even mentioned. \n- Gaps between theory and practice. This paper studies the worst-case bound of RE. But, it is not clear for which kind of instances and what kind of setting, RE can be better than other algorithms. Though this paper presents some numerical results to answer such questions, these results are not very convincing. Please refer to the detailed comments below. \n\n[1] T. Xu, Z. Li, Y. Yu, and Z.-Q. Luo. On generalization of adversarial imitation learning and beyond. arXiv, 2106.10424, 2022.\n\n[2] Z. Liu, Y. Zhang, Z. Fu, Z. Yang, and Z. Wang. Provably efficient generative adversarial imitation learning for online and offline setting with linear function approximation. arXiv, 2108.08765, 2021b.\n\n[3] Xu, T., Li, Z., & Yu, Y. (2020). Error bounds of imitating policies and environments. *Advances in Neural Information Processing Systems*, *33*, 15737-15749.\n\n[4] Y. Wang, T. Liu, Z. Yang, X. Li, Z. Wang, and T. Zhao. On computation and generalization of generative adversarial imitation learning. In Proceedings of the 8th International Conference on Learning Representations, 2020.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "W2qkKGB8IGq", "reviewer": "Reviewer_Jnmr", "summary": "This paper studies the imitation learning problem with general function approximation. The authors utilized the replay estimation (RE) technique to improve the moment matching (MM) method. To analyze RE with general function approximation, they first introduced the membership oracle M and proved a general performance guarantee as a function of M. Then they instantiated the membership oracle with an offline classification oracle and proved a meta-theorem under some assumptions on this offline classification oracle. Based on this result, they argued that when the parameter error matches the generalization error (i.e., $\\mathcal{E} _ {\\Theta, n, \\delta} \\asymp \\mathcal{E}_{\\Theta, n, \\delta}^{\\text {class}}$), the imitation gap of RE is better than that of MM. Finally, they showed the empirical performance of RE on some tasks.", "questions": "1. In lines 56-57, the authors argued that they achieved the best-known imitation gap under significantly weaker assumptions compared with the prior work [Rajaraman et al. (2021a)]. In my opinion, these two works impose two different types of assumptions and it is hard to say which one is weaker. Can the authors elaborate more on this point?\n\n2. In eq.(6), it seems to be $P (s _ { 1 \\cdots t-1 }) \\rightarrow P (s _ { 1 \\cdots t })$ since we should also consider the weight on $s_t$ for estimating the moment in time step $t$?\n\n3. The weak margin condition in assumption 3 is hard to understand. Can the authors explain more about the meaning of this condition and why this condition is required in the analysis? \n", "limitations": "The authors discussed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. This paper conducts a rigorous theoretical study for online imitation learning with general function approximation. To handle the challenge brought by function approximation, the author introduced the membership oracle, which is novel in imitation learning literature.\n2. To instantiate the membership oracle, they proposed an offline classification oracle that can achieve a small statistical error for parameter estimation. With this offline classification oracle, they proved a performance guarantee for the resultant algorithm, which contributes new analysis.  \n\n\nWeakness:\n\n1. The idea of replay estimation is not novel in imitation learning. In the tabular and known transition setting, [1] leverages the transition function to obtain a better estimation of the expert’s moments. Furthermore, [2] extends this idea to the setting where the transition function is unknown but environment interactions are allowed. In [2], they employ BC’s policy to interact with the environment and utilize the collected trajectories to establish a better estimation. This is exactly the replay estimation technique used in this paper. However, the discussion about the relation to the existing work [2] is missing in this work.\n\n2. Due to the first point, the main contribution of this paper is that they applied the replay estimation technique with general function approximation. However, in the setting of general function approximation, the proposed algorithm and analysis require access to a strong offline classification oracle (assumption 2). This offline classification oracle is required to have a small parameter estimation error. Note that this parameter estimation error does not consider the data distribution, in stark contrast to the generalization error used in supervised learning. Therefore, this offline classification oracle asks a much stronger requirement than learning a classifier with a small generalization error in supervised learning. Besides, the RE method is proved to have a better imitation gap than MM only when the parameter error matches the generalization error. However, the key question of when the parameter error matches the generalization error is not answered in this paper.        \n\n     \n3. For the experiment part, I have the following two concerns.\n\n   a). In experiments, the classical moment matching method has poor performance while BC enjoys superior performance. This observation largely conflicts with the well-accepted empirical observation that moment matching based methods can match expert’s performance even with few expert trajectories on locomotion tasks [3, 4, 5]. Besides, in [6], they evaluated a similar MM approach named AdIRL on locomotion tasks from the same Py-Bullet suite. They showed that AdIRL can match expert’s performance across all environments, which conflicts with the experimental observation in this paper. This contradiction raises a doubt about whether the MM method is implemented correctly in experiments.\n   \n   b). Besides, as far as I am concerned, the performance gap established in this paper cannot explain the superior performance of RE in experiments. As the bound $\\min (H^{3 / 2} / N_{\\exp }, H / \\sqrt{N_{\\exp }} )$ suggests, RE has a better performance gap than MM in the large data regime $N_{\\exp} \\succsim H $. Nevertheless, in experiments, they considered the low data regime where $H \\approx 1000$ and $N_{\\exp} \\leq 20$. Therefore, I believe that the developed theory cannot explain the experimental results in this paper.\n\n\n4. Some theoretical results in this paper have been established in the existing work [2]. In particular, for moment matching, the performance gap upper bound and a tight lower bound (Theorems 3.1 and 3.2) have been proved in [2]; see Theorem 1 and Proposition 2 in [2] for details. However, the discussion on [2] is missing in this work.\n\nReferences:\n\n[1]. Rajaraman, N., Yang, L. F., Jiao, J., and Ramachandran, K. Toward the fundamental limits of imitation learning. arXiv preprint arXiv:2009.05990, 2020.\n\n[2]. T. Xu, Z. Li, and Y. Yu. On generalization of adversarial imitation learning and beyond. arXiv, 2106.10424, 2021.\n\n[3] J. Ho and S. Ermon, “Generative adversarial imitation learning,” in Advances in Neural Information Processing Systems 29, 2016, pp. 4565–4573.\n\n[4] I. Kostrikov, K. K. Agrawal, D. Dwibedi, S. Levine, and J. Tompson, “Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning,” in Proceedings of the 7th International Conference on Learning Representations, 2019.\n\n[5] K. Brantley, W. Sun, and M. Henaff, “Disagreement-regularized imitation learning,” in Proceedings of the 8th International Conference on Learning Representations, 2020.\n\n[6] G. Swamy, S. Choudhury, J. A. Bagnell, and S. Wu, “Of moments and matching: A game-theoretic framework for closing the imitation gap,” in Proceeding of the 38th International Conference on Machine Learning, 2021, pp.10 022–10032.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "JIvFem2x2K", "reviewer": "Reviewer_kTcU", "summary": "\nThey propose a reply estimation way that uses a combination of simulated and empirical rollouts to optimally estimate expert moments. They show the finite sample guarantee that can enjoy both of the worlds of BC and moment matching. ", "questions": "I wrote down the question above. ", "limitations": "Not applied. ", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper's result is arguably strong. But at the same time, I feel some of the results are currently hard to understand. Many details are in the appendix currently, and hard to tell what are real assumptions and caveats from the main text.  \n\n(Strengths) \n\n* Their proposed algorithm is sufficiently novel as far as I know.  \n \n* Their theoretical results are also very interesting. At least, I didn't expect we can enjoy both of the worlds of BC and moment matching in the function approximation setting. \n\n(Weakness) \nIt looks like many important stuff are in the appendix. I can understand it might be due to the page constraint. The presentation would be much more improved by cutting some of the stuff. \n\n* I personally feel Section 2 is too long? I guess it is introduced for readers to get some intuition. But, personally, I cannot get so much good intuition from these special instances, figure 2 and figure 3. (other reviewers might feel in a different way though) I would like to suggest authors shorten this section (just introduce BC/moment matching/replace estimation and what are known results in the tabular setting) without introducing these special instances. \n\n* Some of the theorems are written in expectation. Some of the theorems are written in high probability.  This is confusing since comparisons would be not apple to apple sometimes.  Could you try to be consistent? If the authors cannot,  that’s fine. But, could you explain why? \n\n* Some assumptions of theorems are hidden in the appendix. I think they should be in the main text as much as possible. (I think this can be done by shortening section 2 a lot) \n\n\t* For example, in Theorem 3.1, “instantiated with an appropriate discriminator F”… This should be in the man text. And It should be explained if it is not trivial stuff. \n\t* Theorem 2.  Did you assume the realizability of the reward function? But then, assumptions are stronger than BC? In this sense, you cannot directly compare with BC since the assumptions are different?  Does the same stuff apply to moment matching since moment matching might not require the reazaibility of the policy class? \n\t* Assumption 2 sounds strong. I am not sure you can really claim this result handles general function approximation. It looks like it only holds in the simple parametric case. How about more complete classes such as neural networks or random forests? (If not, “function approximation with parametric models” sounds more reasonable as a tile of the paragraph) \n\t* Theorem 3:“Appropriately instantiating —-“ This should be in the main text. \n\t* Line 321: “Whenever we are able to establish —- up to ‘’: How reasonable it is? When does it happen/ \n\n\n\n\t\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept