DataMIL: Selecting Data for Robot Imitation Learning with Datamodels

TL;DR

DataMIL presents a data-driven framework for data selection in robot imitation learning by extending the datamodels paradigm to robotics. It builds tractable, end-to-end estimators (regression and metagradient) to predict each datapoint's influence on policy performance, using a differentiable proxy objective to avoid expensive real-world rollouts. Through clustering, proxy-based evaluation, and co-training on target data, DataMIL achieves consistent performance gains across diverse benchmarks (MetaWorld, LIBERO, OXE) and even transfers to new embodiments. The work demonstrates the value of performance-aware data curation for leveraging large, heterogeneous prior datasets in practical robotic manipulation tasks.

Abstract

Recently, the robotics community has amassed ever larger and more diverse datasets to train generalist robot policies. However, while these policies achieve strong mean performance across a variety of tasks, they often underperform on individual, specialized tasks and require further tuning on newly acquired task-specific data. Combining task-specific data with carefully curated subsets of large prior datasets via co-training can produce better specialized policies, but selecting data naively may actually harm downstream performance. To address this, we introduce DataMIL, a policy-driven data selection framework built on the datamodels paradigm that reasons about data selection in an end-to-end manner, using the policy itself to identify which data points will most improve performance. Unlike standard practices that filter data using human notions of quality (e.g., based on semantic or visual similarity), DataMIL directly optimizes data selection for task success, allowing us to select data that enhance the policy while dropping data that degrade it. To avoid performing expensive rollouts in the environment during selection, we use a novel surrogate loss function on task-specific data, allowing us to use DataMIL in the real world without degrading performance. We validate our approach on a suite of more than 60 simulation and real-world manipulation tasks - most notably showing successful data selection from the Open X-Embodiment datasets-demonstrating consistent gains in success rates and superior performance over multiple baselines. Our results underscore the importance of end-to-end, performance-aware data selection for unlocking the potential of large prior datasets in robotics. More information at https://robin-lab.cs.utexas.edu/datamodels4imitation/

Figures (9)

• Figure 1: Data selection with datamodels. (left) Similarity-based methods select close samples (yellow), but these aren’t always beneficial for learning. DataMIL evaluates data based on its impact on policy performance, selecting the samples that lead to policy improvement. (center) We estimate datamodels that score each sample by its influence on policy performance and select the highest-scoring samples for training. (right) DataMIL explores two datamodel estimation methods, adapted to robotics: regression and metagradient-based estimation (see Sec. \ref{['subsec:method_datamodels']}).
• Figure 2: Comparing true rollout success ($\mathcal{M}$) vs. proxy metric ($\mathcal{\widehat{M}}$)
• Figure 3: Performance of policy trained on selected datasets in sim. environments.
• Figure 4: Results for data selection on OXE. We test the performance of policies trained on data selected from the Open X-Embodiment dataset using different selection strategies. Data selected using DataMIL achieves the highest performance across all tasks, highlighting the need for end-end policy-aware data selection techniques. (Droid-Multitask shows the average success rate across all its tasks. Individual task success rates shown in the appendix)
• Figure 5: Distribution of datasets selected by different methods for Tiago-Sink task