PaperlayerView exact revision on arXiv

Source-linked AI summary

Robot Data Curation with Mutual Information Estimators

Joey Hejna, Suvir Mirchandani, Ashwin Balakrishna, Annie Xie, Ayzaan Wahid, Jonathan Tompson, Pannag Sanketi, Dhruv Shah, Coline Devin, Dorsa Sadigh

arXiv:2502.08623v3cs.RO

TL;DR

Robot demonstration quality is important, but it has received little measurement attention, and common mutual-information estimators can require more data than robotics settings provide. DemInf uses VAE embeddings and k-nearest-neighbor estimation to score demonstrations, and its filtering improves policy outcomes across reported simulation and real-world settings.

  • Problem

    Robotics has limited methods for assessing individual demonstration quality, despite its relevance to imitation learning and the difficulty of collecting large datasets.

    Source: §I · p. 1

  • Method

    DemInf uses VAE embeddings of states and action chunks with k-nearest-neighbor mutual-information estimates, averaged over time to rank demonstrations.

    Source: §I · p. 2, A. Representation Learning · p. 5

  • Results

    Across robot platforms and environments, DemInf partitions expert-rated data quality better than baselines, and filtering improves policy performance in RoboMimic and real ALOHA setups.

    Source: §I · p. 2, §VII · p. 9

  • Takeaways & Limitations

    The authors present mutual information as a useful basis for ranking and selecting individual robot demonstrations across simulation and real-world datasets.

    Source: §VII · p. 9

    The aggregate state-action mutual information analysis largely ignores sequential links through environment dynamics, and disentangling them by using action sequences would reduce estimation data by a factor of T.

    Source: A. Limitations · p. 10

Abstract

from arXiv · show

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.

I. INTRODUCTION

The paper argues that robotics needs demonstration-quality measures that assess action diversity and predictability, then introduces DemInf to rank individual demonstrations using mutual information. It situates this approach among data-curation methods that do not directly assess individual robot action quality.

I. INTRODUCTION

  • The paper motivates curation by noting that maximum-likelihood learning benefits from high-quality data, while robot data collection has expanded in response to successes with larger vision and language datasets.The introduction cites work on supervised learning and scaling in those fields [62],, [33] [74], alongside increased robotics data collection [34] [69].Source: §I · p. 1
  • Because each robot demonstration costs labor, time, and capital, the authors argue that data quality deserves attention alongside dataset size.They cite DROID’s 76K demonstrations across 13 institutions being down-weighted for OpenVLA training after hurting performance.Source: §I · p. 1
  • Demonstration quality is undermeasured in robotics, and methods from vision and language do not directly handle continuous robot states and actions or action quality.The paper proposes state-action mutual information as a measure of relative predictability that affects how well a policy fits the expert.Source: §I · p. 1
  • DemInf embeds states and action chunks with VAEs, estimates mutual information using k-nearest neighbors, and averages estimates over time to rank demonstrations.The authors report greater stability than common estimators on robotics datasets of 50–300 demonstrations.Source: §I · p. 2
  • Across robot platforms and environments, DemInf partitions expert-rated high- and low-quality data better than contemporary baselines and alternative mutual-information estimators, while filtering improves policy performance.The reported policy gains span RoboMimic and real ALOHA setups.Source: §I · p. 2

II. RELATED WORK

  • Unlike coarse vision and language filters, group-level robot data mixing, and similarity retrieval, DemInf scores individual demonstrations without labels; it uses k-NN estimation rather than sample-hungry parametric MI methods.The related work discusses n-gram filtering, CLIP-based curation [23], group mixing, retrieval [53] [21] [44], and parametric MI methods, which may require about a million samples; DemInf uses KSG [40].Source: §II · p. 2

A. Imitation Learning

Imitation learning commonly trains a policy by behavior cloning on expert demonstrations, although real demonstrations can reflect varied collectors and strategies. The paper frames curation as offline, individual-demonstration scoring to shape the empirical expert distribution.

A. Imitation Learning

  • Behavior cloning turns imitation learning into supervised learning that can be trained offline from expert demonstrations.Imitation learning seeks to match expert state visitation, while directly optimizing that objective is challenging because it requires sampling from the learned policy.Source: A. Imitation Learning · p. 3

