SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziniu Li\*1,2, Tian Xu\*3,4, Zeyu Qin\*5, Yang Yu\*3,4, and Zhi-Quan Luo\*1,2
<sup>1</sup>The Chinese University of Hong Kong, Shenzhen
2Shenzhen Research Institute of Big Data
$^{3}$ National Key Laboratory for Novel Software Technology, Nanjing University
4Polixir.ai
$^{5}$ Hong Kong University of Science and Technology
Imitation learning (IL) algorithms excel in acquiring high-quality policies from expert data for sequential decision-making tasks. But, their effectiveness is hampered when faced with limited expert data. To tackle this challenge, a novel framework called (offline) IL with supplementary data has been proposed [25, 61], which enhances learning by incorporating an additional yet imperfect dataset obtained inexpensively from sub-optimal policies. Nonetheless, learning becomes challenging due to the potential inclusion of out-of-expert-distribution samples. In this work, we propose a mathematical formalization of this framework, uncovering its limitations. Our theoretical analysis reveals that a naive approach—applying the behavioral cloning (BC) algorithm concept to the combined set of expert and supplementary data—may fall short of vanilla BC, which solely relies on expert data. This deficiency arises due to the distribution shift between the two data sources. To address this issue, we propose a new importance-sampling-based technique for selecting data within the expert distribution. We prove that the proposed method eliminates the gap of the naive approach, highlighting its efficacy when handling imperfect data. Empirical studies demonstrate that our method outperforms previous state-of-the-art methods in tasks including robotic locomotion control, Atari video games, and image classification.<sup>1</sup> Overall, our work underscores the potential of improving IL by leveraging diverse data sources through effective data selection.
Imitation learning (IL) [3, 34] is an essential technique in artificial intelligence, allowing machines to enhance their performance by imitating expert behaviors. This technique has showcased remarkable achievements across various domains such as robotics [27], self-driving cars [35], and language models [55, 8]. Among the IL approaches, behavioral cloning (BC) [40] stands out as a popular method. In particular, BC leverages temporally-correlated state-action pairs extracted from expert demonstrations, employing them as training samples to learn a mapping from states to actions via maximum likelihood estimation. In this manner, IL expands upon the traditional supervised learning framework, enabling the acquisition of sequential decision-making capabilities.
The quantity of expert trajectories plays a crucial role in achieving satisfactory performance. Previous studies have shown that BC works well when the dataset contains a large number of expert-level
*: Equal contribution. Author ordering is determined by coin flip. Emails: ziniuli@link.cuhk.edu.cn and xut@lamda.nju.edu.cn.
†: Email: zeyu.qin@connect.ust.hk.
: Corresponding authors. Emails: yuy@nju.edu.cn and luozq@cuhk.edu.cn.
<sup>1</sup>The code is available at https://github.com/liziniu/ISWBC.
trajectories [48]. However, the well-known compounding errors issue [44] renders any offline IL algorithm, including BC, ineffective when the number of expert trajectories is small [42, 60]. Furthermore, collecting more trajectories from the expert is costly and impractical in many domains, such as robotics and healthcare.
To overcome the challenge of scarce expert data, we propose to use an additional yet imperfect dataset to supplement the expert data; see Figure 1 for illustration. In particular, this additional dataset can be cheaply obtained by executing sub-optimal policies. However, the incorporation of supplementary dataset introduces a distribution shift issue due to the presence of out-of-expert-distribution trajectories. The distribution shift issue may hamper the model's performance in utilizing the supplementary data, as we will argue later.
We realize that a number of empirical successes have been reported in this direction [25, 24, 61, 32]. Most algorithms rely on a discriminator to distinguish between expert-style and sub-optimal samples, followed by optimization of a weighted BC objective to learn a good policy. For example, DemoDICE [25] uses a regularized state-action distribution matching objective to train the discriminator, while DWBC [61] employs a cooperative training framework for the policy and discriminator. Despite the empirical successes in certain scenarios, there remains a notable absence of systematic theoretical studies, particularly in terms of imitation gap (i.e., performance difference between the expert and learner), which may hinder deep understanding and impede future algorithmic advances.
Figure 1: Compared with the conventional IL framework (shown in cyan), the supplementary data helps address the expert data scarcity issue, and the data selection technique helps address the distribution shift issue in model training.
Table 1: The theoretical guarantees of three methods: (1) BC, which relies solely on expert data; (2) NBCU, which directly uses the union of expert data and supplementary data without selection; and (3) ISW-BC, a new method that employs importance sampling for data selection. Compared with BC, NBCU suffers a nonvanishing error due to the distribution shift between two datasets while ISW-BC does not.
Imitation Gap
BC	O( |S|H2/NE)
NBCU	O((1 - η)(V(πE) - V(πβ)) + |S|H2/NE + NS)
ISW-BC	O( |S|H2/NE + NS/μ)
We aim to bridge the gap between theory and practice in the (offline) IL with supplementary data framework by developing effective algorithms and providing rigorous theoretical guarantees. To the best of our knowledge, only [9] provided imitation gap bounds for a model-based adversarial imitation learning approach in a similar problem. However, our focus lies on the widely used and simpler BC and its variants, which are model-free in nature. Our contributions are summarized below.
- We develop a formal mathematical framework of the IL with supplementary data framework and conduct corresponding theoretical analysis. Our findings highlight the impact of the distribution shift between expert and supplementary data. Our results are summarized in Table 1.4
In particular, our analysis shows that naively applying BC on the union of expert and supplementary data (referred to as the NBCU method in this paper) has a non-vanishing error term in the imitation gap bound (see the second row of Table 1). This means that the direct use of additional data might yield inferior performance compared with solely using the expert data with BC. This necessitates the development of more advanced algorithms.
For instance, web texts and images are easily available for language and vision models, respectively; robotics can benefit from previously collected datasets.
This issue is separate from the intrinsic distribution shift problem that IL already faces, where the training and evaluation distributions differ [44]. It is important to note that these are distinct issues.
4We briefly explain notations in Table 1:  $|\mathcal{S}|$  indicates the state space size,  $H$  is the horizon length,  $N_{\mathrm{E}}$  means the expert data size, and  $N_{\mathrm{S}}$  refers to the supplementary data size. In addition, we define  $\eta = N_{\mathrm{E}} / (N_{\mathrm{E}} + N_{\mathrm{S}})$ , and  $V(\pi^{\mathrm{E}}) - V(\pi^{\beta})$  as the performance gap between expert policy  $V^{\mathrm{E}}$  and supplementary data collection policy  $\pi^{\beta}$ . For ISW-BC's bound,  $\mu >0$  represents the state-action distribution density ratio bound between these two policies (to be detailed in Theorem 3).
- To address the distribution shift issue, we propose a new approach called ISW-BC, which uses the importance sampling technique to select data within the expert distribution. In contrast to prior methods [25, 61], ISW-BC re-weights data in an unbiased way. We develop a new imitation gap bound for ISW-BC (see the last row of Table 1), revealing that it not only eliminates the gap exhibited by the naive approach but also offers a superior guarantee over BC.
- Our theoretical analysis has been validated through extensive experiments, including robotic locomotion control, Atari video games, and image classification. The results affirm the superiority of ISW-BC over existing methods, thus demonstrating the potential of our method in addressing the distribution shift issue in IL with supplementary data.
We review briefly relevant studies in the main text and provide a detailed discussion in Appendix B.
Behavioral Cloning. Behavioral cloning (BC) is a popular algorithm in the offline setting, where the learner cannot interact with the environment. According to the learning theory in [42], only using an expert dataset, BC has an imitation gap of  $\mathcal{O}(|S|H^2 /N_{\mathrm{E}})$ , where  $|\mathcal{S}|$  is the state space size,  $H$  is the planning horizon, and  $N_{\mathrm{E}}$  is the number of expert trajectories. Our work investigates the use of a supplementary dataset to enhance the dependence on the data size.
Adversarial Imitation Learning. In contrast to BC, adversarial imitation learning (AIL) methods, such as GAIL [21], perform imitation through state-action distribution matching. It has been demonstrated both empirically and theoretically that AIL methods do not suffer from the compounding errors issue when the expert data is limited [21, 16, 26, 59, 60, 28, 63]. Under mild conditions, [62] provided a horizon-free bound of  $\mathcal{O}(\min\{1, \sqrt{|S|/N_{\mathrm{E}}}\})$ , which is much better than BC in terms of dependence on  $H$ . However, AIL methods work naturally in the online setting (i.e., the interaction is allowed), which is not directly applicable in the offline setting that we study in this paper. Although the proposed method has a discriminator and a policy like AIL, our discriminator and policy are not designed to compete with each other adversially, as we will explain in detail later.
Imitation Learning with Supplementary Data. Our theoretical study is motivated by recent empirical successes in IL with supplementary data [25, 61, 32, 24]. Compared with [25, 61], a related setting, learning from observation, is studied in [32, 24]. In this setting, expert actions are absent, and only expert states are observed. The importance sampling technique used in our method for addressing distribution shift is also studied in (semi-)supervised learning [49, 11, 30, 15]. Our contribution is to show this technique is also effective in the imitation learning, where data has a Markovian structure.
Markov Decision Process. In this paper, we consider the episodic Markov decision process (MDP) framework [41]. An MDP is defined by the tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{P},r,H,\rho)$ , where  $\mathcal{S}$  and  $\mathcal{A}$  are the state and action space, respectively.  $H$  is the maximum length of a trajectory, and  $\rho$  is the initial state distribution. The non-stationary transition function is specified by  $\mathcal{P} = \{P_1,\dots ,P_H\}$ , where  $P_{h}(s_{h + 1}|s_{h},a_{h})$  determines the probability of transiting to state  $s_{h + 1}$  given the current state  $s_h$  and action  $a_{h}$  in time step  $h$ , for  $h\in [H]$ . Here the symbol  $[x]$  means the set of integers from 1 to  $x$ . Similarly, the reward function  $r = \{r_1,\dots ,r_H\}$  specifies the reward received at each time step, where  $r_h:\mathcal{S}\times \mathcal{A}\to [0,1]$  for  $h\in [H]$ . A policy in an MDP is a function that maps each state to a probability distribution over actions. We consider time-dependent policies  $\pi_h:\mathcal{S}\rightarrow \Delta (\mathcal{A})$ , where  $\Delta (\mathcal{A})$  is the probability simplex. The policy at each time step  $h$  is denoted as  $\pi_h$ , and we use  $\pi$  to denote the collection of time-dependent policies  $\{\pi_h\}_{h = 1}^H$  when the context is clear.
We measure the quality of a policy  $\pi$  by the policy value (i.e., environment-specific long-term return):  $V(\pi) = \mathbb{E}\big[\sum_{h = 1}^{H}r_{h}(s_{h},a_{h})\mid s_{1}\sim \rho ;a_{h}\sim \pi_{h}(\cdot |s_{h}),s_{h + 1}\sim P_{h}(\cdot |s_{h},a_{h}),\forall h\in [H]\big]$ . To facilitate later analysis, we need to introduce the state-action distribution  $d_h^\pi (s,a) = \mathbb{P}(s_h = s,a_h = a|\pi)$ . We use the convention that  $d^{\pi}$  is the collection of time-dependent state-action distributions.
Imitation Learning. Imitation learning (IL) aims to learn a policy that mimics an expert policy based on expert demonstrations. In this paper, we assume that there exists a good expert policy  $\pi^{\mathrm{E}}$  that generates a dataset  $\mathcal{D}^{\mathrm{E}}$  consisting of  $N_{\mathrm{E}}$  trajectories of length  $H$ .
$$
\mathcal {D} ^ {\mathrm {E}} = \big \{\mathrm {t r} = (s _ {1}, a _ {1}, s _ {2}, a _ {2}, \dots , s _ {H}, a _ {H}); s _ {1} \sim \rho , a _ {h} \sim \pi_ {h} ^ {\mathrm {E}} (\cdot | s _ {h}), s _ {h + 1} \sim P _ {h} (\cdot | s _ {h}, a _ {h}), \forall h \in [ H ] \big \}.
$$
The learner aims to imitate the expert using the expert dataset  $\mathcal{D}^{\mathrm{E}}$ . The quality of the imitation is measured by the imitation gap, defined as  $\mathbb{E}\left[V(\pi^{\mathrm{E}}) - V(\pi)\right]$ , where the expectation is taken over the randomness of data collection. It is worth noting that in the training phase, the IL learner does not have access to reward information. A good learner should closely mimic the expert, resulting in a small imitation gap. We assume that the expert policy is deterministic, a common assumption in the literature [42, 43, 60], and applicable to tasks such as MuJoCo locomotion control.
Behavioral Cloning. Behavioral cloning (BC) is a popular imitation learning algorithm that aims to learn a policy from an expert dataset  $\mathcal{D}^{\mathrm{E}}$  via supervised learning. Specifically, BC seeks to find a policy  $\pi^{\mathrm{BC}}$  that maximizes the log-likelihood of the expert actions in the dataset:
$$
\pi^ {\mathrm {B C}} \in \operatorname * {a r g m a x} _ {\pi} \sum_ {h = 1} ^ {H} \sum_ {(s, a) \in \mathcal {S} \times \mathcal {A}} \widehat {d _ {h} ^ {\mathrm {E}}} (s, a) \log \pi_ {h} (a | s), \tag {1}
$$
where  $\widehat{d_h^{\mathrm{E}}}(s,a)$  is the empirical state-action distribution in the expert dataset. By the maximum likelihood estimation (MLE), BC can make good decisions by duplicating expert actions on states visited in  $\mathcal{D}^{\mathrm{E}}$ . However, BC may take sub-optimal actions on non-visited states, resulting in compounding errors and a large imitation gap. This issue is significant when the expert data is limited.
4 Imitation Learning with Supplementary Data
To address the challenge of limited availability of expert data, we consider an IL with a supplementary dataset framework. We assume that a supplementary yet imperfect dataset  $\mathcal{D}^{\mathrm{S}} = \left\{\mathrm{tr} = (s_1,a_1,s_2,a_2,\dots ,s_H,a_H)\right\}$  is collected by a behavior policy  $\pi^{\beta}$ . A naive approach is to perform MLE on the union of the expert and supplementary dataset  $\mathcal{D}^{\mathrm{U}} = \mathcal{D}^{\mathrm{E}}\cup \mathcal{D}^{\mathrm{S}}$ :
$$
\pi^ {\mathrm {N B C U}} \in \operatorname * {a r g m a x} _ {\pi} \sum_ {h = 1} ^ {H} \sum_ {(s, a)} \widehat {d _ {h} ^ {\mathrm {U}}} (s, a) \log \pi_ {h} (a | s), \tag {2}
$$
where  $\widehat{d_h^{\mathrm{U}}}(s,a)$  is the empirical state-action distribution in  $\mathcal{D}^{\mathrm{U}}$ . We refer to this approach as NBCU (naive BC with the union dataset). NBCU treats these two datasets equally and is brittle to distribution shift, as we will demonstrate later. For theoretical analysis purpose, we assume expert data represents a  $\eta \in [0,1]$  fraction of the total union data.
Assumption 1. The expert dataset  $\mathcal{D}^{\mathrm{E}}$  and supplementary dataset  $\mathcal{D}^{\mathrm{S}}$  are collected in the following way: each time, we roll-out a behavior policy  $\pi^{\beta}$  with probability  $1 - \eta$  and the expert policy with probability  $\eta$ . Such an experiment is independent and identically conducted  $N_{\mathrm{tot}}$  times.
Under Assumption 1, we slightly overload our notations: we use  $N_{\mathrm{E}}$  to denote the expected number of expert trajectories, which is given by  $N_{\mathrm{E}} = \eta N_{\mathrm{tot}}$ , and  $N_{\mathrm{S}}$  to denote the expected number of supplementary trajectories, which is given by  $N_{\mathrm{S}} = (1 - \eta)N_{\mathrm{tot}}$ . Note that the probabilistic sampling procedure does not change the nature of our theoretical insights. In practice, one may collect a fixed number of expert and supplementary trajectories, respectively.
To establish a common ground, we begin by specifying the policy representations. Here, we adopt tabular representations, which assume that the parameterized value functions can take any possible form. Specifically, we define  $\pi_h(a|s;\theta) = \langle \phi (s,a),\theta \rangle$  , where  $\phi (s,a)\in \mathbb{R}^d$  is the feature representation and  $\theta \in \mathbb{R}^d$  is the parameter to optimize. In tabular representations, we use one-hot features for  $\phi (s,a)$  . For a discussion on general function approximation schemes, please refer to Appendix E.
Imitation Gap of BC. In order to evaluate the usefulness of the supplementary dataset, we use BC with only the expert dataset as a baseline. The analysis of this approach has been done in the conventional IL set-up in [42], and we re-state their results in our setting.
Theorem 1. Under Assumption 1, if we apply BC only on the expert dataset, we have that  $\mathbb{E}\left[V(\pi^{\mathrm{E}}) - V(\pi^{\mathrm{BC}})\right] = \mathcal{O}\left(\frac{|S|H^{2}}{N_{\mathrm{E}}}\right)$ , where the expectation is taken over the randomness in the dataset collection (same as other expectations).
Proofs of Theorem 1 and other theoretical results are deferred to the Appendix. The proof of Theorem 1 builds on [42], with the main difference being that the number of expert trajectories is a random variable in our set-up. We handle this difficulty by using Lemma 3 in the Appendix. The quadratic dependence on the planning horizon  $H$  indicates the compounding errors issue of BC.
If the expert data is limited (i.e.,  $N_{\mathrm{E}}$  is small), the performance gap can be large, suggesting poor performance of BC in this scenario.
Imitation Gap of NBCU. Guarantees of naively using the supplementary data are presented below.
Theorem 2. Under Assumption 1, if we apply BC on the union dataset, we have  $\mathbb{E}\left[V(\pi^{\mathrm{E}}) - V(\pi^{\mathrm{NBCU}})\right] = \mathcal{O}((1 - \eta)(V(\pi^{\mathrm{E}}) - V(\pi^{\beta})) + \frac{|S|H^{2}\log(N_{\mathrm{tot}})}{N_{\mathrm{tot}}})$
Remark 1. In case when the behavior policy is inferior to the expert policy, we have  $V(\pi^{\mathrm{E}}) - V(\pi^{\beta}) > 0$ . In this case, even if  $N_{\mathrm{tot}}$  is large enough to make the second term negligible, there is still a non-vanishing gap of  $V(\pi^{\mathrm{E}}) - V(\pi^{\beta})$  due to the behavior policy's potential to collect non-expert actions. In other words, the recovered policy may select wrong actions even on expert states, leading to the sub-optimal performance of NBCU.
It is worthy to note that this non-vanishing error term can also be interpreted from the viewpoint of the distribution shift. Specifically, using the analysis in [62], we can show that
$$
V (\pi^ {\mathrm {E}}) - V (\pi^ {\beta}) = \mathcal {O} (H \varepsilon_ {d}) = \mathcal {O} (H ^ {2} \varepsilon_ {\pi}),
$$
where  $\varepsilon_{d} = \max_{h}\mathrm{TV}(d_{h}^{\pi^{\mathrm{E}}},d_{h}^{\pi^{\beta}})$  is the state-action distribution total variation (TV) distance and  $\varepsilon_{\pi} = \max_{h}\max_{s}\mathrm{TV}(\pi_{h}^{\mathrm{E}}(\cdot |s),\pi_{h}^{\beta}(\cdot |s))$  is the policy distribution TV distance. Hence, we can also view Theorem 2 in the context of state-action or policy distribution shifts.
Note also that we do not claim that NBCU is always worse than BC. Instead, Theorem 2 implies that if the distribution shift between two data datasets is small, NBCU could be better than BC by leveraging more data, which we will also show in experiments later.
The next proposition establishes the inevitability of the gap  $V(\pi^{\mathrm{E}}) - V(\pi^{\beta})$  in the worst case.
Proposition 1. Under Assumption 1, there exists an MDP  $\mathcal{M}$ , an expert policy  $\pi^{\mathrm{E}}$  and a behavior policy  $\pi^{\beta}$ , such that we have  $\mathbb{E}\left[V(\pi^{\mathrm{E}}) - V(\pi^{\mathrm{NBCU}})\right] = \Omega((1 - \eta)(V(\pi^{\mathrm{E}}) - V(\pi^{\beta})))$ .
The hard instance in Proposition 1 builds on the following idea: NBCU considers all action labels in the union dataset equally important. Therefore, we can build an instance where the expert  $\pi$  selects a good action with an one-step reward of 1, while the behavior policy  $\pi^{\beta}$  chooses a bad action with an one-step reward of 0. The noise introduced by  $\pi^{\beta}$  results in incorrect learning goals, causing NBCU to make a mistake with probability  $1 - \eta$ , which is the fraction of the noise in the union dataset. By a carefully designed transition probability, we can obtain the expected bound in Proposition 1.
5 Addressing Distribution Shift with Importance Sampling
In this section, we propose a data selection approach to alleviate the distribution shift issue between expert data and supplementary data. Our approach is inspired by recent works [25, 61], where a discriminator is trained to re-weight samples, and a weighted BC objective is used for policy optimization. Specifically, we define the weighted BC objective as follows:
$$
\pi^ {\mathrm {I S W - B C}} \in \operatorname * {a r g m a x} _ {\pi} \sum_ {h = 1} ^ {H} \sum_ {(s, a) \in \mathcal {S} \times \mathcal {A}} \left\{\widehat {d _ {h} ^ {\mathrm {U}}} (s, a) \times \left[ w _ {h} (s, a) \log \pi_ {h} (a | s) \right] \times \mathbb {I} \left[ w _ {h} (s, a) \geq \delta \right] \right\}, \tag {3}
$$
where  $\widehat{d_h^{\mathrm{U}}}(s,a)$  is the empirical state-action distribution of the union dataset, and  $w_{h}(s,a)\in [0,\infty)$  is the weight decided by the discriminator. We introduce a hyper-parameter  $\delta \in [0,\infty)$  to better select samples.
We propose using the importance sampling technique [46, Chapter 9] to transfer samples in the union dataset to the expert policy distribution, which is the key idea behind our method. This technique helps address the failure mode of NBCU. In an ideal scenario where there are infinite samples (i.e., at the population level),  $\widehat{d_h^{\mathrm{U}}}$  would equal  $d_h^{\mathrm{U}}$ . By setting  $w_{h}(s,a) = d_{h}^{\mathrm{E}}(s,a) / d_{h}^{\mathrm{U}}(s,a)$ , we obtain  $\widehat{d_h^{\mathrm{U}}}(s,a)w_h(s,a) = d_h^{\mathrm{E}}(s,a)$ , and the objective (3) enables the learning of a policy as if samples were collected by the expert policy. However, in practice,  $d_h^{\mathrm{E}}(s,a)$  and  $d_h^{\mathrm{U}}(s,a)$  are unknown, and we only have a finite number of samples from each of these distributions. Therefore, we must estimate the grounded importance sampling ratio  $d_h^{\mathrm{E}}(s,a) / d_h^{\mathrm{U}}(s,a)$  from the expert data and union data.
We emphasize that estimating the probability densities of high-dimensional distributions separately for expert and union data and then calculating their quotient can be a challenging task. We take a different
approach. Inspired by [17], we directly train a discriminator to estimate the importance sampling ratio  $d_h^{\mathrm{E}}(s,a) / d_h^{\mathrm{U}}(s,a)$ . To this end, we introduce time-dependent parameterized discriminators  $\{c_{h}:\mathcal{S}\times \mathcal{A}\to [0,1]\}_{h = 1}^{H}$ , each of which is optimized according to the objective function
$$
\max  _ {c _ {h}} \sum_ {(s, a) \in \mathcal {S} \times \mathcal {A}} \widehat {d _ {h} ^ {\mathrm {E}}} (s, a) [ \log \left(c _ {h} (s, a)\right) ] + \sum_ {(s, a) \in \mathcal {S} \times \mathcal {A}} \widehat {d _ {h} ^ {\mathrm {U}}} (s, a) [ \log \left(1 - c _ {h} (s, a)\right) ]. \tag {4}
$$