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
分类: cs.RO
发布日期: 2025-02-12 (更新: 2025-04-22)
备注: Videos and code at https://jhejna.github.io/demonstration-info. Published at RSS 2025
💡 一句话要点
提出基于互信息估计的机器人数据质量评估方法，提升模仿学习性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 机器人 模仿学习 数据质量 互信息估计 变分自编码器
📋 核心要点
- 模仿学习依赖高质量数据集，但机器人数据质量评估研究较少，现有方法难以有效评估数据质量。
- 提出一种基于互信息估计的数据质量评估方法，通过计算轨迹对状态-动作互信息的贡献来衡量数据质量。
- 实验表明，该方法能有效区分数据质量，并能通过过滤低质量数据提升模仿学习策略的性能。
📝 摘要（中文）
模仿学习策略的性能通常取决于训练数据集的质量。本文旨在解决机器人数据质量评估问题，针对给定的演示数据集，评估其中每个演示在动作多样性和可预测性方面的相对质量。具体而言，通过估计轨迹对整个数据集中状态和动作之间互信息的平均贡献来衡量数据质量，该互信息同时捕捉了边缘动作分布的熵和状态条件下的动作熵。针对机器人领域数据量不足的问题，提出了一种基于k近邻互信息估计的新技术，并结合状态和动作的VAE嵌入。实验结果表明，该方法能够根据人类专家评分对演示数据集进行质量划分，并且使用该方法过滤后的数据训练策略，在RoboMimic上性能提升5-10%，并在真实的ALOHA和Franka机器人平台上表现更好。
🔬 方法详解
问题定义：模仿学习的性能高度依赖于训练数据的质量，然而，机器人领域的数据质量评估研究相对滞后。现有的互信息估计方法通常需要大量数据，这在机器人领域中往往难以满足。因此，如何有效地评估机器人演示数据的质量，并利用高质量数据提升模仿学习性能，是一个亟待解决的问题。
核心思路：本文的核心思路是通过估计每个轨迹对整个数据集中状态和动作之间互信息的贡献来衡量数据质量。互信息能够同时捕捉动作的多样性（边缘动作分布的熵）和动作的可预测性（状态条件下的动作熵）。高质量的轨迹应该既包含丰富的动作变化，又能根据状态预测出合理的动作。
技术框架：该方法主要包含以下几个步骤：1) 使用变分自编码器（VAE）对状态和动作进行嵌入，将高维的状态和动作数据映射到低维的潜在空间；2) 利用k近邻（k-NN）方法在潜在空间中估计状态和动作之间的互信息；3) 计算每个轨迹对整体互信息的贡献，作为该轨迹的质量评分；4) 根据质量评分对数据集进行过滤，选择高质量的数据用于训练模仿学习策略。
关键创新：本文的关键创新在于提出了一种适用于机器人数据量较少情况下的互信息估计方法。传统的互信息估计方法需要大量数据才能保证准确性，而本文提出的基于k-NN的互信息估计方法，结合VAE嵌入，能够在数据量有限的情况下，有效地估计状态和动作之间的互信息。这使得该方法能够应用于机器人数据质量评估，并指导数据选择。
关键设计：在VAE的训练过程中，使用了标准的高斯先验和重参数化技巧。k-NN互信息估计器的k值是一个重要的超参数，需要根据数据集的大小和特性进行调整。轨迹对互信息的贡献可以通过多种方式计算，例如，可以计算移除该轨迹后互信息的减少量。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法能够有效区分不同质量的机器人演示数据，并与人类专家评分具有较高的一致性。使用该方法过滤后的数据训练模仿学习策略，在RoboMimic基准测试中性能提升5-10%，并在真实的ALOHA和Franka机器人平台上取得了更好的性能。这些结果验证了该方法在机器人数据质量评估和提升模仿学习性能方面的有效性。
🎯 应用场景
该研究成果可广泛应用于机器人模仿学习领域，帮助研究人员和工程师筛选高质量的训练数据，提升模仿学习策略的性能和泛化能力。此外，该方法还可以用于评估不同数据收集策略的有效性，指导机器人数据采集过程，降低数据收集成本，加速机器人技术的研发和应用。
📄 摘要（原文）
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.