Robot Data Curation with Mutual Information Estimators (2502.08623v3)

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 action diversity and 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 captures both the entropy of the marginal action distribution and the state-conditioned action entropy. 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. Moreover, training policies based on data filtered by our method leads to a 5-10% improvement in RoboMimic and better performance on real ALOHA and Franka setups.

• The paper proposes a novel Demonstration Information Estimation (DemInf) method to score robot demonstrations based on mutual information.
• It combines VAEs for structured representation learning with k-NN-based estimators to compute state-action quality effectively.
• Empirical results show a 5-10% improvement in imitation learning performance on benchmarks after filtering by data quality.

Robot Data Curation Using Mutual Information Estimators

The paper "Robot Data Curation with Mutual Information Estimators" by authors affiliated with Google DeepMind Robotics and Stanford University details an innovative approach to improving the quality of datasets used in robot imitation learning. This research addresses a critical gap in the field where the focus has traditionally been on increasing the volume of data collected for training imitation learning policies, with less attention given to the quality of this data. The authors propose a novel method to estimate the quality of individual robot demonstrations by analyzing both state diversity and action predictability using mutual information (MI) estimators.

Key Contributions and Methodology

The seminal contribution of this paper is the development of a methodology to quantify the relative quality of individual robot demonstrations by estimating their contribution to the mutual information between states and actions in the dataset. This is achieved using a combination of $k$-nearest neighbor ($k$-NN) based MI estimators and Variational Autoencoders (VAEs) for embedding state and action spaces into low-dimensional representations. The choice of $k$-NN for mutual information estimation is particularly noteworthy because traditional MI estimators demand large-scale datasets, which are often not available in robotics due to data collection constraints.

The proposed Demonstration Information Estimation (DemInf) method involves several steps:

1. Representation Learning: Utilizing VAEs to obtain structured low-dimensional embeddings of states and actions.
2. Mutual Information Estimation: Applying $k$-NN-based estimators to compute mutual information on these embeddings.
3. Scoring and Filtering: Averaging mutual information estimates across individual trajectories to partition the dataset based on demonstration quality.

Empirical Evaluation and Results

Empirically, the authors showcase that their approach is capable of segmenting datasets by quality according to human expert evaluations across simulation and real-world benchmarks, demonstrating significant improvements in the training of imitation learning policies. Specifically, training with datasets filtered using their method led to a 5-10% performance enhancement on the RoboMimic benchmark and yielded superior results on practical setups like ALOHA and Franka.

Furthermore, the research emphasizes the importance of not only collecting large datasets but ensuring that high-quality data is identified and leveraged effectively. This paradigm shift from data quantity to data quality could have profound implications for enhancing the performance and generalization capabilities of robotic learning systems.

Implications and Future Developments

The implications of this work are extensive, both practically and theoretically. Practically, the DemInf method equips researchers and engineers with a tool to refine and optimize their dataset for imitation learning, potentially reducing the cost and resource demands associated with collecting high-quality data. Theoretically, the use of mutual information as a measure of data quality paves the way for new insights into the nature of effective motor learning in robotics.

As the field of AI and robotics progresses, the adoption of such quality-driven data approaches could enhance the robustness and adaptability of robotic systems across varied and dynamic real-world environments. Future developments may investigate integrating the proposed mutual information framework with online data collection processes, adapting dynamically to new input data, and refining robotic learning models accordingly. Furthermore, as datasets grow in size and diversity, more scalable and efficient MI estimation methods or alternative representational techniques might be explored to improve the fidelity of state-action correlations in robotic datasets.