Paper title: DataMIL: Selecting Data for Robot Imitation Learning with Datamodels (arXiv 2505.09603, May 2025) Venue: ICLR 2026 Category: Data · Imitation learning Trend tag: Policy-aware data curation Affiliations: UT Austin (Dass, Martín-Martín), MIT (Khaddaj, Engstrom, Mądry), Stanford (Ilyas)

flowchart LR
  Big[Prior dataset D<br/>e.g. OXE, LIBERO-90, MetaWorld replay] --> Cluster[Cluster into sub-trajectories<br/>length 15 / 30 / full traj.]
  Cluster --> DM{Datamodel estimator}
  DM -->|regression| Reg[Train m subset policies,<br/>regress proxy loss on subset masks]
  DM -->|metagradient| Meta[Differentiable through SGD<br/>Engstrom et al. 2025]
  Reg --> Tau[Per-cluster scores τzi]
  Meta --> Tau
  Spec[Target demos D_target<br/>≈ 5–50 per task] --> Proxy[Surrogate loss<br/>negative BC loss on held-out target]
  Proxy --> DM
  Tau --> Sel[Select top-x% positive-influence clusters → D_sel]
  Sel --> Co[Co-train policy:<br/>α from D_target, 1−α from D_sel]
  Co --> Pol[Octo / MLP policy]

Pretraining on large robot datasets helps generalist policies but often hurts specialized-task performance when naively combined with task-specific data. Standard heuristics — semantic, visual, motion-flow, or state-action similarity filters — operate on human notions of quality and do not reliably identify which prior-data samples actually help the downstream policy. The paper formulates this as the optimization

arg max_{D' ⊂ D}  M(A(D'))

where A is a fixed imitation-learning algorithm and M is a target metric (typically rollout success rate). Exhaustive search is infeasible — every evaluation requires a full training run plus real-world rollouts.

DataMIL adapts the datamodels paradigm (Ilyas et al., 2022; Park et al., 2023) to robotics. A datamodel is a function f̂: 2^D → R that predicts model performance given a training subset, without actually running training. The paper considers linear datamodels

f̂(D') = Σ_{zi ∈ D'} τ(zi)

so each training example gets a scalar attribution τ(zi). Two estimators are explored:

• Regression estimator (Eq. 3 in paper): sample m random subsets Dj ⊂ D, train a policy on each, evaluate the metric to obtain Mj, then solve a linear least-squares regression of Mj on subset-membership masks.
• Metagradient estimator (Engstrom et al., 2025): parameterize the dataset by a continuous weight vector w ∈ [0,1]^n, Taylor-expand M(A(w)) ≈ M(A(w0)) + ∇w M(A(w0))^T (w − w0), and compute the influence I = ∇w M(A(w0)) exactly by differentiating through SGD using step-wise auto-differentiation. Each Ii becomes a coefficient.

The true target M is success rate via rollouts — non-differentiable and expensive. DataMIL replaces it with a differentiable proxy computed entirely on a held-out target set D_target:

M̂(π, D_target) = (1/|D_target|) Σ_{(s,a) ∈ D_target} −L_BC(π(s), a)

For Gaussian MLP policies (MetaWorld) L_BC is NLL or L1 loss; for Octo's diffusion head (LIBERO, OXE) it is the standard denoising score-matching loss. Empirically L1 works best with the regression estimator and NLL with the metagradient estimator.

A state-weighted variant doubles the loss weight w(s,a) = 2 for "pre-grasp" states in LIBERO; default is w=1.