B. Demonstration Curation

  • Real demonstrations need not come from one absolute expert: collectors and strategies can differ, making the empirical expert distribution a curation choice.Novice collectors may produce less predictable actions, and pooling demonstrators can complicate the action distribution.Source: B. Demonstration Curation · p. 3, A. Imitation Learning · p. 3
  • The paper scores individual demonstrations offline, assuming they all complete the task, and filters the dataset by retaining demonstrations above a quality threshold.This individual-level approach is finer-grained than group mixing and requires neither environment access nor oracle feedback; the proposed score uses unsupervised objectives such as mutual information.Source: B. Demonstration Curation · p. 3

IV. MUTUAL INFORMATION AS A QUALITY METRIC

The paper treats state-action mutual information as a natural quality metric because it reflects how well states predict actions. Its conditional-entropy account argues that simpler, more predictable actions can make behavior cloning a closer fit to the expert.

IV. MUTUAL INFORMATION AS A QUALITY METRIC

  • State-action mutual information is a natural curation metric because behavior cloning aims to predict actions from states.The paper interprets its two components as minimizing conditional action entropy and maximizing marginal action entropy.Source: §IV · p. 3

A. Minimizing Conditional Action Entropy

  • Lower conditional action entropy may improve behavior cloning because simpler action distributions can be easier to fit and make forward and reverse KL behaviors more similar.The authors also argue that lower-entropy distributions often have fewer modes, where the divergences’ mode-seeking and mode-covering behaviors are more likely to align.Source: A. Minimizing Conditional Action Entropy · pp. 3–4
  • Reducing conditional action entropy can limit the effect of unobserved information on actions when data collectors have privileged sensory information unavailable to the policy.The paper gives the bound H(A|S) ≥ I(A; Z|S) for an unobserved variable Z.Source: A. Minimizing Conditional Action Entropy · p. 4

B. Maximizing Marginal Action Entropy

Mutual information balances diverse marginal actions with predictable state-conditioned actions, while DemInf estimates demonstration contributions using k-NN mutual information on learned representations. VAE embeddings make distances suitable for k-NN estimation in high-dimensional robotics data.

B. Maximizing Marginal Action Entropy

  • Mutual information pairs high marginal action entropy with low state-conditioned action entropy, discouraging constant actions and favoring state-aware predictions.DemInf averages mutual-information estimates across trajectories, so it does not explicitly optimize initial-state diversity and instead targets data predictability.Source: B. Maximizing Marginal Action Entropy · p. 4

V. METHOD

  • DemInf scores individual episodes by their contributions to dataset mutual information, using k-NN estimation to address the sample demands of neural estimators.Its pipeline comprises representation learning, mutual-information estimation, and scoring.Source: §V · p. 4

A. Representation Learning

  • Separate VAEs embed states and actions in low-dimensional spaces, where their isotropic Gaussian prior makes distances meaningful for k-NN estimates.The latent dimension is chosen to be as small as possible while sufficiently capturing each variable.Source: A. Representation Learning · pp. 4–5

B. k-NN MI Estimation

DemInf uses KSG k-NN estimates of mutual information to score state-action pairs, then averages their contributions to rank demonstrations. It uses batching for computation and avoids filtering individual pairs so that harder but necessary task phases are retained.

B. k-NN MI Estimation

  • KSG estimates mutual information from the k-NN distances of state-action pairs in joint and separate latent spaces.The estimator uses nearby-sample density to estimate entropy and extends that estimate to mutual information.Source: B. k-NN MI Estimation · p. 5
  • In Square MH, estimated mutual information is highest at trajectory start and after grasp, but low during grasp at 50–75 steps.The figure reports averages per timestep for high-quality demonstrations.Source: Figure 2 · p. 5
  • DemInf averages state-action mutual-information contributions within batches across repeated dataset shuffles and multiple k values.This randomized approach addresses the computational cost of k-NN as dataset size increases.Source: B. k-NN MI Estimation · p. 5

C. Scoring

  • Each demonstration’s score is its average contribution to the dataset’s mutual-information estimate, and filtering retains demonstrations contributing positively.The score is used for relative ordering; state-action estimates are clipped to the 1st and 99th percentiles to limit outlier influence.Source: C. Scoring · p. 5
  • Filtering individual state-action pairs could discard necessary but less predictable phases, such as grasping, in favor of easier motion.The authors therefore score and filter demonstrations rather than directly filtering their component pairs.Source: C. Scoring · p. 5

VI. EXPERIMENTS

The experiments evaluate demonstration curation across simulated and real robot datasets with expert quality labels. They compare prior data-quality estimators and alternative mutual-information estimators, with figures covering filtering thresholds across RoboMimic and RoboCrowd datasets.

