#  Minimax Optimal Online Imitation Learning via Replay Estimation

Gokul Swamy
Carnegie Mellon University
<EMAIL_ADDRESS>
&Nived Rajaraman∗
UC Berkeley
<EMAIL_ADDRESS>
&Matthew Peng
UC Berkeley
&Sanjiban Choudhury
Cornell University
&J. Andrew Bagnell
Aurora Innovation and Carnegie Mellon University
&Zhiwei Steven Wu
Carnegie Mellon University &Jiantao Jiao
UC Berkeley
&Kannan Ramchandran
UC Berkeley Equal contribution. Correspondence to<EMAIL_ADDRESS>and
<EMAIL_ADDRESS>
###### 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 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
a performance gap that scales with $H^{2}/N_{\mathrm{exp}}$ for behavioral
cloning and $H/\sqrt{N_{\mathrm{exp}}}$ for online moment matching, where $H$
is the horizon and $N_{\mathrm{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 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_{\mathrm{exp}}},{H}/{\sqrt{N_{\mathrm{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.

##  1 Introduction

In online imitation learning (IL), one is given access to (a) a fixed 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 & Ermon (2016),

$\min_{\pi\in\Pi}\sup_{f\in\mathcal{F}}\mathbb{E}_{\pi}[f(s,a)]-\mathbb{E}_{\pi^{E}}[f(s,a)],$
(1)

where $\mathbb{E}_{\pi}[\cdot]$ denotes the expectation over a random
trajectory $\\{(s_{1},a_{1}),\cdots,(s_{H},a_{H})\\}$ generated by rolling out
$\pi$. Swamy et al. (2021) show that for an appropriate choice of
$\mathcal{F}$, approximate solutions 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
($\mathbb{E}_{\pi^{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 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 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. (2011). 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?"

Figure 1: Top: Attempting to exactly match a finite-sample approximation of
expert moments can cause a 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 (e.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. $s_{1}\rightarrow
s_{1}\rightarrow s_{2}$ on the right) that are consistent with expert actions,
one can augment the original demonstration set and compute expert moments more
accurately.

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 $\propto{H^{2}}/{N_{\mathrm{exp}}}$,
online empirical moment matching, $\propto{H}/{\sqrt{N_{\mathrm{exp}}}}$, and
our approach of replay estimation
$\propto\min\\{{H^{3/2}}/{N_{\mathrm{exp}}},{H}/{\sqrt{N_{\mathrm{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:

  1. 1.

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 prefix weights.

  2. 2.

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 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
$\widetilde{O}\left(H^{3/2}d^{5/4}/N_{\mathrm{exp}}\right)$ under
significantly weaker assumptions compared to prior work Rajaraman et al.
(2021a).

  3. 3.

We give multiple practical options for constructing performant membership
oracles. We then use these approximate oracles to significantly 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 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) Pomerleau
(1989) and online moment matching (MM) Swamy et al. (2021). 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 $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.

$s_{1}$$s_{2}$$\text{Unif}(\\{s_{1},s_{2}\\})$

Figure 2: An MDP where the expert always takes the green action that puts them
in the uniform distribution over $s_{1}$ and $s_{2}$. Because of full expert
support, BC will learn to always take this action at both states. However, if
the empirical state distribution is more tilted towards $s_{2}$, MM will take
the incorrect red action.

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 state-action 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-filled 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 $\propto H/\sqrt{N_{\mathrm{exp}}}$.

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. 3.

$s_{1}$$s_{2}$$s_{n}$$\ldots$$s_{x}$$\text{Unif}(\\{s_{1},\dots,s_{n}\\})$

Figure 3: The expert always takes the green action, which places it in a
uniform distribution over $s_{1},\dots,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
$\propto H^{2}/N_{\mathrm{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},\cdots,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 sufficient to establish statistical
optimality in the tabular setting and with linear function approximation, with
an error rate
$\propto\min({H^{3/2}}/{N_{\mathrm{exp}}},{H}/{\sqrt{N_{\mathrm{exp}}}})$. In
short, replay estimation is a practical technique for reducing some of the
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 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 $\mathcal{M}:\mathcal{S}\rightarrow\\{0,1\\}$. Explicitly,
$\mathcal{M}(s)=1$ for states where we know the expert’s action well (e.g.
states where we have lots of similar demonstrations) and $\mathcal{M}(s)=0$
otherwise.

We can then compute expert moments by splitting on the output of the
membership oracle:

$\displaystyle\mathbb{E}_{\pi^{E}}[f(s,a)]=\underbrace{\mathbb{E}_{\pi^{E}}\left[f(s,a)\mathbbm{1}(\mathcal{M}(s)=1)\right]}_{(i)}+\underbrace{\mathbb{E}_{\pi^{E}}\left[f(s,a)\mathbbm{1}(\mathcal{M}(s)=0)\right]}_{(ii)}$
(2)

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 $\mathcal{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 $\mathcal{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 first analyzing the
statistical properties of applying MM to this bipartite estimator before
discussing practical constructions of performant membership oracles.

##  3 Theoretical Analysis

Figure 4: In the tabular setting, RE inherits the superior low-data
performance of moment-matching approaches and is able to perform better than
both MM and BC with enough data.

The proofs of all results in this section are deferred to appendix A. We begin
by introducing some notation.

Notation. Let $\Delta(X)$ denote the probability simplex over set $X$ and let
$\gtrsim,\lesssim,\asymp$ respectively denote greater than, lesser than and
equality up to constants. We study the IL problem in the episodic MDP setting
with state space $\mathcal{S}$, action space $\mathcal{A}$ and horizon $H$. We
assume that the transition, reward function and policies can be non-
stationary. The MDP transition is denoted $P=\\{\rho,P_{1},\cdots,P_{H-1}\\}$,
where $\rho$ is the initial state distribution and
$P_{t}:\mathcal{S}\times\mathcal{A}\to\Delta(\mathcal{S})$, while the reward
function is denoted $r=\\{r_{1},\cdots,r_{H}\\}$ where
$r_{t}:\mathcal{S}\times\mathcal{A}\to[0,1]$. In the online imitation learning
setting, the learner has access to a finite dataset $D$ of $N_{\mathrm{exp}}$
trajectories (i.e. the sequences of states visited and actions played)
generated by rolling out expert policy $\pi^{E}$. Importantly, the learner
does not observe rewards during rollouts. The fundamental goal of the learner
is to learn a policy $\widehat{\pi}$ such that the imitation gap,
$J(\pi^{E})-J(\widehat{\pi})$ is small. Here, $J(\pi)$ denotes the expected
value of the policy $\pi$,
$\mathbb{E}_{\pi}\left[\sum_{t=1}^{H}r_{t}(s_{t},a_{t})\right]$.
$\mathbb{E}_{D}[\cdot]$ 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 classifier from expert states to
expert actions. More formally, for the $0$-$1$ loss $\ell$ (or a continuous
proxy) BC minimizes the empirical classification error,

$\displaystyle\pi^{\texttt{BC}}\leftarrow\operatorname*{arg\,min}_{\pi\in\Pi}\mathbb{E}_{D}\left[\frac{1}{H}\sum\nolimits_{t=1}^{H}\ell(\pi_{t}(s_{t}),a_{t})\right].$
(3)

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
$\lesssim|\mathcal{S}|H^{2}/N_{\mathrm{exp}}$. The fast $1/N_{\mathrm{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_{\mathrm{exp}}$. The $H^{2}$ dependence 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
Section A.3).

######  Theorem 1 (Theorem 6.1 of Rajaraman et al. (2020)).

Then, there exists a tabular MDP instance such that the BC incurs,
$\mathbb{E}\left[J(\pi^{E})-J(\pi^{\texttt{BC}})\right]\gtrsim\min\left\\{H,{|\mathcal{S}|H^{2}}/{N_{\mathrm{exp}}}\right\\}$.

Moment matching. Many standard algorithms in the IL literature fall under the
empirical moment matching framework (e.g. GAIL Ho & Ermon (2016), MaxEnt IRL
Ziebart et al. (2008)); see Table 3 of Swamy et al. (2021) for more examples.
Reward moment matching corresponds to finding a policy which best matches the
state-action visitation measure of $\pi^{E}$, in the sense of minimizing an
Integral Probability Metric (IPM) Müller (1997). In the finite sample setting,
the empirical moment matching learner $\pi^{\texttt{MM}}$ attempts to best
match the empirical state-visitation measure. Namely,

$\displaystyle\pi^{\texttt{MM}}\in\operatorname*{arg\,min}_{\pi\in\Pi}\sup_{f\in\mathcal{F}}\mathbb{E}_{\pi}\\!\left[\frac{\sum_{t=1}^{H}f_{t}(s_{t},a_{t})}{H}\right]\\!{-}\mathbb{E}_{D}\\!\left[\frac{\sum_{t=1}^{H}f_{t}(s_{t},a_{t})}{H}\right].$
(4)

First we show an upper bound on the imitation gap incurred by empirical moment
matching.

######  Theorem 3.1.

Consider the empirical moment matching learner $\pi^{\texttt{MM}}$ (eq. 4),
instantiated with an appropriate discriminator class $\mathcal{F}$. The
imitation gap satisfies
$\mathbb{E}\left[J(\pi^{E})-J(\pi^{\texttt{MM}})\right]\lesssim
H\sqrt{|\mathcal{S}|/N_{\mathrm{exp}}}$.

The proof of this result can be found in Section A.2. It turns out that this
guarantee is essentially tight for empirical moment matching, answering an
open question from Rajaraman et al. (2020).

######  Theorem 3.2.

If $H\geq 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(\pi^{E})-J(\pi^{\texttt{MM}})\gtrsim H/\sqrt{N_{\mathrm{exp}}}$.

The proof of this result is deferred to section 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 modification of the MDP instance shown in
fig. 2, where the reward function is $0$ for $t=1$. For $t\geq 2$, the
transition function is absorbing at both states; the reward function equals
$1$ at the state $s_{1}$ for any action and is $0$ everywhere 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-\delta,1/2+\delta\\}$ for $|\delta|\approx\pm
1/\sqrt{N_{\mathrm{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(\pi^{E})-J(\pi^{\texttt{BC}})\gtrsim|\mathcal{S}|H^{2}/N_{\mathrm{exp}}$
and $J(\pi^{E})-J(\pi^{\texttt{MM}})\gtrsim H/\sqrt{N_{\mathrm{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 $\Pi$, moment class
$\mathcal{F}=\bigoplus_{t=1}^{H}\mathcal{F}_{t}$, simulator SIM, ALG which
returns a membership oracle given a dataset;

1

2Partition the dataset $D$ into $D_{1}$ and $D_{2}$;

3

4Using ALG, train the membership oracle $\mathcal{M}$ on $D_{1}$;

5

6Train $\pi^{\texttt{BC}}$ using behavior cloning on $D_{1}$;

7

8Roll out $\pi^{\texttt{BC}}$ in SIM $N_{\mathrm{replay}}$ times to construct
a new dataset, $D_{\mathrm{replay}}$;

9

10Define prefix weights $\mathcal{P}(s_{1\dots
t-1})=\prod\nolimits_{t^{\prime}=1}^{t-1}\mathcal{M}(s_{t^{\prime}},t^{\prime})$;

11

12Define,  $\displaystyle\widehat{E}(f)$
$\displaystyle=\mathbb{E}_{D_{\mathrm{replay}}}\left[\frac{1}{H}\sum\nolimits_{t=1}^{H}f_{t}(s_{t},a_{t})\left(\mathcal{P}(s_{1\dots
t})\right)\right]+\mathbb{E}_{D_{2}}\left[\frac{1}{H}\sum\nolimits_{t=1}^{H}f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\dots
t})\right)\right].$

Output: $\pi^{\texttt{RE}}$, a solution to the moment-matching problem:
$\displaystyle\operatorname*{arg\,min}_{\pi\in\Pi}\sup_{f\in\mathcal{F}}\mathbb{E}_{\pi}\left[\frac{1}{H}\sum\nolimits_{t=1}^{H}f_{t}(s_{t},a_{t})\right]-\widehat{E}(f)$
(5)

Algorithm 1 Replay Estimation (RE )

To handle this challenge, we introduce the notion of a soft membership oracle
Zadeh (1965), $\mathcal{M}:\mathcal{S}\times[H]\to[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 definition, if
the membership oracle is consistently large at all the states visited in a
trajectory, we can be confident 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,\cdots,H$, we have the decomposition,

$\displaystyle\mathbb{E}_{\pi^{E}}[g(s_{t},a_{t})]$
$\displaystyle=\underbrace{\mathbb{E}_{\pi^{E}}\left[g(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t})\right]}_{(i)}+\underbrace{\mathbb{E}_{\pi^{E}}\left[g(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t})\right)\right]}_{(ii)}$ (6)

where $\mathcal{P}(s_{1\cdots t})$ is defined as the prefix weight
$\prod_{t^{\prime}=1}^{t}\mathcal{M}(s_{t^{\prime}},t^{\prime})$. 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, $(i)$ can be
approximated by replacing the expectation over $\pi^{E}$ by that over
$\pi^{\texttt{BC}}$, 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 to 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 5, 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 $\mathcal{M}$.

######  Theorem 2.

Consider the policy $\pi^{\texttt{RE}}$ returned by Algorithm 1. Assume that
$\pi^{E}\in\Pi$ and the ground truth reward function
$r_{t}\in\mathcal{F}_{t}$, which is assumed to be symmetric
($f_{t}\in\mathcal{F}_{t}\iff-f_{t}\in\mathcal{F}_{t}$) and bounded (For all
$f_{t}\in\mathcal{F}_{t}$, $\|f_{t}\|_{\infty}\leq 1$). Choose
$|D_{1}|,|D_{2}|=\Theta(N_{\mathrm{exp}})$ and suppose
$N_{\mathrm{replay}}\to\infty$. With probability $\geq 1-3\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim$ $\displaystyle\
\mathcal{L}_{1}+\mathcal{L}_{2}+\frac{\log\left(F_{\mathrm{max}}H/\delta\right)}{N_{\mathrm{exp}}}$
(7)

where $F_{\mathrm{max}}\triangleq\max_{t\in[H]}|\mathcal{F}_{t}|$, and,

$\displaystyle\mathcal{L}_{1}\triangleq H^{2}\
\mathbb{E}_{\pi^{E}}\\!\left[\frac{\sum_{t=1}^{H}\mathcal{M}(s_{t},t)\mathsf{TV}\left(\pi^{E}_{t}(\cdot|s_{t}),\pi^{\texttt{BC}}_{t}(\cdot|s_{t})\right)}{H}\right],$
(8) $\displaystyle\mathcal{L}_{2}\triangleq
H^{3/2}\sqrt{\frac{\log\left(F_{\mathrm{max}}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]}{H}}.$

We discuss a proof of this result in Appendix B and include bounds when
$N_{\mathrm{replay}}$ is finite.

######  Remark 2.

Note that Theorem 2 can be extended to infinite function families using the
standard technique of replacing $|\mathcal{F}_{t}|$ by,
$\mathcal{N}(\mathcal{F}_{t},\xi,\|\cdot\|_{\infty})$ the $\xi$ log-covering
number (metric entropy) of $\mathcal{F}_{t}$ in the $L_{\infty}$ norm, for
$\xi=\frac{1}{N_{\mathrm{exp}}H}$. Likewise, we may appropriately replace
$F_{\mathrm{max}}$ by
$\mathcal{N}_{\mathrm{max}}\triangleq\max_{t\in[H]}\mathcal{N}(\mathcal{F}_{t},\frac{1}{N_{\mathrm{exp}}H},\|\cdot\|_{\infty})$.
For ease of exposition here, we stick to the case where $\mathcal{F}_{t}$ is
finite.

The term $\mathcal{L}_{1}$ measures how accurate BC is on states from expert
trajectories where $\mathcal{M}(s_{t},t)$ is large. Intuitively, if we set
$\mathcal{M}(s_{t},t)=1$ on states where BC is accurate and
$\mathcal{M}(s_{t},t)=0$ elsewhere, we would expect this term to be small.
$\mathcal{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-\mathcal{M}(s_{t},t)$ is small
on expert trajectories), we expect this term to be small.

Prima facie, one might think that because $\mathcal{L}_{1}$ resembles the
imitation gap of BC and $\mathcal{L}_{2}$ resembles that of MM , RE can only
perform as well as the best of BC ($\propto H^{2}/N_{\mathrm{exp}}$) and MM
($\propto H/\sqrt{N_{\mathrm{exp}}}$) on a given instance. However, with a
careful choice of $\mathcal{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\left\\{\frac{|\mathcal{S}|H^{3/2}}{N_{\mathrm{exp}}},H\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}\right\\}\log\left(\frac{|\mathcal{S}|H}{\delta}\right)$
and strictly improves over both BC and MM .

We now show how to instantiate the membership oracle for parametric function
approximation and provide a statistical guarantee under a particular margin
assumption.

##  4 Parametric Function Approximation: Reduction to Offline Classification

While BC can be thought of as a reduction of IL to 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, also known
as error compounding Ross et al. (2011). This problem is certainly encountered
in practice Pomerleau (1989), and Rajaraman et al. (2020); Swamy et al. (2021)
show that this issue is a problem for any offline algorithm (i.e. which
doesn’t interact with the MDP, or know its transition). This begs the question
how can we enable algorithms such as BC which reduce IL to offline
classification to benefit from interaction with the environment.

In this section, under the known transition model, we study a novel reduction
of IL with parametric function approximation to parameter estimation in
offline classification, which we define formally. We provide a provable
guarantee, assuming the learner has access to a classification oracle and the
underlying function class admits a Lipschitz parameterization.

######  Definition 1 (IL with function-approximation).

In this setting, for each $t\in[H]$, there is a parameter class
$\Theta_{t}\subseteq\mathbb{B}_{2}^{d}$, the unit $L_{2}$ ball in $d$
dimensions, and an associated function class
$\\{f_{\theta_{t}}:\theta_{t}\in\Theta_{t}\\}$. For each $t\in[H]$ there
exists an unknown $\theta^{E}_{t}\in\Theta_{t}$ such that $\forall
s\in\mathcal{S}$,

$\displaystyle\pi^{E}_{t}(s)=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta^{E}_{t}}(s,a).$
(9)

######  Definition 2 (Lipschitz parameterization).

A function class $\mathcal{G}=\\{g_{\theta}:\theta\in\Theta\\}$ where
$g_{\theta}(\cdot):\mathcal{X}\to\mathbb{R}$ is said to satisfy $L$-Lipschitz
parameterization if,
$\|g_{\theta}(\cdot)-g_{\theta^{\prime}}(\cdot)\|_{\infty}\leq
L\|\theta-\theta^{\prime}\|_{2}$. In other words, for each $x\in\mathcal{X}$,
$g_{\theta}(x)$ is an $L$-Lipschitz function in $\theta$, in the $L_{2}$ norm.

The Lipschitz parameterization condition can be interpreted as saying that a
small change to the underlying classifier does not all of a sudden change the
label of a large mass of points. In particular, points which are classified
with a large enough “margin” continue to stay in the same class even if the
underlying classifier/parameter is perturbed.

######  Assumption 1.

For each $t$, the class $\\{f_{\theta_{t}}:\theta_{t}\in\Theta_{t}\\}$ is
$L$-Lipschitz in its parameterization, $\theta_{t}\in\Theta_{t}$.

Recall that BC is a reduction from IL to offline classification, and assumes
that the learner has access to an oracle which, given samples from an
underlying ground truth classifier, returns a classifier with small
generalization error. In particular, treating the expert dataset at each time
$t$ as a collection of $N_{\mathrm{exp}}$ samples from a classifier
$\pi_{t}(\cdot|\cdot)$ mapping states to actions, the BC learner uses the
oracle to train $H$ classifers on these datasets. Ross et al. (2011) show that
the imitation gap of the policy induced by these classifiers is bounded by the
average generalization error of the $H$ learned classifiers, up to a factor of
$H^{2}$.

In order to extend RE (which uses the transition of the MDP) to deal with
parametric function approximation, and show error guarantees which surpass
that achieved by BC , we assume that the learner has access to a slightly
stronger offline classification oracle, which, given access to a dataset of
classification examples, approximately returns the underlying ground truth
parameter. More formally,

######  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^{i},a^{i})$ where
$s^{i}\overset{\text{i.i.d.}}{\sim}\mathcal{D}$ and
$a^{i}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta^{*}}(s^{i},a)$,
returns a $\hat{\theta}\in\Theta$ such that, with probability $\geq 1-\delta$,
$\|\hat{\theta}-\theta^{*}\|_{2}\leq\mathcal{E}_{\Theta,n,\delta}$.

We assume that this classification 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 finding a classifier with small generalization error
(which need not be close to the ground truth $\theta^{*}$). The generalization
error, i.e.
$\mathbb{Pr}_{s\sim\mathcal{D}}\left[\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta^{*}}(s,a)\neq\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\hat{\theta}}(s,a)\right]$
in the notation of Assumption 2, was previously studied in Daniely et al.
(2013) for multi-class classification. The authors show that up to log-factors
in the number of classes (i.e. the number of actions), the Natarajan dimension
characterizes the generalization error of the best learner, which scales as
$\Theta(1/n)$ given $n$ classification examples. For example, in linear
classification, which we study later in Section B.4, under some mild
regularity assumptions on the input distribution $\mathcal{D}$, we show that
the generalization guarantee carries over to approximately learning the
parameter $\theta$. In particular, we show that under said regularity
conditions, for linear classification (i.e.
$f_{\theta}=\langle\theta,\cdot\rangle$),
$\mathcal{E}_{\Theta,n,\delta}\lesssim\frac{d+\log(1/\delta)}{n}$.

Recall that we assume that the learner trains a sequence of $H$ classifiers on
the state-actions observed in the expert dataset at each time $t$, using the
offline classification oracle in Assumption 2. We denote the resulting offline
classifiers as
$\widehat{\theta}^{\textsf{BC}}=(\widehat{\theta}^{\textsf{BC}}_{1},\cdots,\widehat{\theta}^{\textsf{BC}}_{H})$.

We are now ready to define the membership oracle under which we study RE
below,

$\displaystyle\mathcal{M}(s,t){=}\begin{cases}+1\quad\\!\\!&\text{if }\exists
a\in\mathcal{A}\text{ such that, }\forall a^{\prime}\in\mathcal{A},\
f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a){-}f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a^{\prime})\geq
2L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}\\\
0&\text{otherwise.}\end{cases}$ (10)

The intuitive interpretation of $\mathcal{M}$ is that, state which are
classified by BC as some action with a significant margin are assigned as
$+1$, and the remaining states are assigned as $0$ by the membership oracle.

Finally, we impose an assumption on the IL instances we study. We assume that
the classification problems solved by BC at each $t\in[H]$ satisfy a margin
condition.

######  Assumption 3 (Weak margin condition).

For $t\in[H]$ and $\theta\in\Theta_{t}$, define
$a^{\theta}_{s}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta}(s,a)$ as
the classifier output on the state $s$. The weak margin condition with
parameter $\mu>0$ assumes that for each $t$, there is no classifier
$\theta\in\Theta_{t}$ such that for a large mass of states,
$f_{\theta}(s_{t},a^{\theta}_{s_{t}})-\max_{a\neq
a^{\theta}_{s_{t}}}f_{\theta}(s_{t},a)$, i.e. the “margin” from the nearest
classification boundary, is small. Formally, the weak-margin condition with
parameter $\mu$ states that for any $\eta\leq 1/\mu$,

$\displaystyle\forall\theta\in\Theta_{t},\ \
\mathrm{Pr}_{\pi^{E}}\left(f_{\theta}(s_{t},a^{\theta}_{s_{t}})-\max_{a\neq
a^{\theta}_{s_{t}}}f_{\theta}(s_{t},a)\geq\eta\right)\geq e^{-\mu\eta}.$ (11)

The weak margin condition only assumes that there is at least an exponentially
small (in $\eta$) mass of states with margin at least $\eta$. A smaller $\mu$
indicates a larger mass away from any decision boundary.

######  Remark 3.

It suffices to assume that for each $t$, eq. 11 is only true for the singular
choice $\theta=\widehat{\theta}^{\textsf{BC}}_{t}\in\Theta_{t}$, for our main
guarantee (Theorem 3) to hold.

Under the previously discussed assumptions, we provide a strong guarantee for
RE which uses the classification oracle in Assumption 2 to define BC , and the
membership oracle as defined in eq. 10.

######  Theorem 3.

For IL with parametric function approximation, under Assumptions 1 to 3,
appropriately instatiating RE ensures that with probability $\geq 1-4\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim H^{3/2}\sqrt{\frac{\mu
L\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}}{H}}+\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}.$
(12)

where $F_{\mathrm{max}}=\max_{t\in[H]}|\mathcal{F}_{t}|$ is as defined in
Theorem 2.

Note that we impose the same assumptions on the policy and discriminator
classes employed by RE as in Theorem 2. Namely, $(i)$ $\pi^{E}\in\Pi$, and for
each $t\in[H]$, $(ii)$ the ground truth reward function
$r_{t}\in\mathcal{F}_{t}$, $(iii)$ $\mathcal{F}_{t}$ is symmetric, i.e.
$f_{t}\in\mathcal{F}_{t}\iff-f_{t}\in\mathcal{F}_{t}$, and $(iv)$
$\mathcal{F}_{t}$ is $1$-bounded, i.e. for all $f_{t}\in\mathcal{F}_{t}$,
$\|f_{t}\|_{\infty}\leq 1$.

As discussed in Remark 2, this result can be extended to infinite
discriminator families by replacing $|\mathcal{F}_{t}|$ by the appropriate
log-covering number of $\mathcal{F}_{t}$ in $L_{\infty}$ norm.

The intuition behind the result is as follows. By Assumption 2, the learner is
able to approximately learn
$\widehat{\theta}^{\textsf{BC}}_{t}\approx\theta^{*}_{t}$ at each time $t$.
Since the discriminator functions are Lipschitz (Assumption 1) and there are
not too many states classified with small margin (Assumption 3), this means
that the states classified with large margin by
$\widehat{\theta}^{\textsf{BC}}_{t}$ are correctly classified by
$\theta^{E}_{t}$, while the states which are close to a decision boundary
(induced by BC ) may be misclassified by BC . Therefore, we may set the
membership oracle as $+1$ on states classified by BC with a large margin, and
$0$ on states classified with small margin. In particular, the membership
oracle considered in eq. 10 ensures that on states at which
$\mathcal{M}(s,t)>0$, $\pi^{E}_{t}(\cdot|s)=\pi^{\texttt{BC}}_{t}(\cdot|s)$.
Likewise, the states at which $\mathcal{M}(s,t)=0$ correspond to the states
which are classified by BC with a small margin (i.e. are close to a decision
boundary), the probability mass of which is bounded by the weak margin
condition, Assumption 3. All in all, in the language of Theorem 2, these
results ensure that $\mathcal{L}_{1}=0$ and $\mathcal{L}_{2}$ is significantly
smaller than $H^{3/2}\sqrt{\log(F_{\mathrm{max}}H/\delta)/N_{\mathrm{exp}}}$
which essentially results in the proof of Theorem 3.

In order to interpret Theorem 3, we draw the connection back with BC and MM .
As discussed earlier, Ross et al. (2011) prove the best known general
statistical guarantee for BC in terms of the generalization error of the
underlying classifiers. In particular, with probability $\geq 1-\delta$,

$\displaystyle
J(\pi^{E})-J(\pi^{\texttt{BC}})\leq\texttt{Gap}(\pi^{\texttt{BC}})\triangleq
H^{2}\frac{\sum_{t=1}^{H}\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}^{\text{class}}}{H}$
(13)

Here, $\mathcal{E}_{\Theta,n,\delta}^{\text{class}}$ denotes the best
achievable $0$-$1$ generalization error for multi-class classification: in the
notation of Assumption 2, there exists a learner $\hat{\theta}$, such that on
any classification instance
$\mathbb{E}_{s\sim\mathcal{D}}[\max_{a\in\mathcal{A}}f_{\hat{\theta}}(s,a)\neq\max_{a\in\mathcal{A}}f_{\theta^{*}}(s,a)]\leq\mathcal{E}_{\Theta,n,\delta}^{\text{class}}$
with probability $\geq 1-\delta$. On the other hand, the best general
statistical guarantee for MM is,

$\displaystyle
J(\pi^{E})-J(\pi^{\texttt{MM}})\leq\texttt{Gap}(\pi^{\texttt{MM}})\triangleq
H\sqrt{\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}}$ (14)

(note that we can extend to infinite classes using the covering number
argument in Remark 2).

Now, whenever the statistical error for parameter estimation matches with the
best statistical error for generalization in offline classification, namely,
$\mathcal{E}_{\Theta,n,\delta}\asymp\mathcal{E}_{\Theta,n,\delta}^{\text{class}}$
up to problem dependent constants, the guarantee in Theorem 3 can be
reinterpreted as,

$\displaystyle
J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim\texttt{Gap}(\pi^{\texttt{MM}})\sqrt{\frac{\texttt{Gap}(\pi^{\texttt{BC}})}{H}}.$
(15)

The interpretation of this result is that, whenever BC admits a non-trivial
gap on the imitation gap, namely $\texttt{Gap}(\pi^{\texttt{BC}})\ll H$, the
performance gap in eq. 15 is $\ll\texttt{Gap}(\pi^{\texttt{MM}})$.
Furthermore, from Daniely et al. (2013), for multi-class classification,
$\mathcal{E}_{\Theta,n,\delta}^{\text{class}}\lesssim\frac{\left(\mathfrak{n}_{\Theta}+\log(1/\delta)\right)\log(n)\log|\mathcal{A}|}{n}$
where $\mathfrak{n}_{\Theta}$ denotes the Natarajan dimension of the function
class $\\{f_{\theta}:\theta\in\Theta\\}$ and $\mathcal{A}$ denotes the set of
labels (which here, are the set of actions). Therefore, from eq. 15, we get
the guarantee,

$\displaystyle
J(\pi^{E})-J(\pi^{\texttt{RE}})\leq\widetilde{O}\left(\frac{H^{3/2}}{N_{\mathrm{exp}}}\left(\log(F_{\mathrm{max}})\times\frac{\sum_{t=1}^{H}\mathfrak{n}_{\Theta_{t}}}{H}\right)^{1/2}\right),$
(16)

whenever the underlying classification problem allows
$\mathcal{E}_{\Theta,n,\delta}^{\text{class}}\asymp\mathcal{E}_{\Theta,n,\delta}$
up to problem dependent constants. Note that the polylogarthmic factors in eq.
16 are in $|\mathcal{A}|$, $N_{\mathrm{exp}}$ and $1/\delta$. Under these
conditions, we essentially recover a performance guarantee for RE which scales
as $H^{3/2}/N_{\mathrm{exp}}$. This improves on the quadratic $H$ dependence
incurred by BC , and is optimal in the worst case, even in the tabular setting
with just $3$ states, as shown in Rajaraman et al. (2021b). This guarantee
also suggests a natural measure of complexity for IL - the average Natarajan
dimension, $\frac{\sum_{t=1}^{H}\mathfrak{n}_{\Theta_{t}}}{H}$ multiplied by
the maximum log-covering number of the discriminator (or reward) class,
$\log(F_{\mathrm{max}})$.

From this discussion, an important problem stands out: For what kind of
classification problems is
$\mathcal{E}_{\Theta,n,\delta}\asymp\mathcal{E}_{\Theta,n,\delta}^{\text{class}}$?
While we do not answer this question in its full generality, focusing on the
special case of linear classification and provide a set of sufficient
conditions under which
$\mathcal{E}_{\Theta,n,\delta}\asymp\mathcal{E}_{\Theta,n,\delta}^{\text{class}}$
up to problem dependent constants. In conjunction with eq. 16, this results in
a novel guarantee for IL with linear function approximation, under
significantly weaker conditions compared to prior work. We defer the
discussion of these results to Section B.4.

Our guarantees for RE in the presence of parametric function approximation and
linear function approximation give us reasons to expect RE to improve the
performance of MM . We now turn our attention to confirming this holds true in
practice.

##  5 Practical Algorithm

When considering the implementation of RE (Alg. 1) in practice, two main
questions arise:

  1. 1.

How does one construct a membership oracle in practice, especially when action
spaces may be continuous?

  2. 2.

How does one design good solvers for the moment matching problem in eq. 5?

We now provide potential answers to both of these questions.

Membership Oracle.  Recall from the interpretation of Theorem 2 that the
membership oracle $\mathcal{M}$ intuitively should capture how uncertain BC is
about the expert’s action at a state. For continuous action spaces, ideally
one would assign the membership oracle at that state based on an appropriate
notion of distance between the action played by BC and that played by the
expert. For example, for a sigmoid function $\sigma$ and constants
$\mu,\beta$,

$\mathcal{M}_{\texttt{EXP}}(s,t)=\sigma\left(\frac{\mu-\left\lVert\pi^{\texttt{BC}}(s)-\pi^{E}(s)\right\rVert_{2}}{\beta}\right),$
(17)

However, in the non-interactive setting where the demonstrator cannot be
queried at states, we can only approximate this quantity. The first
approximation we propose 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 to use RND as a measure of the epistemic
uncertainty of BC about expert actions. That is,

$\mathcal{M}_{\texttt{RND}}(s,t)=\sigma\left(\frac{\mu-\left\lVert\pi^{\texttt{BC}}(s)-\widehat{\pi^{\texttt{BC}}}(s)\right\rVert_{2}}{\beta}\right),$
(18)

where $\widehat{\pi^{\texttt{BC}}}$ is a network trained to imitate the output
of the classifier $\pi^{\texttt{BC}}$ on the states observed in the expert
dataset. To train $\widehat{\pi^{\texttt{BC}}}$, we evaluate
$\pi^{\texttt{BC}}$ on states observed in the expert dataset, and plug this
new dataset into the standard BC pipeline.

We can also utilize other uncertainty measures like the disagreement of an
ensemble to measure epistemic uncertainty, which has previously shown success
on various simulated sequential decision making tasks (Pathak et al., 2019).
Past work by Brantley et al. (2019) proposes to regularize the standard BC
(classification) error, by the variance of an ensemble of independently
trained BC learners at the states visited by the learner’s policy. This
encourages the learner to mimic BC on the states where the action predicted by
all the policies in the ensemble are similar, and avoid states where they are
different (i.e. the variance at these states is high). In contrast, we define,

$\mathcal{M}_{\texttt{VAR}}(s,t)=\sigma\left(\frac{\mu-\text{Var}(\pi^{\texttt{BC}(1)}(s),\dots,\pi^{\texttt{BC}(k)}(s))}{\beta}\right),$
(19)

where $\\{\pi^{\texttt{BC}(1)},\cdots,\pi^{\texttt{BC}(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 across the ensemble as a
measure of uncertainty:

$\mathcal{M}_{\texttt{MAX}}(s,t)=\sigma\left(\frac{\mu-\max_{i,j\in[k]}\left\lVert\pi^{\texttt{BC}}_{i}(s)-\pi^{\texttt{BC}}_{j}(s)\right\rVert_{2}}{\beta}\right),$
(20)

as suggested by the work of Kidambi et al. (2021). We compare these four
choices below. 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.

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 $\mathcal{M}_{\texttt{EXP}}$ and the other
oracles and see $\mathcal{M}_{\texttt{MAX}}$ correlate most.

Figure 6: We see RE with $\mathcal{M}_{\texttt{MAX}}$ improve the performance
of MM on the Noisy Walker2DBulletEnv and HopperBulletEnv tasks. We see RE (and
MM ) out-perform BC on the initial-state-perturbed Walker2DBulletEnv and
HopperBulletEnv tasks .

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 & 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 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 Syrgkanis et
al. (2015) and in practice Daskalakis et al. (2017). Together, these changes
lead to an implementation which significantly 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.

##  6 Experimental Results

We now quantify the empirical benefits of RE on several continuous control
tasks from the the PyBullet suite Coumans & Bai (2016–2019). All the task we
consider have long horizons ($H\approx 1000$) and we use relatively few
demonstrations. ($N_{\mathrm{exp}}\leq 20$). We set $N_{\mathrm{replay}}$ as
100 BC rollouts (Line 4 of Algorithm 1). We test all four membership oracles
from the previous section ($\mathcal{M}_{\texttt{EXP}}$ as an idealized
target, $\mathcal{M}_{\texttt{RND}}$, $\mathcal{M}_{\texttt{VAR}}$, and
$\mathcal{M}_{\texttt{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 sufficient for replay estimation to
generate a 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 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 significantly more challenging than
the standard version of the Walker task. RE is still able to match expert
performance, with $\mathcal{M}_{\texttt{MAX}}$ working notably well. The
correlation plot in Fig. 5 (right) shows us $\mathcal{M}_{\texttt{MAX}}$
appears to be best correlated with the idealized prefix weights,
$\mathcal{M}_{\texttt{EXP}}$ under the state distribution induced by BC .
Because of its superior performance, we use $\mathcal{M}_{\texttt{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. 111 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.

##  7 Discussion and open problems

In this paper we propose and analyze a new meta-algorithm for imitation
learning called Replay Estimation (RE ). The key technique behind RE is to
improve the quality of moment estimates by rolling out artificial trajectories
by using the knowledge of the expert policy at states. This approach can be
implemented on top of any online moment-matching algorithm (e.g. Ziebart et
al. (2008), Ho & Ermon (2016), Swamy et al. (2021)) to improve performance. We
prove a general meta theorem bounding the imitation gap of the policies
produced by the algorithm in terms of the parameter estimation error in
offline classification. This results in a guarantee for IL with the optimal
dependence on the horizon, $H$, and number of expert rollouts available,
$N_{\mathrm{exp}}$, under the assumption that the parameter estimation
guarantee matches the generalization guarantee for the underlying offline
classification problem. Under these conditions, the analysis also suggests a
natural measure of complexity for IL depending on the average Natarajan
dimension and log-covering number (metric entropy) of the discriminator class.
It is a significant open question to extend the analysis of RE to depend on
less stringent classification oracles, and only require constructing learners
with bounded generalization error, as required by BC . This is motivated by a
briefly discussion on the drawbacks of the offline classification oracle we
consider in Assumption 2.

##### Drawbacks of the offline classification oracle (Assumption 2).

Note that the offline classification oracle we consider in Assumption 2
requires a learner with bounded parameter estimation error, compared to one
with bounded generalization error which is typically studied in practice
(Daniely et al., 2013). While under certain conditions, both measures of error
have the same asymptotic scaling in $n$ and the Natarajan dimension
$\mathfrak{n}_{\Theta}$, in general they can be different. For instance,
consider the case of linear (binary) classification under very poor coverage:
denoting the true classifier as $\theta^{*}\in\mathbb{R}^{d}$, and with the
input distribution as $s\sim\mathrm{Unif}(\\{e_{1},-e_{1}\\})$ where
$e_{1}=(1,0,\cdots,0)\in\mathbb{R}^{d}$. Then, given
$n=\Omega(\log(1/\delta))$ samples, there exists a learner which with
probability $1-\delta$ learns a classifier with $0$ generalization error.
However, regardless of the number of samples $n$, no learner can guarantee to
learn $\theta^{*}$, since the remaining $d-1$ coordinates of $\theta^{*}$
aside from $\theta_{1}^{*}$ do not affect the labels of the inputs, which are
what are observed by a learner. Thus, under poor coverage, the parameter
$\theta^{*}$ cannot be learned consistently even though there exists a trivial
classifier with bounded generalization error.

##### Designing practical algorithms which do not require parameter tuning.

The membership oracle resulting from the analysis in theory (eq. 10) as well
as those implemented in practice (eqs. 17, 18, 19 and 20) require tuning
either the margin threshold
$2L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}$, or the scale
parameters $\mu$ and $\beta$. An interesting next direction would be to
develop an algorithm which uses data dependent scales (e.g. Haarnoja et al.
(2018)) to reduce the effort required to fine-tuning these parameters.

## References

  * Audibert & Tsybakov (2007) Audibert, J.-Y. and Tsybakov, A. B.  Fast learning rates for plug-in classifiers.  _The Annals of Statistics_ , 35(2):608 – 633, 2007.  doi: 10.1214/009053606000001217.  URL https://doi.org/10.1214/009053606000001217.
  * Brantley et al. (2019) Brantley, K., Sun, W., and Henaff, M.  Disagreement-regularized imitation learning.  In _International Conference on Learning Representations_ , 2019.
  * Burda et al. (2018) Burda, Y., Edwards, H., Storkey, A., and Klimov, O.  Exploration by random network distillation.  _arXiv preprint arXiv:1810.12894_ , 2018.
  * Coumans & Bai (2016–2019) Coumans, E. and Bai, Y.  Pybullet, a python module for physics simulation for games, robotics and machine learning.  http://pybullet.org, 2016–2019.
  * Daniely & Shalev-Shwartz (2014) Daniely, A. and Shalev-Shwartz, S.  Optimal learners for multiclass problems.  _CoRR_ , abs/1405.2420, 2014.  URL http://arxiv.org/abs/1405.2420.
  * Daniely et al. (2013) Daniely, A., Sabato, S., Ben-David, S., and Shalev-Shwartz, S.  Multiclass learnability and the ERM principle.  _CoRR_ , abs/1308.2893, 2013.  URL http://arxiv.org/abs/1308.2893.
  * Daskalakis et al. (2017) Daskalakis, C., Ilyas, A., Syrgkanis, V., and Zeng, H.  Training gans with optimism.  _arXiv preprint arXiv:1711.00141_ , 2017.
  * Gulrajani et al. (2017) Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A.  Improved training of wasserstein gans, 2017.
  * Haarnoja et al. (2018) Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.
  * Han et al. (2015) Han, Y., Jiao, J., and Weissman, T.  Minimax estimation of discrete distributions under $\ell_{1}$ loss, 2015\.
  * Ho & Ermon (2016) Ho, J. and Ermon, S.  Generative adversarial imitation learning, 2016.
  * Kidambi et al. (2021) Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T.  Morel : Model-based offline reinforcement learning, 2021.
  * Mammen & Tsybakov (1999) Mammen, E. and Tsybakov, A. B.  Smooth discrimination analysis.  _The Annals of Statistics_ , 27(6):1808–1829, 1999.  ISSN 00905364.  URL http://www.jstor.org/stable/2673938.
  * McAllester & Schapire (2000) McAllester, D. A. and Schapire, R. E.  On the convergence rate of good-turing estimators.  In _Proceedings of the Thirteenth Annual Conference on Computational Learning Theory_ , COLT ’00, pp. 1–6, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc.  ISBN 155860703X.
  * Müller (1997) Müller, A.  Integral probability metrics and their generating classes of functions.  _Advances in Applied Probability_ , 29(2):429–443, 1997.
  * Pathak et al. (2019) Pathak, D., Gandhi, D., and Gupta, A.  Self-supervised exploration via disagreement, 2019.
  * Pomerleau (1989) Pomerleau, D. A.  Alvinn: An autonomous land vehicle in a neural network.  Technical report, CARNEGIE-MELLON UNIV PITTSBURGH PA ARTIFICIAL INTELLIGENCE AND PSYCHOLOGY …, 1989.
  * Raffin (2020) Raffin, A.  Rl baselines3 zoo.  https://github.com/DLR-RM/rl-baselines3-zoo, 2020.
  * Raffin et al. (2021) Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., and Dormann, N.  Stable-baselines3: Reliable reinforcement learning implementations.  _Journal of Machine Learning Research_ , 22(268):1–8, 2021.  URL http://jmlr.org/papers/v22/20-1364.html.
  * Rajaraman et al. (2020) Rajaraman, N., Yang, L. F., Jiao, J., and Ramachandran, K.  Toward the fundamental limits of imitation learning.  _arXiv preprint arXiv:2009.05990_ , 2020.
  * Rajaraman et al. (2021a) Rajaraman, N., Han, Y., Yang, L., Liu, J., Jiao, J., and Ramchandran, K.  On the value of interaction and function approximation in imitation learning.  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021a.  URL https://openreview.net/forum?id=Xa9Ba6NsJ6.
  * Rajaraman et al. (2021b) Rajaraman, N., Han, Y., Yang, L. F., Ramchandran, K., and Jiao, J.  Provably breaking the quadratic error compounding barrier in imitation learning, optimally, 2021b.
  * Reddy et al. (2019) Reddy, S., Dragan, A. D., and Levine, S.  Sqil: Imitation learning via reinforcement learning with sparse rewards, 2019.
  * Ross & Bagnell (2010) Ross, S. and Bagnell, D.  Efficient reductions for imitation learning.  In Teh, Y. W. and Titterington, M. (eds.), _Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics_ , volume 9 of _Proceedings of Machine Learning Research_ , pp. 661–668, Chia Laguna Resort, Sardinia, Italy, 13–15 May 2010. PMLR.  URL https://proceedings.mlr.press/v9/ross10a.html.
  * Ross et al. (2011) Ross, S., Gordon, G. J., and Bagnell, J. A.  A reduction of imitation learning and structured prediction to no-regret online learning, 2011.
  * Schulman et al. (2017) Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347, 2017.  URL http://arxiv.org/abs/1707.06347.
  * Swamy et al. (2021) Swamy, G., Choudhury, S., Bagnell, J. A., and Wu, Z. S.  Of moments and matching: A game-theoretic framework for closing the imitation gap.  2021\.  URL https://arxiv.org/abs/2103.03236.
  * Syrgkanis et al. (2015) Syrgkanis, V., Agarwal, A., Luo, H., and Schapire, R. E.  Fast convergence of regularized learning in games.  _arXiv preprint arXiv:1507.00407_ , 2015.
  * Wang et al. (2019) Wang, R., Ciliberto, C., Amadori, P. V., and Demiris, Y.  Random expert distillation: Imitation learning via expert policy support estimation.  In _International Conference on Machine Learning_ , pp. 6536–6544. PMLR, 2019.
  * Zadeh (1965) Zadeh, L.  Fuzzy sets.  _Information and Control_ , 8(3):338–353, 1965\.  ISSN 0019-9958.  doi: https://doi.org/10.1016/S0019-9958(65)90241-X.  URL https://www.sciencedirect.com/science/article/pii/S001999586590241X.
  * Ziebart et al. (2008) Ziebart, B. D., Maas, A. L., Bagnell, J. A., Dey, A. K., et al.  Maximum entropy inverse reinforcement learning.  In _Aaai_ , volume 8, pp. 1433–1438. Chicago, IL, USA, 2008.

##  Appendix A Proofs

###  A.1 Notation

In this appendix, we use the notation $d_{t}^{\pi}(\cdot,\cdot)$ to indicate
the state-action visitation measure induced by the policy $\pi$ at time $t$.
We overload the notation $d_{t}^{\pi}(\cdot)$ to denote the state-visitation
measure induced by the policy $\pi$ at time $t$. Likewise, the notations
$d_{t}^{D}(\cdot,\cdot)$ and $d_{t}^{D}(\cdot)$ indicate the empirical
visitation measures in the dataset $D$. For a function
$g:\mathcal{X}\to\mathbb{R}$, the norm
$\|g\|_{\infty}\triangleq\sup_{x\in\mathcal{X}}|g(x)|$.

Before discussing the proofs of the results, we also explain the instantiation
of the function class in the tabular setting below.

######  Remark 4.

In the tabular setting, we instantiate the discriminator class as
$\mathcal{F}_{t}=\\{f_{t}:\|f_{t}\|_{\infty}\leq 1\\}$ for each $t$, as the
set of all $1$-bounded functions, and the policy class $\Pi$ as the set of all
policies. eq. 4 corresponds to finding a policy which best matches the
empirical state-action visitation measure observed in the dataset $D$ in total
variation (TV) distance (see Section A.2 for a proof).

###  A.2 Imitation gap upper bound on empirical moment matching (3.1)

Below we restate 3.1 and provide a proof of this result. The key observation
is that since the learner $\pi^{\texttt{MM}}$ best matches the empirical
distribution in the dataset, which is in turn close to the population
visitation measure induced by $\pi^{E}$, we can expect the visitation measure
induced by $\pi^{E}$ and $\pi^{\texttt{MM}}$ to be close. This in turns
implies that both policies will collect a similar value under any reward
function. Precisely characterizing the rates at which these distributions
converge to one another results in the final bound.

See 3.1

###### Proof.

Recall that the learner $\pi^{\texttt{MM}}$ is the solution to the following
optimization problem,

$\displaystyle\operatorname*{arg\,min}_{\pi}\sup_{f\in\mathcal{F}}\left\\{\mathbb{E}_{\pi}\left[\frac{\sum_{t=1}^{H}f_{t}(s_{t},a_{t})}{H}\right]-\mathbb{E}_{D}\left[\frac{\sum_{t=1}^{H}f_{t}(s_{t},a_{t})}{H}\right]\right\\}$
(21)

Exchanging the summation and maximization operators and recalling from Remark
4 that in the tabular setting, the discriminator class $\mathcal{F}$ is
instantiated as the set of all $1$-bounded functions
$\bigoplus_{t=1}^{H}\\{f_{t}:\|f_{t}\|_{\infty}\leq 1\\}$, $\pi^{\texttt{MM}}$
is a solution to

$\operatorname*{arg\,min}_{\pi}\frac{1}{H}\sum_{t=1}^{H}\Big{(}\sup_{f:\|f\|_{\infty}\leq
1}\mathbb{E}_{\pi}\left[f_{t}(s_{t},a_{t})\right]-\mathbb{E}_{D}\left[f_{t}(s_{t},a_{t})\right]\Big{)}=\operatorname*{arg\,min}_{\pi}\frac{1}{H}\sum_{t=1}^{H}\mathsf{TV}\left(d_{t}^{\pi},d_{t}^{D}\right)$
(22)

where the equation follows by the variational definition of the total
variation distance, and where $d_{t}^{\pi}$ is the state-action visitation
measure induced by $\pi^{E}$ and $d_{t}^{D}$ is the empirical state-action
visitation measure in the dataset $D$. The imitation gap of this policy can be
upper bounded by,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{MM}})$
$\displaystyle=\mathbb{E}_{\pi^{E}}\left[\sum_{t=1}^{H}r_{t}(s_{t},a_{t})\right]-\mathbb{E}_{\pi^{\texttt{MM}}}\left[\sum_{t=1}^{H}r_{t}(s_{t},a_{t})\right]$
(23)
$\displaystyle\overset{(i)}{\leq}\sum_{t=1}^{H}\sup_{r_{t}:\|r_{t}\|_{\infty}\leq
1}\Big{(}\mathbb{E}_{\pi^{E}}\left[r_{t}(s_{t},a_{t})\right]-\mathbb{E}_{\pi^{\texttt{MM}}}\left[r_{t}(s_{t},a_{t})\right]\Big{)}$
(24)
$\displaystyle\overset{(ii)}{=}\sum_{t=1}^{H}\mathsf{TV}\left(d_{t}^{\pi^{E}}(\cdot,\cdot),d_{t}^{\pi^{\texttt{MM}}}(\cdot,\cdot)\right)$
(25)

where $(i)$ maximizes over the reward function which is assumed to lie in the
interval $[0,1]$ pointwise. $(ii)$ again follows from the variational
definition of total variation distance. This goes to show that in the tabular
setting, MM is equivalent to finding the policy which best matches (in TV-
distance) the empirical state-action distribution observed in the dataset.

By an application of triangle inequality,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{MM}})$
$\displaystyle\leq\sum_{t=1}^{H}\mathsf{TV}\left(d_{t}^{\pi^{E}}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)+\mathsf{TV}\left(d_{t}^{D}(\cdot,\cdot),d_{t}^{\pi^{\texttt{MM}}}(\cdot,\cdot)\right)$
(26) $\displaystyle\leq
2\sum_{t=1}^{H}\mathsf{TV}\left(d_{t}^{\pi^{E}}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)$
(27)

where $(i)$ follows from eq. 22 which shows that $\pi^{\texttt{MM}}$ is the
policy which best approximates the empirical state-action visitation measure
in total variation distance, and therefore
$\mathsf{TV}\left(d_{t}^{\pi^{\texttt{MM}}}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)\leq\mathsf{TV}\left(d_{t}^{\pi^{E}}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)$.
The final element is to identify the rate of convergence of the empirical
visitation measure $d_{t}^{D}$, to the population $d_{t}^{\pi^{E}}$ in total
variation distance. This result is known from Theorem 1 of Han et al. (2015),
which shows that
$\mathbb{E}\left[\mathsf{TV}\left(d_{t}^{\pi^{E}}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)\right]\lesssim\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}$
noting that $d_{t}^{\pi^{E}}$ is a distribution with support size
$|\mathcal{S}|$ since $\pi^{E}$ is deterministic. Putting it together with eq.
27 after taking expectations on both sides gives,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{MM}})$
$\displaystyle\lesssim\sum_{t=1}^{H}\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}=H\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}.$
(28)

This completes the proof of the result. ∎

###  A.3 Lower bounding the Imitation gap of Behavior Cloning

Since BC is an offline algorithm - namely, the learner does not interact with
the MDP, any lower bound against offline algorithms applies for behavior
cloning as well. The lower bound instance in Rajaraman et al. (2020) is one
such example. One state in the MDP is labelled as the “bad” state, $b$, which
is absorbing and offers no reward. The remaining states each have a single
“good” action which re-initializes the policy in a particular distribution
$\rho$ on the set of good states, and offering a reward of $1$. The other
actions at these states are “bad” and transition the learner to the bad state
$b$ with probability $1$.

The key idea in the lower bound is that any offline algorithm does not know
(i) which action the expert would have chosen at states unvisited in the
dataset, and (ii) which action does what at these states. At best, the learner
can correctly guess the good actions at a state with probability
$1/|\mathcal{A}|\leq 1/2$. So, at each state unvisited in the dataset, the
learner has a constant probability of getting stuck at the bad state in the
MDP and collecting no reward then on. On the other hand, the expert would
never choose bad actions at states and collects the maximum possible reward.
By carefully counting the probability mass on the unvisited states, the
expected imitation gap of any offline IL algorithm (such as BC ) can be lower
bounded by $\Omega(|\mathcal{S}|H^{2}/N_{\mathrm{exp}})$ on these instances.
Thus we have the following theorem,

######  Theorem 4 (Theorem 6.1 of Rajaraman et al. (2020)).

Consider any learner $\widehat{\pi}$ which carries out an offline imitation
learning algorithm (such as behavior cloning). Then, there exists an MDP
instance such that the expected imitation gap is lower bounded by,

$\displaystyle\mathbb{E}\left[J(\pi^{E})-J(\widehat{\pi})\right]\gtrsim\min\left\\{H,\frac{|\mathcal{S}|H^{2}}{N_{\mathrm{exp}}}\right\\}.$
(29)

###  A.4 Lower bounding the imitation gap of Empirical Moment Matching

In this section, we show that in the tabular setting, empirical moment
matching is suboptimal compared for imitation learning in the worst-case. The
main result we prove in this section is,

See 3.2

######  Remark 5.

It is known from Rajaraman et al. (2020) that the Mimic-MD algorithm achieves
an imitation gap of
$\min\left\\{H,\frac{\mathcal{S}|H^{3/2}}{N_{\mathrm{exp}}},H\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}\right\\}$.
This is always better than the worst case error bound incurred by TV
distribution matching from 3.2. In fact when $N_{\mathrm{exp}}\gtrsim\sqrt{H}$
the bound $\frac{H^{3/2}}{N_{\mathrm{exp}}}$ is significantly better than
$\frac{H}{\sqrt{N_{\mathrm{exp}}}}$ which decays as
$1/\sqrt{N_{\mathrm{exp}}}$. This is illustrated in Figure 4 .

First note that the learner $\pi^{\texttt{MM}}$ carries out empirical moment
matching (eq. 4), with the discriminator class $\mathcal{F}$ as initialized in
Remark 4. As shown in eq. 22, the empirical moment matching learner can be
redefined as the solution to a distribution matching problem,

$\operatorname*{arg\,min}_{\pi}\frac{1}{H}\sum_{t=1}^{H}\mathsf{TV}\left(d_{t}^{\pi}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot)\right)$
(30)

