Resumen Técnico: DataMIL – Selección de Datos para el Aprendizaje por Imitación de Robots con Datamodels

1. Declaración del Problema

La comunidad de la robótica se ha desplazado recientemente hacia el entrenamiento de políticas fundacionales generalistas utilizando conjuntos de datos masivos y diversos (por ejemplo, Open-X Embodiment). Si bien estas políticas logran un sólido rendimiento medio en tareas variadas, frecuentemente rinden por debajo de lo esperado en tareas individuales especializadas, lo que requiere un ajuste fino (fine-tuning) con nuevos datos específicos de la tarea.

El desafío central abordado por este trabajo es la selección de datos: identificar qué subconjuntos de grandes conjuntos de datos previos, al combinarse con demostraciones limitadas de la tarea específica, producen las mejores políticas especializadas.

• Limitaciones de los Métodos Actuales: La selección ingenua basada en nociones humanas de calidad (por ejemplo, similitud semántica, similitud visual o proximidad estado-acción) a menudo falla. Estos heurísticos asumen que los datos más "similares" son los más "útiles", ignorando el impacto real de un punto de datos en el rendimiento de la política descendente (downstream).
• La Barrera de Optimización: Idealmente, uno seleccionaría los datos mediante la prueba exhaustiva de subconjuntos a través del reentrenamiento y la evaluación de la política resultante. Sin embargo, en robótica, la evaluación de la política requiere ejecuciones (rollouts) en el mundo real, lo cual es lento, peligroso e computacionalmente inviable para una búsqueda a gran escala.
• La Brecha: Los marcos de atribución de datos existentes como los datamodels (que predicen el rendimiento del modelo basándose en los datos de entrenamiento) han tenido éxito en NLP y Visión Computacional, pero no han sido directamente aplicables a la robótica debido a la naturaleza no diferenciable y costosa de la evaluación de políticas en el mundo real.

2. Metodología: DataMIL

Los autores presentan DataMIL (Datamodels para Imitación de Aprendizaje), un marco que extiende el paradigma de los datamodels a la robótica. DataMIL razona sobre la selección de datos de manera de extremo a extremo (end-to-end), utilizando la propia política para identificar los puntos de datos que mejoran el rendimiento, en lugar de depender de heurísticos estáticos.

2.1 Marco Central

DataMIL trata el algoritmo de aprendizaje de la política como una "caja negra" y entrena un estimador (un datamodel) para predecir el rendimiento de una política entrenada en un subconjunto específico de datos.

• Objetivo: Encontrar un subconjunto $D' \subset D$ que maximice la métrica objetivo $M(A(D'))$, donde $A$ es el algoritmo de aprendizaje y $M$ es la métrica de rendimiento.
• Formulación del Datamodel: En lugar de entrenar la política, DataMIL aprende una función $\hat{f}(D') \approx M(A(D'))$. Para datamodels lineales, esto se reduce a asignar una puntuación escalar $\tau(z_i)$ a cada punto de datos $z_i$, donde el rendimiento predicho es la suma de las puntuaciones para el subconjunto seleccionado.

2.2 Adaptaciones Clave para la Robótica

Para hacer que los datamodels sean tratables en robótica, los autores introducen tres modificaciones críticas:

1. Función de Pérdida Sustituta (Evitando Rollouts):
   Las ejecuciones en el mundo real son no diferenciables y costosas. DataMIL reemplaza la métrica real de tasa de éxito por una métrica proxy $\tilde{M}$, definida como la pérdida negativa de clonación de comportamiento (BC) sobre un conjunto de demostraciones objetivo retenido:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   Esta métrica es totalmente diferenciable y no requiere interacción con el entorno, lo que permite el uso de métodos de estimación basados en gradientes.

2. Estimadores de Datamodel:
   El artículo emplea dos métodos para estimar las puntuaciones de influencia $\tau(z_i)$:

• Estimador de Regresión: Muestrea subconjuntos aleatorios de datos, entrena políticas, evalúa la métrica proxy y ajusta un modelo lineal para predecir el rendimiento basado en la presencia de los datos. Es preciso pero computacionalmente pesado.
• Estimador Basado en Metagradiente: Utiliza funciones de influencia y metagradientes para calcular la derivada de la métrica objetivo con respecto a los pesos de los datos. Este método es más eficiente para modelos grandes (como Octo) y evita la necesidad de reentrenar políticas para cada subconjunto.

3. Agrupación (Clustering) y Mitigación del Cambio de Distribución:

• Clustering: Para reducir el ruido en la estimación de la influencia, los pares estado-acción individuales se agrupan en clusters temporales (sub-trayectorias o trayectorias completas). La granularidad óptima del cluster depende del tamaño del conjunto de datos; los conjuntos de datos más grandes (como OXE) se benefician de la agregación a nivel de trayectoria.
• Cambio de Distribución: Para prevenir estimaciones deficientes causadas por cambios entre los dominios previos y el objetivo, se incluye una pequeña fracción de datos de la tarea objetivo durante la fase de estimación del datamodel para alinear el proceso de aprendizaje con el dominio objetivo.

