Equivariant Diffusion Policy

Autores: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · 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; su figura se muestra con crédito a los autores.

En resumen

Equivariant Diffusion Policy incorpora la simetría de rotación SO(2) en el eliminador de ruido de Diffusion Policy, eleva el éxito medio en simulación 21,9 puntos con 100 demostraciones y aprende tareas reales a partir de 20-60 demostraciones.

Por qué importa

Las demostraciones reales son caras, e incorporar la simetría de rotación sobre la mesa permitió a la política tener éxito en seis tareas reales, incluida una tarea de horizonte largo de hornear bagels, con 20 a 60 demostraciones, donde una referencia de Diffusion Policy con vóxeles fracasó en la mayoría de los casos. Como el análisis cubre acciones completas de la pinza en SE(3), se aplica a tareas que requieren rotación fuera del plano, que una variante SE(2) no podía resolver.

Problema

Diffusion Policy captura el comportamiento multimodal de las demostraciones, pero debe aprender una función de eliminación de ruido sobre observaciones, acciones con ruido y niveles de ruido, lo que es más difícil que aprender una política directa de estado a acción y requiere más datos. Las políticas de manipulación equivariantes anteriores usaban espacios de acciones SE(3) en lazo abierto o SE(2) en lazo cerrado, y ninguna aplicaba la simetría a las políticas de difusión.

Solución propuesta

Los autores demuestran que la función de predicción de ruido es SO(2)-equivariante siempre que lo sea la política experta, y deducen cómo actúa la rotación alrededor del eje de la gravedad sobre las acciones completas de 6-DoF de la pinza, tanto con control de pose absoluto como relativo. El eliminador de ruido usa codificadores equivariantes de observaciones y de acciones (construidos con la biblioteca escnn) que producen embeddings para un grupo de rotación discreto, una U-Net temporal 1D compartida aplicada a cada elemento del grupo y un decodificador de ruido equivariante. Las entradas pueden ser imágenes RGB o una rejilla de vóxeles de 64×64×64.

Ideas clave

• Si la política experta es SO(2)-equivariante, también lo es la función de predicción de ruido de la difusión, lo que convierte a las redes equivariantes en un sesgo inductivo correcto para el eliminador de ruido.
• Las mejoras son mayores con pocas demostraciones y en tareas con poses iniciales de los objetos muy aleatorizadas; en las tareas de Robomimic con poca aleatoriedad, el método fue similar o ligeramente mejor que Diffusion Policy.
• En las ablaciones, la estructura equivariante importó más que la entrada de vóxeles: con 100 demostraciones, quitar la equivarianza redujo el éxito medio 17,6 puntos y quitar los vóxeles, 10,3.
• El aumento de datos por rotación en una CNN estándar ayudó en tareas sencillas como Stack, pero quedó muy por debajo de la red equivariante en las más difíciles.
• Un espacio de acciones SE(3) completo importa: una variante SE(2) obtuvo un 0,0 % en Coffee Preparation frente al 85,3 % de SE(3), y un 12,7 % frente al 40,0 % en Threading.

Resultados principales

• En 12 tareas de MimicGen con control absoluto y 100 demostraciones, la versión con vóxeles promedió un 63,9 % de éxito frente al 42,0 % de Diffusion Policy (DiffPo-C), una mejora de 21,9 puntos; la versión RGB promedió un 53,7 %.
• Con 200 demostraciones, la versión con vóxeles (72,6 %) superó a todas las referencias entrenadas con 1000 demostraciones (la mejor, 71,4 %).
• Con control absoluto, la versión con vóxeles superó a las referencias en 11 de 12 entornos; la excepción fue Hammer Cleanup D1.
• En un brazo Franka real, el éxito frente a una referencia de Diffusion Policy con vóxeles fue: Oven Opening 95 % frente a 60 % (20 demostraciones); Banana in Bowl 95 % frente a 30 %, Letter Alignment 95 % frente a 0 %, Trash Sweeping 90 % frente a 5 % (40 demostraciones cada una); Hammer to Drawer 85 % frente a 5 % (60 demostraciones); Bagel Baking 80 % frente a 10 % (58 demostraciones).
• Todos los fallos de Bagel Baking se debieron a los límites articulares del robot, y la política generalizó zero-shot a rotaciones del horno no vistas, salvo en una pose bloqueada por los límites articulares.

Limitaciones

• Los autores señalan que la equivarianza solo se aprovecha en parte debido a un desajuste de simetría en el sistema de visión; el brazo que aparece en la rejilla de vóxeles y el ruido de la cámara pueden romper la simetría.
• Una "equivarianza incorrecta" puede perjudicar el rendimiento cuando la simetría del modelo entra en conflicto con las demostraciones.
• La teoría de acciones SO(2) solo se muestra con políticas de difusión, no con otros pipelines como BC RNN, que rindió bien con control relativo.
• El método aún no se ha extendido a la navegación, la locomoción ni la manipulación móvil.
• Con entrada RGB, la cámara de vista del agente no es cenital, por lo que la rotación de la imagen no coincide con la rotación real del estado; con control relativo, la versión RGB solo fue marginalmente mejor que las referencias.

Impacto en el campo

Finalista del CoRL 2024 Outstanding Paper Award y citado 107 veces. El artículo remite a una página del proyecto en equidiff.github.io. No nombra un repositorio de código en su texto.

Resumen

Trabajos recientes han mostrado que los modelos de difusión son un enfoque eficaz para aprender las distribuciones multimodales que surgen de los datos de demostración en la clonación de comportamiento. Sin embargo, un inconveniente de este enfoque es la necesidad de aprender una función de eliminación de ruido, que es significativamente más compleja que aprender una política explícita. En este trabajo proponemos Equivariant Diffusion Policy, un nuevo método de aprendizaje de políticas de difusión que aprovecha las simetrías del dominio para obtener una mejor eficiencia muestral y generalización en la función de eliminación de ruido. Analizamos teóricamente la simetría SO(2) del control completo de 6-DoF y caracterizamos cuándo un modelo de difusión es SO(2)-equivariante. Además, evaluamos empíricamente el método en un conjunto de 12 tareas de simulación en MimicGen y mostramos que obtiene una tasa de éxito que, de media, es un 21,9 % superior a la de la referencia Diffusion Policy. También evaluamos el método en un sistema del mundo real para mostrar que se pueden aprender políticas eficaces con relativamente pocas muestras de entrenamiento, mientras que la referencia Diffusion Policy no puede.

Enlaces

• arXiv 2407.01812: https://arxiv.org/abs/2407.01812
• DOI 10.48550/arXiv.2407.01812: https://doi.org/10.48550/arXiv.2407.01812
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/wang25a/wang25a.pdf
• Página del proyecto: https://equidiff.github.io/