Consider an MDP instance with $2$ states and $2$ actions with a non-stationary
transition and reward structure as described in Figure fig. 7. State $1$
effectively has a single action (i.e. two actions, $a_{1}$ and $a_{2}$ with
both inducing the same next-state distribution and reward). One of the actions
at state $2$ induces the uniform distribution over next states. The other
action deterministically keeps the learner at state $2$. The reward function
is $0$ at $t=1$, and the action $a_{2}$ at state $2$ is the only one which
offers a reward of $1$. The initial state distribution is highly skewed toward
the state $s=1$ and places approximately $1/\sqrt{N_{\mathrm{exp}}}$ mass on
$s=2$ and the remaining on $s=1$.

$s_{1}$$\mathrm{Unif}(\mathcal{S})$$s_{2}$ Figure 7: MDP instance which shows
that $L_{1}$ distribution matching is suboptimal. Here the transition
structure is illustrated for $t=1$. Both states have one action which
reinitializes in the uniform distribution. State $2$ has an additional action
which keeps the state the same. The reward function is $0$ for $t=1$. For
$t\geq 2$ the transition function is absorbing at both states; the reward
function equals $1$ at the state $s=1$ for any action and is $0$ everywhere
else.

MDP transition: The state $2$ is the only one with two actions. Action $a_{1}$
induces the uniform distribution over states, while action $a_{2}$ transitions
the learner to state $2$ with probability $1$. Namely,

