Abstracting Robot Manipulation Skills via Mixture-of-Experts Diffusion Policies

Riepilogo Tecnico: Astrazione di Competenze di Manipolazione Robotica tramite Mixture-of-Experts Diffusion Policies

Definizione del Problema

Le policy basate sulla diffusione (diffusion-based policies) hanno dimostrato un successo straordinario nella manipolazione robotica a compito singolo, modellando la generazione di azioni come un processo di denoising. Tuttavia, estendere queste policy a scenari multi-task presenta sfide significative. Le strategie attuali si affidano spesso all'aumento delle dimensioni del modello e dei dataset di dimostrazioni per ottenere la generalizzazione. Questo approccio comporta costi proibitivi: i modelli sovradimensionati soffrono di velocità di inferenza ridotte, non adatte al controllo in tempo reale, e i dati di dimostrazione richiesti crescono esponenzialmente con la diversità dei compiti. Inoltre, le attuali applicazioni di Mixture-of-Experts (MoE) nella robotica spesso non riescono a disincagliare esplicitamente le competenze di manipolazione riutilizzabili, portando a comportamenti intrecciati e scarsa interpretabilità. La sfida centrale è consentire una generalizzazione multi-task efficace con dimensioni moderate del modello e bassa latenza di campionamento, garantendo al contempo che le competenze siano riutilizzabili e trasferibili.

Metodologia: Skill Mixture-of-Experts Policy (SMP)

Gli autori propongono la Skill Mixture-of-Experts Policy (SMP), un framework basato sulla diffusione che apprende una base di competenze compatta e ortogonale e compone le azioni da un sottoinsieme sparso di esperti.

1. Base di Competenze Ortogonale Adattiva allo Stato

Invece di utilizzare una base di azioni fissa, SMP apprende una base ortonormale dipendente dallo stato $B(s) \in \mathbb{R}^{d \times K}$ (dove $K \ll d$).

• Costruzione: Una rete neurale leggera genera una matrice non vincolata $W(s)$, che viene proiettata sulla varietà di Stiefel tramite una ritrazione thin-QR differenziabile con stabilizzazione del segno. Ciò assicura che i vettori della base rimangano ortonormali ed evolvano continuamente con lo stato del robot (ad esempio, posa del braccio, contatto), preservando il disincaglio locale.
• Decomposizione dell'Azione: Le azioni vengono decodificate come $a_t = B(s_t)(g_t \odot z_t)$, dove $g_t$ sono i pesi di gating e $z_t$ sono i coefficienti delle competenze. Questa decomposizione ortogonale assicura che ogni esperto contribuisca a una direzione distinta e non sovrapponibile nello spazio delle azioni.

2. Routing "Sticky" e Utilizzo Globale

Per affrontare il problema dello switching rapido e instabile tra gli esperti (chatter), SMP impiega una dinamica di Dirichlet-Markov "sticky" per il meccanismo di gating.

• Consistenza Temporale: Il gate $g_t$ è modellato come un processo di Markov in cui lo stato corrente dipende dal gate precedente $g_{t-1}$ e da un vettore di utilizzo globale $\vartheta$. Ciò incoraggia attivazioni di competenze a tratti costanti, allineandosi con le fasi quasi stazionarie dei compiti di manipolazione (ad esempio, afferrare, muovere, rilasciare).
• Utilizzo Globale: Un vettore globale $\vartheta$ cattura la prevalenza complessiva delle competenze, impedendo alla policy di collassare in un singolo esperto e mantenendo un utilizzo bilanciato tra i vari compiti.

3. Obiettivo di Addestramento Variazionale

SMP è addestrata tramite inferenza variazionale per massimizzare l'evidence lower bound (ELBO), che si decompone in tre componenti:

• Errore di Ricostruzione ($L_{recon}$): Assicura che l'azione sintetizzata nella base bianca corrisponda alla verità di base (ground truth). Fondamentalmente, l'addestramento utilizza due target di coefficienti: un target con stop-gradient per il surrogato di diffusione e un target con gradiente per la ricostruzione dell'azione, permettendo alla base delle competenze $B(s)$ di essere aggiornata senza destabilizzare il processo di diffusione.
• Regolarizzazione del Gate ($L_{gate}$): Allinea la posterior del gate ammortizzato con il prior di Dirichlet "sticky", imponendo vincoli di persistenza temporale e di utilizzo globale.
• Regolarizzazione dei Coefficienti ($L_{coeff}$): Allinea i coefficienti delle competenze con i loro prior di diffusione.

4. Attivazione Adattiva degli Esperti in Inferenza

