ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Autoren: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

Lizenzen

Open Access
Paper: CC BY 4.0

Alles Folgende stützt sich auf das vollständige Paper unter seiner Lizenz CC BY 4.0.

Kurz gesagt

ReMix nutzt gruppenweise verteilungsrobuste Optimierung auf dem Excess Loss einer Policy, um zu lernen, wie stark jeder Roboterdatensatz gesampelt werden soll, und schlägt auf realen Robotern gleichmäßige Gewichte und die RT-X-Expertengewichte.

Warum es wichtig ist

Teams, die auf zusammengelegten Roboterdaten vortrainieren, können handabgestimmte Sampling-Gewichte durch aus den Daten berechnete Gewichte ersetzen und in den berichteten Tests 75 % der Daten bei begrenztem Leistungsverlust weglassen. Die gelernten Gewichte zeigen außerdem, welche Datensätze herabgewichtet werden sollten und warum, etwa die autonom gesammelten Kuka-Daten, was bei der Entscheidung hilft, welche Daten gesammelt oder behalten werden.

Problem

Große Roboterdatensätze wie Open X-Embodiment bündeln Daten verschiedener Roboter, Kameras, Aktionsräume und Regler, doch die Mischungen, mit denen generalistische Policies wie RT-X, Octo und OpenVLA trainiert wurden, wurden von Hand gewählt. Datenfiltermethoden aus Bild- und Sprachverarbeitung lassen sich schlecht übertragen, weil n-Gramme nicht anwendbar sind und visuelle Einbettungen Aktionen und den sequenziellen Charakter von Episoden ignorieren.

Lösungsansatz

ReMix fasst Kuratierung als Min-Max-Problem: Eine diskrete Policy minimiert ihren Excess Loss beim Behavior Cloning gegenüber einem Referenzmodell, während die Domänengewichte per exponentiertem Gradientenanstieg aktualisiert werden, um ihn zu maximieren, sodass Gewicht an Domänen geht, auf denen sich die Policy noch verbessern kann. Um Verluste über Roboter hinweg vergleichbar zu machen, werden Aktionen je Domäne gaußnormalisiert und in Intervalle diskretisiert, und der Referenz-Checkpoint wird per Early Stopping gewählt, bevor er sich an eine Domäne überanpasst. Die über das Training gemittelten Gewichte werden anschließend wiederverwendet, um separate Diffusion Policies zu trainieren oder die Daten unterabzutasten.

Wichtigste Erkenntnisse

Die Verwendung des Excess Loss gegenüber einer Referenz-Policy statt des rohen Verlusts gewichtet sowohl Domänen herab, die bereits angepasst sind, als auch solche, die zu schwer anzupassen sind.
Unterschiede in der Aktionsskala verzerren die robuste Optimierung: In einem Spielzeugtest mit Grenzen-Normalisierung ging 94,3 % des Gewichts an eine Domäne aus reinen Zufallsaktionen; mit Gauß-Normalisierung je Domäne erhielten die realen Daten 84,2 %.
Kontinuierliche L1/L2-Aktionsverluste werden von Ausreißern und multimodalen Aktionen beeinflusst, daher liefert die Diskretisierung der Aktionen verlässlichere Gewichte, selbst wenn die finale Policy ein kontinuierliches Diffusionsmodell ist.
Roboterdatensätze (etwa 10k-100k Demonstrationen) sind klein genug, dass das Referenzmodell überangepasst wird; ein Checkpoint 50K Schritte nach dem Early-Stopping-Punkt senkte die Leistung im Mittel um über 15 %.

Hauptergebnisse

