Data Quality in Imitation Learning
Suneel Belkhale Yuchen Cui Dorsa Sadigh
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.