Technical Summary: DataMIL – Selecting Data for Robot Imitation Learning with Datamodels

1. Problem Statement

The robotics community has recently shifted toward training generalist foundation policies using massive, diverse datasets (e.g., Open-X Embodiment). While these policies achieve strong mean performance across varied tasks, they frequently underperform on individual, specialized tasks, necessitating fine-tuning on newly acquired, task-specific data.

The core challenge addressed by this work is data selection: identifying which subsets of large prior datasets, when combined with limited task-specific demonstrations, yield the best specialized policies.

• Limitations of Current Methods: Naive selection based on human notions of quality (e.g., semantic similarity, visual similarity, or state-action proximity) often fails. These heuristics assume that the most "similar" data is the most "useful," ignoring the actual impact of a data point on downstream policy performance.
• The Optimization Barrier: Ideally, one would select data by exhaustively testing subsets via retraining and evaluating the resulting policy. However, in robotics, policy evaluation requires real-world rollouts, which are time-consuming, dangerous, and computationally infeasible for large-scale search.
• The Gap: Existing data attribution frameworks like datamodels (which predict model performance based on training data) have been successful in NLP and Computer Vision but have not been directly applicable to robotics due to the non-differentiable and expensive nature of real-world policy evaluation.

2. Methodology: DataMIL

The authors introduce DataMIL (Datamodels for Imitation Learning), a framework that extends the datamodel paradigm to robotics. DataMIL reasons about data selection in an end-to-end manner, using the policy itself to identify data points that improve performance, rather than relying on static heuristics.

2.1 Core Framework

DataMIL treats the policy learning algorithm as a "black box" and trains an estimator (a datamodel) to predict the performance of a policy trained on a specific subset of data.

• Objective: Find a subset $D' \subset D$ that maximizes the target metric $M(A(D'))$, where $A$ is the learning algorithm and $M$ is the performance metric.
• Datamodel Formulation: Instead of training the policy, DataMIL learns a function $\hat{f}(D') \approx M(A(D'))$. For linear datamodels, this reduces to assigning a scalar score $\tau(z_i)$ to each data point $z_i$, where the predicted performance is the sum of scores for the selected subset.

2.2 Key Adaptations for Robotics

To make datamodels tractable for robotics, the authors introduce three critical modifications:

1. Surrogate Loss Function (Avoiding Rollouts):
   Real-world rollouts are non-differentiable and expensive. DataMIL replaces the true success rate metric with a proxy metric $\tilde{M}$, defined as the negative behavior cloning (BC) loss on a held-out set of target demonstrations:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   This metric is fully differentiable and requires no environment interaction, enabling the use of gradient-based estimation methods.

2. Datamodel Estimators:
   The paper employs two methods to estimate the influence scores $\tau(z_i)$:

   • Regression Estimator: Samples random subsets of data, trains policies, evaluates the proxy metric, and fits a linear model to predict performance based on data presence. This is accurate but computationally heavy.
   • Metagradient-Based Estimator: Uses influence functions and metagradients to compute the derivative of the target metric with respect to data weights. This method is more efficient for large models (like Octo) and avoids the need to retrain policies for every subset.

3. Clustering and Distribution Shift Mitigation:

   • Clustering: To reduce noise in influence estimation, individual state-action pairs are grouped into temporal clusters (sub-trajectories or full trajectories). The optimal cluster granularity depends on dataset size; larger datasets (like OXE) benefit from trajectory-level aggregation.
   • Distribution Shift: To prevent poor estimates caused by shifts between prior and target domains, a small fraction of target task data is included during the datamodel estimation phase to align the learning process with the target domain.

2.3 Training Pipeline

Once influence scores are estimated:

1. The top $x\%$ of prior data (highest positive influence) is selected to form $D_{sel}$.
2. The final policy is trained via co-training on the target data $D_{target}$ and the selected data $D_{sel}$. At each step, a batch is sampled from $D_{target}$ with probability $\alpha$ and from $D_{sel}$ with probability $1-\alpha$.

3. Key Contributions

• Extension of Datamodels to Robotics: DataMIL is the first framework to successfully apply datamodels to robot imitation learning, overcoming the barrier of expensive real-world rollouts via a differentiable surrogate loss.
• End-to-End Performance Awareness: Unlike heuristic methods (visual, motion, or state similarity), DataMIL directly optimizes for the impact on policy success, allowing it to select data that improves performance and discard data that degrades it, even if the data appears visually similar.
• Scalable Estimators: The introduction of the metagradient-based estimator allows the framework to scale to large, complex policies (e.g., Octo) and massive datasets (e.g., OXE) where regression-based approaches would be computationally prohibitive.
• Cross-Embodiment and Multi-Task Selection: The method demonstrates the ability to select useful data for tasks and robot embodiments not present in the prior dataset, and to curate datasets that support multiple downstream objectives simultaneously.

4. Experimental Results

The authors validate DataMIL across 60+ tasks in simulation and real-world settings:

• MetaWorld (50 Tasks): DataMIL achieved a 10% performance boost over state-of-the-art baselines (Behavior Retrieval, Flow Retrieval, STRAP). It successfully filtered out noisy, suboptimal demonstrations from autonomous exploration data, whereas similarity-based baselines often retrieved irrelevant or low-quality data.
• LIBERO (10 Tasks): Using the Octo transformer policy, DataMIL consistently outperformed baselines across diverse long-horizon tasks. While visual similarity baselines performed well in some specific tasks, DataMIL provided the most robust average performance.
• Open-X Embodiment (OXE) (Real-World):
  • Franka-Ball & Franka-Pouch: DataMIL selected data that improved success rates significantly compared to random selection and heuristic baselines.
  • Tiago-Sink: In a challenging cross-embodiment setting (selecting data for a Tiago robot from a dataset containing no Tiago data), DataMIL successfully identified relevant demonstrations from other robots (e.g., RT-1, BC-Z) that shared the "essence" of the task (tabletop manipulation), whereas baselines failed to generalize.
  • Droid-Multitask: DataMIL curated a single dataset that improved performance on three distinct tasks simultaneously, outperforming baselines in average success rate.

Qualitative Findings:

• DataMIL selects a diverse mix of datasets rather than over-relying on a single source (unlike baselines which often retrieve 80%+ from one dataset).
• It effectively distinguishes between helpful and harmful data even when they look visually identical (e.g., same state but different action distributions), a failure mode common in similarity-based retrieval.

5. Significance and Claims

The paper claims that end-to-end, performance-aware data selection is critical for unlocking the potential of large prior datasets in robotics. The authors argue that while large foundation models provide a strong base, their specialization relies heavily on the quality of the fine-tuning data.

• Modest Claims on Scalability: The authors acknowledge that estimating datamodels remains computationally expensive (several times the cost of training a policy on all data), though the metagradient approach mitigates this. They note that hyperparameter tuning (e.g., cluster size, selection ratio) currently lacks strong theoretical intuition.
• Scope: The work focuses primarily on single-task specialization, with the Droid-Multitask setting serving as a preliminary step toward multi-task curation.
• Conclusion: DataMIL demonstrates that moving beyond human-defined heuristics to data-driven, model-aware selection yields consistent gains in success rates, particularly in complex, heterogeneous, and real-world scenarios where traditional retrieval methods break down.