PocketDP3: Efficient Pocket-Scale 3D Visuomotor Policy

Sintesi Tecnica: PocketDP3: Una Politica Visuomotoria 3D Efficiente su Scala "Pocket"

Problematica
I recenti progressi nella manipolazione robotica hanno sfruttato le politiche di diffusione basate sulla visione 3D per apprendere abilità complesse dai dati. Tuttavia, gli autori identificano un disallineamento architettonico critico nei modelli allo stato dell'arte (SOTA) attuali, come DP3. Questi modelli accoppiano tipicamente un encoder di nuvole di punti leggero ed efficiente con un decoder backbone eccessivamente grande, spesso una U-Net condizionale. Sebbene le U-Net siano adatte alla sintesi di immagini ad alta risoluzione dove è richiesta una ricostruzione dei pixel fine, gli autori sostengono che questo design sia inefficiente per il controllo visuomotorio 3D. In questo contesto, l'encoder produce già una rappresentazione della scena compatta e ricca di informazioni, e l'obiettivo della predizione è una traiettoria di azione a bassa frequenza. Di conseguenza, il massiccio decoder introduce una ridondanza di parametri sostanziale, aumenta l'impronta di memoria e ostacola l'implementazione in tempo reale su robot con risorse limitate, senza fornire guadagni proporzionali nelle prestazioni di controllo. Inoltre, il campionamento della diffusione standard richiede molti passaggi iterativi di denoising, creando colli di bottiglia nella latenza che spesso richiedono tecniche complesse di distillazione della coerenza per accelerare l'inferenza.

Metodologia
Il paper propone PocketDP3, una politica di diffusione 3D su scala "pocket" progettata per risolvere lo squilibrio architettonico e accelerare l'inferenza.

• Architettura (Diffusion Mixer - DiM): Invece di una pesante U-Net condizionale, PocketDP3 impiega un decoder leggero denominato Diffusion Mixer (DiM) basato su blocchi MLP-Mixer. L'architettura mantiene l'efficiente encoder di nuvole di punti di DP3 (utilizzando Farthest Point Sampling e MLP per generare una rappresentina 3D compatta), ma sostituisce il decoder con $K$ blocchi DiM impilati.

  • Meccanismo di Fusione: I blocchi DiM eseguono una fusione delle informazioni efficiente alternando tra mixing temporale (aggregando informazioni attraverso la dimensione temporale della traiettoria) e mixing dei canali (aggregando attraverso le dimensioni delle feature).
  • Condizionamento: Il modello incorpora le rappresentazioni della scena e gli stati del robot tramite una proiezione lineare e inietta queste informazioni di condizionamento nei blocchi DiM utilizzando strati FiLM (Feature-wise Linear Modulation) combinati con connessioni residue a gate.
  • Efficienza dei Parametri: Gli autori propongono due configurazioni, "Tiny" (0,53M di parametri) e "Base" (1,73M di parametri), che rappresentano meno dell'1% dei parametri utilizzati da metodi precedenti come DP3 (~250M di parametri).

• Strategia di Inferenza: Il metodo utilizza il campionamento in stile DDIM. Notevolmente, gli autori dimostrano che PocketDP3 supporta l'inferenza a due step (2 valutazioni della funzione) senza richiedere ulteriore distillazione della coerenza o obiettivi di addestramento specializzati (ad esempio, One-Step Diffusion o Modelli di Coerenza). Ciò riduce significativamente la latenza di inferenza mantenendo le prestazioni.

• Addestramento: Il modello viene addestrato utilizzando un obiettivo di predizione $x_0$, minimizzando l'errore quadratico medio tra la traiettoria pulita predetta e la traiettoria esperta reale, condizionata dall'osservazione e dal timestep.

Contributi Chiave