$\displaystyle P_{1}(\cdot|s=1,a)$
$\displaystyle=\mathrm{Unif}(\mathcal{S})\text{ for all }a\in\mathcal{A}$ (31)
$\displaystyle P_{1}(\cdot|s=2,a_{1})$
$\displaystyle=\mathrm{Unif}(\mathcal{S})$ (32) $\displaystyle
P_{1}(\cdot|s=2,a_{2})$ $\displaystyle=\delta_{2}$ (33)

From time $t=2$ onward, the actions are all absorbing. Namely, for all $t\geq
2$, $s\in\mathcal{S}$ and $a\in\mathcal{A}$,

$\displaystyle P_{t}(\cdot|s,a)=\delta_{s}.$ (34)

Initial state distribution: The initial state distribution
$\rho=\left(1-\frac{1}{\sqrt{N_{\mathrm{exp}}}},\frac{1}{\sqrt{N_{\mathrm{exp}}}}\right)$.

MDP reward function: The reward function of the MDP encourages the learner to
stay at the state $s=1$ from time $t=2$ onward. Namely,

$\displaystyle r_{t}(s,a)=\begin{cases}1,\qquad&\text{if }t\geq 2\text{ and
}s=1\\\ 0,&\text{otherwise.}\end{cases}$ (35)