VI. EXPERIMENTS

  • Experiments ask how well DemInf curates data, how estimator choice affects performance, whether mutual-information curation improves downstream policy learning, and what drives DemInf’s performance.The experiments span simulated, real single-arm, and real bi-arm robots.Source: §VI · pp. 5–6, 1) Datasets · p. 6

1) Datasets:

  • Datasets include RoboMimic’s Lift, Can, and Square tasks, Franka’s PenInCup and DishRack tasks, and ALOHA’s candy bin-picking tasks from RoboCrowd.The datasets include expert quality labels and varying quality, with RoboCrowd versions both with and without unstructured play data.Source: 1) Datasets · p. 6, Figure 3 · p. 6

2) Baselines:

  • Baselines include ensemble uncertainty, demonstration compatibility, VIP, InfoNCE, and MINE, alongside alternative mutual-information estimators.The InfoNCE comparison uses the symmetric objective used to train CLIP; MINE estimates mutual information with a learned critic.Source: 2) Baselines · pp. 6–7
  • Figures 4 and 5 compare average quality remaining after filtering at different score thresholds across RoboMimic and RoboCrowd tasks.Figure 4 contrasts state and image inputs; Figure 5 contrasts task-relevant datasets with versions containing additional unstructured play data.Source: Figure 4 · p. 7, Figure 5 · p. 7

3) Architectures:

DemInf is evaluated as a data-curation method across simulation and real-world settings, with results covering expert-label alignment, estimator stability, and downstream policy performance. The section also specifies the architectures and training setup used in these experiments.

3) Architectures:

  • State-based experiments use two-layer 512-unit MLPs; image-based experiments use ResNet-18 encoders with spatial softmax and two-layer 1024-unit MLPs.Models use Adam at learning rate 0.0001, batch size 256, and three seeds; state-based models train for 50,000 steps and image-based models for 100,000.Source: 3) Architectures · p. 7

B. How well does DemInf curate data?

  • DemInf matches or exceeds baselines across state-based RoboMimic environments and surpasses all methods on image-based RoboMimic and Franka, matching the oracle on DishRack.For RoboCrowd, DemInf most consistently filters low-quality data, while the limited, diverse datasets challenge all baselines.Source: B. How well does DemInf curate data? · pp. 7–8, Figure 6 · p. 7

C. Mutual Information Estimators

  • InfoNCE and MINE perform worse on image-based settings and data-limited RoboCrowd than DemInf, which shows lower variance across seeds.The parametric estimators had one or two runs that performed far worse than others.Source: C. Mutual Information Estimators · p. 8, Figure 7 · p. 8

D. Does demonstration curation affect policy performance?

E. Ablations

DemInf is evaluated through design ablations, benchmark comparisons, and discussion of its scope and future extensions. The authors report that it outperforms prior methods at measuring demonstration quality across real-world and simulated datasets.

E. Ablations

  • DemInf is generally robust to k across HersheyKiss and Square MH, while KSG outperforms BiKSG and KL; latent-space and estimator quality matter for downstream performance.The authors also report that changing latent dimensions or VAE β has relatively little effect on RoboMimic Image performance, based on two seeds for each experiment.Source: E. Ablations · p. 9

VII. CONCLUSION

  • Across real-world and simulated datasets, DemInf outperforms several prior methods at measuring demonstration quality.The conclusion describes results across several datasets of human-teleoperated demonstrations.Source: §VII · p. 9

A. Limitations

  • DemInf aggregates state-action mutual information and may favor predictable pauses; the authors recommend filtering pauses and ensuring demonstrations complete the task.Curation is greedy rather than globally optimal, and DemInf assumes all supplied demonstrations succeed.Source: A. Limitations · p. 10

B. Future Work

  • Future directions include multi-task curation, scaling to larger datasets, and integrating DemInf into online data-collection interfaces.Multi-task use requires disentangling task conditioning from mutual-information estimation, especially when tasks are not enumerable.Source: B. Future Work · p. 10

APPENDIX

The appendix adds results and ablations omitted from the main text, including estimator, policy-learning, parameter, action-space, and baseline comparisons.