1. Intuizione Architetturale: Il paper identifica che l'accoppiamento di encoder 3D compatti con massicci decoder U-Net è un'inefficienza diffusa nelle politiche di diffusione 3D, che porta a uno spreco di parametri.
2. Architettura PocketDP3: L'introduzione di una politica su scala pocket che sostituisce la U-Net con un Diffusion Mixer (DiM), consentendo un efficace mixing temporale e dei canali con un budget di parametri minimo.
3. Campionamento Veloce Senza Distillazione: Dimostrazione che un'inferenza ad alte prestazioni in due step è ottenibile senza la distillazione della coerenza, semplificando la pipeline di addestramento e consentendo l'implementazione a bassa latenza.
4. Validazione Empirica: Valutazione completa che mostra come PocketDP3 eguagli o superi le prestazioni SOTA riducendo drasticamente la dimensione del modello e i tempi di inferenza.

Risultati Sperimentali
Il metodo è stato valutato su tre benchmark di simulazione (RoboTwin2.0, Adroit e MetaWorld) ed esperimenti nel mondo reale.

• Prestazioni in Simulazione:

  • RoboTwin2.0: PocketDP3-Base ha raggiunto un tasso di successo medio del 71,6%, superando il precedente SOTA (DP3) del 20,8%. Anche PocketDP3-Tiny ha superato DP3 del 15,2%.
  • Adroit & MetaWorld: PocketDP3-Base ha raggiunto un tasso di successo medio del 77,4%, superando DP3 (73,0%) e FlowPolicy (72,6%).
  • Studi di Ablazione: I confronti con un "Vanilla-MLP" (che non è riuscito a convergere) e un "Vanilla-U-Net" ridotto (che ha ottenuto prestazioni inferiori) hanno confermato che il meccanismo di fusione specifico del blocco DiM è critico per le prestazioni sotto rigorosi budget di parametri.
  • Step di Inferenza: L'ablatore sul numero di valutazioni della funzione (NFE) ha mostrato che le prestazioni si stabilizzano dopo 2 step; aumentare a 5 o 10 step non ha prodotto guadagni significativi.

• Efficienza:

  • Dimensione del Modello: PocketDP3-Base utilizza 1,73M di parametri, rispetto ai 255,1M di DP3 (una riduzione di >99%).
  • Latenza: Con 2 step di inferenza, PocketDP3-Base raggiunge una latenza di 4,80ms su una RTX 5880 Ada, significativamente più veloce di DP3 (51,4ms con 10 step) e paragonabile a FlowPolicy (7,04ms con 1 step), ma con due ordini di grandezza in meno di parametri.

• Implementazione nel Mondo Reale:

  • Testato su un robot AgileX Piper con una telecamera RealSense D455.
  • PocketDP3-Base ha dimostrato un forte trasferimento sim-to-real, superando il baseline DP3 del 15,3% in media su task come il posizionamento di oggetti e l'impilamento di blocchi.

Significato e Rivendicazioni
Il paper sostiene che PocketDP3 offra una via pratica verso un controllo visuomotorio 3D efficiente. Sfidando l'assunto che grandi decoder siano necessari per le politiche basate sulla diffusione quando la rappresentazione dell'input è già compressa, gli autori dimostrano che "più piccolo è meglio" in questo specifico contesto. Il metodo raggiunge prestazioni allo stato dell'arte con meno dell'1% dei parametri dei metodi precedenti e consente l'inferenza in tempo reale senza complessi addestramenti di distillazione. Gli autori posizionano questo lavoro come un passo verso l'implementazione di sofisticate politiche robotiche su hardware con risorse limitate.

Limitazioni e Lavoro Futuro
Gli autori riconoscono che l'architettura DiM potrebbe affrontare un "collo di bottiglia spaziale" per task che richiedono un ragionamento geometrico locale fine (ad esempio, inserimento sub-millimetrico) a causa della mancanza di bias spaziali allineati ai pixel tipici delle U-Net. Notano inoltre che il meccanismo sottostante che permette un'inferenza a due step di alta qualità rimane poco chiaro e richiede ulteriori indagini. Il lavoro futuro suggerito si concentra sulla progettazione di encoder di condizionamento più efficienti e ricchi di informazioni e sullo studio sistematico dei compromessi tra capacità di condizionamento, complessità del task e scala del decoder.