Expert policy: At both states in the MDP, the expert picks the action $a_{1}$
to play, which induces the uniform distribution over actions at the next
state. Namely, for each $t\in[H]$ and $s\in\mathcal{S}$,

$\displaystyle\pi^{E}_{t}(\cdot|s)=\delta_{a_{1}}$ (36)

The intuition behind the lower bound is as follows. The only action which
affects the value of a policy is the choice made at $s=2$ at time $t=1$. At
all other states, we may assume that there is effectively only a single
action.

By the absorbing nature of states for $t\geq 2$, it turns out that if the
observed empirical distribution in the dataset at time $2$ is skewed toward
state $2$ (which is possible because of the inherent randomness in the data
generation process), the learner’s behavior at time $1$ may be to ignore the
expert’s action observed at state $s=2$, and instead pick the action $a_{2}$
which moves the learner to the state $s=2$ deterministically. The learner is
willing to choose a different action because the loss function eq. 30
encourages the state-action distribution at time $t=2$ also to be well matched
with what is observed in the dataset. Even if it comes at the cost of picking
an action different from what the expert plays. By exploiting this fact, we
are able to show that the error incurred by a learner which solves eq. 30 in
this simple $2$ state example must be $\Omega(H/\sqrt{N_{\mathrm{exp}}})$.

Formally, we define $3$ events,

  1. (i)

$\mathcal{E}_{1}$: All states in the MDP are visited in the dataset at each
time $t=1,2,\cdots,H$.

  2. (ii)

$\mathcal{E}_{2}$: State $2$ is visited at most $\sqrt{N_{\mathrm{exp}}}$
times at time $1$ in the dataset $D$. In other words,
$d_{1}^{D}(s=2)=\delta^{\prime}$ where
$\delta^{\prime}\leq\frac{1}{\sqrt{N_{\mathrm{exp}}}}$.

  3. (iii)

$\mathcal{E}_{3}$: At time $2$ in the dataset $D$, the empirical distribution
over states is of the form
$\left(\frac{1}{2}-\delta,\frac{1}{2}+\delta\right)$ for some
$\delta\geq\frac{2}{\sqrt{N_{\mathrm{exp}}}}$.

######  Lemma 1.

Jointly, the events $\mathcal{E}_{1},\mathcal{E}_{2}$ and $\mathcal{E}_{3}$
occur with at least constant probability.

$\displaystyle\mathrm{Pr}(\mathcal{E}_{1}\cap\mathcal{E}_{2}\cap\mathcal{E}_{3})\geq
C,$ (37)