A. Extended Results

  • The appendix supplies additional results and ablations that could not fit in the main text.It includes figures on Franka estimator performance and RoboMimic policy-learning performance from state.Source: A. Extended Results · p. 15, Figure 12 · p. 15, Figure 13 · p. 15
  • Figures 14–16 examine how k affects RoboCrowd and RoboMimic Image and how latent dimensions affect RoboMimic Image.The latent-dimension experiment used two seeds rather than three; its caption reports performance as relatively robust to that parameter.Source: Figure 14 · p. 15, Figure 15 · p. 15, Figure 16 · p. 16
  • Figure 17 examines VAE β on RoboMimic Image, while Figure 18 examines action-space choices for RoboCrowd.The β experiment used two seeds rather than three, and its caption reports performance as relatively robust to β for RoboMimic.Source: Figure 17 · p. 16, Figure 18 · p. 16
  • Figures 19–21 place all methods on shared plots for direct comparison on RoboMimic, RoboCrowd, and Franka.The extended plots also include Policy Loss, defined as a behavioral-cloning policy's loss.Source: A. Extended Results · p. 16, Figure 19 · p. 16, Figure 20 · p. 17

B. Method Details

The appendix describes how curation scores are aggregated across demonstrations and summarizes the scoring procedures used by DemInf and comparison methods.

B. Method Details

  • All methods aggregate scores over each demonstration, and state-action scores are clipped at the 1st and 99th percentiles in practice.The general scoring form applies a function h to state-action pairs in data composed of trajectories.Source: B. Method Details · p. 17
  • DemInf fits state and action VAEs, estimates mutual information in random batches, and uses the resulting pair scores as its curation signal.It iterates over the dataset four times; batches contain 1024 examples, and its estimator uses k-NN counts in the latent spaces.Source: B. Method Details · p. 17
  • BiKSG changes DemInf's mutual-information estimator and uses L2 distance over the state-action latent spaces without the L∞ norm.The KL alternative estimates H(S), H(A), and H(S, A) separately using the differentiable Kozachenko–Leonenko entropy estimator.Source: B. Method Details · p. 17
  • MINE scores state-action pairs with a learned critic trained using joint and marginal samples, and uses an exponential moving average to debias its gradient denominator.The appendix sets MINE's α parameter to 0.9.Source: B. Method Details · p. 18
  • InfoNCE trains state and action encoders with a symmetric objective, while VIP estimates transition progress toward each demonstration's final state.VIP uses an ensemble of five policies for compatibility and uncertainty scores; Policy Loss is negative L2 loss, so lower-loss demonstrations score higher.Source: B. Method Details · p. 18

C. Implementation Details

The implementation appendix specifies model architectures, training resources, hyperparameter reporting, and checkpoint and scoring procedures.

C. Implementation Details

  • State-based experiments use two-layer 512-unit ReLU MLPs, while image experiments use ResNet18 encoders with spatial softmax and camera representations combined with state.RoboMimic uses a three-layer 512-unit MLP; Franka and RoboCrowd use two-layer 1024-unit MLPs, with corresponding architectures for state encoders and VAEs.Source: C. Implementation Details · p. 18
  • Models were trained on TPU v4-8 VMs; 100K training steps took about 30 minutes for VAEs and one hour for other methods.Table I lists hyperparameters, which the authors chose consistently across settings where possible.Source: C. Implementation Details · pp. 18–19, Table I · p. 19
  • k-NN scores use random batches of 1024 over four dataset passes, and checkpoint choices account for faster overfitting by InfoNCE and MINE.State-based models trained for 50K steps and image-based models for 100K; InfoNCE checkpoints were taken at 20K and 40K, and MINE at 50K and 60K, respectively.Source: C. Implementation Details · p. 19

D. Evaluation Details

The evaluation details specify training and scoring protocols for RoboMimic, RoboCrowd, and Franka. The setups differ in trial counts, task scoring, architectures, and training duration.

RoboMimic

  • RoboMimic policies use an MLP action head with L2 loss, train for 100K timesteps, and undergo 200 evaluation episodes truncated at 400 timesteps.The architecture matches the image-based data quality experiments.Source: D. Evaluation Details · p. 19

RoboCrowd

  • RoboCrowd uses ten ALOHA trials per method, scored from 0 to 3 by candy grasping and return outcomes, with ACT-style policies trained for 200K timesteps.The ACT-style setup uses an encoder-decoder transformer, L1 loss, and action chunks of size 100.Source: D. Evaluation Details · p. 19

Franka

  • Franka evaluation uses 15 trials per method, binary task success criteria, and Diffusion Policy with action chunks of 16 and execution size of 8.Compared with Khazatsky et al. [34], the setup uses scratch-initialized ResNet-34s with GroupNorm and trains for 200K rather than 50K steps.Source: D. Evaluation Details · p. 19

Related papers

Loading 2502.08623v3…