ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Autori: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · 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.

In breve

ReMix usa l'ottimizzazione distribuzionalmente robusta per gruppi sulla loss in eccesso di una policy per apprendere quanto campionare ciascun dataset robotico, battendo sui robot reali i pesi uniformi e i pesi degli esperti di RT-X.

Perché è importante

I team che fanno pre-addestramento su dati robotici aggregati possono sostituire i pesi di campionamento tarati a mano con pesi calcolati dai dati e, nelle prove riportate, possono scartare il 75 % dei dati con una perdita di prestazioni limitata. I pesi appresi mostrano anche quali dataset sottopesare e perché, per esempio i dati Kuka raccolti in modo autonomo, il che aiuta a decidere quali dati raccogliere o conservare.

Problema

Grandi dataset robotici come Open X-Embodiment riuniscono dati di robot, telecamere, spazi delle azioni e controllori diversi, eppure le miscele usate per addestrare policy generaliste come RT-X, Octo e OpenVLA sono state scelte a mano. I metodi di filtraggio dei dati della visione e del linguaggio si trasferiscono male, perché gli n-grammi non sono applicabili e gli embedding visivi ignorano le azioni e la natura sequenziale degli episodi.

Soluzione proposta

ReMix imposta la cura dei dati come un problema min-max: una policy discreta minimizza la propria loss di behavior cloning in eccesso rispetto a un modello di riferimento, mentre i pesi di dominio vengono aggiornati con un'ascesa del gradiente esponenziata per massimizzarla, così il peso va ai domini in cui la policy può ancora migliorare. Per rendere le loss confrontabili tra robot, le azioni vengono normalizzate con una gaussiana per dominio e discretizzate in intervalli, e il checkpoint di riferimento viene scelto con early stopping prima che vada in overfitting su qualche dominio. I pesi mediati durante l'addestramento vengono poi riutilizzati per addestrare policy di diffusione separate o per sottocampionare i dati.

Idee chiave

• Usare la loss in eccesso rispetto a una policy di riferimento, anziché la loss grezza, riduce il peso sia dei domini già appresi sia dei domini troppo difficili da apprendere.
• Le differenze di scala delle azioni distorcono l'ottimizzazione robusta: in un test giocattolo con normalizzazione sui limiti, il 94,3 % del peso è andato a un dominio di azioni puramente casuali; con la normalizzazione gaussiana per dominio, i dati reali hanno ricevuto l'84,2 %.
• Le loss continue L1/L2 sulle azioni sono influenzate da valori anomali e azioni multimodali, quindi discretizzare le azioni dà pesi più affidabili anche quando la policy finale è un modello di diffusione continuo.
• I dataset robotici (circa 10k-100k dimostrazioni) sono abbastanza piccoli da portare il modello di riferimento all'overfitting; usare un checkpoint 50K passi dopo l'early stopping ha ridotto le prestazioni di oltre il 15 % in media.

Risultati principali

• Sulla miscela di dati RT-X, le policy addestrate con i pesi di ReMix hanno superato i pesi uniformi del 38 % e i pesi degli esperti umani di RT-X del 32 % in media, su quattro compiti reali con WidowX e due con Franka; i pesi degli esperti hanno battuto quelli uniformi del 6 %.
• Su sottoinsiemi del 25 % della miscela RT-X, ReMix ha mantenuto le prestazioni mentre i pesi umani sono calati di oltre il 10 %; ReMix ha conservato il 72 % dei dati UR5 e il 12 % di quelli Kuka, contro il 30 % e il 24 % dei pesi umani.
• ReMix ha portato il piccolo dataset Toto dal 3,42 % (uniforme) al 16,3 % e ha ridotto Kuka dal 24,9 % al 12,1 %.
• Sul dataset Bridge V2 completo non c'è stata una differenza rilevante tra ReMix e i pesi uniformi; su sottoinsiemi del 25 % di Bridge, ReMix ha superato i pesi uniformi e la baseline SSP basata su embedding di immagini.
• In un compito WidowX con un cubo e un oggetto distrattore, i pesi di ReMix sono riusciti 7/10 contro 0/10 sia per i pesi umani sia per quelli uniformi.
• In simulazione (Robomimic NutAssemblySquare suddiviso per operatore), ReMix ha ottenuto 77/100 contro 53/100 dei pesi uniformi con sottocampionamento al 50 %, e 59/100 contro 39/100 al 25 %.

Limiti

• Gli autori notano che la valutazione nel mondo reale ha riguardato solo due bracci, WidowX e Franka, e lasciano altri embodiment a lavori futuri.
• ReMix aumenta il peso di dataset con distribuzioni di azioni anomale e multimodali come Toto, cosa che secondo gli autori non è necessariamente desiderabile.
• Il calcolo dei pesi richiede di addestrare due volte sull'intero insieme di dati, una per il modello di riferimento e una per l'ottimizzazione robusta.
• Provato su Bridge V2 e sul sottoinsieme RT-X di OpenX, non sull'intero dataset OpenX di oltre 2M episodi. La sezione 4.1 descrive Bridge V2 come 50k dimostrazioni, mentre l'appendice B.2 indica che sono state usate circa 45.000 traiettorie delle 60K del dataset completo.
• Le policy addestrate solo sui dati di pre-addestramento non hanno generalizzato zero-shot alle configurazioni degli autori, quindi ogni valutazione si basa su un co-addestramento con 25 dimostrazioni del dominio specifico per compito.

Impatto sul settore

Finalista del CoRL 2024 Outstanding Paper Award e citato 74 volte. Gli autori rilasciano il codice su github.com/jhejna/remix. L'articolo pubblica inoltre i pesi di dominio appresi per la miscela RT-X (tabella 2) e per i 32 domini di Bridge V2 (tabella 5).

Abstract

Vengono raccolti dataset di apprendimento per imitazione sempre più grandi con l'obiettivo di addestrare foundation model per la robotica. Tuttavia, nonostante la selezione dei dati sia stata di fondamentale importanza nella visione e nell'elaborazione del linguaggio naturale, pochi lavori in robotica si sono chiesti su quali dati tali modelli debbano effettivamente essere addestrati. In questo lavoro studiamo come pesare diversi sottoinsiemi o “domini” dei dataset robotici per il pre-addestramento di foundation model per robot. Concretamente, usiamo l'ottimizzazione distribuzionalmente robusta (DRO) per massimizzare le prestazioni nel caso peggiore su tutti i possibili domini a valle. Il nostro metodo, Re-Mix, affronta l'ampia gamma di sfide che emergono applicando la DRO ai dataset robotici, tra cui la variabilità degli spazi delle azioni e della dinamica tra dataset. Re-Mix impiega early stopping, normalizzazione delle azioni e discretizzazione per contrastare questi problemi. Attraverso un'ampia sperimentazione sul più grande dataset open source di manipolazione robotica, il dataset Open X-Embodiment, dimostriamo che la cura dei dati può avere un impatto sproporzionato sulle prestazioni a valle. In particolare, i pesi di dominio appresi da Re-Mix superano i pesi uniformi del 38 % in media e i pesi scelti da persone del 32 % sui dataset usati per addestrare policy robotiche generaliste esistenti, in particolare i modelli RT-X.

Link

• arXiv 2408.14037: https://arxiv.org/abs/2408.14037
• DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
• Codice: https://github.com/jhejna/remix