for some constant $C>0$.

###### Proof.

By the absorbing nature of states for $t\geq 2$, it suffices for both states
of the MDP to be visited in the dataset at time $t=1,2$. At time $t=2$, the
marginal state distribution under $\pi^{E}$ is the uniform distribution. By
binomial concentration, both states are observed in the dataset at time $t=2$
with probability $\geq 1-e^{-C_{1}N_{\mathrm{exp}}}$ for some constant
$C_{1}>0$. On the other hand, at time $t=1$, the marginal state distribution
is
$\rho=\left(1-\frac{1}{\sqrt{N_{\mathrm{exp}}}},\frac{1}{\sqrt{N_{\mathrm{exp}}}}\right)$.
Yet again, by binomial concentration, both states are observed with
probability $\geq 1-e^{-C_{2}\sqrt{N_{\mathrm{exp}}}}$ for some constant
$C_{2}>0$. By union bounding,

$\mathrm{Pr}(\mathcal{E}_{1})\geq
1-e^{-C_{1}N_{\mathrm{exp}}}-e^{-C_{2}\sqrt{N_{\mathrm{exp}}}}.$ (38)

Next we study $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$ together. First of all,
note that the state observed at $t=1$ and $t=2$ in a rollout of the expert
policy are independent. This is because at both states at $t=1$, the next
state distribution under $\pi^{E}$ is uniform. Because of this fact,
$\mathcal{E}_{2}$ and $\mathcal{E}_{3}$ are independent. Next we individually
bound the probability of the two events.

$\mathcal{E}_{2}$: The number of times $s=2$ is the initial state in
trajectories the dataset $D$ is distributed as a binomial random variable with
distribution $\mathrm{Bin}(N_{\mathrm{exp}},q)$ with
$q=\rho(s=2)=\frac{1}{\sqrt{N_{\mathrm{exp}}}}$. A median of a binomial random
variable is $N_{\mathrm{exp}}q=\sqrt{N_{\mathrm{exp}}}$ (in fact any number in
the interval $[\lfloor N_{\mathrm{exp}}q\rfloor,\lceil
N_{\mathrm{exp}}q\rceil]$ is a median). Therefore, the probability that $s=2$
is visited $\leq\sqrt{N_{\mathrm{exp}}}$ times in the dataset at time $1$ is
at least $1/2$. In summary,

$\mathrm{Pr}(\mathcal{E}_{2})\geq\frac{1}{2}$ (39)

$\mathcal{E}_{3}$: The marginal distribution over states at time $2$ in the
dataset is uniform. Therefore, we expect the states $1$ and $2$ to be visited
roughly $N_{\mathrm{exp}}/2$ times each in the dataset, but with a random
variation of $\approx\sqrt{N_{\mathrm{exp}}}$ around this average. In other
words, the empirical distribution fluctuates as
$\left(\frac{1}{2}-\delta,\frac{1}{2}+\delta\right)$ with
$\delta\geq\frac{2}{\sqrt{N_{\mathrm{exp}}}}$ with constant probability.

By the independence of $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$ and union
bounding to account for $\mathcal{E}_{1}$, the statement of the lemma follows.
∎

######  Lemma 2.

For each $t\geq 2$,

$\mathsf{TV}(d_{t}^{\pi}(\cdot,\cdot),d_{t}^{D}(\cdot,\cdot))\geq\mathsf{TV}(d_{2}^{\pi}(\cdot),d_{2}^{D}(\cdot)).$
(40)

The RHS is the TV distance between the state-visitation measure at time $t=2$
under $\pi$ and that empirically observed in the dataset $D$. Conditioned on
the events $\mathcal{E}_{1}$, $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$
occuring, equality is met in eq. 40 if any only if
$\pi_{t}(\cdot|s)=\pi^{E}_{t}(\cdot|s)$ for all states $s\in\mathcal{S}$.

###### Proof.

For any state $s\in\mathcal{S}$ and $t\geq 2$, observe that,

$\displaystyle\sum_{a\in\mathcal{A}}\left|d_{t}^{\pi}(s,a)-d_{t}^{D}(s,a)\right|$
(41)
$\displaystyle=d_{t}^{\pi}(s)(1-\pi_{t}(a^{*}|s))+\left|d_{t}^{\pi}(s)\pi_{t}(a^{*}|s)-d_{t}^{D}(s,a^{*})\right|,\text{
where }a^{*}=\pi^{E}_{t}(s),$ (42)
$\displaystyle\overset{(i)}{=}d_{2}^{\pi}(s)(1-\pi_{t}(a^{*}|s))+\left|d_{2}^{\pi}(s)\pi_{t}(a^{*}|s)-d_{2}^{D}(s,a^{*})\right|$
(43)
$\displaystyle\overset{(ii)}{\geq}\left|d_{2}^{\pi}(s)-d_{2}^{D}(s)\right|,$
(44)

where $(i)$ follows by the fact that the states of the MDP are absorbing under
$\pi$ for $t\geq 2$. $(ii)$ follows by triangle inequality and using the fact
that $\pi^{E}$ is deterministic, so $d_{t}^{D}(s,a^{*})=d_{t}^{D}(s)$.
Equality is met only if $\pi_{t}(a^{*}|s)=1$ (since $d_{t}^{D}(s,a^{*})>0$
conditioned on $\mathcal{E}_{1}$). ∎

The above lemma asserts the behavior of $\pi^{\texttt{MM}}$ in eq. 30 for
$t\geq 2$. Namely, conditioned on the event $\mathcal{E}_{1}$ which happens
with very high probability, all states are visited in the MDP and therefore,
$\pi^{\texttt{MM}}_{t}(\cdot|s)=\pi^{E}_{t}(\cdot|s)$ for each state
$s\in\mathcal{S}$ and time $t\geq 2$.

The only thing left to study is the MM learner’s behavior at $t=1$. We wish to
show that with constant probability, the learner may choose to deviate from
the expert policy in order to better match empirical state-action visitation
measures. Conditioned on $\mathcal{E}_{1}$, the learner’s policy at time $t=1$
can be computed by solving the following optimization problem,

$\displaystyle\mathsf{TV}(d_{1}^{\pi}(\cdot,\cdot),d_{1}^{D}(\cdot,\cdot))+(H-1)\textsf{TV}(d_{2}^{\pi}(\cdot),d_{2}^{D}(\cdot)).$
(45)

This follows directly by simplifying the learner’s objective using Lemma 2.

Now, conditioned on the event $\mathcal{E}_{1}$, at time $t=1$, the learner
policy only needs to be optimized at the state $s=2$. At the state $s=1$, we
may assume that the learner picks the expert’s action $\pi^{E}_{1}(s=1)$. To
this end, suppose the learner picks the action $a_{1}$ with probability $p$
and the action $a_{2}$ with probability $1-p$.

$\displaystyle\textsf{TV}(d_{1}^{\pi}(\cdot,\cdot),d_{1}^{D}(\cdot,\cdot))$
$\displaystyle=\sum_{a\in\mathcal{A}}\left|d_{1}^{\pi^{E}}(s=2,a)-d_{1}^{D}(s=2,a)\right|$
(46)
$\displaystyle=\left|\rho(2)p-\delta^{\prime}\right|+\left|\rho(2)(1-p)-0\right|$
(47)
$\displaystyle=\left|\frac{p}{\sqrt{N_{\mathrm{exp}}}}-\delta^{\prime}\right|+\frac{1-p}{\sqrt{N_{\mathrm{exp}}}}.$
(48)

which follows by plugging in $\rho(2)=\frac{1}{\sqrt{N_{\mathrm{exp}}}}$. And,

$\displaystyle\textsf{TV}(d_{2}^{\pi}(\cdot),d_{2}^{D}(\cdot))$
$\displaystyle=\left|\left(\frac{1}{2}-\delta\right)-\frac{\rho(1)}{2}-\rho(2)\frac{p}{2}\right|+\left|\left(\frac{1}{2}+\delta\right)-\frac{\rho(1)}{2}-\rho(2)\left((1-p)+\frac{p}{2}\right)\right|.$
(49)

Plugging in $\rho(2)=\frac{1}{\sqrt{N_{\mathrm{exp}}}}$ and
$\rho(1)=1-\frac{1}{\sqrt{N_{\mathrm{exp}}}}$, we get,

$\displaystyle\textsf{TV}(d_{2}^{\pi}(\cdot),d_{2}^{D}(\cdot))$
$\displaystyle=\left|\frac{1}{2\sqrt{N_{\mathrm{exp}}}}-\delta-\frac{p}{2\sqrt{N_{\mathrm{exp}}}}\right|+\left|\frac{p}{2\sqrt{N_{\mathrm{exp}}}}-\frac{1}{2\sqrt{N_{\mathrm{exp}}}}+\delta\right|.$
(50)

Summing up eqs. 48 and 50, $p$ minimizes,

$\underbrace{\left|\frac{p}{\sqrt{N_{\mathrm{exp}}}}-\delta^{\prime}\right|+\frac{1-p}{\sqrt{N_{\mathrm{exp}}}}}_{(i)}+\underbrace{(H-1)\left(\left|\frac{p}{2\sqrt{N_{\mathrm{exp}}}}+\delta-\frac{1}{2\sqrt{N_{\mathrm{exp}}}}\right|+\left|\frac{1}{2\sqrt{N_{\mathrm{exp}}}}-\delta-\frac{p}{2\sqrt{N_{\mathrm{exp}}}}\right|\right)}_{(ii)}.$
(51)

Intuitively, term $(i)$ captures the error incurred by the learner in the loss
eq. 30 by deviating from $\pi^{E}$ at the first time step. Term $(ii)$
captures the decrease in the error at every subsequent time step because of
the same deviation, since the learner is able to better match the state
distribution at future time steps. In the next lemma we show that under events
that hold with at least constant probability, the empirical moment matching
learner chooses to play the wrong action at time $t=1$ at the state $s=2$.

######  Lemma 3.

Conditioned on the events $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$, for $H\geq
4$, the unique minimizer of eq. 51 for $p\in[0,1]$ is $p=0$.

###### Proof.

The first term of eq. 51 is
$\left|\frac{p}{\sqrt{N_{\mathrm{exp}}}}-\delta^{\prime}\right|+\frac{1-p}{\sqrt{N_{\mathrm{exp}}}}$,
the error from not picking the expert’s action at state $1$ at time $1$
decreases at most linearly with a slope of
$\frac{2}{\sqrt{N_{\mathrm{exp}}}}$.

Conditioned on the event $\mathcal{E}_{3}$,
$\delta\geq\frac{2}{\sqrt{N_{\mathrm{exp}}}}$. Therefore,
$\left|\frac{p}{2\sqrt{N_{\mathrm{exp}}}}+\delta-\frac{1}{2\sqrt{N_{\mathrm{exp}}}}\right|=\frac{p}{2\sqrt{N_{\mathrm{exp}}}}+\delta-\frac{1}{2\sqrt{N_{\mathrm{exp}}}}$.
Therefore, the decrease in error at future steps by deviating from $\pi^{E}$
at the time $t=1$, term $(ii)$ in eq. 51 is,

$\displaystyle
2(H-1)\left(\frac{p}{2\sqrt{N_{\mathrm{exp}}}}+\delta-\frac{1}{2\sqrt{N_{\mathrm{exp}}}}\right)$
(52)

which is an increasing function of $p$ with slope
$\frac{H-1}{\sqrt{N_{\mathrm{exp}}}}$. For $H\geq 4$ and the argument from the
previous paragraph, this implies that term $(ii)$ increases more rapidly in
$p$ than the rate at which term $(i)$ decreases. Therefore, the minimizer must
be $p=0$. ∎

Thus, we conclude from Lemmas 2 and 3 that conditioned on the events
$\mathcal{E}_{1},\mathcal{E}_{2}$ and $\mathcal{E}_{3}$, the learner
$\pi^{\texttt{MM}}$ perfectly mimics $\pi^{E}$ at each time $t\geq 2$, but
deviates from the action played by $\pi^{E}$ at the state $s=1$ at time $t=1$.

Finally, we bound the difference in value between $\pi^{E}$ and
$\pi^{\texttt{MM}}$ induced because of this deviation under the reward eq. 35.

######  Lemma 4.

Under the events $\mathcal{E}_{1},\mathcal{E}_{2}$ and $\mathcal{E}_{3}$,
under the reward eq. 35, the empirical moment matching learner
$\pi^{\texttt{MM}}$ incurs imitation gap,

$J(\pi^{E})-J(\pi^{\texttt{MM}})=\frac{H}{2\sqrt{N_{\mathrm{exp}}}}.$ (53)

###### Proof.

Recall that under the events $\mathcal{E}_{1},\mathcal{E}_{2}$ and
$\mathcal{E}_{3}$, the learner $\pi^{\texttt{MM}}$ is identical to $\pi^{E}$
except at the state $s=2$ where they perfectly deviate from each other. The
state distribution induced by $\pi^{E}$ at each time $t\geq 2$ is the uniform
distribution over states $\left(\frac{1}{2},\frac{1}{2}\right)$. On the other
hand, for $t\geq 2$, the state distribution induced by $\pi^{\texttt{MM}}$ at
each time $t\geq 2$ is
$\left(\rho(1)\frac{1}{2},\rho(1)\frac{1}{2}+\rho(2)\right)=\left(\frac{1-1/\sqrt{N_{\mathrm{exp}}}}{2},\frac{1+1/\sqrt{N_{\mathrm{exp}}}}{2}\right)$.
Since the reward function is $1$ on state $1$, the difference in value between
the expert and learner policy is,

$J(\pi^{E})-J(\pi^{\texttt{MM}})=\frac{H}{2}-H\left(\frac{1-1/\sqrt{N_{\mathrm{exp}}}}{2}\right)=\frac{H}{2\sqrt{N_{\mathrm{exp}}}}.$
(54)

This completes the proof. ∎

Since $\mathcal{E}_{1}$, $\mathcal{E}_{2}$ and $\mathcal{E}_{3}$ jointly occur
with constant probability by Lemma 1, this completes the proof of 3.2.

##  Appendix B Imitation gap of RE : Proof of Theorem 2

In this section, we discuss a proof of a more general version of Theorem 2,
where $N_{\mathrm{replay}}$ can be finite. We prove the following result,

######  Theorem 5.

Consider the policy $\pi^{\texttt{RE}}$ returned by Algorithm 1. Assume that
$\pi^{E}\in\Pi$ and the ground truth reward function
$r_{t}\in\mathcal{F}_{t}$, which is assumed to be symmetric
($f_{t}\in\mathcal{F}_{t}\iff-f_{t}\in\mathcal{F}_{t}$) and bounded (For all
$f_{t}\in\mathcal{F}_{t}$, $\|f_{t}\|_{\infty}\leq 1$). Choose
$|D_{1}|,|D_{2}|=\Theta(N_{\mathrm{exp}})$ and suppose
$N_{\mathrm{replay}}\to\infty$. With probability $\geq 1-3\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim$ $\displaystyle\
\mathcal{L}_{1}+\mathcal{L}_{2}+\mathcal{L}_{3}+\frac{\log\left(F_{\mathrm{max}}H/\delta\right)}{N_{\mathrm{exp}}}$
(55)

where $F_{\mathrm{max}}\triangleq\max_{t\in[H]}|\mathcal{F}_{t}|$, and,

