SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Suneel Belkhale
Stanford University belkhale@stanford.edu
Yuchen Cui
Stanford University yuchenc@stanford.edu
Dorsa Sadigh
Stanford University dorsa@stanford.edu
In supervised learning, the question of data quality and curation has been overshadowed in recent years by increasingly more powerful and expressive models that can ingest internet-scale data. However, in offline learning for robotics, we simply lack internet scale data, and so high quality datasets are a necessity. This is especially true in imitation learning (IL), a sample efficient paradigm for robot learning using expert demonstrations. Policies learned through IL suffer from state distribution shift at test time due to compounding errors in action prediction, which leads to unseen states that the policy cannot recover from. Instead of designing new algorithms to address distribution shift, an alternative perspective is to develop new ways of assessing and curating datasets. There is growing evidence that the same IL algorithms can have substantially different performance across different datasets. This calls for a formalism for defining metrics of "data quality" that can further be leveraged for data curation. In this work, we take the first step toward formalizing data quality for imitation learning through the lens of distribution shift: a high quality dataset encourages the policy to stay in distribution at test time. We propose two fundamental properties that shape the quality of a dataset: i) action divergence: the mismatch between the expert and learned policy at certain states; and ii) transition diversity: the noise present in the system for a given state and action. We investigate the combined effect of these two key properties in imitation learning theoretically, and we empirically analyze models trained on a variety of different data sources. We show that state diversity is not always beneficial, and we demonstrate how action divergence and transition diversity interact in practice.
Supervised learning methods have seen large strides in recent years in computer vision (CV), natural language processing (NLP), and human-level game playing [18, 26, 43, 8, 51, 13, 44]. These domains have benefited from large and complex models that are trained on massive internet-scale datasets. Despite their undeniable power, biases present in these large datasets can result in the models exhibiting unexpected or undesirable outcomes. For example, foundation models such as GPT-3 trained on uncurated datasets have resulted in instances of racist behavior such as associating Muslims with violence [1, 5]. Thus offline data curation is immensely important for both safety and cost-effectiveness, and it is gaining in prominence in training foundation models [22, 10, 40].
Data curation is even more important in the field of robotics, where internet-scale data is not readily available and real-world datasets are small and uncurated. Noise or biases present in the data can lead to dangerous situations in many robotics tasks, for example injuring a person or damaging equipment. In such scenarios, deciding which data to collect and how best to collect it are especially important [37, 34, 4]. Of course, the quality of data depends on the algorithm that uses the data. A common paradigm in robot learning from offline datasets is imitation learning (IL), a data-driven, sample efficient framework for learning robot policies by mimicking expert demonstrations. However, when learning from offline data using IL, estimating data quality becomes especially difficult, since the "test set" the robot is evaluated on is an entirely new data distribution due to compounding errors
incurred by the model, i.e., action prediction errors take the model to unseen states. This phenomenon is studied and referred to as the distribution shift problem, and prior work has viewed and addressed it through several angles [46, 49].
Broadly, prior work address distribution shift either by taking an algorithm-centric approach to account for biases in the dataset, or by directly modifying the dataset collection process in a data-centric manner. Algorithm-centric approaches learn robust policies by imposing task-specific assumptions [23, 28, 25], acquiring additional environment data [21, 41], or leveraging inductive biases in representing states [54, 30, 38, 12, 12, 37] and actions [53, 48, 12]. What these algorithm-centric works overlook is that changing the data can be as or more effective for policy learning than changing the algorithm. In prior data-centric methods, the goal is usually to maximize state diversity through shared control [46, 45, 16, 27, 32, 36, 47], noise injection [34], or active queries [15, 6, 42, 29]. By only focusing on state coverage, these works are missing the role that actions (i.e., the expert) plays in the quality of data. A more complete understanding of data quality that integrates both state and action quality would not only improve performance but also save countless hours of data collection.
To better understand the role of both states and actions in the data for learning good policies, consider a single state transition — the three core factors that affect distribution shift are: the policy action distribution, the stochasticity of transitions, and the previous state distribution. Note that the previous states are also just a function of the policy and dynamics back through time. Through this we extract two fundamental properties that can influence data quality in IL: action divergence and transition diversity. Action divergence captures how different the learned policy action distribution is from the expert's actions at a given state: for example, if the expert is very consistent but the policy has high variance, then action divergence will be high and distribution shift is likely. Transition diversity captures the inherent variability of the environment at each state, which determines how the state distribution evolves for a given policy: for example, noisy dynamics in the low data regime can reduce overlap between the expert data and the learned policy distribution. Importantly, these factors can compound over time, thus greatly increasing the potential for distribution shift. While state coverage has been discussed in prior work, we are the first work to formalize the roles of both state and action distributions in data quality, along with how they interact through time: this new data-focused framing leads to insights about how to curate data to learn more effective policies.
To validate and study these properties empirically, we conduct two sets of experiments: (1) Data Noising, where we ablate these properties in robotics datasets to change the policy success rates, and (2) Data Measuring, where we observe human and machine generated datasets and approximately measure these properties in relationship to the policy success rates. Both experiments show that state diversity, a commonly used heuristic for quality in prior work, is not always correlated with success. Furthermore, we find that in several human generated datasets, less consistent actions at each state is often associated with decreases in policy performance.
Data quality research dates back to the early time of computing and the full literature is out of the scope of this paper. Existing literature in machine learning has proposed different dimensions of data quality including accuracy, completeness, consistency, timeliness, and accessibility [33, 14, 20]. We could draw similarities between action divergence with the concept of data consistency, and state diversity with completeness. However, instantiation of these metrics is non-trivial and domain-specific.
Our work is closely related to the imitation learning literature that explicitly addresses the distribution shift problem in various ways. Prior work can be divided into two camps: ones that take an algorithm-centric approach to account for biases in the dataset, and ones that employ data-centric methods for modifying the data collection process.
Algorithm-centric. Robust learning approaches including model-based imitation learning methods learn a dynamics model of the environment and therefore can plan to go back in distribution when the agent visits out of distribution states [21, 41]. Data augmentation is a useful post-processing technique when one could impose domain or task specific knowledge [23, 28, 25]. Prior work has also investigated how to explicitly learn from sub-optimal demonstrations by learning a weighting function over demonstrations either to guide BC training [9, 3] or as a reward function for RL algorithms [7, 11]. These methods either require additional information such as rankings of trajectories [9], demonstrator identity [3], or collecting additional environment data [7, 11]. Recent efforts on demonstration retrieval augment limited task-specific demonstrations with past robot experiences [39]. A recent
body of work build better inductive biases into the model for learning robust state features, like pretrained vision or language representations [54, 30, 38, 12]. Some approaches modify the action representation to be more expressive to capture all the expert's actions, for example using Gaussian or mixture model action spaces [12, 37]. Others consider temporally abstracted action spaces like waypoints or action sequences to reduce the effective task length and thus mitigate compounding errors [53, 48, 12]. What these algorithm-centric works overlook is that changing the data can be as or more effective for policy learning than changing the algorithm. However, we still lack a comprehensive understanding of what properties in the data matter in imitation learning.
Data-centric. In more data-centric prior works that discuss data quality, the primary goal is often just to maximize state diversity. A large body of research focuses on modifying the data collection process such that the expert experience a diverse set of states. Ross et al. [46] proposed to iteratively collect on-policy demonstration data with shared-control between the expert and the robot, randomly switching with a gradually decreasing weight on the expert's input, such that the training data contains direct samples of the expert policy at states experienced by the learned policy. However, randomly switching the control can make it unnatural for the human demonstrator and leads to noisier human control. To mitigate this issue, methods have been proposed to gate the control more effectively by evaluating metrics such as state uncertainty and novelty [16, 27, 50]. Other methods allow the human to gate the control and effectively correct the robot's behavior only when necessary [32, 36, 24, 47]. Closely related to our work, Laskey et al. [34] takes an optimal control perspective and showed that injecting control noise during data collection can give similar benefits as DAgger-style iterative methods. Active learning methods have also been developed to guide data collection towards more informative samples [15, 6, 42, 29]. By only focusing on state coverage, these works are missing the role that actions (i.e., the expert) plays in the quality of data.
In imitation learning (IL), we assume access to a dataset  $\mathcal{D}_N = \{\tau_1, \dots, \tau_N\}$  of  $N$  successful demonstrations. Each demonstration  $\tau_i$  consists of a sequence of continuous state-action pairs of length  $T_i$ ,  $\tau_i = \{(s_1, a_1), \dots, (s_{T_i}, a_{T_i})\}$ , with states  $s \in S$  and actions  $a \in \mathcal{A}$ . Demonstrations are generated by sampling actions from the demonstration policy  $\pi_E$  through environment dynamics  $\rho(s'|s, a)$ . The objective of imitation learning is to learn a policy  $\pi_\theta : S \to \mathcal{A}$  that maps states to actions. Standard behavioral cloning optimizes a supervised loss maximizing the likelihood of the state-action pairs in the dataset:
$$
\mathcal {L} (\theta) = - \frac {1}{| \mathcal {D} _ {N} |} \sum_ {(s, a) \in \mathcal {D} _ {N}} \log \pi_ {\theta} (a | s), \tag {1}
$$
which is optimizing the following objective under finite samples from the demonstration policy:
$$
\mathbb {E} _ {s \sim \rho_ {\pi_ {E}} (\cdot)} [ D _ {\mathrm {K L}} (\pi_ {E} (\cdot | s), \pi_ {\theta} (\cdot | s)) ] = - \mathbb {E} _ {s \sim \rho_ {\pi_ {E}} (\cdot), a \sim \pi_ {E} (\cdot | s)} [ \log \pi_ {\theta} (a | s) ] + C \tag {2}
$$
The  $C$  term here captures the entropy of the demonstration state-action distribution, which is constant with respect to  $\theta$  and thus it does not affect optimality of  $\theta$ .  $\rho_{\pi_E}(s)$  is the state visitation of the demonstration policy, defined for any policy  $\pi$  as follows:
$$
\rho_ {\pi} (s) = \frac {1}{T} \sum_ {t = 1} ^ {T} \rho_ {\pi} ^ {t} (s) \tag {3}
$$
$$
\rho_ {\pi} ^ {t} \left(s ^ {\prime}\right) = \int_ {s, a} \rho_ {\pi} ^ {t} \left(s, a, s ^ {\prime}\right) d s d a = \int_ {s, a} \pi (a | s) \rho \left(s ^ {\prime} \mid s, a\right) \rho_ {\pi} ^ {t - 1} (s) d s d a \tag {4}
$$
3.1 Distribution Shift in IL
Behavioral cloning methods as in Eq. (2) often assume that the dataset distribution is a good approximation of the demonstration policy distribution. In most applications, however, the learned policy is bound to experience novel states that were not part of the training data due to stochasticity of the environment dynamics and the learned policy. Herein lies the fundamental challenge with imitation learning, i.e., state distribution shift between training and test time. Consider the training sample  $(s,a,s^{\prime})$  at timestep  $t$  in demonstration  $\tau_{i}$ . If the learned policy outputs  $\tilde{a}\sim \pi (\cdot |s)$  which has a small action error  $\epsilon = \tilde{a} -a$ , the new state at the next time step will also deviate:  $\tilde{s}^{\prime}\sim \rho (s^{\prime}|s,a + \epsilon)$ , which in turn affects the policy output at the next step. In practice, the change in next state can be
highly disproportionate to  $||\epsilon||$ , so small errors can quickly lead the policy out of the data distribution. Stochastic dynamics, for example system noise, can compound with the distribution shift caused by policy errors, and as we continue to execute for  $T - t$  steps, both these factors can compound over time to pull the policy out of distribution, often leading to task failure.
We can address this distribution shift problem by minimizing the mismatch between state visitation distribution of a learned policy  $\pi$  and the state visitation distribution of an demonstration policy  $\pi_{E}$  under some  $f$ -divergence  $D_{f}$ :
$$
J (\pi , \pi_ {E}) = D _ {f} \left(\rho_ {\pi} (s), \rho_ {\pi_ {E}} (s)\right) \tag {5}
$$
Next, we connect this distribution shift problem to the question of data quality.
3.2 Formalizing Data Quality
How can we define and measure the quality of a dataset  $\mathcal{D}$  in IL? In existing literature, data quality has been used interchangeably with either the proficiency level of the demonstrator at the task or the coverage of state space. However, these notions of quality are loose and incomplete: for example, they do not explain what concretely makes a demonstrator better or worse, nor do they consider how stochasticity in the transitions impacts state space coverage and thus downstream learning. Our goal is to more formally define a measure of quality so that we can then optimize these properties in our datasets.
We posit that a complete notion of dataset quality is one that minimizes distribution shift in Eq. (5). This suggests that we cannot discuss data quality in isolation. The notion of data quality is heavily tied to the algorithm  $A$  that leads to the learned policy  $\pi_A$  as well as the demonstration policy  $\pi_E$ . We thus formalize the quality of a dataset  $\mathcal{D}_N$  conditioned on the demonstration policy  $\pi_E$  that generates it along with a policy learning algorithm  $A$ , as the negative distribution shift of a learned policy under some  $f$ -divergence  $D_f$ :
$$
Q \left(\mathcal {D} _ {N}; \pi_ {E}, A\right) = - D _ {f} \left(\rho_ {\pi_ {A}} (s), \rho_ {\pi_ {E}} (s)\right), \text {w h e r e} \pi_ {A} = A \left(\mathcal {D} _ {N}\right) \tag {6}
$$
Here,  $\pi_A$  is the policy learned by algorithm  $A$  using dataset  $\mathcal{D}_N$  of size  $N$ , which is generated from demonstrator  $\pi_E$ . Note that quality is affected by several components: the choice in demonstrator, the choice of dataset size, and the algorithm. The choice of demonstration policy  $\pi_E$  changes the trajectories in  $\mathcal{D}_N$  (which in turn affects policy learning) along with the desired state distribution  $\rho_{\pi_E}(s)$ . The dataset size controls the amount of information about the demonstration state distribution present in the dataset: note that  $\rho \pi_A(s)$  should match  $\rho_{\pi_E}(s)$ , but  $A$  learns only from  $\mathcal{D}_N$  not from the full  $\rho_{\pi_E}(s)$ . The algorithm controls how the data is processed to produce the final policy  $\pi_A$  and thus the visited distribution  $\rho_{\pi_A}(s)$ .
To optimize this notion of quality, we might toggle any one of these factors. Prior work has studied algorithm modifications at length [37, 54, 4], but few works study how  $\pi_E$  and the dataset  $D_N$  should be altered to perform best for any given algorithm  $A$ . We refer to this as data curation. In practice we often lack full control over the choice of  $\pi_E$ , since this is usually a human demonstrator; still, we can often influence  $\pi_E$ , for example through prompting or feedback [24], or we can curate the dataset derived from  $\pi_E$  through filtering [19].
4 Properties of Data Quality
To identify properties that affect data quality in imitation learning, we can study how each state transition  $(s,a,s^{\prime})\in \mathcal{D}$  in the dataset affects the dataset quality in detail. As we defined in Eq. (6), the dataset quality relies on distribution shift given an demonstration policy  $\pi_E$  and an imitation learning policy  $\pi_A$ . This relies on the state visitation distribution  $\rho_{\pi_A}^t (s)$ , which based on Eq. (4) depends on three terms:  $\pi_A(a|s)$ ,  $\rho (s'|s,a)$ , and  $\rho_{\pi_A}^{t - 1}(s)$ . Intuitively, three clear factors emerge for managing distribution shift: how different the policy  $\pi_A(a|s)$  is from the demonstrator  $\pi_E(a|s)$  - which we refer to as action divergence, how diverse the dynamics  $\rho (s'|a,s)$  are - which we refer to as transition diversity, and how these factors interact over time to produce past state visitation distribution  $\rho_{\pi_A}^{t - 1}(s)$ . Importantly, the past state visitation can be controlled through action divergence and transition diversity at previous time steps, so in this section, we will formalize these two properties and discuss their implication on data curation.
4.1 Action Divergence
Action divergence is a measure of distance between the learned policy and the demonstration policy  $D_{f}(\pi_{A}(\cdot |s),\pi_{E}(\cdot |s))$ . This can stem from biases in the algorithm or dataset such as mismatched action representations or lack of samples. While this is typically viewed in prior work as a facet of the algorithm or dataset size, importantly action divergence and thus data quality can also be influenced by the demonstration policy itself. For example, if the demonstrator knows the action representation used by the learning agent a priori, then the policy mismatch can be reduced by taking actions that are consistent with that action space. In Theorem 4.1<sup>1</sup>, we illustrate the importance of action divergence by showing that our notion of data quality is lower bounded by the action divergence under the visited state distribution.
Theorem 4.1. Given a policy  $\pi_A$  and demonstrator  $\pi_E$  and environment horizon length  $H$ , the distribution shift:
$$
D _ {K L} \left(\rho_ {\pi_ {A}}, \rho_ {\pi_ {E}}\right) \leq \frac {1}{H} \sum_ {t = 0} ^ {H - 1} (H - t) \mathbb {E} _ {s \sim \rho_ {\pi_ {A}} ^ {t}} \left[ D _ {K L} \left(\pi_ {A} (\cdot | s), \pi_ {E} (\cdot | s)\right) \right]
$$
When using KL divergence for distribution shift in Eq. (5), we can see how the quality is lower bounded by the policy action divergence from the optimal policy under the visited state distribution, weighted at each step by time-to-go. Many prior works have noted the compounding error problem [46, 45], and here too action divergence at earlier timesteps has an out-sized effect on the overall distribution shift. Thus to optimize for the quality of a dataset, we should reduce the aggregated policy mismatch across visited trajectories.
**Optimality:** Many prior ideas around the "optimality" of a demonstration are intimately related to this action divergence property. Suboptimalities that are naturally present in demonstrations like pauses or other action noise have been a source of difficulty when learning imitation policies [35, 4]. These factors make it harder for the model to learn the demonstration action distributions, thus increasing action divergence. Suboptimality can also come from multi-modal demonstration policy distributions, and more expressive policy representations (i.e., reducing action divergence between the demonstrator and the learned policy) has been shown to help [17, 12, 24]. The speed of demonstration is another common notion of optimality, which decreases total action divergence by reducing horizon length  $H$ , provided it does not increase the per time step action divergence to compensate.
State Visitation: Critically, the visitation distribution in Theorem 4.1 determines where action divergence should be low, whereas in the standard BC objective in Eq. (2), the divergence is only minimized under samples from the demonstration distribution  $\rho_{\pi_E}^t$ . To better understand how the visitation distribution evolves, we now analyze state transitions in greater detail.
4.2 Transition Diversity
Transition diversity encompasses the diversity of next state transitions seen at a given state for a certain policy. What role does transition diversity play in minimizing distribution shift? Intuitively, if we consider the upper bound in Theorem 4.1, the demonstrator's state coverage should overlap as much as possible with the visited state distribution, but without increasing the action divergence.
Lemma 4.2. Given a learned policy  $\pi_A$  and an demonstration policy  $\pi_E$ , assume that the policy is learned such that when  $s \in \mathrm{supp}(\rho_{\pi_E}^t)$ ,  $D_{KL}(\pi_A(\cdot | s), \pi_E(\cdot | s)) \leq \beta$ . Then:
$$
\begin{array}{l} \mathbb {E} _ {s \sim \rho_ {\pi_ {A}} ^ {t}} \left[ D _ {K L} \left(\pi_ {A} (\cdot | s), \pi_ {E} (\cdot | s)\right) \right] \leq \mathbb {E} _ {s \sim \rho_ {\pi_ {A}} ^ {t}} \left[ \beta \mathbb {1} \left(s \in \operatorname {s u p p} \left(\rho_ {\pi_ {E}} ^ {t}\right)\right) \right. \\ + \mathbb {1} (s \notin \operatorname {s u p p} \left(\rho_ {\pi_ {E}} ^ {t}\right)) D _ {K L} \left(\pi_ {A} (\cdot | s), \pi_ {E} (\cdot | s)\right) ] \\ \end{array}
$$
Lemma 4.2 shows that with a good learning algorithm (i.e., when  $\beta$  is small), it is important to maximize the overlap between the visited policy and the demonstration data (minimizes the contribution of the unbounded right term), but as shown in Theorem 4.1 this should not come at the expense of increasing the action divergence of the policy. Rather than broadly maximizing state diversity, as prior works do, a more sensible approach is for the demonstrator to maximize the two other factors that affect  $\rho_{\pi_A}^t$ : system noise and initial state variance in the demonstrator data.
Proof for all theorems and lemmas in Appendix A
Finite Data and Low Coverage: While the above analysis holds for infinite data, maximizing transition diversity can lead to thinly spread coverage in the finite data regime. To analyze the finite data regime in isolation, we consider Gaussian system noise for simplicity. First, we formulate an empirical estimate of state "coverage" of the demonstrator (i.e. the deviation from states in the demonstration dataset). We then show that higher system noise can cause the learned policy  $\pi_A$  to deviate more from states sampled in the dataset.