SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Definition 1. For a given start state  $s$ , define the next state coverage probability over  $N$  samples from  $\pi_E$  for tolerance  $\epsilon$  as  $P_S(s; N, \epsilon) \coloneqq P(\min_{i \in \{1 \dots N\}} || s' - s'_*, i||_\infty \leq \epsilon)$ , where  $\{s'_*, i \sim \rho_{\pi_E}^t(\cdot | s)\}$  is a set of  $N$  sampled next state transitions under the demonstration policy  $\pi_E$  starting at  $s$ , and  $s' \sim \rho_{\pi_A}^t(\cdot | s)$  is a sampled next state under the policy  $\pi_A$  starting at  $S$ .
We can think of  $P_{S}(s;N,\epsilon)$  as the probability under potential datasets for  $\pi_E$  of seeing a next state at test time (under  $\pi_A$ ) that was nearby the next states in the dataset, conditioned at a starting state  $s$  (i.e., coverage). This is related to a single step of distribution shift, but measured over a dataset rather than a distribution. Defining coverage as the distance to the dataset being below some threshold is reasonable for function approximators like neural networks [2].
Theorem 4.3. Given a policy  $\pi_A$  and demonstrator  $\pi_E$ , assume that for state  $s$ , if  $\rho_{\pi_E}^{t-1}(s) > 0$ , then  $\pi_A(a|s) = \pi_E(a|s)$ . Assume that transitions are normally distributed with fixed and diagonal variance,  $\rho(s'|s,a) = \mathcal{N}(\mu(s,a), \sigma^2 I)$ , then the next state coverage probability is  $P_S(s; N, \epsilon) \geq 1 - (1 - (\frac{c\epsilon}{\sigma})^d \exp(-\alpha^2 d))^N - \exp(-(\alpha - 1)^2 d)$ , where  $d$  is the dimensionality of the state,  $c$  is a constant, and  $\alpha \geq 1$  is a parameter chosen to maximize the bound.
In Theorem 4.3, we see that even under a policy that is perfect when in distribution, the lower bound of next state coverage probability increases as the variance of system noise  $\sigma^2$  decreases, for a fixed sampling budget  $N$  (and fixed  $\alpha$ ). However, increasing  $N$  has a much stronger positive effect on this lower bound of state coverage than decreases in  $\sigma$ , suggesting that system noise is only an issue when the dataset size is sufficiently small. Furthermore,  $\epsilon$  here represents some simplistic form of the generalization capacity of the policy, and we see here that increasing  $\epsilon$  makes us more robust to proportional increases in  $\sigma$ .
Finite Data and Good Coverage: However, if we assume  $N$  is large enough to mitigate the coverage effect of system noise, are there any other effects of transition diversity in the dataset? The answer lies in generalization properties of the algorithm, which as provided in Theorem 4.3 are fairly conservative (i.e., in the definition of  $P_{S}(s;N,\epsilon)$ ). In Theorem A.1 in Appendix A, we use a more loose generalization definition and relax the assumption that the learned policy is perfect. We show as system noise increases, the resulting boost in sample coverage can actually replicate and overcome the effects of high learned policy noise, suggesting that more system noise can actually be beneficial for learning. This finding sheds light on results from prior work and our own experiments in Section 5 that show the benefits of collecting data in the presence of high system noise [34].
4.3 Implications for Data Curation
We have shown how action divergence and transition diversity are both tied to distribution shift and thus data quality. Based on action divergence and transition diversity, we now examine downstream implications of these properties on what matters for data curation, where the goal is to collect and then select high quality demonstrations in our dataset  $\mathcal{D}_N$  for good policy learning.
Action Consistency: To minimize action divergence, the algorithm action representation should align with the demonstrator's action distribution, for the given dataset size. One algorithmic solution is to improve the expressiveness of the policy action space so it can capture the exact actions at every state that was demonstrated. However, in practice, individual states might only be visited a few times in the dataset, and so if the entropy of the demonstration policy at those states is high, the learned policy will find it difficult to perfectly match the actions even for the most expressive action spaces. Instead, we argue that demonstration data should be curated to have more consistent actions, e.g., reducing the entropy of the demonstration policy:  $\mathbb{E}_{s\sim \rho_{\pi_A}(\cdot)}[\mathcal{H}(\pi_E(\cdot |s))]$ . Since knowing the visited state distribution is impossible beforehand, the best we can do is to encourage low entropy in the demonstration data distribution:
$$
\min  _ {\pi_ {E}} \mathbb {E} _ {s \sim \rho_ {\pi_ {E}} (\cdot)} [ \mathcal {H} (\pi_ {E} (\cdot | s)) ] \tag {7}
$$
Figure 1: Case Study: Trajectories and action variance for scripted (left two plots) compared to human demonstration data (right two plots). Even though the human data (right) has high state coverage, the action variance is high, leading to high action divergence, and vice versa.
State Diversity: As discussed previously, the state visitation of the policy depends on the trajectory action distribution, the initial state distribution, and the system noise distribution. Many prior works seek to improve the state coverage of the dataset, using some metric similar to Definition 1 [6, 29]. However, the required state diversity is a function of the learned policy mismatch (action divergence) and the noise present in the system (transition diversity), and is thus a coarse representation of data quality. Uniformly increasing coverage over the state space is not necessarily a good thing — as shown in Section 4.1, if state coverage comes at the expense of action consistency, the result might be worse than a policy trained on less state coverage but more consistent actions. Instead, we argue that we should optimize for system noise.
System Noise: Although system noise at a given state cannot be controlled, demonstrators can control system noise at the trajectory level (e.g. visiting states with more or less system noise). Should system noise be encouraged or discouraged in a dataset? Based on the discussion in Section 4.2, we hypothesize that when the dataset size is fixed, increasing the system noise can increase the coverage of the demonstrations and thus improve robustness in the learned policy (Theorem A.1), but only up to a certain point, when the data becomes too sparse to generalize (Theorem 4.3). Thus, in addition to consistent actions, we posit that demonstrators should encourage paths with high system entropy for learning better policies, such that the overall state entropy stays below some threshold  $\gamma$  to avoid coverage that is too sparse.
$$
\max  _ {\pi_ {E}} \mathbb {E} _ {s \sim \rho_ {\pi_ {E}} (\cdot), a \sim \pi_ {E} (\cdot | s)} [ \mathcal {H} (\rho (\cdot | a, s)) ]
$$
$$
s. t. H \left(\rho_ {\pi_ {E}} (s)\right) \leq \gamma \tag {8}
$$
Horizon Length: The length of trajectories in the demonstration data also can have a large effect on the demonstration and visited state distributions, and thus on the quality of the dataset. However, while horizon length certainly plays a role, like state diversity, it is a downstream effect of action divergence and transition diversity. Based on previous analysis, what really matters is minimizing the aggregated action divergence and transition diversity produced by the demonstrator across time. Horizon length alone only crudely measures this, but is often correlated in practice as we show in Section 5.2.
While prior work in data curation primarily aim to maximize state diversity alone, our analysis of data quality reveals several new properties that matter for data curation, such as action consistency, system noise, and their combined effect over time. As we demonstrate in the next section, understanding these properties and the inherent tradeoffs between them is vital for data curation in imitation learning.
To empirically analyze how different data properties affect imitation learning, we conduct two sets of experiments. In the first set, we generate various quality demonstration datasets by adding different types of noise to scripted policies, i.e., policies that are hand-designed to be successful at the task. An example scripted policy compared to human demonstrations is shown in Fig. 1. We study both noise in the system – considering transition diversity – and noise added to the demonstration policy – considering both transition diversity and action divergence – as a function of dataset size. In the second set of experiments, we evaluate the empirical properties (see Appendix B) for real human collected datasets.
Figure 2: BC Success rates in PMObstacle (top row) for 1000 and 10 episodes of data, and in Square (bottom row) for 200 and 50 episodes of data (error bars over 3 datasets). X-axis corresponds to injected system or policy noise in the dataset  $(\sigma_s^d$  or  $\sigma_p^d)$  and each line corresponds to injected noise  $(\sigma_s$  or  $\sigma_p$  ) during evaluation. System Noise (left): for large datasets (a)(c), higher system noise during evaluation decreases performance, but more system noise during training does the best. For small datasets (b)(d), we note a similar but exaggerated effect. Policy Noise (mid) For large datasets (e)(g), unlike system noise, more demonstration policy noise often hurts performance despite similar state coverage. For small datasets (f)(h), adding policy noise exaggerates this effect and produces large variance in performance. For Square, the dotted lines mark comparable values of noise in terms of state coverage. System + Policy Noise (i-l): Adding some system noise can make the policy more robust to action divergence.
5.1 Data Noising
We train Behavior Cloning (BC) with data generated with system noise and policy noise in two environments: PMObstacle, a 2D environment where a point mass agent must reach a target point without hitting an obstacle in its way, where demonstrations are multimodal (left and right around the obstacle); and Square, a 3D environment where a robot arm must pick up a square nut and insert it onto a peg (shown in Fig. 1). In both, system and policy noise are Gaussian random variables added to the dynamics and scripted policy, respectively, at each time step, and BC uses an MLP architecture. Fig. 2 shows results in PMObstacle (top row) and Square (bottom row). We additionally show Transformer architecture results in Fig. 3.
State Diversity through Transition Diversity improves performance, up to a point. The left plots in Fig. 2 (a)-(d) show policy success rates as we vary system noise  $(\sigma_s)$  (more noise for lighter color shades), where (a)(c) show the high-data regime and the (b)(d) show the low data regime. Higher system noise tends to improve policy performance in the high data regime, but only up to a point for Square — after  $\sigma_s = 0.3$  in the demonstration data (c), performance starts to drop off, which we hypothesize is due to the low coverage, based on the analysis of transition diversity in Section 4.2. In the low data regime, the story is similar but even more exaggerated, with increasing system noise leading to comparable performance as the high data regime. Once again for Square, increasing transition diversity (d) helps until the state coverage is too thin. See Table 2, Table 5, and Table 6 in Appendix C for the full sweep of system noise values in each environment.
State Diversity at the cost of Action Divergence hurts performance. Plots in Fig. 2 (e)(f) and Fig. 2 (i)(j) show policy success rates for increasing the policy noise  $(\sigma_{p})$  in the dataset, where the  $(\mathrm{e})(\mathrm{g})$  show the high data regime and the  $(\mathrm{f})(\mathrm{h})$  show the low data regime. Note that each value of policy noise yields the same demonstration state distribution as the corresponding amount of system noise. Since this noise is zero-mean and the learned policy is deterministic, policy noise in the high data regime  $(\mathrm{e})(\mathrm{g})$  is only moderately worse as compared to equivalent amounts of system noise  $(\mathrm{a})(\mathrm{c})$ , suggesting that the action divergence is minor. In fact, due to the added state diversity, adding policy noise in the high data regime helps up to a certain point. However in the low data regime  $(\mathrm{f})(\mathrm{h})$ , performance is substantially worse compared to comparable system noise, since the policy can not recover the unbiased demonstration policy from just a few noisy examples (note that in (c)(d), the x-axes are not aligned with those in  $(\mathrm{g})(\mathrm{h})$ , and  $\sigma_{s} = 0.2$  corresponds to  $\sigma_{p} = 0.02$  as shown by the dotted line). This illustrates that state diversity coming from the noise in the demonstration policy
can increase action divergence, and thus high state diversity is not universally desirable. See Table 3, Table 7, and Table 8 in Appendix C for the full sweep of policy noise values in each environment.
Transition Diversity can counteract the effects of policy noise. In the right two plots (i)(j) for PMObstacle in Fig. 2, the dataset combines both system and policy noise. The system noise is fixed  $(\sigma_{s} = 0.03)$  and the policy noise is varied in the same range as the middle two plots (just policy noise). Given the analysis in Section 4.2 and Theorem A.1 in Appendix A, we would expect that having some system noise could actually make the policy much more robust to policy noise, since the state coverage provided in the dataset can help learn a more robust policy (generalization). Indeed we find that just by adding system noise, the learned policy becomes very robust to added policy noise in both the high (i) and low (j) regimes. This suggests that adding transition diversity can help mitigate the affects of action divergence incurred by noisy or sub-optimal demonstrators. See Table 4 in Appendix C for the full sweep of policy noise values for fixed system noise in each environment.
Figure 3: Transformer-BC Success rates in PMObstacle for 1000 and 10 episodes. X-axis corresponds to injected system or policy noise in the dataset and each line corresponds to injected noise  $(\sigma_s$  or  $\sigma_p)$  during evaluation. System Noise (left): for large datasets (a), higher system noise during evaluation decreases performance, but more system noise during training does the best. For small datasets (b), we note a similar but exaggerated effect. Policy Noise (mid) For large datasets (c), unlike system noise, more demonstration policy noise often hurts performance despite Transformers being more expressive than simple MLPs. For small datasets (d), adding policy noise exaggerates this effect and produces large variance in performance. Overall, we find that Transformers can suffer from similar data quality issues as MLPs despite being more expressive.
5.2 Data Measuring
Next we empirically instantiate several metrics from Section 4.3 that capture various facets of data quality, and measure different datasets: (1) action variance (higher means less action consistency) measured through clustering nearby states (2) horizon length  $T$  (higher means longer trajectories), and (3) state similarity (opposite of state diversity) measured as the average cluster size. See Appendix B for the exact empirical definitions. We leave out system noise since these human datasets were collected in deterministic environments.
In Table 1, we consider single and multi-human datasets from the Square and Can tasks from robotic mimic [37]. We see that higher action variance and horizon length are often accompanied by decreases in success rates. The Worse multi-human datasets have seemingly lower action variance but less state similarity, and yet qualitatively we observed highly multi-modal behaviors in these datasets (e.g., grasping different parts of the nut instead of the handle). We suspect that this type of multi-modality is not measured well by the single time step action variance metric, but rather requires a longer term view of the effects of that action on the scene – as indicated in Theorem 4.1. Additionally, state diversity once again is not correlated with performance on the task.
We emphasize that these metrics likely do not provide a complete picture of data quality, but they do provide us insights into why performance is low and how data collection can be improved in future datasets. Through future efforts from the community, we envision a comprehensive set of data metrics that practitioners can use to quickly evaluate the quality of datasets without needing to first train and evaluate models on those datasets.
Square	Can
PH	Better	Okay	Worse	PH	Better	Okay	Worse
Success Rate	58	36	12	2	96	56	40	22
Dataset Size (N)	200	100	100	100	200	100	100	100
Action Variance	0.073	0.062	0.099	0.061	0.051	0.066	0.079	0.063
Horizon Length (T)	150	190	250	350	115	140	180	300
State Similarity	8.2e-5	1.8e-4	1.7e-4	1.2e-4	1.0e-4	2.1e-4	2.4e-4	2.0e-4
Table 1: Data Quality metrics evaluated on Square (left) and Can (right) for proficient human (PH), and multi human (Better, Okay, Worse) datasets. Note that state similarity (defined in Appendix B) is lower for PH than Better, however importantly Better actually represents overlapping data with the PH dataset (Better consists of 2 proficient demonstrators). We use this as calibration between different dataset sizes, and we see that there is some effect of dataset size on our chosen metrics. Note that dataset size corresponds to number of episodes.
Data curation is incredibly important for maximizing the potential of any given algorithm, and especially so in imitation learning. To curate data, one needs to establish a formalism for assessing data quality. The predominant notions of data quality today are solely centered around maximizing state diversity in a dataset, which ignore the quality of the demonstrated actions and how these factors interplay. In this work, we propose a holistic view of data curation centered around minimizing distribution shift. We demonstrate that action divergence and transition diversity are crucial factors in data quality that can often be controlled or measured, and we draw several valuable insights from this shedding light on effective strategies for curating demonstration datasets. We find that making actions more consistent tends to increase policy success. Furthermore, we observe a fundamental tradeoff between state diversity and action divergence — increasing state diversity can often come at the expense of action divergence. Instead of uniformly maximizing state diversity, we show that increasing transition diversity improves performance until the data coverage becomes too sparse. While the metrics presented in our experiments are often good measures of final performance, we find that comprehensively measuring data quality in practice for real world datasets can be quite challenging. Finally, we envision a broader set of metrics informed by our formalism for data quality that practitioners can use to curate datasets, and we believe our work is an important first step in formulating and evaluating these metrics.
Limitations and Future Work: We hope our analysis will inspire future research in data quality for imitation learning. Firstly, we hope to relax several theoretical assumptions like Gaussian noise in future work. The metrics defined in Section 5.2 do not necessarily capture the full scope of factors that comprise data quality. Future work should refine these metrics or find alternate ways of measuring data quality without requiring running the algorithm first, or potentially with only limited algorithm training. While we show the performance of several simple metrics on two interesting control settings, future work is needed in determining how best to use these metrics for guiding data collection. One option is to use the metrics as a score function that demonstrators might use as feedback. Another is to explicitly filter or weight data using the metrics. Furthermore, it would be beneficial to test the insights in this paper on real world settings, for example under more realistic models of system noise than simply Gaussian noise. Nonetheless, we believe our work to be an important first step in defining and measuring data quality.
Fairness in robotics: As discussed in our introduction, foundation models like GPT-3 trained on uncurated datasets have learned harmful biases and stereotypes from uncurated datasets. Efforts in the broader ML community have started to focus on curation as a means to solve this problem. In the policy learning setting in robotics, fairness and safety are less clearly defined. Our primary focus in this paper is on data curation for learning high performing policies, but we believe future work should find analogous notions of fairness in the robotics setting, and to consider how metrics of data quality that we define will affect this notion of fairness. For example, we might want to ensure data is collected over a diverse range of institutions rather than in just a select few locations. Our action consistency metric might artificially select only demonstrators from a small range of institutions, and so this should also be considered during data curation.
This research was supported by ONR, DARPA YFA, Ford, and NSF Awards #1941722, #2218760, and #2125511. We also give special thanks to Stephen Tu for helping to streamline and improve several proofs for theorems in the text.