$\displaystyle\mathcal{L}_{1}\triangleq H^{2}\
\mathbb{E}_{\pi^{E}}\\!\left[\frac{\sum_{t=1}^{H}\mathcal{M}(s_{t},t)\mathsf{TV}\left(\pi^{E}_{t}(\cdot|s_{t}),\pi^{\texttt{BC}}_{t}(\cdot|s_{t})\right)}{H}\right],$
(56) $\displaystyle\mathcal{L}_{2}\triangleq
H^{3/2}\sqrt{\frac{\log\left(F_{\mathrm{max}}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]}{H}},$
(57)

And,

$\displaystyle\mathcal{L}_{3}\triangleq
H\sqrt{\frac{\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{replay}}}}+\frac{H\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{replay}}}.$
(58)

Recall that the learner carrying out replay estimation returns the policy
which minimizes the loss $\sup_{f\in\mathcal{F}}J_{f}(\pi)-\widehat{E}(f)$
over policies $\pi$, where
$J_{f}(\pi)\triangleq\mathbb{E}_{\pi}\left[\frac{1}{H}\sum_{t=1}^{H}f_{t}(s_{t},a_{t})\right]$
and abbreviating the notation $f=(f_{1},\cdots,f_{H})$. Note that,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})$
$\displaystyle\overset{(i)}{\leq}\sup_{f\in\mathcal{F}}J_{f}(\pi^{E})-J_{f}(\pi^{\texttt{RE}})$
(59)
$\displaystyle\leq\sup_{f\in\mathcal{F}}\left|J_{f}(\pi^{E})-\widehat{E}(f)\right|+\sup_{f\in\mathcal{F}}\left|\widehat{E}(f)-J_{f}(\pi^{\texttt{RE}})\right|$
(60)
$\displaystyle\overset{(ii)}{\leq}2\sup_{f\in\mathcal{F}}\left|J_{f}(\pi^{E})-\widehat{E}(f)\right|.$
(61)

where $(i)$ uses the realizability assumption that the ground truth reward
lies in $\mathcal{F}$, and $(ii)$ uses the fact that $\pi^{\texttt{RE}}$ is a
minimizer of eq. 5 and the fact that $\mathcal{F}$ is symmetric (this implies
that
$\sup_{f\in\mathcal{F}}J_{f}(\pi^{E})-\widehat{E}(f)=\sup_{f\in\mathcal{F}}-J_{f}(\pi^{E})+\widehat{E}(f)=\sup_{f\in\mathcal{F}}\left|J_{f}(\pi^{E})-\widehat{E}(f)\right|$).

Note that $\widehat{E}(f)$ can be decomposed into a sum of two parts,

$\displaystyle\widehat{E}^{(1)}(f)=\mathbb{E}_{D_{\mathrm{replay}}}\left[\frac{1}{H}\sum_{t=1}^{H}f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\dots
t-1})\right)\right],\text{ and,}$ (62)
$\displaystyle\widehat{E}^{(2)}(f)=\mathbb{E}_{D_{2}}\left[\frac{1}{H}\sum_{t=1}^{H}f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\dots
t-1})\right)\right]$ (63)

Likewise, we can decompose $J_{f}(\pi^{E})$ into two terms,

$\displaystyle J_{f}^{(1)}(\pi^{E})$
$\displaystyle\triangleq\mathbb{E}_{\pi^{E}}\left[\sum_{t=1}^{H}f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\dots
t-1})\right],\text{ and}$ (64) $\displaystyle J_{f}^{(2)}(\pi^{E})$
$\displaystyle\triangleq\mathbb{E}_{\pi^{E}}\left[\sum_{t=1}^{H}f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\dots
t-1})\right)\right]$ (65)

Then, from eq. 61,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})$ $\displaystyle\leq
2\sup_{f\in\mathcal{F}}\left|J_{f}(\pi^{E})-\widehat{E}(f)\right|$ (66)
$\displaystyle\leq
2\underbrace{\sup_{f\in\mathcal{F}}\left|J_{f}^{(1)}(\pi^{E})-\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]\right|}_{\text{(I)}}+2\underbrace{\sup_{f\in\mathcal{F}}\left|\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]-\widehat{E}^{(1)}(f)\right|}_{\text{(II)}}$
$\displaystyle\hskip
160.00024pt+2\underbrace{\sup_{f\in\mathcal{F}}\left|J_{f}^{(2)}(\pi^{E})-\widehat{E}^{(2)}(f)\right|}_{\text{(III)}}.$
(67)

where the last line follows by triangle inequality. We bound each of these
terms in the next 3 lemmas, starting with (I).

######  Lemma 5.

$\displaystyle\sup_{f\in\mathcal{F}}\left|J_{f}^{(1)}(\pi^{E})-\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]\right|$
$\displaystyle\leq
H\sum_{h=1}^{H}\mathbb{E}_{\pi^{E}}\left[\mathcal{M}(s_{h},h)\textsf{TV}\left(\pi^{E}_{h}(\cdot|s_{h}),\pi^{\texttt{BC}}_{h}(\cdot|s_{h})\right)\right]$
(68)

###### Proof.

The proof of this result closely follows the supervised learning reduction of
BC (cf. Ross & Bagnell (2010)). Note that,

$\displaystyle\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]-V^{(1)}_{f}(\pi^{E})$
$\displaystyle=\sum_{t=1}^{H}\mathbb{E}_{\pi^{\texttt{BC}}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right].$ (69)

Define $\pi^{(h)}$ as the policy which plays $\pi^{E}$ until (and including)
time $h$ and $\pi^{\texttt{BC}}$ after time $h$. Then, by cascading,

$\displaystyle\mathbb{E}_{\pi^{\texttt{BC}}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]$
$\displaystyle=\sum_{h=0}^{t-1}\mathbb{E}_{\pi^{(h)}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]-\mathbb{E}_{\pi^{(h+1)}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]$ (70)

Define, the uncertainty weighted state visitation measure $d^{\mathcal{M}}$
and the uncertainty weighted look-forward reward $\rho^{\mathcal{M}}$ as
follows,

$\displaystyle
d^{\mathcal{M}}_{h+1}(s^{\prime})\triangleq\mathbb{E}_{\pi^{E}}\left[\mathbbm{1}(s_{h+1}=s^{\prime})\prod_{t^{\prime}=1}^{h}\mathcal{M}(s_{t^{\prime}},t^{\prime})\right]$
(71)
$\displaystyle\rho^{\mathcal{M}}_{h+1}(s^{\prime},a^{\prime})\triangleq\mathbb{E}_{\pi^{\texttt{BC}}}\left[f_{t}(s_{t},a_{t})\prod_{t^{\prime}=h+2}^{t}\mathcal{M}(s_{t^{\prime}},t^{\prime})\middle|s_{h+1}=s^{\prime},a_{h+1}=a^{\prime}\right]$
(72)

By decomposing expectations along trajectories, using the fact that
$\mathcal{P}(s_{1\cdots
t-1})=\prod_{t^{\prime}=1}^{H}\mathcal{M}(s_{t^{\prime}},t^{\prime})$ some
simplification results in the following equation,

$\displaystyle\left|\mathbb{E}_{\pi^{(h)}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]-\mathbb{E}_{\pi^{(h+1)}}\left[f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots
t-1})\right]\right|$ (73)
$\displaystyle=\left|\sum_{s^{\prime}\in\mathcal{S}}\sum_{a^{\prime}\in\mathcal{A}}d^{\mathcal{M}}_{h+1}(s^{\prime})\mathcal{M}(s^{\prime},h+1)\left(\pi^{E}_{h+1}(a^{\prime}|s^{\prime})-\pi^{\texttt{BC}}_{h+1}(a^{\prime}|s^{\prime})\right)\rho^{\mathcal{M}}_{h+1}(s^{\prime},a^{\prime})\right|$
(74)
$\displaystyle\overset{(i)}{\leq}\sum_{s^{\prime}\in\mathcal{S}}d^{\mathcal{M}}_{h+1}(s^{\prime})\mathcal{M}(s^{\prime},h+1)\textsf{TV}\left(\pi^{E}_{h+1}(\cdot|s^{\prime}),\pi^{\texttt{BC}}_{h+1}(\cdot|s^{\prime})\right)$
(75)
$\displaystyle=\mathbb{E}_{\pi^{E}}\left[\mathcal{M}(s_{h+1},h+1)\textsf{TV}\left(\pi^{E}_{h+1}(\cdot|s_{h+1}),\pi^{\texttt{BC}}_{h+1}(\cdot|s_{h+1})\right)\right].$
(76)

where $(i)$ uses the fact that the membership oracle is a function $\in[0,1]$
and $f$ is bounded and lies in the interval $[0,1]$ (which implies that
$\rho^{\mathcal{M}}$ also lies in $[0,1]$ pointwise). Plugging into eq. 70 and
subsequently into eq. 69 completes the proof. ∎

Next we bound the 3rd term, (III). This follows by an application of
Bernstein’s inequality.

######  Lemma 6.

With probability $\geq 1-\delta$,

$\displaystyle\sup_{f\in\mathcal{F}}\left|J_{f}^{(2)}(\pi^{E})-\widehat{E}^{(2)}(f)\right|\leq
H\sqrt{\frac{\log(F_{\mathrm{max}}H/\delta)\sum_{t=1}^{H-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]}{N_{\mathrm{exp}}}}+\frac{H\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{exp}}}$
(77)

###### Proof.

First observe that,

$\displaystyle J_{f}^{(2)}(\pi^{E})-\widehat{V}^{(2)}_{f}$
$\displaystyle=\sum_{t=1}^{H}\mathbb{E}_{\textsf{tr}\sim\mathrm{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]$ (78)

For each $t$, note that $f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)$ is bounded in the range $[0,1]$. Therefore, invoking Bernstein’s
inequality, with probability $\geq 1-\delta$,

$\displaystyle\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (79)
$\displaystyle\lesssim\sqrt{\frac{\text{Var}_{\pi^{E}}(f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right))\log(1/\delta)}{N_{\mathrm{exp}}}}+\frac{\log(1/\delta)}{N_{\mathrm{exp}}}$
(80)
$\displaystyle\leq\sqrt{\frac{\mathbb{E}_{\pi^{E}}[(f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right))^{2}]\log(1/\delta)}{N_{\mathrm{exp}}}}+\frac{\log(1/\delta)}{N_{\mathrm{exp}}}$
(81)
$\displaystyle\overset{(i)}{\leq}\sqrt{\frac{\mathbb{E}_{\pi^{E}}[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)]\log(1/\delta)}{N_{\mathrm{exp}}}}+\frac{\log(1/\delta)}{N_{\mathrm{exp}}}$
(82)
$\displaystyle\overset{(ii)}{\leq}\sqrt{\frac{\mathbb{E}_{\pi^{E}}\left[1-\mathcal{P}(s_{1\cdots
t-1})\right]\log(1/\delta)}{N_{\mathrm{exp}}}}+\frac{\log(1/\delta)}{N_{\mathrm{exp}}}$
(83)

where $(i)$ uses the fact that
$f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots t-1})\right)$ is bounded in
the range $[0,1]$, and $(ii)$ uses the fact that $0\leq f_{t}(s_{t},a_{t})\leq
1$. Assuming $0\leq x_{i}\leq 1$ for all $i\in[n]$, we have the inequality,

$1-\prod_{i=1}^{n}x_{i}\leq\sum_{i=1}^{n}1-x_{i}$ (84)

Applying this to eq. 83 for $1-\mathcal{P}(s_{1\cdots
t-1})=1-\prod_{t^{\prime}=1}^{t-1}\mathcal{M}(s_{t^{\prime}},t^{\prime})$, we
have,

$\displaystyle\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (85)
$\displaystyle\leq\sqrt{\frac{\sum_{t^{\prime}=1}^{t-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t^{\prime}},t^{\prime})\right]\log(1/\delta)}{N_{\mathrm{exp}}}}+\frac{\log(1/\delta)}{N_{\mathrm{exp}}}$
(86)

Therefore, by union bounding, with probability $\geq 1-\delta/H$,
simultaneously for every $f_{t}\in\mathcal{F}_{t}$,

$\displaystyle\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (87)
$\displaystyle\lesssim\sqrt{\frac{\log(|\mathcal{F}_{t}|H/\delta)\sum_{t^{\prime}=1}^{t-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t^{\prime}},t^{\prime})\right]}{N_{\mathrm{exp}}}}+\frac{\log(|\mathcal{F}_{t}|H/\delta)}{N_{\mathrm{exp}}}.$
(88)

This implies that the maximum over $f_{t}$ of the LHS is upper bounded by the
RHS. Union bounding over $t=1,\cdots,H$ and plugging into eq. 78, we have that
with probability $\geq 1-\delta$,

$\displaystyle\sup_{f\in\mathcal{F}}\left|J_{f}^{(2)}(\pi^{E})-\widehat{V}^{(2)}_{f}\right|$
(89)
$\displaystyle\leq\sum_{t=1}^{H}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (90) $\displaystyle\lesssim
H\sqrt{\frac{\log(F_{\mathrm{max}}H/\delta)\sum_{t=1}^{H-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]}{N_{\mathrm{exp}}}}+\frac{H\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{exp}}}.$
(91)

∎

######  Lemma 7.

With probability $\geq 1-\delta$,

$\displaystyle\sup_{f\in\mathcal{F}}\left|\mathbb{E}\left[\widehat{V}_{f}^{(1)}\middle|D_{1}\right]-\widehat{V}_{f}^{(1)}\right|\lesssim
H\sqrt{\frac{\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{replay}}}}+\frac{H\log(F_{\mathrm{max}}H/\delta)}{N_{\mathrm{replay}}}.$
(92)

###### Proof.

The proof follows essentially the same structure as Lemma 6 by decomposing
$\widehat{V}_{f}^{(1)}$ into a sum of $H$ terms of the form
$f_{t}(s_{t},a_{t})\mathcal{P}(s_{1\cdots t-1})$, applying Bernstein’s
inequality to bound the deviation of each term from its mean and finally union
bounding over the rewards $f_{t}\in\mathcal{F}_{t}$ to get the uniform bound
over all discriminators $f\in\mathcal{F}$. ∎

Putting together Lemmas 5, 6 and 7 with eq. 67 completes the proof of Theorem
2.

###  B.1 Bound in the tabular setting (Theorem 6)

In this section, we provide an upper bound on the imitation gap of RE in the
tabular setting when the expert is a deterministic policy. This recovers the
bound on the imitation gap for RE proved in Rajaraman et al. (2020).

######  Theorem 6.

Consider an appropriately initialized version of RE , and let the size of the
replay dataset $N_{\mathrm{replay}}\to\infty$. For any tabular IL instance
with $H\geq 10$, with probability $\geq 1-3\delta$,

$J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim\min\left\\{\frac{|\mathcal{S}|H^{3/2}}{N_{\mathrm{exp}}},H\sqrt{\frac{|\mathcal{S}|}{N_{\mathrm{exp}}}}\right\\}\log\left(\frac{|\mathcal{S}|H}{\delta}\right).$
(93)

Below we describe the implementation of RE corresponding to Theorem 6 in more
detail.

The membership oracle we use in this setting for RE is defined below,

$\mathcal{M}(s,t)=\begin{cases}1\quad&\text{if }s\text{ is visited in
}D_{1}\text{ at time }t\\\ 0&\text{otherwise.}\end{cases}$ (94)

The function class $\mathcal{F}$ which we use is identical to that for
empirical moment matching, which is described in Remark 4.