• Clustering. Per-sample influences are noisy (each example seen few times during training). DataMIL groups examples into temporal clusters — sub-trajectories of horizon length 15 for LIBERO; full trajectories for OXE; and full trajectories for MetaWorld (App. E.1). Optimal granularity scales inversely with dataset size (coarser for larger datasets). (Separately, the similarity baselines BR/SR/AR slide a fixed horizon H = 50 window over MetaWorld data for their retrieval — this is a baseline detail, not DataMIL's clustering.)
• Distribution shift. Datamodels trained purely on prior data degrade when the target domain differs in lighting, camera, or embodiment. For OXE (real world) DataMIL splits D_target in half — one half mixed into the datamodel-estimation training set, the other half reserved for the proxy. Skipped in MetaWorld and LIBERO since prior and target distributions match.

Top-x% of clusters by attribution form D_sel. The downstream policy is co-trained: each minibatch draws from D_target with probability α and from D_sel with 1−α. Authors use α = 0.5 everywhere and x = 10% for MetaWorld and LIBERO; for the much larger OXE pool the per-task budget is far smaller (Table 2: 1% for Franka-Pick / Droid-Multitask, 0.75% for Franka-Pouch, 0.5% for Tiago-Sink).

Prior dataset = 4,000 scripted-expert episodes (350K env steps) + multi-task SAC replay (1M env steps, ~3× the scripted data, from a SAC agent averaging ~21% success). Each task uses 5 expert demos as D_target and selects 10% of the prior data with the regression estimator. (Per App. E.1, MetaWorld trains on D_sel only — co-training with D_target gave negligible gains here and is excluded.)

In the proxy validation experiment on pick-place-wall, DataMIL with proxy metric achieves ~7× higher success than All-Data while Target-Only fails almost entirely. Without goal-conditioning, DataMIL still beats the best baseline by 10% but absolute scores drop across all methods.

Prior = LIBERO-90 (4,500 human teleop demos, 50 per task). Target = 10 LIBERO-10 long-horizon tasks. Sub-trajectory clustering (length 15), metagradient estimator, top-10% selection.

(All numbers are mean ± std over 5 seeds in the paper; see Table 3 of the paper.) Note that on LIBERO's clean single-view structure, similarity baselines are competitive — STRAP with the authors' relaxed retrieval (top sub-trajectories until 10% of prior data, vs. STRAP's original 100 segments) jumps from 24.72% to 34.96%. DataMIL still wins on average.

4 tasks across 2 embodiments: Franka-Ball, Franka-Pouch, Tiago-Sink, Droid-Multitask (3 sub-tasks: bread-in-bowl, napkin-in-drawer, open-drawer). Subsets OXE-13 / OXE-23 / OXE-24 from Octo's training collections. Tiago is never seen in prior data.

(Paper Table 3; for OXE the "All-Data" column is replaced by a Random baseline matched to DataMIL's selection budget. Real-robot evaluations use a single seed.) DataMIL is best on average (61.0 vs. the next-best 40.0 for Flow) and dominates on the harder cross-embodiment / multi-task settings (Franka-Pouch 70.6, Tiago-Sink 64.1, Droid-Multitask 59.4). The only task where a similarity baseline edges it out is Franka-Ball (STRAP 57.1 vs. DataMIL 50.0), where the simple single-source target distribution favors visual retrieval.

• For Tiago-Sink (unseen embodiment), DataMIL's most-selected datasets are RT-1, BC-Z, and Bridge — visually different but sharing table-top ego-perspective semantics.
• DataMIL spreads selection across many source datasets; AR / Flow / BR concentrate on a single source (e.g., AR pulls 80%+ from RT-1 on Tiago-Sink).
• High- and low-ranked samples often look similar — supporting the data-attribution intuition that harmful samples differ from helpful ones in subtle action/label structure rather than appearance.

• Datamodel estimator quality (Fig. 2 on pick-place-wall). DM-rollouts (true success target) ≈ DataMIL-rg (proxy + regression) > DataMIL-meta (proxy + metagradient). The proxy incurs only a marginal drop versus rollouts; metagradient incurs a further small drop but is ~8× faster to train. All three selected-data policies achieve up to 7× higher success than All-Data, while Target-Only fails almost entirely — the paper's central evidence that the differentiable proxy is a valid stand-in for rollout success and that curation (not just more data) is what helps.
• Co-training ratio. α = 0.5 is reported to "perform well across all experiments" (Sec. 4.3); the reachable version gives no α sweep figure.
• Clustering granularity is chosen per dataset size (sub-trajectory length 15 for LIBERO, full trajectories for OXE/MetaWorld); the paper motivates this qualitatively rather than via a published cluster-size ablation.

Audit note: A prior revision of this page contained an "Ablation Studies" section citing Figure 13a/b/c (selection-percentage / cluster-size / α sweeps), a "LIBERO-5 ablation suite (5 seeds)", and an "All-Data longer-training" table with values 43.36 / 44.13 / 38.40. None of these appear in the reachable source (arXiv 2505.09603 v1 and the OpenReview submission PDF are both 21 pages; figures cap at Figure 9 and the numbers do not occur anywhere in the text). They have been removed as unverifiable/likely fabricated. Restore only if a longer camera-ready version is confirmed to contain them.

1. Compute cost. Even with the metagradient estimator, datamodel estimation costs several times more than training the policy on all data; the authors suggest accelerating it with a scaled-down version of the model or prior dataset as future work.
2. Hyperparameter intuition. Target-set size, cluster granularity, and selection percentage lack principled selection rules; the paper's choices are empirical.
3. Single-task focus. Most experiments target single tasks; Droid-Multitask is the only multi-task evaluation. Truly large-scale multi-task selection is open.

Data selection becomes policy-aware rather than feature-aware — the policy itself defines which samples are useful. This is a different axis from:

• LBM Co-training Study (TRI) (asks how much prior data to mix, not which subset).
• Visual / motion / state-action similarity baselines like STRAP, FlowRetrieval, BehaviorRetrieval which all encode hand-crafted notions of relevance.
• Concurrent CUPID (Agia et al., 2025) which uses a policy-gradient influence measure but requires online rollouts; DataMIL is fully offline.

DataMIL's main provenance is the datamodels line (Ilyas et al., 2022; Park et al., 2023; Engstrom et al., 2025) — this is the first paper to bring exact metagradient datamodel estimation to imitation learning. It's a practical recipe for the increasingly common workflow of fine-tuning a generalist policy (Octo, π0, OpenVLA) on a narrow target task using a giant heterogeneous prior dataset (OXE-scale): instead of "use all of OXE" or "use a similarity-filtered slice," score each cluster by its differentiable contribution to your held-out target loss.

• arXiv: https://arxiv.org/abs/2505.09603
• Project page: https://robin-lab.cs.utexas.edu/datamodels4imitation/
• Code: https://github.com/UT-Austin-RobIn/datamil
• OpenReview: https://openreview.net/forum?id=AcTsKglDdh
• PDF: https://openreview.net/pdf?id=AcTsKglDdh

• LBM Co-training Study (TRI)
• Hybrid Training

← Back to ICLR-2026