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

Licenças

Acesso aberto
Artigo: CC BY 4.0

Tudo o que segue se baseia no artigo completo sob a licença CC BY 4.0; sua figura é mostrada com crédito aos autores.

Em resumo

A Equivariant Diffusion Policy incorpora a simetria de rotação SO(2) ao removedor de ruído da Diffusion Policy, elevando o sucesso médio em simulação em 21,9 pontos com 100 demonstrações e aprendendo tarefas reais a partir de 20-60 demonstrações.

Por que importa

Demonstrações reais são caras, e incorporar a simetria de rotação sobre a mesa permitiu que a política tivesse sucesso em seis tarefas reais, incluindo uma tarefa de horizonte longo de assar bagels, com 20 a 60 demonstrações, enquanto uma baseline de Diffusion Policy com voxels falhou na maioria dos casos. Como a análise cobre ações completas da garra em SE(3), ela se aplica a tarefas que exigem rotação fora do plano, que uma variante SE(2) não conseguia resolver.

Problema

A Diffusion Policy captura o comportamento multimodal das demonstrações, mas precisa aprender uma função de remoção de ruído sobre observações, ações ruidosas e níveis de ruído, o que é mais difícil do que aprender uma política direta de estado para ação e exige mais dados. As políticas de manipulação equivariantes anteriores usavam espaços de ações SE(3) em malha aberta ou SE(2) em malha fechada, e nenhuma aplicava simetria a políticas de difusão.

Solução proposta

Os autores provam que a função de predição de ruído é SO(2)-equivariante sempre que a política especialista o é, e derivam como a rotação em torno do eixo da gravidade atua sobre as ações completas de 6-DoF da garra, tanto com controle de pose absoluto quanto relativo. O removedor de ruído usa codificadores equivariantes de observações e de ações (construídos com a biblioteca escnn) que produzem embeddings para um grupo de rotação discreto, uma U-Net temporal 1D compartilhada aplicada a cada elemento do grupo e um decodificador de ruído equivariante. As entradas podem ser imagens RGB ou uma grade de voxels de 64×64×64.

Principais ideias

Se a política especialista é SO(2)-equivariante, a função de predição de ruído da difusão também é, o que torna as redes equivariantes um viés indutivo correto para o removedor de ruído.
Os ganhos são maiores com poucas demonstrações e em tarefas com poses iniciais dos objetos amplamente aleatorizadas; nas tarefas do Robomimic com pouca aleatoriedade, o método foi semelhante ou ligeiramente melhor que a Diffusion Policy.
Nas ablações, a estrutura equivariante importou mais que a entrada em voxels: com 100 demonstrações, remover a equivariância reduziu o sucesso médio em 17,6 pontos e remover os voxels, em 10,3.
O aumento de dados por rotação em uma CNN padrão ajudou em tarefas simples como Stack, mas ficou bem abaixo da rede equivariante nas mais difíceis.
Um espaço de ações SE(3) completo importa: uma variante SE(2) obteve 0,0 % em Coffee Preparation contra 85,3 % da SE(3), e 12,7 % contra 40,0 % em Threading.

Principais resultados

Em 12 tarefas do MimicGen com controle absoluto e 100 demonstrações, a versão com voxels teve média de 63,9 % de sucesso contra 42,0 % da Diffusion Policy (DiffPo-C), um ganho de 21,9 pontos; a versão RGB teve média de 53,7 %.
Com 200 demonstrações, a versão com voxels (72,6 %) superou todas as baselines treinadas com 1000 demonstrações (a melhor, 71,4 %).
Com controle absoluto, a versão com voxels superou as baselines em 11 de 12 ambientes, sendo a exceção Hammer Cleanup D1.
Em um braço Franka real, o sucesso contra uma baseline de Diffusion Policy com voxels foi: Oven Opening 95 % contra 60 % (20 demonstrações); Banana in Bowl 95 % contra 30 %, Letter Alignment 95 % contra 0 %, Trash Sweeping 90 % contra 5 % (40 demonstrações cada); Hammer to Drawer 85 % contra 5 % (60 demonstrações); Bagel Baking 80 % contra 10 % (58 demonstrações).
Todas as falhas em Bagel Baking vieram dos limites das juntas do robô, e a política generalizou zero-shot para rotações do forno não vistas, exceto uma pose bloqueada pelos limites das juntas.

Limitações

Os autores observam que a equivariância só é parcialmente explorada por causa de um descompasso de simetria no sistema de visão; o braço que aparece na grade de voxels e o ruído da câmera podem quebrar a simetria.
Uma "equivariância incorreta" pode prejudicar o desempenho quando a simetria do modelo entra em conflito com as demonstrações.
A teoria de ações SO(2) só é demonstrada com políticas de difusão, e não com outros pipelines como o BC RNN, que teve bom desempenho com controle relativo.
O método ainda não foi estendido para navegação, locomoção ou manipulação móvel.
Com entrada RGB, a câmera de visão do agente não é de cima para baixo, então a rotação da imagem não corresponde à rotação real do estado; com controle relativo, a versão RGB foi apenas marginalmente melhor que as baselines.

Impacto na área

Finalista do CoRL 2024 Outstanding Paper Award e citado 107 vezes. O artigo indica uma página do projeto em equidiff.github.io. Ele não cita um repositório de código no texto.

Resumo

Trabalhos recentes mostraram que modelos de difusão são uma abordagem eficaz para aprender as distribuições multimodais que surgem dos dados de demonstração na clonagem de comportamento. No entanto, uma desvantagem dessa abordagem é a necessidade de aprender uma função de remoção de ruído, que é significativamente mais complexa do que aprender uma política explícita. Neste trabalho, propomos a Equivariant Diffusion Policy, um novo método de aprendizado de políticas de difusão que aproveita simetrias do domínio para obter melhor eficiência amostral e generalização na função de remoção de ruído. Analisamos teoricamente a simetria SO(2) do controle completo de 6-DoF e caracterizamos quando um modelo de difusão é SO(2)-equivariante. Além disso, avaliamos empiricamente o método em um conjunto de 12 tarefas de simulação no MimicGen e mostramos que ele obtém uma taxa de sucesso que é, em média, 21,9 % maior que a da baseline Diffusion Policy. Também avaliamos o método em um sistema do mundo real para mostrar que políticas eficazes podem ser aprendidas com relativamente poucas amostras de treinamento, ao passo que a baseline Diffusion Policy não consegue.

Open Source Robotics