Equivariant Diffusion Policy

Autori: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · CoRL 2024

Licenze

• Open access
• Articolo: CC BY 4.0

Tutto ciò che segue si basa sull'articolo completo secondo la sua licenza CC BY 4.0; la figura è riportata citando gli autori.

In breve

Equivariant Diffusion Policy integra la simmetria di rotazione SO(2) nel denoiser di Diffusion Policy, aumentando il successo medio in simulazione di 21,9 punti con 100 dimostrazioni e apprendendo compiti reali da 20-60 dimostrazioni.

Perché è importante

Le dimostrazioni reali sono costose, e integrare la simmetria di rotazione sul piano del tavolo ha permesso alla policy di riuscire in sei compiti reali, tra cui un compito a lungo orizzonte di cottura di bagel, con 20-60 dimostrazioni, dove una baseline Diffusion Policy a voxel ha per lo più fallito. Poiché l'analisi copre azioni complete della pinza in SE(3), si applica a compiti che richiedono rotazioni fuori dal piano, che una variante SE(2) non riusciva a risolvere.

Problema

Diffusion Policy cattura il comportamento multimodale delle dimostrazioni, ma deve apprendere una funzione di denoising su osservazioni, azioni rumorose e livelli di rumore, il che è più difficile che apprendere una policy diretta da stato ad azione e richiede più dati. Le precedenti policy di manipolazione equivarianti usavano spazi delle azioni SE(3) ad anello aperto o SE(2) ad anello chiuso, e nessuna applicava la simmetria alle policy di diffusione.

Soluzione proposta

Gli autori dimostrano che la funzione di predizione del rumore è SO(2)-equivariante ogni volta che lo è la policy esperta, e ricavano come la rotazione attorno all'asse di gravità agisce sulle azioni complete a 6-DoF della pinza, sia con controllo di posa assoluto sia relativo. Il denoiser usa encoder equivarianti delle osservazioni e delle azioni (costruiti con la libreria escnn) che producono embedding per un gruppo di rotazioni discreto, una U-Net temporale 1D condivisa applicata a ciascun elemento del gruppo e un decoder del rumore equivariante. Gli input possono essere immagini RGB o una griglia di voxel 64×64×64.

Idee chiave

• Se la policy esperta è SO(2)-equivariante, lo è anche la funzione di predizione del rumore della diffusione, il che rende le reti equivarianti un bias induttivo corretto per il denoiser.
• I guadagni sono maggiori con poche dimostrazioni e nei compiti con pose iniziali degli oggetti fortemente randomizzate; nei compiti Robomimic a bassa casualità il metodo è risultato simile o leggermente migliore di Diffusion Policy.
• Nelle ablazioni la struttura equivariante ha contato più dell'input a voxel: con 100 dimostrazioni, togliere l'equivarianza ha ridotto il successo medio di 17,6 punti e togliere i voxel di 10,3.
• La data augmentation per rotazione su una CNN standard ha aiutato in compiti semplici come Stack, ma è rimasta molto al di sotto della rete equivariante in quelli più difficili.
• Uno spazio delle azioni SE(3) completo conta: una variante SE(2) ha ottenuto lo 0,0 % su Coffee Preparation contro l'85,3 % di SE(3), e il 12,7 % contro il 40,0 % su Threading.

Risultati principali

• Su 12 compiti MimicGen con controllo assoluto e 100 dimostrazioni, la versione a voxel ha ottenuto in media il 63,9 % di successo contro il 42,0 % di Diffusion Policy (DiffPo-C), un guadagno di 21,9 punti; la versione RGB ha ottenuto in media il 53,7 %.
• Con 200 dimostrazioni la versione a voxel (72,6 %) ha battuto tutte le baseline addestrate con 1000 dimostrazioni (la migliore 71,4 %).
• Con controllo assoluto la versione a voxel ha superato le baseline in 11 ambienti su 12, con l'eccezione di Hammer Cleanup D1.
• Su un braccio Franka reale, il successo rispetto a una baseline Diffusion Policy a voxel è stato: Oven Opening 95 % contro 60 % (20 dimostrazioni); Banana in Bowl 95 % contro 30 %, Letter Alignment 95 % contro 0 %, Trash Sweeping 90 % contro 5 % (40 dimostrazioni ciascuno); Hammer to Drawer 85 % contro 5 % (60 dimostrazioni); Bagel Baking 80 % contro 10 % (58 dimostrazioni).
• Tutti i fallimenti di Bagel Baking sono dovuti ai limiti dei giunti del robot, e la policy ha generalizzato zero-shot a rotazioni del forno mai viste, tranne una posa bloccata dai limiti dei giunti.

Limiti

• Gli autori notano che l'equivarianza è sfruttata solo in parte a causa di un disallineamento di simmetria nel sistema di visione; il braccio che compare nella griglia di voxel e il rumore della telecamera possono rompere la simmetria.
• Un'"equivarianza errata" può peggiorare le prestazioni quando la simmetria del modello è in conflitto con le dimostrazioni.
• La teoria delle azioni SO(2) è mostrata solo con policy di diffusione, non con altre pipeline come BC RNN, che ha dato buoni risultati con il controllo relativo.
• Il metodo non è ancora esteso a navigazione, locomozione o manipolazione mobile.
• Con input RGB la telecamera dalla vista dell'agente non è dall'alto, quindi la rotazione dell'immagine non corrisponde alla vera rotazione dello stato; con controllo relativo la versione RGB è stata solo marginalmente migliore delle baseline.

Impatto sul settore

Finalista del CoRL 2024 Outstanding Paper Award e citato 107 volte. L'articolo rimanda a una pagina del progetto su equidiff.github.io. Nel testo non indica un repository di codice.

Abstract

Lavori recenti hanno mostrato che i modelli di diffusione sono un approccio efficace per apprendere le distribuzioni multimodali che emergono dai dati di dimostrazione nel behavior cloning. Tuttavia, uno svantaggio di questo approccio è la necessità di apprendere una funzione di denoising, che è significativamente più complessa dell'apprendimento di una policy esplicita. In questo lavoro proponiamo Equivariant Diffusion Policy, un nuovo metodo di apprendimento di policy di diffusione che sfrutta le simmetrie del dominio per ottenere una migliore efficienza campionaria e generalizzazione nella funzione di denoising. Analizziamo teoricamente la simmetria SO(2) del controllo completo a 6-DoF e caratterizziamo quando un modello di diffusione è SO(2)-equivariante. Valutiamo inoltre empiricamente il metodo su un insieme di 12 compiti di simulazione in MimicGen e mostriamo che ottiene un tasso di successo in media superiore del 21,9 % rispetto alla baseline Diffusion Policy. Valutiamo il metodo anche su un sistema reale per mostrare che si possono apprendere policy efficaci con relativamente pochi campioni di addestramento, mentre la baseline Diffusion Policy non ci riesce.

Link

• 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
• Pagina del progetto: https://equidiff.github.io/