Auf der RT-X-Datenmischung übertrafen mit ReMix-Gewichten trainierte Policies gleichmäßige Gewichte um 38 % und die RT-X-Gewichte menschlicher Experten um 32 %, gemittelt über vier reale WidowX- und zwei reale Franka-Aufgaben; Expertengewichte schlugen gleichmäßige um 6 %.
Auf 25-%-Teilmengen der RT-X-Mischung behielt ReMix seine Leistung, während menschliche Gewichte um über 10 % abfielen; ReMix behielt 72 % der UR5-Daten und 12 % von Kuka, gegenüber 30 % und 24 % bei den menschlichen Gewichten.
ReMix hob den kleinen Toto-Datensatz von 3,42 % (gleichmäßig) auf 16,3 % an und senkte Kuka von 24,9 % auf 12,1 %.
Auf dem vollständigen Bridge V2-Datensatz gab es keinen nennenswerten Unterschied zwischen ReMix und gleichmäßigen Gewichten; auf 25-%-Teilmengen von Bridge übertraf ReMix gleichmäßige Gewichte und die SSP-Baseline mit Bildeinbettungen.
Bei einer WidowX-Würfelaufgabe mit Ablenkobjekt waren ReMix-Gewichte in 7/10 Fällen erfolgreich, gegenüber 0/10 sowohl für menschliche als auch für gleichmäßige Gewichte.
In der Simulation (Robomimic NutAssemblySquare, nach Bediener aufgeteilt) erzielte ReMix 77/100 gegenüber 53/100 für gleichmäßige Gewichte bei 50 % Teilmenge und 59/100 gegenüber 39/100 bei 25 %.

Einschränkungen

Die Autoren merken an, dass die reale Evaluierung nur zwei Arme abdeckte, WidowX und Franka, und überlassen weitere Embodiments künftiger Arbeit.
ReMix gewichtet Datensätze mit abnormalen, multimodalen Aktionsverteilungen wie Toto höher, was laut den Autoren nicht unbedingt wünschenswert ist.
Die Berechnung der Gewichte erfordert zweimaliges Training auf den vollständigen Daten, einmal für das Referenzmodell und einmal für die robuste Optimierung.
Getestet auf Bridge V2 und der RT-X-Teilmenge von OpenX, nicht auf dem vollständigen OpenX-Datensatz mit über 2M Episoden. Abschnitt 4.1 beschreibt Bridge V2 als 50k Demonstrationen, während Anhang B.2 angibt, dass etwa 45.000 der 60K Trajektorien des vollständigen Datensatzes verwendet wurden.
Nur auf den Vortrainingsdaten trainierte Policies generalisierten nicht zero-shot auf die Aufbauten der Autoren, sodass jede Evaluierung auf Co-Training mit 25 domäneninternen Demonstrationen je Aufgabe beruht.

Wirkung auf das Fachgebiet

Finalist für den CoRL 2024 Outstanding Paper Award und 74-mal zitiert. Die Autoren veröffentlichen Code unter github.com/jhejna/remix. Das Paper veröffentlicht außerdem seine gelernten Domänengewichte für die RT-X-Mischung (Tabelle 2) und für die 32 Bridge V2-Domänen (Tabelle 5).

Zusammenfassung

Zunehmend große Datensätze für Imitationslernen werden mit dem Ziel gesammelt, Foundation Models für die Robotik zu trainieren. Doch obwohl die Datenauswahl in der Bildverarbeitung und der Verarbeitung natürlicher Sprache von größter Bedeutung war, haben nur wenige Arbeiten in der Robotik hinterfragt, auf welchen Daten solche Modelle eigentlich trainiert werden sollten. In dieser Arbeit untersuchen wir, wie verschiedene Teilmengen oder „Domänen“ von Robotikdatensätzen für das Vortraining von Roboter-Foundation-Models gewichtet werden sollten. Konkret verwenden wir verteilungsrobuste Optimierung (DRO), um die Leistung im ungünstigsten Fall über alle möglichen nachgelagerten Domänen zu maximieren. Unsere Methode Re-Mix geht die vielfältigen Herausforderungen an, die bei der Anwendung von DRO auf Robotikdatensätze entstehen, darunter die Variabilität von Aktionsräumen und Dynamik zwischen Datensätzen. Re-Mix setzt Early Stopping, Aktionsnormalisierung und Diskretisierung ein, um diesen Problemen entgegenzuwirken. Durch umfangreiche Experimente auf dem größten Open-Source-Datensatz zur Robotermanipulation, dem Open X-Embodiment-Datensatz, zeigen wir, dass Datenkuratierung einen überproportionalen Einfluss auf die nachgelagerte Leistung haben kann. Insbesondere übertreffen die von Re-Mix gelernten Domänengewichte gleichmäßige Gewichte im Mittel um 38 % und von Menschen gewählte Gewichte um 32 % auf Datensätzen, mit denen bestehende generalistische Roboter-Policies trainiert wurden, namentlich die RT-X-Modelle.

Open Source Robotics