SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Solving the optimization problem in (4) is equivalent to training a binary classifier that assigns positive labels to expert data and negative labels to union data. We can obtain the optimal discriminator at the population level, from which we can derive the importance sampling ratio formula:
$$
c _ {h} ^ {\star} (s, a) = \frac {d _ {h} ^ {\mathrm {E}} (s , a)}{d _ {h} ^ {\mathrm {E}} (s , a) + d _ {h} ^ {\mathrm {U}} (s , a)}, \quad w _ {h} (s, a) = \frac {c _ {h} ^ {\star} (s , a)}{1 - c _ {h} ^ {\star} (s , a)}. \tag {5}
$$
Based on the previous discussion, we present the implementation of our proposed method, named ISW-BC (importance-sampling-weighted BC), in Algorithm 1. It is worth noting that ISW-BC employs an unbiased weighting rule since it directly estimates the importance sampling ratio. In contrast, previous approaches such as [25, 61] use regularized weighting rules that may fail to recover the expert policy even with infinite samples. For further details on the differences between our method and previous ones, please refer to Appendix B.
5.1 Negative Result of ISW-BC with Tabular Representations of Discriminator
We have not yet specified the representations of the discriminator. One natural choice is to use tabular representations, which correspond to linear function approximation with one-hot features. Tabular representations have a strong representation power since they can span all possible functions. However, surprisingly, we show that tabular representations can fail when considering generalization.
Proposition 2. If the discriminator uses the one-hot feature with  $\delta = 0$ , we have  $\pi^{\mathrm{ISW - BC}} = \pi^{\mathrm{BC}}$ .
Proposition 2 suggests that even if we have a large number of supplementary samples and use importance sampling, ISW-BC is not guaranteed to outperform BC based on tabular representations. To illustrate, suppose we have a sample  $(s,a)$  that is an expert-style sample but only appears in the supplementary dataset, meaning that  $d_h^{\mathrm{E}}(s,a) > 0$ ,  $\widehat{d_h^{\mathrm{E}}}(s,a) = 0$  and  $\widehat{d_h^{\mathrm{U}}}(s,a) > 0$ . Using tabular representations, we can compute the closed-form solution  $c_h^\star(s,a) = \widehat{d_h^{\mathrm{E}}}(s,a) / (\widehat{d_h^{\mathrm{E}}}(s,a) + \widehat{d_h^{\mathrm{U}}}(s,a)) = 0$ . This implies that the importance sampling ratio  $w_h(s,a) = c_h^\star(s,a) / (1 - c_h^\star(s,a)) = 0$ , so this good sample does not contribute to the learning objective (3). The failure of tabular representations is due to their discrete treatment of data, ignoring internal correlations. Consequently, although they work well in minimizing the empirical loss, they are not good at generalization. This kind of failure mode is also mentioned in the GAN literature [4].
Algorithm 1 ISW-BC
Input: Expert data  $\mathcal{D}^{\mathrm{E}}$  and supplementary data  $\mathcal{D}^{\mathrm{S}}$ .
1:  $\mathcal{D}^{\mathrm{U}}\gets \mathcal{D}^{\mathrm{E}}\cup \mathcal{D}^{\mathrm{S}}$
2:  $c\gets$  Train a binary classifier with positive labels for  $\mathcal{D}^{\mathrm{E}}$  and negative labels for  $\mathcal{D}^{\mathrm{U}}$
3:  $w \gets$  Compute importance sampling ratio by Eq. (5).
4:  $\pi^{\mathrm{ISW - BC}}\gets$  Apply weighted BC to learn a policy by Eq. (3) with  $\mathcal{D}^{\mathrm{U}}$  and  $w$
Output: Policy  $\pi^{\mathrm{ISW - BC}}$
Figure 2: Illustration for ISW-BC. Please refer to the text below Assumption 2 for a detailed explanation.
5.2 Positive Result of ISW-BC with Function Approximation of Discriminator
In this section, we address the issue raised in the previous section by investigating ISW-BC with a specific function approximation. To avoid the limitations of tabular representations, we consider that the discriminator is parameterized by  $c_{h}(s,a;\theta_{h}) = \frac{1}{1 + \exp(-\langle\phi_{h}(s,a),\theta_{h}\rangle)}$ , where  $\phi_h:\mathcal{S}\times \mathcal{A}\to \mathbb{R}^d$  is a fixed feature mapping and  $\theta_h\in \mathbb{R}^d$  is the parameter to be trained. Note that we require  $d < |\mathcal{S}||\mathcal{A}|$  to avoid the tabular representations. Let  $g(x) = \log (1 + \exp (x))$ . Then, the optimization problem of
the discriminator becomes:
$$
\min  _ {\theta_ {h}} \mathcal {L} _ {h} \left(\theta_ {h}\right) \triangleq \sum_ {(s, a)} \widehat {d _ {h} ^ {\mathrm {E}}} (s, a) g \left(- \left\langle \phi_ {h} (s, a), \theta_ {h} \right\rangle\right) + \sum_ {(s, a)} \widehat {d _ {h} ^ {\mathrm {U}}} (s, a) g \left(\left\langle \phi_ {h} (s, a), \theta_ {h} \right\rangle\right) \tag {6}
$$
Let  $\theta^{\star} = \{\theta_{1}^{\star},\dots ,\theta_{H}^{\star}\}$  be the optimal solution obtained from Eq. (6). With the feature vector, samples are no longer treated independently, and the discriminator can perform structured estimation. To be consistent with the previous results, the policy is still based on tabular representations.
In the context of general linear function approximation, it is no longer possible to obtain a closed-form solution for  $c^{\star}$  as in Eq. (5). This raises the question: what can we infer about  $c^{\star}$ ? Our intuition can be described as follows. We can envision the supplementary dataset containing two types of samples: some that were in-expert distribution, and others that were out-of-expert distribution. We expect that  $w_{h}(s,a)$  is large in the former case and small in the latter case. Note that  $w_{h}(s,a)$  is monotonic with respect to the inner product  $\langle \phi_h(s,a),\theta \rangle$ . Therefore, we conclude that a larger value of  $\langle \phi_h(s,a),\theta \rangle$  implies a more significant contribution to the learning objective (3). In the following part, we demonstrate that the aforementioned intuition can be achieved under mild assumptions.
Assumption 2. Let  $\mathcal{D}_h^{\mathrm{S}}$  denote the set of state-action pairs in  $\mathcal{D}^{\mathrm{S}}$  in  $h$ . Define  $\mathcal{D}_h^{\mathrm{S},1} = \{(s,a)\in$ $\mathcal{D}_h^{\mathrm{S}}:d_h^{\pi^{\mathrm{E}}}(\bar{s}) > 0,a = \pi_h^{\mathrm{E}}(s)\}$  as the in-expert-distribution dataset in  $\mathcal{D}_h^{\mathrm{S}}$  and  $\mathcal{D}_h^{\mathrm{S},2} = \mathcal{D}_h^{\mathrm{S}}\setminus \mathcal{D}_h^{\mathrm{S},1}$  as the out-of-expert-distribution dataset. There exists a ground truth parameter  $\bar{\theta}_h\in \mathbb{R}^d$ , for any  $(s,a)\in \mathcal{D}_h^{\mathrm{E}}\cup \mathcal{D}_h^{\mathrm{S},1}$  and  $(s',a')\in \mathcal{D}_h^{\mathrm{S},2}$ , it holds that  $\langle \bar{\theta}_h,\phi_h(s,a)\rangle >0$  and  $\langle \bar{\theta}_h,\phi_h(s',a')\rangle < 0$ .
Readers may realize that Assumption 2 is closely related to the notion of "margin" in the classification problem. Define  $\Delta_h(\theta) \triangleq \min_{(s,a) \in \mathcal{D}_h^{\mathrm{E}} \cup \mathcal{D}_h^{\mathrm{S},1}} \langle \theta, \phi_h(s,a) \rangle - \max_{(s',a') \in \mathcal{D}_h^{\mathrm{S},2}} \langle \theta, \phi_h(s',a') \rangle$ . From Assumption 2, we have  $\Delta_h(\bar{\theta}_h) > 0$ . This means that there exists a classifier that recognizes samples from both  $\mathcal{D}_h^{\mathrm{E}}$  and  $\mathcal{D}_h^{\mathrm{S},1}$  as in-expert-distribution samples and samples from  $\mathcal{D}_h^{\mathrm{S},2}$  as out-of-expert-distribution samples. Note that such a nice classifier is assumed to exist, which is not identical to what is learned via Eq. (6). Before further discussion, we note that  $\bar{\theta}_h$  is not unique if it exists. Without loss of generality, we define  $\bar{\theta}_h$  as that can achieve the maximum margin (among all unit vectors).
Let us delve into the technical challenge that arises from Assumption 2. Although we assume two modes in the supplementary dataset, the learner is not aware of them beforehand. To gain a better understanding, refer to Figure 2, where the "star" corresponds to the expert data and the "triangle" corresponds to the supplementary data. The green and red parts of the triangle represent  $\mathcal{D}^{\mathrm{S},1}$  and  $\mathcal{D}^{\mathrm{S},2}$ , respectively. While training the discriminator, we assign positive labels (shown in "+" to the expert data and negative labels (shown in "-") to the union data. Consequently, it becomes challenging to determine the learned decision boundary theoretically. To address this challenge, we develop the landscape properties, Lipschitz continuity and quadratic growth conditions, in Lemma 1 and Lemma 2, respectively. These terminologies are from the optimization literature; see [23, 13]. Incorporating these properties will enable us to infer the learned decision boundary.
Lemma 1. For any  $\theta \in \mathbb{R}^d$ , the margin function is  $L_{h}$ -Lipschitz continuous in the sense that  $\Delta_h(\bar{\theta}_h) - \Delta_h(\theta) \leq L_h\|\bar{\theta}_h - \theta\|$ , where  $L_h = \|\phi_h(s^1, a^1) - \phi_h(s^2, a^2)\|$  with  $(s^1, a^1) \in \mathrm{argmin}_{(s,a) \in \mathcal{D}_h^{\mathrm{E}} \cup \mathcal{D}_h^{\mathrm{S},1}}\langle \theta, \phi_h(s,a) \rangle$  and  $(s^2, a^2) \in \mathrm{argmax}_{(s,a) \in \mathcal{D}_h^{\mathrm{S},2}}\langle \theta, \phi_h(s,a) \rangle$ .
Lemma 2. For any  $h$ , let  $A_{h} \in \mathbb{R}^{N_{\mathrm{tot}} \times d}$  be the matrix that aggregates the feature vectors of samples in  $\mathcal{D}_h^{\mathrm{U}}$ . Assume that  $\operatorname{rank}(A_h) = d$ , then  $\mathcal{L}_h$  (defined in Eq. (6)) has a (one-sided) quadratic growth condition. That is, there exists  $\tau_h > 0$  such that  $\mathcal{L}_h(\bar{\theta}_h) \geq \mathcal{L}_h(\theta_h^\star) + \frac{\tau_h}{2} \left\| \bar{\theta}_h - \theta_h^\star \right\|^2$ .
Using Lemma 1 and Lemma 2, we are ready to obtain the imitation gap bound of ISW-BC.
Theorem 3. Under Assumptions 1 and 2, let  $\mu = \max_{(s,h)\in S\times [H]}d_h^{\pi^{\mathrm{E}}} (s,\pi_h^{\mathrm{E}}(s)) / d_h^{\pi^\beta}(s,\pi_h^{\mathrm{E}}(s)) < \infty$ , if the feature is designed such that  $\sqrt{\frac{2\left(\mathcal{L}_h(\bar{\theta}_h) - \mathcal{L}_h(\theta_h^\star)\right)}{\tau_h}} < \frac{\Delta_h(\bar{\theta}_h)}{L_h}$  holds, then we have  $\Delta_h(\theta_h^\star) > 0$ . Furthermore, we have the imitation gap bound  $\mathbb{E}[V(\pi^{\mathrm{E}}) - V(\pi^{\mathrm{ISW - BC}})] = \mathcal{O}(\frac{|S|H^2}{N_{\mathrm{E}} + N_{\mathrm{S}} / \mu})$ .
In order to interpret Theorem 3, it is important to note that  $\Delta_h(\theta_h^\star) > 0$  means that there exists a  $\delta > 0$  such that  $w_h(s,a;\theta_h^\star) > \delta$  for  $(s,a) \in \mathcal{D}_h^{\mathrm{E}} \cup \mathcal{D}_h^{\mathrm{S},1}$  and  $w_h(s,a;\theta_h^\star) < \delta$  for  $(s,a) \in \mathcal{D}_h^{\mathrm{S},2}$ . As a result, all samples from  $\mathcal{D}_h^{\mathrm{E}}$  and  $\mathcal{D}_h^{\mathrm{S},1}$  are assigned with large weights, which allows ISW-BC to make use of additional samples and outperform BC.
We remark that the imitation gap bound of ISW-BC is dependent on the number of expert-style state-action pairs presented in the union of  $\mathcal{D}_h^{\mathrm{E}}$  and  $\mathcal{D}_h^{\mathrm{S},1}$ . This number is represented as  $N_{\mathrm{E}} + N_{\mathrm{S}} / \mu$ , where  $\mu$  is a state-action coverage parameter. It is important to mention that a similar notation is used in the literature of offline RL, as seen in [33, 10]. Additionally, ISW-BC has the ability to eliminate the gap of NBCU, meaning there is no non-vanishing error in Theorem 3. Moreover, ISW-BC can perform well even when  $\mathcal{D}_h^{\mathrm{S},2}$  has noisy action labels, a scenario where NBCU may fail.
Although Theorem 3 produces desirable outcomes, it does have some limitations. First, the theoretical analysis necessitates knowledge of  $\delta$ , which is typically challenging to determine beforehand. However, our empirical findings in Section 6 demonstrate that setting  $\delta = 0$  is effective in practice. Second, Theorem 3 mandates the use of good smooth features to ensure the required inequality holds, thereby avoiding the undesirable case presented in Proposition 2. Our paper does not offer a solution for finding such feature representations. Nevertheless, our experiments indicate that neural networks can usually learn suitable features. We present a simple mathematical example corresponding to Theorem 3 in Appendix D.5. We leave more general results of ISW-BC to future work.
To validate the theoretical claims, we perform numerical experiments. We provide a brief overview of the experiment set-up below, and the details can be found in Appendix H.
6.1 Robotic Locomotion Control
In this section, we present our experiment on locomotion control, where we train a robot to run like a human in four environments from the Gym MuJoCo suite [14]: Ant, Hopper, Halfcheetah, and Walker. We adopt online SAC [19] to train an agent for each environment with 1M steps, and consider the resultant policy as the expert. For each environment, the expert data contains 1 trajectory collected by the expert policy. We consider two types of supplementary datasets:
- Full Replay (small distribution shift): the supplementary dataset (1 million samples) is directly sampled from the experience replay buffer of the online SAC agent, which is suggested by [25]. This setting has a small distribution shift as the online agent quickly converges to the expert policy (see Figure 5 in the Appendix), resulting in abundant expert trajectories in the replay buffer.
- Noisy Expert (large distribution shift): the supplementary dataset consists of 10 clean expert trajectories and 5 noisy expert trajectories where the action labels are corrupted (i.e., replaced by random actions). This introduces a large state-action distribution shift, as discussed in Remark 1. For further discussion on dataset corruption and distribution shift, please refer to Appendix E.2.
Table 2: Environment return of algorithms on 4 robotic locomotion control tasks. Digits correspond to the mean performance over 5 random seeds and the subscript  $\pm$  indicates the standard deviation. "Avg" computes the normalized score over environments. Same as the other tables.
Ant	HalfCheetah	Hopper	Walker	Avg
Random	-326	-280	-20	2	0%
Expert	5229	11115	3589	5082	100%
BC	1759±287	931±273	2468±164	1738±311	38%
Full Replay	NBCU	4932±148	10566±86	3241±276	4462±105	92%
DemoDICE	5000±124	10781±67	3394±93	4537±125	94%
DWBC	2951±155	1485±377	2567±88	1572±225	44%
ISW-BC	4933±110	10786±56	3434±38	4475±164	94%
Noisy Expert	NBCU	3259±159	5561±539	558±23	518±56	35%
DemoDICE	2523±244	6020±346	1990±90	1685±160	49%
DWBC	3270±238	5688±557	3317±59	1985±175	62%
ISW-BC	3075±268	9284±346	2624±249	2859±407	69%
Besides our proposed methods, we also evaluate two state-of-the-art methods in the locomotion control domain: DemoDICE [25] and DWBC [61]. We report the experiment results in Table 2.
We observe that BC suffers since the amount of expert data is limited. In the full replay task, NBCU performs well due to the small distribution shift. However, in the noisy expert task, NBCU's average
performance is inferior to that of  $\mathrm{BC}^5$ , while ISW-BC outperforms NBCU in average significantly, demonstrating the robustness of ISW-BC to distribution shift. It is worth noting that among all evaluated methods, only our proposed method ISW-BC performs well in both settings. Prior methods such as DemoDICE and DWBC only perform well in one of the two settings.
6.2 Atari Video Games
In this section, we evaluate algorithms on Atari games [5], which involve video frames as inputs and discrete controls as outputs. Furthermore, environment transitions are stochastic for these games. We consider 5 games, namely Alien, MsPacman, Phoenix, Qbert, and SpaceInvaders. We obtain the offline expert data and supplementary data from the replay buffer of an online DQN agent, as released by [2]. We use the expert data from the buffer with the last index, which only contains  $50\mathrm{k}$  frames, to create a challenging learning setting. To augment this data, we use earlier replay buffer data to obtain supplementary data with approximately  $200\mathrm{k}$  frames. We consider the same baselines as in Section 6.1. All methods build on the classical convolutional neural networks used in DQN.
Similar to Section 6.1, we consider two types of supplementary data. The full replay setting involves supplementary data that is close to the expert data, exhibiting a small distribution shift. The noisy expert setting has noisy action labels, leading to a large distribution shift. Experiment details can be found in Appendix H.1.2. We report the game scores of the trained policies in Table 3.
Table 3: Environment return of algorithms on 5 Atari video games.
Alien	MsPacman	Phoenix	Qbert	SpaceInvaders	Avg
Random	-228	307	761	164	148	0%
Expert	2443	3601	4869	10955	1783	100%
BC	1051±21	1799±27	1520±56	4769±111	472±10	32%
Full Replay	NBCU	1405±28	2089±48	2431±104	8065±109	600±13	50%
DemoDICE	1401±16	2146±52	2192±72	7820±206	558±29	48%
DWBC	122±4	1251±56	583±33	1078±50	287±6	7%
ISW-BC	1452±37	2162±36	2299±76	7848±237	613±16	50%
Noisy Expert	NBCU	944±22	1378±30	1491±55	4366±458	418±14	27%
DemoDICE	1054±38	1604±59	1448±112	5354±295	395±10	31%
DWBC	643±18	656±16	1165±87	3860±104	296±5	16%
ISW-BC	1122±28	1980±51	1618±51	5247±328	497±6	36%
Our observations are consistent with those of the previous experiments. NBCU performs well when the distribution shift is small, while only ISW-BC is robust when the distribution shift is large.
6.3 Image Classification
(a) Clipart
(b) Infograph
(c) Painting
Figure 3: Samples of tiger class from 6 sub-datasets of the DomainNet [36] dataset. Infograph and quickdraw have quite different patterns (i.e., distribution shift) compared with the others.
(d) Quickdraw
(e) Real
(f) Sketch
In our final experiment, we tackle an image classification task. This task is a special type of imitation learning where the planning horizon is 1 and there are no environment transitions. The reward is classification accuracy. Please note that our main purpose here is to use this degraded one-step task to verify the theoretical results.
We use a famous dataset, DomainNet [36], which comprises 6 sub-datasets (clipart, infographic, painting, quickdraw, real, and sketch) that have different feature patterns and hence distribution shifts; see Figure 3 for an illustration. Following [22], our task is to perform 10-class classification (bird, feather, headphones, ice_cream, teapot, tiger, whale, windmill, wine_glass, and
We provide an explanation of why NBCU is better than BC on Ant and Walker in this scenario in Appendix.
zebra) using  $80\%$  of the images for training and  $20\%$  for test. Each sub-dataset has roughly 2000-5000 images.
We build the classifier on the pretrained ResNet-18 [20], as directly training ResNet-18 on the DomainNet dataset failed. We then optimize a 2-hidden-layer neural network, where inputs are from the feature representations extracted by the pretrained and fixed ResNet-18. We create 6 sub-tasks, where one of the 6 sub-datasets is used as the expert data while the other 5 sub-datasets are used as the supplementary datasets. We evaluate the classification accuracy on the expert test data. Note that there is no natural extension of DemoDICE for this task. More details can be found in Appendix H.1.3.
The results of our experiment are presented in Table 4. We observe that due to the presence of distribution shifts, NBCU performs even worse than BC, even though NBCU use more data than BC. On the other hand, ISW-BC can improve the performance over BC on 5 out of 6 tasks by re-weighting the supplementary data. At the same time, ISW-BC is more effective than DWBC.
Table 4: Test classification accuracy (%) of algorithms on 6 types of expert and supplementary data.
Clipart	Infograph	Painting	Quickdraw	Real	Sketch	Avg
BC	89.31±0.01	55.80±0.01	90.14±0.00	85.61±0.01	96.19±0.00	87.58±0.01	84.10
NBCU	89.16±0.02	56.32±0.02	88.29±0.01	84.78±0.03	95.31±0.00	87.57±0.00	83.57
DWBC	90.00±0.09	57.44±0.06	90.89±0.04	85.09±0.09	96.35±0.01	88.86±0.09	84.77
ISW-BC	90.86±0.00	57.52±0.01	91.78±0.01	84.97±0.01	96.56±0.01	89.63±0.06	85.22
This paper introduces a formal mathematical framework for imitation learning with a supplementary yet imperfect dataset, which is designed to address the scarcity of expert data. Within this framework, we present new theoretical insights that illuminate the distribution shift challenge between expert and supplementary data. To deal with this challenge, we devise a new method named ISW-BC, employing the importance sampling technique to select data within the expert distribution. Through both theoretical analysis and empirical evaluations, we show the superiority of the proposed approach.
Our research is closely connected with data-centric artificial intelligence (AI) [39, 54, 64]. Here, the emphasis lies in the quality, availability, and management of data as foundational elements for constructing effective AI models and applications. The importance sampling technique developed in this paper proves valuable for processing imperfect data.
Furthermore, our methods can extend beyond the scope of consideration in this paper. To illustrate, the core concept of data re-weighting and selection can find utility in the realm of large language models. In specific downstream tasks with limited expert data, we can judiciously select a set of pre-training tasks to construct supplementary data and to enhance overall performance of language models; for additional insights, please refer to recent progress in [58, 18]. Other potential avenues for future exploration include the extension to multi-task imitation learning, as well as the unsupervised case where the expert data is not available.<sup>6</sup>
Overall, our findings demonstrate the potential of improving imitation learning performance by leveraging diverse data sources through effective data selection. We aspire for this work to serve as inspiration for future advancements in the field.
We thank Congliang Chen for discussing a technical lemma. The work of Yang Yu is supported by National Key Research and Development Program of China (2020AAA0107200), NSFC (61876077), and Collaborative Innovation Center of Novel Software Technology and Industrialization. The work of Zhi-Quan Luo is supported in part by the National Key Research and Development Project under grant 2022YFA1003900, and in part by the Guangdong Provincial Key Laboratory of Big Data Computing.
In this scenario, we can consider using a small fraction of imperfect data to reconstruct a surrogate expert policy and expert data (e.g., via majority voting). Subsequently, we can regard the remaining imperfect data as supplementary and apply the proposed framework.