Minimax Optimal Online Imitation Learning via Replay Estimation
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 inﬁnite sample regime, exact moment matching achieves value equivalence to the expert policy.
However, in the ﬁnite sample regime, even if one has no optimization error, empirical variance can lead to a performance gap that scales with H 2 /N exp for behavioral cloning and H/ p
N exp for online moment matching, where H is the horizon and N 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 parametric function approximation, we prove a meta theorem reducing the performance gap of our approach to the parameter estimation error for ofﬂine classiﬁcation (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 e
O  min( H 3 / 2 /N exp , H/ p
N exp  dependency, under signiﬁcantly weaker assumptions compared to prior work Rajaraman et al.
We implement multiple instantiations of our approach on several continuous control tasks and ﬁnd that we are able to signiﬁcantly improve policy performance across a variety of dataset sizes.
Introduction
In online imitation learning (IL), one is given access to (a) a ﬁxed set of expert demonstrations and (b) an environment or simulator to perform rollouts in.
Many online IL approaches fall under the umbrella of solving a moment matching problem between learner and expert trajectory distributions
Ziebart et al.
[2008], Ho and Ermon [2016],
[Equation]
where E ⇡ [ · ] denotes the expectation over a random trajectory { ( s 1 , a 1 ) , · · · , ( s H , a H ) } generated by rolling out ⇡ .
Swamy et al.
[2021] show that for an appropriate choice of
F , approximate solutions
[Figure] sampling [Figure] of (1) have a performance gap linear in the horizon, the gold standard for sequential problems.
However, a key assumption in their work is that expert moments ( E ⇡ E [ f ( s, a )] ) can be estimated arbitrarily well from the available demonstrations.
The resulting bounds are therefore purely a function of optimization error .
Moving to the ﬁnite-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 an adversarial optimization procedure like (1), the learner may choose to take incorrect actions in order to match noisy moments estimated from the dataset, leading to policies that perform poorly at test-time.
Ideally, one would solve this problem by querying the expert for more demonstrations, as in the work of Ross et al.
However, when we are unable to do so, we still have to 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.
Effectively, 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 that this technique is at once conceptually simple, practically feasible, and minimax optimal in several settings.
Formally, we prove that in the worst case, behavioral cloning has a performance gap / H 2 /N exp , online empirical moment matching, / H/ p
N exp , and our approach of replay estimation / min { H 3 / 2 /N exp , H/ p
N exp } in the tabular setting as well as with linear function approximation.
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-matching algorithms, for strong test-time policy performance.
More explicitly, our work makes the following three contributions:
We extend replay estimation ( RE ) Rajaraman et al.
[2020] beyond the tabular and deterministic setting by introducing the notion of a soft membership oracle and preﬁx weights .
We show how to instantiate the membership oracles for IL with parametric function approximation and prove a meta-theorem relating the imitation gap of RE to the parameter estimation error for ofﬂine classiﬁcation 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 e O  H 3 / 2 d 5 / 4 /N
exp  under signiﬁcantly weaker assumptions compared to prior work Rajaraman et al.
We give multiple practical options for constructing performant membership oracles.
We then use these approximate oracles to signiﬁcantly improve the performance of online IL on several continuous control tasks across a variety of dataset sizes.
We also investigate the differences between our proposed oracles.
We sketch the beneﬁts of replay estimation before providing theoretical and empirical evidence to support our claims.
The Replay Estimator
We begin with a tabular vignette to illustrate our key insight in greater detail.
We compare two algorithms: ofﬂine behavioral cloning ( BC ) Pomerleau [1989] and online moment matching ( MM )
Swamy et al.
Throughout, we focus on learning policies from ﬁnite 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 s 1 or s 2 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.
[Figure]
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 s 2 more than we see s 1 , the learner might realize that the only way to match the observed state distribution (a prerequisite for matching the observed stateaction distribution) is to occasionally take the red action at s 2 .
In general, this could cause the learner to spend an unnecessary amount of time in s 2 which may be undesirable (e.g. if s 2 corresponds to the tree-ﬁlled area in Fig.
1 (top)).
The core issue we hope to illustrate in this example is that by treating the empirical estimate of the expert’s behavior as perfectly accurate, distribution matching can force the learner to take incorrect actions to minimize training error, leading to test-time performance degradation.
As we will discuss in Sec. 3, this can lead to slow statistical rates / H/
[Equation]
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.
Consider, for example, the MDP in Fig.
[Equation] [Figure]
[Equation]
[Figure]
[Equation] [Figure]
[Equation]
Figure 3: The expert always takes the green action, which places it in a uniform distribution over s 1 , . . .
, s n .
At states where we have demonstrations (e.g. s 1 , s 2 ), both BC and MM will take the same, correct action.
However, at states where we have no demonstrations (e.g. s 3 ), MM will correctly take the green action to get back to states with demonstration support, while BC might not. Let us assume that the expert always takes the green action, dropping them in a state in the top row with uniform probability.
In a small demonstration set, we might not see expert actions at some states in the top row.
At all such states, BC 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 a 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 to errors / H 2 /N exp in the worst case.
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 a single algorithm?
It turns out it is indeed possible to do so, via the technique of replay estimation .
In its simplest form, replay estimation involves playing a cached expert action whenever possible and re-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 s 1 and s 2 .
Similarly, in Fig. 3, replay estimation would bring us toward a uniform distribution over the states { s 1 , · · · , s n } in the expert demonstrations.
We could then plug in this improved distribution estimate into the MM procedure (1).
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 s 2 , It turns out this fact is sufﬁcient to establish statistical optimality in the tabular setting and with linear function approximation, with an error rate / min( H 3 / 2 /N exp , H/ p
N exp ) .
In short, replay estimation is a practical technique for reducing some of the ﬁnite-sample variance in expert demonstrations that enables MM to perform optimally in the ﬁnite 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 to answer the question of “ on what states do we know the expert’s action? "
To enable us to answer this question more generally, we introduce the notion of a membership oracle M : S ! { 0 , 1 } .
Explicitly, M ( s ) = 1 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, M ( s ) = 1 on states where BC , which attempts to directly output expert actions, is accurate and 0 otherwise.
We can then compute expert moments by splitting on the output of the membership oracle:
[Equation]
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 a 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 M to a soft membership oracle Zadeh [1965], in order to handle uncertainty in how well we know the expert’s action at a given state.
We proceed by ﬁrst analyzing the statistical properties of applying MM to this bipartite estimator before discussing practical constructions of performant membership oracles.
Theoretical Analysis
[Figure]
The proofs of all results in this section are deferred to appendix A. We begin by introducing some notation.
Notation.
Let  ( X ) denote the probability simplex over set X and let & , . , ⇣ respectively denote greater than, lesser than and equality up to 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 = { ⇢ , P 1 , · · · , P H  1 } , where ⇢ is the initial state distribu- tion and P t : S ⇥ A !
( S ) , while the reward function is denoted r = { r 1 , · · · , r H } where r t : S ⇥ A !
In the online imitation learning setting, the learner has access to a ﬁnite dataset D of N exp 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 fundamental goal of the learner is to learn a policy b ⇡ such that the imitation gap , J ( ⇡ E )  J ( b ⇡ ) is small.
Here, J ( ⇡ ) denotes the expected value of the policy ⇡ , E ⇡
[Equation]
· ] denotes the empirical expectation computed using trajectories from the dataset D .
Behavior cloning.
A standard approach for imitation learning is behavioral cloning ( BC ) Pomerleau [1989], which trains a classiﬁer from expert states to expert actions.
More formally, for the 0 - 1 loss ` (or a continuous proxy) BC minimizes the empirical classiﬁcation error,
[Equation]
It is known from Ross et al.
[2011] and Rajaraman et al.
[2020] that in the tabular setting, the expected imitation gap for BC is always . |S| H 2 /N exp .
The fast 1 /N exp rate comes from the fact that on states where the learner observes the deterministic expert’s actions, BC simply replays them.
Thus, the performance gap is proportional to the mass of unseen states, which decays as 1 /N exp .
The Hdependence comes from the fact that a single mistake can take the learner out of the distribution of the expert, causing it to make mistakes for the rest of the horizon.
Importantly, this bound is tight – i.e. there exists an MDP instance on which BC incurs this imitation gap (See Appendix A.3).
Theorem 1 (Theorem 6.1 of Rajaraman et al.
Then, there exists a tabular MDP instance such that the BC incurs, E
[Equation]
Moment matching.
Many standard algorithms in the IL literature fall under the empirical moment matching framework (e.g. GAIL Ho and Ermon [2016], MaxEnt IRL Ziebart et al.
[2008]); see Table 3 of Swamy et al.
[2021] for more examples.
Reward moment matching corresponds to ﬁnding a policy which best matches the state-action visitation measure of ⇡ E , in the sense of minimizing an
Integral Probability Metric (IPM) Müller [1997].
In the ﬁnite sample setting, the empirical moment matching learner ⇡ MM attempts to best match the empirical state-visitation measure.
Namely,
[Equation]
First we show an upper bound on the imitation gap incurred by empirical moment matching.
Theorem 3.1.
Consider the empirical moment matching learner ⇡ MM (eq.
(4) ), instantiated with an appropriate discriminator class F .
The imitation gap satisﬁes E
[Equation]
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.
Theorem 3.2.
If H  4 , there is a tabular IL instance with 2 states and actions on which with constant probability, the empirical moment matching learner (eq. (4) ) incurs, J ( ⇡ E )  J ( ⇡ MM ) & H/ p
N exp .
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 a slight modiﬁcation of the MDP instance shown in ﬁg. 2, where the reward function is 0 for t = 1 .
For t  2 , the transition function is absorbing at both states; the reward function equals 1 at the state s 1 for any action and iseverywhere else.
Then, the expert state distribution at time 2 and every time thereon is in uniform across the two states, { 1 / 2 , 1 / 2 } .
However, in the dataset D , the learner sees a noisy realization of this distribution in the dataset of the form { 1 / 2   , 1 / 2 +  } for |  | ⇡ ± 1 / p
N exp .
Because of this noise, the empirical moment matching learner may be encouraged to deviate from the expert’s observed behavior and pick the red action at s 2 as this results in a better match to the empirical state visitation measures at every point in the rest of the episode - a prerequisite to matching the empirical state-action visitation measure.
The learner is willing to pick an action different from what the expert played in order to better match the inherently noisy empirical state-action visitation distribution.
Remark 1.
Theorems 1 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 ( ⇡ E )
J ( ⇡ BC ) & |S| H 2 /N exp and J ( ⇡ E )  J ( ⇡ MM ) & H/ p
N exp .
On this mixture instance, training both
BC and empirical moment matching and choosing the better of the two is also statistically suboptimal.
3.1 Replay Estimation
A natural question at this point is whether there is an algorithm that is better than both worlds , i.e. 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, it is unclear how to extend this approach beyond the tabular setting as the algorithm relies on a large measure of states being visited in the demonstrations.
Input: Expert demonstrations D , policy class ⇧ , moment class F =
L
Ht =1 F t , simulator SIM , ALG
) =
[Equation]
; 6 Deﬁne,
[Equation]
[Equation]
Output: ⇡ RE , a solution to the moment-matching problem:
[Equation]
To handle this challenge, we introduce the notion of a soft membership oracle Zadeh [1965], M :
S ⇥ [ H ] !
[0 , 1] which captures the learner’s inherent uncertainty in the expert’s actions at a state at each point in an episode.
The soft membership oracle assigns high weight to a 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 deﬁnition, if the membership oracle is consistently large at all the states visited in a trajectory, we can be conﬁdent that a trajectory generated by BC is as though it was a rollout from the expert policy .
Formally, for any function g and time t = 1 , · · · , H , we have the decomposition,
[Equation]
where P ( s 1 ··· t ) is deﬁned as the preﬁx weight
Q tt 0 =1 M ( s t 0 , t 0 ) .
We need to use preﬁx 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 preﬁx weights, in eq.
(6), ( i ) can be approximated by replacing the expectation over ⇡ E by that over ⇡ BC , i.e. replay estimation.
On the other hand, since the preﬁx weight is low on the remaining trajectories in ( ii ) , we know that BC is inaccurate, so we resort to using a simple empirical estimate to estimate this term.
While we leave the particular choice of the soft membership oracle ﬂexible, 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 ﬁrst 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 ⇡ RE returned by Algorithm 1.
Assume that ⇡ E 2 ⇧ and the ground truth reward function r t 2 F t , which is assumed to be symmetric ( f t 2 F t ()  f t 2 F t ) and bounded (For all f t 2 F t , k f t k 1  1 ).
Choose | D 1 | , | D 2 | = ⇥ ( N exp ) and suppose N replay !
With probability  1  3  ,
[Equation]
where F max , max t 2 [ H ] |F t | , and,
[Equation]
We discuss a proof of this result in Appendix B and include bounds when N replay is ﬁnite.
Remark 2.
Note that Theorem 2 can be extended to inﬁnite function families using the standard technique of replacing |F t | by the ✏ log-covering number of F t in say, the L 2 norm, for ✏ = 1 N exp .
For ease of exposition here, we stick to the case where F t is ﬁnite.
The term L 1 measures how accurate BC is on states from expert trajectories where M ( s t , t ) is large.
Intuitively, if we set M ( s t , t ) = 1 on states where BC is accurate and M ( s t , t ) = {+
+} elsewhere, we would expect this term to be small.
L 2 can be thought of a measure of BC ’s coverage: it tells us how much of the expert’s visitation distribution we believe BC to be inaccurate on.
If BC has good coverage (i.e. 1  M ( s t , t ) is small on expert trajectories), we expect this term to be small.
Prima facie, one might think that because L 1 resembles the imitation gap of BC and L 2 resembles that of MM , RE can only perform as well as the best of BC ( / H 2 /N exp ) and MM ( / H/ p
N exp ) on a given instance.
However, with a careful choice of M , one can achieve “better than both worlds” statistical rates.
In particular, since RE is a generalization of Mimic-MD of Rajaraman et al.
[2020], in the tabular setting, an appropriately initialized version of RE achieves the optimal imitation gap of min
[Equation]
and strictly improves over both BC and MM .
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.
While BC can be thought of as a reduction of IL to ofﬂine classiﬁcation, the algorithm does not take into account the knowledge of the transition of the MDP. This is reﬂected in the quadratic dependency in the horizon, error compounding Ross et al. In this section, we study a novel reduction of IL with parametric function approximation to parameter estimation in ofﬂine classiﬁcation , which we deﬁne formally.
We provide a provable guarantee, assuming the learner has access to a classiﬁcation oracle and the underlying function class admits a Lipschitz parameterization.
Deﬁnition 1 (IL with function-approximation) .
In this setting, for each t 2 [ H ] , there is a parameter class ⇥ t ✓ B d 2 , the unit L 2 ball in d dimensions, and an associated function class { f ✓ t : ✓ t 2 ⇥ t } .
For each t 2 [ H ] there exists an unknown ✓ Et 2 ⇥ t such that 8 s 2 S , [Equation]
[Equation]
Deﬁnition 2 (Lipschitz parameterization) .
A function class G = { g ✓ : ✓ 2 ⇥ } where g ✓ ( · ) : X !
R is said to satisfy L -Lipschitz parameterization if,
[Equation]
[Equation]
L k ✓  ✓ 0 k 2 .
In other words, for each x 2 X , g ✓ ( x ) is an L -Lipschitz function in ✓ , in the L 2 norm.
Assumption 1.
For each t , the class { f ✓ t : ✓ t 2 ⇥ t } is L -Lipschitz in its parameterization, ✓ t 2 ⇥ t .
To deal with parametric function approximation, we assume that the learner has access to an ofﬂine classiﬁcation oracle which, given a dataset of classiﬁcation examples, approximately returns the underlying ground truth parameter.
More formally,
Assumption 2 (Ofﬂine classiﬁcation oracle) .
We assume that the learner has access to a multiclass classiﬁcation oracle, which given n examples of the form, ( s i , a i ) where s i i.i.d.
⇠ D and a i = arg max a 2 A f ✓ ⇤ ( s i , a ) , returns a ˆ ✓ 2 ⇥ such that, with probability  1   , k ˆ ✓  ✓ ⇤ k 2  E ⇥ ,n,  .
We assume that this classiﬁcation oracle is used by RE to train the BC policy in Line 3 of Algorithm 1.
A careful reader might note that Assumption 2 asks for a slightly stronger requirement than just ﬁnding a classiﬁer with small generalization error (which need not be close to the ground truth ✓ ⇤ ).
The latter problem was studied in Daniely et al.
[2013] who show that the Natarajan dimension , up to log-factors in the number of classes (i.e. number of actions) captures the generalization error of the best learner, which scales as ⇥ (1 /n ) given n classiﬁcation examples.
Under some mild regularity (coverage) assumptions on the input distribution D , we show that the generalization guarantee carries over to learning the parameter ✓ , for example in linear classiﬁcation (Appendix B.4).
In particular, we show that E ⇥ ,n,  .
d +log(1 /  ) n when f ✓ = h ✓ , · i .
The membership oracle we study is deﬁned below,
M ( s, t )= ( + if 9 a 2 A such that,
[Equation]
A , f
[Equation]
otherwise.
Intuitively, M assigns a state as +1 if BC classiﬁes it with a signiﬁcant margin as some action.
Finally, we introduce the main assumption on the IL instances we study.
We assume that the classiﬁcation problems solved by BC at each t {+
+}
[ H ] satisfy a margin condition.
Assumption 3 (Weak margin condition) .
The weak margin condition assumes that for each t , there is no classiﬁer ✓ 2 ⇥ t such that for a large mass of states, f ✓ ( s t , ⇡ Et ( s t ))
max a 6 = ⇡ Et ( s t )
f ✓ ( s t , a ) , i.e. the “margin” from the nearest classiﬁcation boundary, is small.
Formally, the weak-margin condition with parameter µ states that, for any ✓ 2 ⇥ t and ⌘  1 /µ ,
[Equation]
The weak margin condition only assumes that there is at least an exponentially small (in ⌘ ) mass of states with margin at least ⌘ .
Smaller µ indicates a larger mass away from any decision boundary.
It sufﬁces to assume that eq.
(11) is only true for ✓ as the classiﬁer in Assumption 2 for our guarantees (Theorem 3) to hold.
Under these three assumptions and using the membership oracle deﬁned above, we can provide a strong guarantee for RE :
Theorem 3.
For IL with parametric function approximation, under Assumptions 1 to 3, appropriately instatiating RE ensures that with probability  1  4  ,
[Equation]
In Appendix B.4, we explicitly instantiate these guarantees for the special case of where the expert follows a linear classiﬁer 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 e
[Equation]
for RE in the linear expert setting, under the restrictive assumptions of a 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 signiﬁcant weakening of the uniform feature distribution constraint to a multi-class analog of the strong distribution assumption Audibert and Tsybakov [2007].
Under this assumption, we show that the best achievable parameter recovery error, E ⇥ t ,N exp ,  /H and the best achievable 0generalization error ( E s ⇠ D [max a 2 A f ˆ ✓ ( s, a ) 6 = max a 2 A f ✓ ⇤ ( s, a )] , in the notation of Assumption 2) for classiﬁcation, E class ⇥ t ,N exp ,  /H , match up to problem dependent constants .
The best known statistical guarantees on imitation gap for BC (Ross et al.
[2011]) and MM are,
[Equation]
[Equation]
[Equation]
The guarantee in Theorem 3 for RE therefore grows as Gap ( ⇡ MM ) p
Gap ( ⇡ BC ) /H , whenever we are able to establish that E ⇥ t ,N exp ,  /H ⇣ E class ⇥ t ,N exp ,  /H up to problem dependent constants.
Our guarantees for RE in the presence of parametric function approximation thus give us reasons to expect RE to improve the performance of MM .
We now turn our attention to validating this in practice.
Practical Algorithm
When considering RE (Alg. 1)
[Figure] two main questions arise: (i) How does one construct a membership oracle in practice, especially when action spaces may be continuous? [Figure] and (ii) How does one solve the moment matching problem (eq.
We now provide potential answers to both of these questions. [Figure] [Figure]
Membership Oracle.
M represents how uncertain BC is 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  and constants µ,  ,
[Equation]
[Equation]
However, in the non-interactive setting, we can only approximate this quantity. The ﬁrst approximation we propose using is inspired by Random Network Distillation (RND)
[Burda et al., 2018], used by Wang et al. [2019] to estimate the support of the expert policy. We instead propose using RND as a measure of how uncertain BC is about expert actions. That is,
[Equation]
where c ⇡ BC is a network trained to imitate ⇡ BC on expert states. To train c ⇡ BC , we evaluate ⇡ BC on 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. [2019] utilizes the variance of a set of BC learners as a regularizer on top of standard BC error. We instead propose using it to weight the two halves of our estimator.
Instead, we deﬁne, [Equation] where { ⇡ BC (1) , · · · , ⇡ BC ( k ) } are BC policies trained with different initializations, which produces sufﬁcient diversity when using deep networks as function approximators. Lastly, we can also use the maximum difference in the ensemble Kidambi et al.
[2021] as our disagreement measure: [Equation] For computing preﬁx 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 modiﬁcation 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 a best-response counterpart [Swamy et al., 2021].
Our approach is related to the GAIL algorithm of Ho and Ermon [2016] which we improve in 4 ways: (i) we use a general
Integral Probability Metric Müller [1997] instead of the Jensen-Shannon Divergence used in the original paper which improves the representation power of the the discriminator, (ii) we add in gradient penalties to the discriminator, which improves convergence rates Gulrajani et al.
[2017], (iii) we solve the entropy-regularized forward problem via Soft-Actor Critic Haarnoja et al.
[2018] as the policy optimizer, for improved sample efﬁciency, 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 Syrgkanis et al.
[2015] and in practice Daskalakis et al.
Together, these changes lead to an implementation which signiﬁcantly out-performs the original , giving us a strong baseline to compare against.
We include an ablation to conﬁrm 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.
Experimental Results
We now quantify the empirical beneﬁts of RE on several continuous control tasks from the the PyBullet suite Coumans and Bai [2016–2019].
All the task we consider have long horizons ( H ⇡ 1000 ) and we use relatively few demonstrations.
( N exp  20 ).
We set N replay as 100 BC rollouts (Line 4 of Algorithm 1).
We test all four membership oracles from the previous section ( M EXP as an idealized target, M RND , M VAR , and M MAX as practical solutions).
In Fig.
5 (left), we see that with only twelve trajectories, RE is able to 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 sufﬁcient for replay estimation to generate a signiﬁcant 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 a similar result in Fig.
5 (center), where we add in i.i.d. noise to the environment dynamics at each timestep.
This makes the problem signiﬁcantly more challenging than the standard version of the Walker task.
RE is still able to match expert performance, with M MAX working notably well.
The correlation plot in Fig.
5 (right) shows us M MAX appears to be best correlated with the idealized preﬁx weights, M EXP under the state distribution induced by
BC .
Because of its superior performance, we use M MAX for the rest of our experiments.
In the left half of Fig. 6, we see RE improve the performance of MM .
In the right half, we see RE out-perform
BC in responding to an extremely tiny amount of noise added to the initial velocity of the agent (similar to the experiments of Reddy et al.
[2019]) – we 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://github.com/gkswamy98/replay_est .
Acknowledgements
ZSW is supported in part by the NSF FAI Award #1939606, a Google Faculty Research Award, a
J.P. Morgan Faculty Award, a Facebook Research Award, an Okawa Foundation Research Grant, and a Mozilla Research Grant.
GS is supported by a GPU award from NVIDIA.
NR and JJ were partially supported by NSF Grants IIS-1901252, and CCF-1909499.
KR is supported by ARO fund 051242-002.