.3 Pipeline de Entrenamiento

Una vez estimadas las puntuaciones de influencia:

1. Se selecciona el $x\%$ superior de los datos previos (mayor influencia positiva) para formar $D_{sel}$.
2. La política final se entrena mediante co-entrenamiento en los datos objetivo $D_{target}$ y los datos seleccionados $D_{sel}$. En cada paso, se muestrea un lote de $D_{target}$ con probabilidad $\alpha$ y de $D_{sel}$ con probabilidad $1-\alpha$.

3. Contribuciones Clave

• Extensión de los Datamodels a la Robótica: DataMIL es el primer marco que aplica con éxito los datamodels al aprendizaje por imitación de robots, superando la barrera de las ejecuciones costosas en el mundo real mediante una pérdida sustituta diferenciable.
• Conciencia de Rendimiento de Extremo a Extremo: A diferencia de los métodos heurísticos (similitud visual, de movimiento o de estado), DataMIL optimiza directamente el impacto en el éxito de la política, lo que le permite seleccionar datos que mejoran el rendimiento y descartar aquellos que lo degradan, incluso si los datos parecen visualmente similares.
• Estimadores Escalables: La introducción del estimador basado en metagradientes permite que el marco escale a políticas grandes y complejas (como Octo) y conjuntos de datos masivos (como OXE), donde los enfoques basados en regresión serían computacionalmente prohibitivos.
• Selección Cross-Embodiment y Multi-Tarea: El método demuestra la capacidad de seleccionar datos útiles para tareas y embodiments de robots que no están presentes en el conjunto de datos previo, y de curar conjuntos de datos que soportan múltiples objetivos descendentes simultáneamente.

4. Resultados Experimentales

Los autores validan DataMIL a través de más de 60 tareas en entornos de simulación y del mundo real:

• MetaWorld (50 Tareas): DataMIL logró un aumento del 10% en el rendimiento sobre las líneas base de vanguardia (Behavior Retrieval, Flow Retrieval, STRAP). Filtró con éxito las demostraciones ruidosas y subóptimas de los datos de exploración autónoma, mientras que los baselines basados en similitud a menudo recuperaron datos irrelevantes o de baja calidad.
• LIBERO (10 Tareas): Utilizando la política de transformador Octo, DataMIL superó consistentemente a los baselines en diversas tareas de largo horizonte. Mientras que los baselines de similitud visual funcionaron bien en algunas tareas específicas, DataMIL proporcionó el rendimiento promedio más robusto.
• Open-X Embodiment (OXE) (Mundo Real):
  • Franka-Ball & Franka-Pouch: DataMIL seleccionó datos que mejoraron significativamente las tasas de éxito en comparación con la selección aleatoria y los baselines heurísticos.
  • Tiago-Sink: En un entorno desafiante de cross-embodiment (seleccionando datos para un robot Tiago de un dataset que no contiene datos de Tiago), DataMIL identificó con éxito demostraciones relevantes de otros robots (ej. RT-1, BC-Z) que compartían la "esencia" de la tarea (manipulación en mesa), mientras que los baselines fallaron en generalizar.
  • Droid-Multitask: DataMIL curó un único conjunto de datos que mejoró el rendimiento en tres tareas distintas simultáneamente, superando a los baselines en la tasa de éxito promedio.

Hallazgos Cualitativos:

• DataMIL selecciona una mezcla diversa de conjuntos de datos en lugar de depender excesivamente de una sola fuente (a diferencia de los baselines que a menudo recuperan más del 80% de un solo dataset).
• Distingue eficazmente entre datos útiles y dañinos incluso cuando se ven visualmente idénticos (ej. mismo estado pero diferentes distribuciones de acción), un modo de fallo común en la recuperación por similitud.

5. Significancia y Reivindicaciones

El artículo sostiene que la selección de datos consciente del rendimiento y de extremo a extremo es crítica para desbloquear el potencial de los grandes conjuntos de datos previos en robótica. Los autores argumentan que, si bien los modelos fundacionales proporcionan una base sólida, su especialización depende fuertemente de la calidad de los datos de ajuste fino.

• Reivindicaciones Modestas sobre Escalabilidad: Los autores reconocen que estimar los datamodels sigue siendo computacionalmente costoso (varias veces el costo de entrenar una política con todos los datos), aunque el enfoque de metagradiente mitiga esto. Señalan que el ajuste de hiperparámetros (ej. tamaño del cluster, ratio de selección) carece actualmente de una fuerte intuición teórica.
• Alcance: El trabajo se centra principalmente en la especialización de tarea única, siendo el entorno Droid-Multitask un paso preliminar hacia la curación multitarea.
• Conclusión: DataMIL demuestra que ir más allá de los heurísticos definidos por humanos hacia una selección impulsada por datos y consciente del modelo produce ganancias consistentes en las tasas de éxito, particularmente en escenarios complejos, heterogéneos y del mundo real donde los métodos de recuperación tradicionales fallan.