Source-linked AI summary
Data Quality in Imitation Learning
Suneel Belkhale, Yuchen Cui, Dorsa Sadigh
TL;DR
Robotics lacks internet-scale datasets, while imitation-learning policies can leave the demonstration distribution as prediction errors compound. The paper formalizes dataset quality through distribution shift and examines action divergence and transition diversity theoretically and across data sources. It finds that state diversity is not always beneficial and that the two properties interact in practice.
Problem
Robotics has small, uncurated datasets, and imitation learning faces training-to-test state distribution shift as policy errors and stochastic dynamics compound.
Source: §1 · p. 1, §3.1 · p. 3Method
The paper formalizes dataset quality as negative learned-versus-expert state distribution shift and studies action divergence and transition diversity theoretically and across datasets.
Source: §1 · p. 2, §3.2 · p. 4Results
State diversity is not always correlated with success, and less consistent actions at a state are often associated with lower performance in several human-generated datasets.
Source: §1 · p. 2Takeaways & Limitations
Action divergence and transition diversity are important data-quality factors that can inform strategies for curating expert datasets.
Source: §1 · p. 2, §6 · p. 9The proposed metrics likely do not provide a complete picture of data quality.
Source: §5.2 · p. 9
Abstract
from arXiv · showhide
In supervised learning, the question of data quality and curation has been over-shadowed 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.
1 Introduction
The introduction argues that robotics needs especially careful data curation because datasets are small and imitation learning can shift into unseen states. It frames data quality around action divergence and transition diversity, rather than state coverage alone.
1 Introduction
- Robotics has small, uncurated datasets, and imitation-learning errors can move policies into unseen states, making data curation especially important.Prior data-centric approaches often maximize state diversity, but this focus misses the role of expert actions in data quality.Source: §1 · pp. 1–2
- Action divergence and transition diversity shape imitation-learning data quality because policy actions and stochastic transitions affect distribution shift over time.Action divergence measures mismatch between learned and expert actions at a state; transition diversity captures environmental variability for a state and action.Source: §1 · p. 2
- State diversity is not always correlated with success, while less consistent actions at a state are often associated with lower performance in several human-generated datasets.The work studies these relationships through data-noising experiments and approximate measurements of human- and machine-generated datasets.Source: §1 · p. 2
2 Related Work
Prior work addresses imitation-learning distribution shift through algorithm-centric changes or data-centric collection methods. The paper situates its approach against state-coverage-focused data curation and broader, domain-specific notions of data quality.
2 Related Work
- Machine-learning literature names quality dimensions such as accuracy, completeness, consistency, timeliness, and accessibility, but their instantiation is domain-specific.The paper relates action divergence to consistency and state diversity to completeness.Source: §2 · p. 2
- The paper distinguishes algorithm-centric approaches from data-centric methods that modify data collection to address distribution shift.It argues that changing data can be as or more effective for policy learning than changing the algorithm.Source: §2 · pp. 2–3
- Algorithm-centric imitation-learning methods address distribution shift through robust models, augmentation, demonstration weighting, retrieval, and improved state or action representations.Some approaches require rankings, demonstrator identity, or additional environment data.Source: §2 · pp. 2–3
- Data-centric methods commonly seek state diversity through shared control, noise injection, and active sampling, but this emphasis can miss the expert actions’ role in data quality.Random control switching can make demonstrations unnatural and produce noisier human control; gating methods aim to mitigate this issue.Source: §2 · p. 3
3 Preliminaries
The preliminaries describe behavioral cloning from expert demonstrations, explain how policy and dynamics errors compound into distribution shift, and define dataset quality by the learned policy’s shift from the expert.
3 Preliminaries
- Imitation learning uses expert state-action demonstrations, and behavioral cloning fits a policy by maximizing the likelihood of those pairs.Its objective is expressed as expected KL divergence between expert and learned actions over expert state visitation, with the expert entropy term constant in the policy parameters.Source: §3 · p. 3
3.1 Distribution Shift in IL
- Small action errors can alter later states, and policy errors and stochastic dynamics can compound over time, pulling the learned policy outside the training distribution.The resulting training-to-test state distribution shift can lead to task failure.Source: §3.1 · pp. 3–4
3.2 Formalizing Data Quality
- Dataset quality is formalized as negative distribution shift between learned and expert state visitation, conditioned on the demonstrator and learning algorithm.The formalization also notes that dataset size affects how much information about the expert state distribution is available to the learner.Source: §3.2 · p. 4
4 Properties of Data Quality
The paper frames imitation-learning data quality through action divergence and transition diversity, which shape distribution shift over time. Its curation implications are to favor consistent actions and manage system noise and coverage together, rather than maximizing state diversity alone.
4.1 Action Divergence
- Action divergence is the mismatch between learned and expert actions; under the paper’s analysis, mismatches earlier in a trajectory have an outsized effect on distribution shift.The authors connect lower aggregate mismatch with higher data quality and note that expert action entropy can make matching difficult when states are sparsely sampled.Source: §4 · p. 4, §4.1 · p. 5
4 Properties of Data Quality
- Data quality depends on how action mismatch, transition noise, and prior state visitation jointly shape distribution shift, not on state coverage alone.The paper presents action divergence and transition diversity as the two properties to formalize for data curation.Source: §4 · p. 4, §4.3 · p. 6
4.2 Transition Diversity
- Transition diversity can improve next-state coverage, but its benefit depends on sample size and must not come at the cost of increased action divergence.Under the stated Gaussian-transition assumptions, coverage falls as system-noise variance rises at fixed sample count; the paper also argues that added coverage can offset learned-policy noise when generalization is considered.Source: §4.2 · pp. 5–6
4.3 Implications for Data Curation
- Dataset curation should favor consistent expert actions and manage system noise so state coverage does not become too sparse or increase action divergence.The authors also caution that horizon length is only a crude proxy for the aggregated effects of action divergence and transition diversity.Source: §4.3 · pp. 6–7
5 Analysis
The experiments vary system and expert-policy noise in scripted demonstrations, then examine metrics in human-collected datasets. They find that state diversity is not uniformly beneficial: effects depend on action divergence, transition diversity, and dataset size.
5.1 Data Noising
- Adding system noise generally improved behavior-cloning performance in larger datasets, while excessive noise could reduce Square performance when state coverage became too thin.In Square’s high-data setting, performance began to drop after σs = 0.3; the authors hypothesize that low coverage explains the decline.Source: §5.1 · p. 8, Figure 2 · p. 8
- Expert-policy noise could hurt more than comparable system noise in low-data settings, while adding system noise made policies more robust to policy noise.In high-data settings, policy noise was only moderately worse than comparable system noise and could help up to a point; the low-data comparison was substantially worse.Source: §5.1 · p. 8, Figure 2 · p. 8
5.2 Data Measuring
- Across human-collected datasets, higher action variance and longer horizons often accompanied lower success, but state diversity was not consistently correlated with performance.The authors note that single-step action variance may miss longer-term multimodal behavior, such as demonstrations grasping different parts of the nut.Source: §5.2 · p. 9, Table 1 · p. 9
- The evaluated metrics offer insights into low performance and data collection, but the authors say they do not provide a complete picture of data quality.The paper envisions a more comprehensive set of metrics for evaluating datasets without first training and evaluating models.Source: §5.2 · p. 9
6 Discussion
The discussion frames data curation for imitation learning around minimizing distribution shift, with action divergence and transition diversity as key data-quality factors. It reports trade-offs in state diversity and cautions that measuring data quality comprehensively on real-world datasets remains challenging.
6 Discussion
- Making actions more consistent tends to increase policy success, while increasing state diversity can come at the expense of action consistency.The discussion argues against uniformly maximizing state diversity.Source: §6 · p. 9
- Increasing transition diversity improves performance until data coverage becomes too sparse.Source: §6 · p. 9
- The proposed metrics are a first step, but comprehensive data-quality measurement on real-world datasets can be challenging.The authors envision a broader set of metrics informed by their formalism.Source: §6 · p. 9
A Theoretical Results
The theoretical results address distribution shift between learned and expert policies and characterize next-state coverage under specified assumptions. The supplied results include a Gaussian-noise coverage expression and an analysis of how sample size and system noise affect coverage.
A Theoretical Results
- Theorem 4.1 states an upper bound on the distribution shift DKL(ρπA, ρπE) between learned-policy and expert state distributions.Source: A Theoretical Results · p. 13
- Under Gaussian transitions with fixed diagonal variance and the stated policy-matching condition, Theorem 4.3 gives next-state coverage as PS(s; N, ϵ) = 1 − (1 − erf(ϵ / 2σ)^d)^N.Here d is the state dimensionality.Source: A Theoretical Results · p. 14
- The left-side coverage analysis varies system noise and sample size N for PS(s; N, ϵ).The supplied passage identifies the plotted variables but does not report the plotted outcome.Source: A Theoretical Results · p. 14
A.1 Generalization under System Noise
This section studies next-state coverage when system noise is present, including a looser coverage definition that allows interpolation among expert samples. Its figure contrasts the conservative PS measure with PB, under which noise can make coverage likely even when test-time policy noise is higher.
A.1 Generalization under System Noise
- PB(s; N) measures whether the learned policy’s next state falls within the hyper-sphere defined by N expert next-state samples, a proxy for neural-network interpolation among training samples.The sphere is centered at the mean expert next state and sized to fit the sampled expert next states.Source: §A.1 · pp. 14–15
- The theorem considers PB under a deterministic demonstrator, Gaussian transitions with fixed diagonal variance, and simplified linear dynamics.Source: §A.1 · p. 15
- Under the looser PB coverage model, added system noise can make coverage likely even with double or triple the learned policy’s noise; under conservative PS, noise quickly reduces coverage, with greater robustness at larger N.Figure 3 plots coverage probability against system-noise standard deviation; its right panel fixes N = 10.Source: Figure 3 · p. 16, §A.1 · p. 16
B Metrics of Data Quality
The section proposes empirical proxies for action consistency and state similarity based on clustering nearby states. The authors caution that these approximate metrics do not cover the full spectrum of data quality.
B Metrics of Data Quality
- Action variance approximates action consistency by measuring per-dimension action variance within clusters of nearby states.The neighborhood is defined using a coverage distance ϵ because the underlying expert action distribution is unavailable.Source: B Metrics of Data Quality · p. 16
- State similarity is estimated as the expected nearby-state cluster size divided by the overall dataset size.It uses the same clustering process as the action-variance metric.Source: B Metrics of Data Quality · p. 16
- These approximate metrics do not capture the full spectrum of data quality, and comprehensive measurement on real-world datasets can be challenging.Source: B Metrics of Data Quality · p. 16
C Results
The results compare behavioral cloning under system noise, policy noise, and both, across PMObstacle and Square and multiple data sizes. The outcomes differ depending on the noise source and the amount of data.
C Results
- Across PMObstacle and Square, more system noise during training generally produces the best models, while low-coverage policies generalize worst as evaluation noise rises.The PMObstacle results cover 1,000-episode and 10-episode datasets; the Square tables cover 200-episode and 50-episode datasets.Source: Table 2 · p. 17, Table 5 · p. 18, Table 6 · p. 18
- With matched state diversity, policy-noise datasets perform almost universally worse than system-noise datasets in PMObstacle.In the large-data regime, more policy noise often produces the worst models; with little data, performance varies substantially across runs.Source: Table 3 · p. 17
- In PMObstacle, system noise improves model robustness to added policy noise in the combined-noise condition.Source: Table 4 · p. 17
- In Square, policy noise tends to improve performance with 200 episodes, but with 50 episodes it causes a significant performance drop and greater variation.The paper attributes the high-data pattern to the model recovering an unbiased policy, which it says it can no longer do with reduced data.Source: Table 7 · p. 18, Table 8 · p. 18