ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Autores: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

Licencias

• Acceso abierto
• Artículo: CC BY 4.0

Todo lo que sigue se basa en el artículo completo bajo su licencia CC BY 4.0.

En resumen

ReMix usa optimización distribucionalmente robusta por grupos sobre la pérdida excedente de una política para aprender cuánto muestrear cada conjunto de datos robótico, y supera a los pesos uniformes y a los pesos expertos de RT-X en robots reales.

Por qué importa

Los equipos que preentrenan con datos robóticos agrupados pueden sustituir los pesos de muestreo ajustados a mano por pesos calculados a partir de los datos y, en las pruebas presentadas, pueden descartar el 75 % de los datos con una pérdida de rendimiento limitada. Los pesos aprendidos también muestran qué conjuntos de datos conviene ponderar a la baja y por qué, por ejemplo los datos de Kuka recopilados de forma autónoma, lo que ayuda a decidir qué datos recopilar o conservar.

Problema

Los grandes conjuntos de datos robóticos como Open X-Embodiment agrupan datos de distintos robots, cámaras, espacios de acciones y controladores, pero las mezclas usadas para entrenar políticas generalistas como RT-X, Octo y OpenVLA se eligieron a mano. Los métodos de filtrado de datos de visión y lenguaje se transfieren mal, porque los n-gramas no son aplicables y los embeddings visuales ignoran las acciones y la naturaleza secuencial de los episodios.

Solución propuesta

ReMix plantea la curación como un problema min-max: una política discreta minimiza su pérdida excedente de clonación de comportamiento respecto a un modelo de referencia, mientras que los pesos de dominio se actualizan por ascenso de gradiente exponenciado para maximizarla, de modo que el peso va a los dominios en los que la política todavía puede mejorar. Para que las pérdidas sean comparables entre robots, las acciones se normalizan con una gaussiana por dominio y se discretizan en intervalos, y el checkpoint de referencia se elige por parada temprana antes de que se sobreajuste a algún dominio. Los pesos promediados durante el entrenamiento se reutilizan después para entrenar políticas de difusión independientes o para submuestrear los datos.

Ideas clave

• Usar la pérdida excedente respecto a una política de referencia, en lugar de la pérdida bruta, pondera a la baja tanto los dominios que ya están ajustados como los que son demasiado difíciles de ajustar.
• Las diferencias de escala de las acciones sesgan la optimización robusta: en una prueba de juguete con normalización por límites, el 94,3 % del peso fue a un dominio de acciones puramente aleatorias; con normalización gaussiana por dominio, los datos reales recibieron el 84,2 %.
• Las pérdidas continuas L1/L2 sobre las acciones se ven influidas por valores atípicos y acciones multimodales, por lo que discretizar las acciones da pesos más fiables incluso cuando la política final es un modelo de difusión continuo.
• Los conjuntos de datos robóticos (unas 10k-100k demostraciones) son lo bastante pequeños para que el modelo de referencia se sobreajuste; usar un checkpoint 50K pasos después de la parada temprana redujo el rendimiento en más de un 15 % de media.

Resultados principales

• Con la mezcla de datos de RT-X, las políticas entrenadas con los pesos de ReMix superaron a los pesos uniformes en un 38 % y a los pesos de expertos humanos de RT-X en un 32 % de media en cuatro tareas reales con WidowX y dos con Franka; los pesos expertos superaron a los uniformes en un 6 %.
• En subconjuntos del 25 % de la mezcla de RT-X, ReMix mantuvo el rendimiento mientras que los pesos humanos cayeron más de un 10 %; ReMix conservó el 72 % de los datos de UR5 y el 12 % de los de Kuka, frente al 30 % y el 24 % de los pesos humanos.
• ReMix elevó el pequeño conjunto de datos Toto del 3,42 % (uniforme) al 16,3 % y redujo Kuka del 24,9 % al 12,1 %.
• Con el conjunto de datos Bridge V2 completo no hubo una diferencia notable entre ReMix y los pesos uniformes; en subconjuntos del 25 % de Bridge, ReMix superó a los pesos uniformes y a la referencia SSP basada en embeddings de imagen.
• En una tarea de WidowX con un cubo y un objeto distractor, los pesos de ReMix tuvieron éxito 7/10 frente a 0/10 tanto con pesos humanos como uniformes.
• En simulación (Robomimic NutAssemblySquare dividido por operador), ReMix obtuvo 77/100 frente a 53/100 con pesos uniformes al submuestrear al 50 %, y 59/100 frente a 39/100 al 25 %.

Limitaciones

• Los autores señalan que la evaluación en el mundo real abarcó solo dos brazos, WidowX y Franka, y dejan más embodiments para trabajos futuros.
• ReMix pondera al alza conjuntos de datos con distribuciones de acciones anómalas y multimodales como Toto, lo que según los autores no es necesariamente deseable.
• Calcular los pesos exige entrenar dos veces con todos los datos, una para el modelo de referencia y otra para la optimización robusta.
• Se probó con Bridge V2 y con el subconjunto RT-X de OpenX, no con el conjunto de datos OpenX completo de más de 2M episodios. La sección 4.1 describe Bridge V2 como 50k demostraciones, mientras que el apéndice B.2 indica que se usaron unas 45.000 trayectorias de las 60K del conjunto de datos completo.
• Las políticas entrenadas solo con los datos de preentrenamiento no generalizaron zero-shot a las configuraciones de los autores, por lo que toda la evaluación se basa en un coentrenamiento con 25 demostraciones del propio dominio por tarea.

Impacto en el campo

Finalista del CoRL 2024 Outstanding Paper Award y citado 74 veces. Los autores publican el código en github.com/jhejna/remix. El artículo también publica sus pesos de dominio aprendidos para la mezcla de RT-X (tabla 2) y para los 32 dominios de Bridge V2 (tabla 5).

Resumen

Se están recopilando conjuntos de datos de aprendizaje por imitación cada vez más grandes con el objetivo de entrenar modelos fundacionales para robótica. Sin embargo, pese a que la selección de datos ha sido de suma importancia en visión y en procesamiento del lenguaje natural, pocos trabajos en robótica se han preguntado con qué datos deberían entrenarse realmente estos modelos. En este trabajo investigamos cómo ponderar distintos subconjuntos o “dominios” de los conjuntos de datos robóticos para el preentrenamiento de modelos fundacionales para robots. En concreto, usamos optimización distribucionalmente robusta (DRO) para maximizar el rendimiento en el peor caso en todos los posibles dominios de aplicación posteriores. Nuestro método, Re-Mix, aborda la amplia gama de retos que surgen al aplicar DRO a conjuntos de datos robóticos, incluida la variabilidad de los espacios de acciones y de la dinámica entre conjuntos de datos. Re-Mix emplea parada temprana, normalización de acciones y discretización para contrarrestar estos problemas. Mediante una amplia experimentación con el mayor conjunto de datos de código abierto de manipulación robótica, el conjunto de datos Open X-Embodiment, demostramos que la curación de datos puede tener un impacto desproporcionado en el rendimiento posterior. En concreto, los pesos de dominio aprendidos por Re-Mix superan a los pesos uniformes en un 38 % de media y a los pesos elegidos por personas en un 32 % en los conjuntos de datos usados para entrenar políticas robóticas generalistas existentes, en concreto los modelos RT-X.

Enlaces

• arXiv 2408.14037: https://arxiv.org/abs/2408.14037
• DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
• Código: https://github.com/jhejna/remix