Note that in the tabular setting, BC simply mimics the deterministic expert’s
actions at states visited in the dataset $D_{1}$ and plays an arbitrary
deterministic action on the remaining states. As a consequence of this
definition, if
$\mathcal{M}(s,t)=1\iff\pi^{\texttt{BC}}_{t}(\cdot|s)=\pi^{E}_{t}(\cdot|s)$
and $\mathcal{M}(s,t)=0$ otherwise. We instantiate the family of
discriminators as in Remark 4, as
$\mathcal{F}=\bigoplus_{t=1}^{H}\\{f_{t}:\|f_{t}\|_{\infty}\leq 1\\}$ and the
set of policies $\Pi$ optimized over is chosen as the set of all deterministic
policies. While the guarantee of Theorem 2 depends on
$F_{\mathrm{max}}=\max_{t\in[H]}|\mathcal{F}_{t}|$ which is unbounded (or
$\exp(|\mathcal{S}||\mathcal{A}|)$ by using a discretization of the reward
space), note that we can improve the guarantee to effectively have
$F_{\mathrm{max}}\approx\exp(|\mathcal{S}|)$ noting the structure of the set
of discriminators. Looking into the proof of Theorem 2 we bring out this
dependence below. We note that there are many ways of bringing out this
dependence, including a careful net argument directly on top of the guarantee
of Theorem 2. We simply present one such argument below.

The critical step where the finiteness of the set of discriminators
$\mathcal{F}$ is used, is in union bounding the gap between the population and
the empirical estimate of $f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)$ in eq. 86.

$\displaystyle\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (95)

In the next step of the proof of Theorem 2, we union bound over all
$f_{t}\in\mathcal{F}_{t}$. However, note that for
$\mathcal{F}_{t}=\\{f_{t}:\|f_{t}\|_{\infty}\leq 1\\}$, we have that,

$\displaystyle\sup_{f_{t}:\|f_{t}\|_{\infty}\leq
1}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (96)
$\displaystyle\overset{(i)}{=}\sum_{s\in\mathcal{S}}\sum_{a\in\mathcal{A}}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[\mathbb{I}(s_{t}=s,a_{t}=a)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s,a_{t}=a)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (97)
$\displaystyle\overset{(ii)}{=}\sum_{s\in\mathcal{S}}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (98)
$\displaystyle\overset{(iii)}{\leq}\sum_{s\in\mathcal{S}}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (99)

where $(i)$ follows similar to the equivalence between the variational
representation of TV distance
($\mathsf{TV}(P,Q)=\frac{1}{2}\sup_{f:\|f\|_{\infty}\leq
1}\mathbb{E}_{P}[f]-\mathbb{E}_{Q}[f]$) and the relationship to the $L_{1}$
distance, $\mathsf{TV}(P,Q)=\frac{1}{2}L_{1}(P,Q)$. On the other hand, $(ii)$
follows by noting that the expert is a deterministic policy (and $D_{2}$ is
generated by rolling out $\pi^{E}$). $(iii)$ follows by an application of
Holder’s inequality. By subgaussian concentration, for each $s\in\mathcal{S}$,
with probability $\geq 1-\frac{\delta}{|\mathcal{S}|H}$,

$\displaystyle\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (100)
$\displaystyle\lesssim\sqrt{\frac{\text{Var}_{\pi^{E}}\left(\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right)\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}+\frac{\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(101)
$\displaystyle\overset{(i)}{\leq}\sqrt{\frac{\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}+\frac{\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(102)

where $(i)$ uses the fact that
$0\leq\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots t-1})\right)\leq 1$.
Combining with eq. 99, union bounding and applying Cauchy Schwarz inequality,
with probability $\geq 1-\frac{\delta}{H}$,

$\displaystyle\sup_{f_{t}:\|f_{t}\|_{\infty}\leq
1}\left|\mathbb{E}_{\text{Unif}(D_{2})}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]-\mathbb{E}_{\pi^{E}}\left[f_{t}(s_{t},a_{t})\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\right|$ (103)
$\displaystyle\lesssim\sqrt{|\mathcal{S}|}\sqrt{\sum_{s\in\mathcal{S}}\frac{\mathbb{E}_{\pi^{E}}\left[\mathbb{I}(s_{t}=s)\left(1-\mathcal{P}(s_{1\cdots
t-1})\right)\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}+\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(104)
$\displaystyle=\sqrt{|\mathcal{S}|}\sqrt{\frac{\mathbb{E}_{\pi^{E}}\left[1-\mathcal{P}(s_{1\cdots
t-1})\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}+\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(105)
$\displaystyle\overset{(i)}{\leq}\min\left\\{\sqrt{\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}},\sqrt{|\mathcal{S}|\frac{\sum_{t=1}^{H-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}\right\\}+\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(106)

where $(i)$ follows by the same simplification as in eq. 84. Comparing with
eq. 88, this roughly corresponds to setting
$F_{\mathrm{max}}\approx\exp(|\mathcal{S}|)$. All in all, summing eq. 106 over
$t\in[H]$ and plugging into eq. 90, with probability $\geq 1-\delta$,

$\displaystyle\sup_{f\in\mathcal{F}}\left|J_{f}^{(2)}(\pi^{E})-\widehat{V}^{(2)}_{f}\right|$
(107) $\displaystyle\lesssim
H\left\\{\sqrt{\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}},\sqrt{|\mathcal{S}|\frac{\sum_{t=1}^{H-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}\right\\}+\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(108)

Finally, we plug this into eq. 67, which is restated below,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})$ $\displaystyle\leq
2\sup_{f\in\mathcal{F}}\left|J_{f}(\pi^{E})-\widehat{E}(f)\right|$ (109)
$\displaystyle\leq
2\underbrace{\sup_{f\in\mathcal{F}}\left|J_{f}^{(1)}(\pi^{E})-\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]\right|}_{\text{(I)}}+2\underbrace{\sup_{f\in\mathcal{F}}\left|\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]-\widehat{E}^{(1)}(f)\right|}_{\text{(II)}}$
(110) $\displaystyle\hskip
160.00024pt+2\underbrace{\sup_{f\in\mathcal{F}}\left|J_{f}^{(2)}(\pi^{E})-\widehat{E}^{(2)}(f)\right|}_{\text{(III)}}.$
(111)

For the chosen membership oracle in eq. 94, the term (I) is $0$, since by
Lemma 5 it is upper bounded by
$H\sum_{h=1}^{H}\mathbb{E}_{\pi^{E}}\left[\mathcal{M}(s_{h},h)\textsf{TV}\left(\pi^{E}_{h}(\cdot|s_{h}),\pi^{\texttt{BC}}_{h}(\cdot|s_{h})\right)\right]$.
This is equal to $0$ since $\mathcal{M}(s,t)=0$ wherever
$\pi^{E}_{t}(\cdot|s)\neq\pi^{\texttt{BC}}_{t}(\cdot|s)$. On the other hand,
$N_{\mathrm{replay}}\to\infty$ ensures that the term (III) goes to $0$ by the
strong law of large numbers. Therefore, with probability $\geq 1-2\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})$ (112) $\displaystyle\leq
2\sup_{f\in\mathcal{F}}\left|\mathbb{E}\left[\widehat{E}^{(1)}(f)\middle|D_{1}\right]-\widehat{E}^{(1)}(f)\right|$
(113) $\displaystyle\lesssim
H\left\\{\sqrt{\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}},\sqrt{|\mathcal{S}|\frac{\sum_{t=1}^{H-1}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}}\right\\}+\frac{|\mathcal{S}|H\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{2}|}$
(114)

Finally, we bound $\mathbb{E}_{\pi^{E}}[1-\mathcal{M}(s_{t},t)]$ for the
membership oracle defined in eq. 94. By definition, this quantity is the same
as $\mathrm{Pr}_{\pi^{E}}\left(s_{t}\text{ not visited in }D_{1}\text{ at time
t }\right)$. This is the probability that given $N_{\mathrm{exp}}$ samples
from a distribution (the state visited at time $t$ in an expert rollout), the
probability that a new sample from the same distribution is not in the support
of the observed samples. This is known as the missing mass McAllester &
Schapire (2000). In Lemma A.3 Rajaraman et al. (2020) it is shown that with
probability $\geq 1-\delta$,

$\displaystyle\sum_{t=1}^{H-1}\mathrm{Pr}_{\pi^{E}}\left(s_{t}\text{ not
visited in }D_{1}\text{ at time t
}\right)\lesssim\frac{|\mathcal{S}|H}{|D_{1}|}+\frac{\sqrt{|\mathcal{S}|}H\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{|D_{1}|}$
(115)

Finally, combining with eq. 114 and using the fact that that
$|D_{1}|,|D_{2}|=\Theta(N_{\mathrm{exp}})$, with probability $\geq 1-3\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})$
$\displaystyle\lesssim\min\left\\{H\sqrt{\frac{|\mathcal{S}|\log\left(\frac{|\mathcal{S}|H}{\delta}\right)}{N_{\mathrm{exp}}}},\frac{|\mathcal{S}|H^{3/2}}{N_{\mathrm{exp}}}\log\left(\frac{|\mathcal{S}|H}{\delta}\right)\right\\}.$
(116)

This completes the proof of Theorem 6.

###  B.2 Bound with parametric function approximation under Lipschitzness

In this section, we provide an upper bound on the imitation gap of RE in the
presence of parametric function approximation under a Lipschitzness assumption
on the function classes, and assuming access to a parameter estimation oracle
for offline classification.

######  Definition (IL with function-approximation).

In this setting, for each $t\in[H]$, there is a parameter class
$\Theta_{t}\subseteq\mathbb{B}_{2}^{d}$, the unit $L_{2}$ ball in $d$
dimensions, and an associated function class
$\\{f_{\theta_{t}}:\theta_{t}\in\Theta_{t}\\}$. For each $t\in[H]$ there
exists an unknown $\theta^{E}_{t}\in\Theta_{t}$ such that $\forall
s\in\mathcal{S}$,

$\displaystyle\pi^{E}_{t}(s)=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta^{E}_{t}}(s,a).$
(117)

######  Definition 3 (Policy induced by a classifier).

Consider a set of parameters $\theta=\\{\theta_{1},\cdots,\theta_{H}\\}$ where
$\theta_{t}\in\Theta_{t}$ for each $t$. A policy $\pi^{\theta}$ is said to be
induced by the set of classifiers defined by $\theta$ if for all
$s\in\mathcal{S}$ and $t\in[H]$,

$\pi^{\theta}_{t}(s)=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta_{t}}(s,a).$
(118)

By this definition, $\pi^{E}=\pi^{\theta^{E}}$ where
$\theta^{E}=\\{\theta^{E}_{1},\cdots,\theta^{E}_{H}\\}$.

######  Definition 4 (Lipschitz parameterization).

A function class $\mathcal{G}=\\{g_{\theta}:\theta\in\Theta\\}$ where
$g_{\theta}(\cdot):\mathcal{X}\to\mathbb{R}$ is said to satisfy $L$-Lipschitz
parameterization if,
$\|g_{\theta}(\cdot)-g_{\theta^{\prime}}(\cdot)\|_{\infty}\leq
L\|\theta-\theta^{\prime}\|_{2}$ for all $\theta,\theta^{\prime}\in\Theta$. In
other words, for each $x\in\mathcal{X}$, $g_{\theta}(x)$ is an $L$-Lipschitz
function in $\theta$, in the $L_{2}$ norm.

######  Assumption 4 (Assumption 1 restated).

For each $t$, the class $\\{f_{\theta_{t}}:\theta_{t}\in\Theta_{t}\\}$ is
$L$-Lipschitz in its parameterization, $\theta_{t}\in\Theta_{t}$.

To deal with parametric function approximation, we assume that the learner has
access to a proper offline classification oracle, which given a dataset of
classification examples, guarantees to approximately return the underlying
ground truth parameter. Namely,

######  Assumption 5 (Assumption 2 restated).

We assume that the learner has access to a multi-class classification oracle,
which given $n$ examples of the form, $(s^{i},a^{i})$ where
$s^{i}\overset{\text{i.i.d.}}{\sim}\mathcal{D}$ and
$a^{i}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta^{*}}(s^{i},a)$,
returns a $\hat{\theta}\in\Theta$ such that, with probability $\geq 1-\delta$,
$\|\hat{\theta}-\theta^{*}\|_{2}\leq\mathcal{E}_{\Theta,n,\delta}$.

This assumption implies that the parameter class $\Theta_{t}$ (and the
associated function class $\\{f_{\theta_{t}}:\theta_{t}\in\Theta_{t}\\}$)
admits finite sample complexity guarantees for learning the parameter
$\theta_{t}^{*}$ given classification examples from the underlying ground
truth function $f_{\theta_{t}^{*}}$. As we discuss in more detail later, we
will assume that this classification oracle is used by RE to train the BC
policy in Line 3 of Algorithm 1.

Finally, we introduce the main assumption on the IL instances we study. We
assume that the classification problems solved by BC at each $t\in[H]$ satisfy
a margin condition.

######  Assumption 6 (Assumption 3 restated).

For $\theta\in\Theta_{t}$, define
$a^{\theta}_{s}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta}(s,a)$ as
the classifier output on the state $s$. The weak margin condition assumes that
for each $t$, there is no classifier $\theta\in\Theta_{t}$ such that for a
large mass of states, $f_{\theta}(s_{t},a^{\theta}_{s_{t}})-\max_{a\neq
a^{\theta}_{s_{t}}}f_{\theta}(s_{t},a)$, i.e. the “margin” from the nearest
classification boundary, is small. Formally, the weak-margin condition with
parameter $\mu$ states that, for any $\theta\in\Theta_{t}$ and $\eta\leq
1/\mu$,

$\displaystyle\mathrm{Pr}_{\pi^{E}}\left(f_{\theta}(s_{t},a^{\theta}_{s_{t}})-\max_{a\neq
a^{\theta}_{s_{t}}}f_{\theta}(s_{t},a)\geq\eta\right)\geq e^{-\mu\eta}.$ (119)

The weak margin condition only assumes that there is at least an exponentially
small (in $\eta$) mass of states with margin at least $\eta$. Smaller $\mu$
indicates a larger mass away from any decision boundary. It suffices to assume
that eq. 119 is only true for $\theta$ as the classifier in Assumption 5 for
our guarantees (Theorem 7) to hold.

######  Remark 6.

Note that the weak margin condition is the multi-class extension of the
Tsybakov margin condition of Mammen & Tsybakov (1999); Audibert & Tsybakov
(2007) defined for the binary case. In particular, in eq. 119, we may replace
the RHS by $1-\mu\eta$, or $1-(\mu\eta)^{\alpha}$ for $\alpha>0$ to get
different analogs of the margin condition and the main guarantee, Theorem 3,
as we discuss in Section B.2.

######  Theorem 7.

For IL with parametric function approximation, under Assumptions 4 to 6,
appropriately instantiating RE ensures that with probability $\geq 1-4\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim H^{3/2}\sqrt{\frac{\mu
L\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}}{H}}+\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}.$
(120)

Note that we assume the same conditions on $\mathcal{F}$ as required in
Theorem 2.

######  Remark 7.

The classification oracle in Assumption 5 asks for a stronger condition than
just finding a classifier with small generalization error, which need not be
close to the ground truth $\theta^{*}$ in the parameter space. Learning
classifiers with small generalization error is 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) is the right statistical complexity
measure which characterizes the generalization error of the best learner. In
the realizable case, the optimal generalization error guarantee scales as
$\widetilde{O}(1/n)$ where $n$ is the number of classification examples. Under
certain assumptions on the input distribution $\mathcal{D}$ and the function
family (e.g. for linear families), we later show that the generalization error
guarantee can be extended to approximately learning the parameter as well (up
to problem dependent constants). Generally, under two conditions,

  1. 1.