Per ridurre il costo computazionale, SMP non valuta tutti gli esperti ad ogni step.

• Router Basato Solo sullo Stato: Un router leggero predice i pesi del gate basandosi esclusivamente sullo stato $s_t$.
• Selezione Sparsa: Gli esperti sono classificati in base alla loro "massa" ($m_i = \bar{g}_{t,i}^2$). Il sistema attiva solo il top-$k$ esperti o il sottoinsieme più piccolo che cattura una soglia (ad esempio, il 95%) della massa totale. Ciò produce un'attivazione sparsa e dipendente dallo stato, riducendo significativamente i parametri attivi e la latenza.

Contributi Chiave

1. Framework SMP: Una policy MoE basata sulla diffusione che astrae esplicitamente competenze di manipolazione riutilizzabili tramite una base di azioni ortonormale dipendente dallo stato con routing "sticky".
2. Strategia di Attivazione Adattiva: Un meccanismo di selezione dinamica che attiva un sottoinsieme compatto di esperti durante l'inferenza, riducendo il costo computazionale pur mantenendo l'accuratezza del campionamento delle azioni.
3. Validazione Empirica: Valutazione completa in simulazione (RoboTwin-2, RLBench-2) e su una piattaforma reale a doppio braccio, dimostrando tassi di successo multi-task superiori e costi di inferenza inferiori rispetto a forti baseline di diffusione.

Risultati Sperimentali

Gli autori hanno valutato SMP contro vanilla diffusion policies (DP, DP3), policy basate su transformer (ACT), grandi modelli fondazionali (RDT) e altre baseline di astrazione delle competenze o MoE.

• Apprendimento Multi-Task: SMP ha ottenuto tassi di successo costantemente più elevati in compiti bimanuali (ad esempio, sollevare vassoi, passaggi di oggetti, spazzare) rispetto alle baseline. Mentre la scalabilità di RDT (un modello da 1.2B di parametri) ha migliorato il successo di solo il 19% rispetto a una baseline più piccola, SMP ha ottenuto prestazioni superiori con un footprint significativamente minore.
• Efficienza: SMP ha ridotto il tempo di inferenza e i parametri attivi. Ha attivato circa il 30% dei propri parametri (circa il 7% dei parametri totali di RDT) mantenendo un alto successo. Al contrario, le baseline dense e le varianti FFN-MoE hanno sofferto di una latenza maggiore a causa della necessità di sincronizzare o valutare tutti gli esperti.
• Apprendimento per Trasferimento (Transfer Learning):
  • Adattamento Few-Shot: Quando perfezionata (fine-tuned) su nuovi compiti con sole 10 dimostrazioni, SMP ha superato le baseline riutilizzando efficacementmente le competenze rilevanti senza dover riaddestrare l'intero backbone.
  • Composizione delle Competenze: In compiti che richiedono la ricomposizione di competenze apprese (ad esempio, mettere le patatine in una padella), il fine-tuning del solo router di SMP ha permesso di ricomporre con successo le competenze, mentre le baseline hanno mostrato ruoli sfumati e fasi inconsistenti.
• Prestazioni su Robot Reale: Su una piattaforma dual-arm PiPER, SMP ha ottenuto i punteggi di progresso più alti in quattro compiti bimanuali complessi (ad esempio, rimuovere tappi di penne, versare fagioli) con una latenza inferiore rispetto alle baseline di diffusione.

Significato e Rivendicazioni

Il paper sostiene che SMP fornisca una via pratica verso una manipolazione multi-task scalabile e trasferibile. Imparando le competenze riutilizzabili una sola volta e attivando solo ciò che è necessario, SMP affronta il compromesso tra capacità del modello ed efficienza di inferenza. Gli autori sottolineano che le competenze apprese si allineano con ruoli di manipolazione intuitivi (braccio sinistro/destro, fasi di pick/move/place) e che il meccanismo di attivazione adattiva offre un semplice controllo per scambiare accuratezza contro latenza. Il lavoro suggerisce che il disincaglio esplicito delle competenze e il routing coerente con le fasi siano critici per un controllo bimanuale robusto, offrendo un'alternaria più efficiente rispetto alla scalabilità brutale dei modelli.

Gli autori riconoscono i limiti, notando che lo studio attuale si concentra sulla manipolazione bimanuale con backbone relativamente piccoli. Si propone come lavoro futuro la scalabilità di SMP a modelli più grandi, l'estensione della valutazione alla manipolazione a braccio singolo e mobile, e la conduzione di studi più estesi su robot reali per valutare la robustezza rispetto al rumore di sensing e al domain shift.