Robot Data Curation with Mutual Information Estimators

Joey Hejna1,2, Suvir Mirchandani2, Ashwin Balakrishna1, Annie Xie1, Ayzaan Wahid1, Jonathan Tompson1, Pannag Sanketi1, Dhruv Shah1, Coline Devin1,†, Dorsa Sadigh1,†
Google Deepmind1, Stanford University2, Equal Advising†
ArXiv Code

Abstract

The performance of imitation learning policies often hinges on the datasets with which they are trained. Consequently, investment in data collection for robotics has grown across both industrial and academic labs. However, despite the marked increase in the quantity of demonstrations collected, little work has sought to assess the quality of said data despite mounting evidence of its importance in other areas such as vision and language. In this work, we take a critical step towards addressing the data quality in robotics. Given a dataset of demonstrations, we aim to estimate the relative quality of individual demonstrations in terms of both state diversity and action predictability. To do so, we estimate the average contribution of a trajectory towards the mutual information between states and actions in the entire dataset, which precisely captures both the entropy of the state distribution and the state-conditioned entropy of actions. Though commonly used mutual information estimators require vast amounts of data often beyond the scale available in robotics, we introduce a novel technique based on k-nearest neighbor estimates of mutual information on top of simple VAE embeddings of states and actions. Empirically, we demonstrate that our approach is able to partition demonstration datasets by quality according to human expert scores across a diverse set of benchmarks spanning simulation and real world environments.

Method Figure

We tested data filtering on various robot platforms and datasets with ground-truth quality labels, and found that filtering data with our method often lead to better downstream policy performance.

Robot Results

Robot Results

RoboCrowd HersheysKiss

DemInf 50% Filter
Random 50% Filter
All Data

Put Dishes Away

Select one of the examples below to view a set of trials.

DemInf 50% Filter
Random 50% Filter
All Data

Pen In Cup

Select one of the examples below to view a set of trials.

DemInf 50% Filter
Random 50% Filter
All Data

BibTeX

@article{hejna2025robotdatacurationmutual,
    title={Robot Data Curation with Mutual Information Estimators}, 
    author={Joey Hejna and Suvir Mirchandani and Ashwin Balakrishna and Annie Xie and Ayzaan Wahid and Jonathan Tompson and Pannag Sanketi and Dhruv Shah and Coline Devin and Dorsa Sadigh},
    year={2025},
    eprint={2502.08623},
    archivePrefix={arXiv},
    primaryClass={cs.RO},
    url={https://arxiv.org/abs/2502.08623}, 
}