Generalization error guarantees which scale as
$\widetilde{O}(1/N_{\mathrm{exp}})$ can be extended to parameter learning,

  2. 2.

The dependence of the generalization error on the failure probability $\delta$
scales as $\text{polylog}(1/\delta)$,

the optimality gap for RE which we prove in Theorem 7 scales as
$\widetilde{O}\left(H^{3/2}/N_{\mathrm{exp}}\right)$. The constants in
$\widetilde{O}(\cdot)$ here depend on the Natarajan dimension and the covering
number of the function classes $\mathcal{F}_{t}$ among problem dependent
constants.

In proving this result, we first discuss the implementation of RE .

##### Implementation of RE (Algorithm 1)

We discuss the instantiation of RE in the Lipschitz setting below. The
underlying function class $\mathcal{F}$ is chosen arbitrarily (note that the
guarantee we prove depends on this function class, and the only constraints on
$\mathcal{F}$ are those in Theorem 2 \- the ground truth reward must belong in
$\mathcal{F}=\otimes_{t=1}^{H}\mathcal{F}_{t}$, the function class is
symmetric, i.e., $f_{t}\in\mathcal{F}_{t}\iff-f_{t}\in\mathcal{F}_{t}$ for
each $t$ and for all $f_{t}\in\mathcal{F}_{t}$, $\|f_{t}\|_{\infty}\leq 1$)
This requires specifying the choice of the membership oracle $\mathcal{M}$ and
describing the instantiation of BC .

##### Implementation of BC :

Recall that in Algorithm 1, the learner trains BC on the dataset $D_{1}$. In
particular, under the offline classification oracle condition, Assumption 5,
the learner trains $H$ classifiers, one for each $t$, trained on the state-
action pairs (i.e. state is the input, and the action at this state is the
corresponding class) observed in the expert dataset at time $t$ using the
offline classifier in Assumption 5. We assume that each classifier is trained
with the failure probability chosen as $\delta/H$. Denoting this set of $H$
classifiers as
$\widehat{\theta}^{\textsf{BC}}=\left\\{\widehat{\theta}^{\textsf{BC}}_{1},\cdots,\widehat{\theta}^{\textsf{BC}}_{H}\right\\}$,
this corresponds to the to the policy
$\pi^{\texttt{BC}}=\pi^{\widehat{\theta}^{\textsf{BC}}}$ induced by the
classifier $\widehat{\theta}^{\textsf{BC}}$ (Definition 3).

In particular, by union bounding, the classifiers
$\widehat{\theta}^{\textsf{BC}}$ satisfy with probability $\geq 1-\delta$
simultaneously for each time $t\in[H]$,

$\|\theta^{E}_{t}-\hat{\theta}_{t}\|_{2}\leq\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}.$
(121)

##### Membership oracle:

Fix a time-step $t\in[H]$. The membership oracle $\mathcal{M}$ is defined in
eq. 10 as,

$\displaystyle\mathcal{M}(s,t)=\begin{cases}+1\qquad&\text{if }\exists
a\in\mathcal{A}\text{ such that, }\forall a^{\prime}\in\mathcal{A},\
f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a)-f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a^{\prime})\geq
2L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}\\\
0&\text{otherwise.}\end{cases}$ (122)

We first show that on the states such that the membership oracle is $1$, the
expert policy perfectly matches the learner’s policy.

######  Lemma 8.

At every state $s$ such that $\mathcal{M}(s,t)=+1$,
$\pi^{E}_{t}(s)=\pi^{\texttt{BC}}_{t}(s)$.

###### Proof.

Note that $\theta^{E}_{t}$ satisfies
$\|\theta_{t}^{E}-\widehat{\theta}^{\textsf{BC}}_{t}\|_{2}\leq\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}$
with probability $1-\delta$. Consider the action $a$ played by the learner,
for any $a^{\prime}\in\mathcal{A}$,

$\displaystyle f_{\theta^{E}_{t}}(s,a)-f_{\theta^{E}_{t}}(s,a^{\prime})$
$\displaystyle\geq
f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a)-\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}L-f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a^{\prime})-\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}L$
(123) $\displaystyle\geq 0$ (124)

where the first inequality follows by Lipschitzness of $f_{\cdot}(s,a)$ and
the last inequality follows by definition of the set of states where
$\mathcal{M}(s,t)=+1$: $\forall a^{\prime}\in\mathcal{A},\
f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a)-f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a^{\prime})\geq
2\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}L$.

Since for this action $a$,
$f_{\theta_{t}^{E}}(s,a)-f_{\theta_{t}^{E}}(s,a^{\prime})\geq 0$ for all other
actions $a^{\prime}\in\mathcal{A}$, $a$ must be the action played by the
expert policy. This completes the proof. ∎

Note that $\pi^{\texttt{BC}}$ always matches $\pi^{E}$ wherever the membership
oracle $\mathcal{M}$ is non-zero. We run Algorithm 1. Therefore, from Theorem
2, with probability $\geq 1-4\delta$, the imitation gap of the learner is
bounded by,

$\displaystyle J_{r}(\pi^{E})-J_{r}(\pi^{\texttt{RE}})\lesssim
H^{3/2}\sqrt{\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]}{H}}+\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}.$
(125)

To complete the proof, we must bound
$\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]$, which is the
measure of states $s$ such that $\forall a\in\mathcal{A},\exists
a^{\prime}\in\mathcal{A}:f_{\hat{\theta}_{t}}(s,a)-f_{\hat{\theta}_{t}}(s,a^{\prime})\leq
2L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}$, i.e. the mass of states
which are very close to a decision boundary. The probability of this set of
states is upper bounded by the weak margin condition. Indeed, for each
$t\in[H]$, defining
$a^{*}_{s}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s,a)$,

$\displaystyle\mathrm{Pr}_{\pi^{E}}\left(f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s_{t},a^{*}_{s_{t}})-\max_{a\neq
a^{*}_{s_{t}}}f_{\widehat{\theta}^{\textsf{BC}}_{t}}(s_{t},a)\geq
2L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}\right)$
$\displaystyle\geq e^{-\mu
L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}}$ (126) $\displaystyle\geq
1-\mu L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}.$ (127)

Therefore,

$\displaystyle\mathbb{E}_{\pi^{E}}\left[1-\mathcal{M}(s_{t},t)\right]\lesssim\mu
L\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}.$ (128)

Putting it together with eq. 125, and simplifying, with probability $\geq
1-4\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim H^{3/2}\sqrt{\frac{\mu
L\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H}}{H}}+\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}.$
(129)

Note that in Equation 127, we only use the fact that the probability mass of
states which are $\eta$-close to any decision boundary is not too high.
Similar to Audibert & Tsybakov (2007), we may consider relaxations of the weak
margin condition, as below.

######  Assumption 7 ($\alpha$-weak margin condition).

Consider any $t\in[H]$ and $\theta\in\Theta_{t}$. For each $s\in\mathcal{S}$,
define $a^{*}_{s}=\operatorname*{arg\,max}_{a\in\mathcal{A}}f_{\theta}(s,a)$
as the classifier output under $f_{\theta}$. The $\alpha$ weak margin
condition with parameter $\mu$ assumes that, for any $\eta\leq 1/\mu$,

$\displaystyle\forall\theta\in\Theta_{t},\ \
\mathrm{Pr}_{\pi^{E}}\left(f_{\theta}(s_{t},a^{*}_{s_{t}})-\max_{a\neq
a^{*}_{s_{t}}}f_{\theta}(s_{t},a)\geq\eta\right)\geq 1-(\mu\eta)^{\alpha}.$
(130)

When $\alpha=1$, this condition is effectively equivalent to the weak margin
condition in Assumption 6.

Following the proof of Theorem 7, we may obtain the following result under the
$\alpha$ weak margin condition for $\alpha\neq 1$.

######  Theorem 8.

For IL with parametric function approximation, under Assumptions 4, 5 and 7,
appropriately instatiating RE ensures that with probability $\geq 1-4\delta$,

$\displaystyle J(\pi^{E})-J(\pi^{\texttt{RE}})\lesssim H^{3/2}\sqrt{\frac{(\mu
L)^{\alpha}\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}\frac{\sum_{t=1}^{H}(\mathcal{E}_{\Theta_{t},N_{\mathrm{exp}},\delta/H})^{\alpha}}{H}}+\frac{\log\left(F_{\max}H/\delta\right)}{N_{\mathrm{exp}}}.\noindent$
(131)

Once again, we assume the same conditions on $\mathcal{F}$ as required in
Theorem 2.

###  B.3 Extension to unbounded discriminator families

Note that when the family of discriminators $\mathcal{F}$ does not have finite
cardinality, it in fact suffices to just bound the imitation gap against a
finite covering of $\mathcal{F}$. We spell out the details explicitly below.

In particular, we can replace $F_{\mathrm{max}}$ by
$\max_{t\in[H]}\mathcal{N}(\mathcal{F}_{t},1/N_{\mathrm{exp}},\|\cdot\|_{\infty})$,
where $\mathcal{N}(\mathcal{G},\xi,\|\cdot\|)$ denotes the covering number of
$\mathcal{G}$ in the norm $\|\cdot\|$ as defined below.

######  Definition 5 (Covering number).

For a function class $\mathcal{G}$, tolerance $\xi$ and norm $\|\cdot\|$, the
covering number $\mathcal{N}(\mathcal{G},\xi,\|\cdot\|)$ is defined as the
cardinality of the smallest set of functions $\mathcal{G}^{\xi}$ such that for
each $g\in\mathcal{G}$, there exists a $g^{\prime}\in\mathcal{G}^{\xi}$,

$\displaystyle\|g-g^{\prime}\|\leq\xi.$ (132)

######  Corollary 1.

When $\mathcal{G}$ is chosen as the set of $1$-bounded linear functions,
$\mathcal{G}=\\{\\{\langle
x,\theta\rangle:x\in\mathbb{B}_{2}^{d}\\}:\theta\in\mathbb{B}_{2}^{d}\\}$,
where $\mathbb{B}_{2}^{d}$ denotes the $L_{2}$ unit ball in $\mathbb{R}^{d}$,
$\mathcal{N}(\mathcal{G},\xi,\|\cdot\|_{\infty})\leq\left(\frac{2\sqrt{d}}{\xi}+1\right)^{d}$.

###### Proof.

For any $g,g^{\prime}\in\mathcal{G}$, where $g$ and $g^{\prime}$ correspond to
parameters $\theta,\theta^{\prime}\in\mathbb{B}_{2}^{d}$,

$\displaystyle\|g-g\|_{\infty}$
$\displaystyle\leq\max_{x\in\mathcal{X}}\langle x,\theta-\theta^{\prime}\\}$
(133) $\displaystyle\leq\|x\|_{2}\|\theta-\theta^{\prime}\|_{2}$ (134)
$\displaystyle\leq\|\theta-\theta^{\prime}\|_{2}.$ (135)

Since the $L_{2}$ covering number of $\mathbb{B}_{2}^{d}$ is bounded by
$\left(\frac{2\sqrt{d}}{\xi}+1\right)^{d}$, the result immediately follows by
defining the covering of $\mathcal{G}$ as
$\\{\langle\theta,\cdot\rangle:\theta\in\mathcal{K}\\}$ where $\mathcal{K}$ is
the optimal covering of $\mathbb{B}_{2}^{d}$ in $L_{2}$ norm. ∎

######  Definition 6 (Discretization of discriminator space).

Define $\mathcal{F}_{t}^{\xi}$ as the optimal covering of $\mathcal{F}_{t}$
under the $L_{\infty}$ norm in the sense of Definition 5. The discretized
family of discriminators we consider is,
$\mathcal{F}^{\xi}=\otimes_{t=1}^{H}\mathcal{F}^{\xi}_{t}$.

######  Lemma 9.

Suppose for all functions $f^{\prime}\in\mathcal{F}_{t}^{\xi_{1}/H}$,
simultaneously
$J_{f^{\prime}}(\pi^{E})-J_{f^{\prime}}(\pi^{\texttt{RE}})\leq\xi_{2}$. Then,
for all discriminators $f\in\mathcal{F}$,
$J_{f}(\pi^{E})-J_{f}(\pi^{\texttt{RE}})\leq 2\xi_{1}+\xi_{2}$.

###### Proof.

Consider any discriminator $f\in\mathcal{F}$. By construction, there exists an
$f^{\prime}\in\mathcal{F}^{\xi_{1}/H}$ such that,

$\displaystyle\|f-f^{\prime}\|_{\infty}\leq\xi_{1}/H.$ (136)

Since for any policy $\pi$, the value $J_{f}(\pi)$ under a discriminator
$f\in\mathcal{F}$ is an $H$-Lipschitz function of $f$, we can make a statement
about how well $J_{f^{\prime}}(\pi)$ approximates $J_{f}(\pi)$ for an
appropriately chosen $f^{\prime}\in\mathcal{F}^{\xi_{1}/H}$. In particular,
the nearest (in $L_{\infty}$ norm) $f^{\prime}\in\mathcal{F}^{\xi_{1}/H}$ to
$f\in\mathcal{F}$ satisfies that for any policy $\pi$,

$\displaystyle|J_{f}(\pi)-J_{f^{\prime}}(\pi)|\leq H\times\frac{\xi_{1}}{H}.$
(137)

As a consequence, for any discriminator $f\in\mathcal{F}$,

$\displaystyle J_{f}(\pi^{E})-J_{f}(\pi^{\texttt{RE}})$
$\displaystyle\leq|J_{f}(\pi^{E})-J_{f^{\prime}}(\pi^{E})|+J_{f^{\prime}}(\pi^{E})-J_{f^{\prime}}(\pi^{\texttt{RE}})+|J_{f^{\prime}}(\pi^{\texttt{RE}})-J_{f}(\pi^{\texttt{RE}})|$
(138) $\displaystyle\leq\xi_{1}+\xi_{2}+\xi_{1}=2\xi_{1}+\xi_{2}.$ (139)

∎

In particular, this means that if we minimize
$J_{f^{\prime}}(\pi^{E})-J_{f^{\prime}}(\pi^{\texttt{RE}})\leq\xi_{2}$ for all
$f^{\prime}\in\mathcal{F}^{1/N_{\mathrm{exp}}H}$, then we can ensure that for
all $f\in\mathcal{F}$,

$\displaystyle
J_{f}(\pi^{E})-J_{f}(\pi^{\texttt{RE}})\leq\frac{2}{N_{\mathrm{exp}}}+\xi_{2}.$
(140)

This implies the following theorem,

######  Theorem 9.

Consider the policy $\pi^{\texttt{RE}}$ returned by Algorithm 1 where
$\mathcal{F}$ is instead chosen as $\mathcal{F}^{\frac{1}{HN_{\mathrm{exp}}}}$
(as defined in Definition 6). Assume that $\pi^{E}\in\Pi$, the ground truth
reward function $r_{t}\in\mathcal{F}_{t}$ which is assumed to be bounded (For
all $f_{t}\in\mathcal{F}_{t}$, $\|f_{t}\|_{\infty}\leq 1$). Choose
$|D_{1}|,|D_{2}|=\Theta(N_{\mathrm{exp}})$ and suppose