Technisches Resümee: RoboSSM – Skalierbares In-Context Imitation Learning mittels State-Space Models

1. Problemstellung

In-Context Imitation Learning (ICIL) ermöglicht es Robotern, sich durch wenige Demonstrationsschritte (Few-Shot) an neue Aufgaben anzupassen, ohne dass während des Einsatzes Parameteraktualisierungen erforderlich sind. Während jüngste ICIL-Methoden erfolgreich Transformer-basierte Architekturen eingesetzt haben, um Imitationslernen als Sequenzmodellierung zu formulieren, stehen sie vor erheblichen Skalierbarkeitsproblemen. Transformer weisen eine quadratische Zeitkomplexität ($O(L^2)$) in Bezug auf die Sequenzlänge auf und haben Schwierigkeiten, über die während des Trainings gesehenen Prompt-Längen hinauszuextrapolieren. Folglich verschlechtert sich die Leistung bestehender ICIL-Methoden drastisch, wenn sie längere Test-Prompts verarbeiten müssen (z. B. mehr Demonstrationen oder Aufgaben mit langem Horizont) als während des Trainings verwendet wurden, was ihre Anwendbarkeit in komplexen, realen Szenarien, die umfangreichen Kontext erfordern, einschränkt.

2. Methodik: RoboSSM

Die Autoren schlagen RoboSSM vor, ein skalierbares ICIL-Framework, das Transformer durch State-Space Models (SSMs) ersetzt und dabei die Longhorn-Backbone-Architektur nutzt.

Architektur und Kodierung

• Input-Verarbeitung: RoboSSM kodiert multimodale Beobachtungen (visuelle Daten via CNNs und propriozeptive Daten via MLPs) in Embeddings. Bemerkenswerterweise schließt die Eingangsrepräsentation keine Aktions-Token und Aufgaben-Sprachinstruktionen ein, um das Modell dazu zu zwingen, die Aufgabenabsicht ausschließlich aus den In-Context-Demonstrationen abzuleiten.
• Longhorn-Backbone: Die Sequenz der Beobachtungs-Embeddings wird von Longhorn verarbeitet, welches eine Speicherzustandsmatrix $s_t$ mittels einer linearen Rekurrenz pflegt. Im Gegensatz zum Self-Attention-Mechanismus von Transformern verwendet Longhorn eine rekursive Aktualisierungsregel, die aus dem Online-Konvexen Programmieren abgeleitet ist:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Diese Formulierung balanciert die Beibehaltung früherer Zustandsinformationen mit der Integration neuer Eingaben aus.
• $\beta$-Skalierung: Um die Aufmerksamkeit speziell auf die Demonstrations-Prompts im ICIL-Setting zu fördern, führen die Autoren einen Testzeit-Skalierungsmechanismus ein, bei dem der Gewichtungsvektor $\beta_t$ mit einem Faktor $\gamma \in (0, 1]$ skaliert wird. Dies moduliert den Kompromiss zwischen Zustandsretention und der Integration neuer Informationen.

Training und Inferenz

• Training: Das Modell wird auf einem Multi-Task-Datensatz trainiert, wobei der Input aus einem Prompt (einem Satz von $N_{train}$ Demonstrations-Trajektorien) und einer Query-Trajektorie besteht. Das Modell sagt die Aktionen für die Query-Trajektorie unter Berücksichtigung des Prompts voraus.
• Inferenz: Zur Testzeit passt sich die Policy an ungesehene Aufgaben unter Verwendung eines Prompts von $N_{test}$ Demonstrationen an. Das Modell sagt iterativ Aktionen basierend auf der initialen Beobachtung und dem durch den Prompt bereitgestellten Kontext voraus, ohne dass Parameteraktualisierungen stattfinden.

3. Zentrale Beiträge

1. Skalierbares ICIL-Framework: Die Einführung von RoboSSM, dem ersten Ansatz, der SSMs für In-Context Imitation Learning nutzt und somit eine lineare Inferenzzeit ($O(L)$) sowie starke Extrapolationsfähigkeiten ermöglicht.
2. Long-Context Extrapolation: Der Nachweis, dass RoboSSM in der Lage ist, Test-Prompts zu verarbeiten, die bis zu 16-mal länger sind als die während des Trainings verwendeten, eine Fähigkeit, bei der Transformer-basierte Methoden versagen.
3. Effiziente Inferenz: Erzielung signifikant niedrigerer Inferenz-Laufzeiten im Vergleich zu Transformer-Baselines bei zunehmender Prompt-Länge, bedingt durch die lineare Komplexität von SSMs gegenüber der quadratischen Komplexität von Transformern.
4. Generalisierung auf Aufgaben mit langem Horizont: Aufzeigen, dass das Modell von kurzfristigen Trainingsverteilungen auf langfristige, mehrstufige Aufgaben und zeitlich dilatierte Szenarien generalisieren kann.

4. Experimentelle Ergebnisse

Die Experimente wurden auf dem LIBERO-Benchmark durchgeführt, der diverse Manipulationsaufgaben (Object, Goal, Spatial, Long und 90 Suites) abdeckt.

• Generalisierung der Prompt-Länge:
  • Wenn mit nur 2 Demonstrationen trainiert wurde ($N_{train}=2$), konnte RoboSSM die Erfolgsraten beibehalten oder sogar verbessern, während die Anzahl der Test-Demonstrationen ($N_{test}$) auf bis zu 32 anstieg.
  • Im Gegensatz dazu erlebte die Transformer-basierte Baseline (ICRT) einen starken Leistungsabfall und kollabierte auf eine Erfolgsrate von 0 %, wenn $N_{test} > N_{train}$.
• Long-Horizon und Zeit-Dilatation:
  • RoboSSM löste erfolgreich Aufgaben mit langem Horizont (z. B. „Lege sowohl Objekt A als auch Objekt B in Objekt C“) sowie Aufgaben mit zeitlich dilatierten Demonstrationen (bis zu 16-fache zeitliche Streckung), während ICRT vollständig versagte.
• In-Distribution-Performance:
  • Selbst wenn $N_{train} = N_{test}$ galt, übertraf RoboSSM ICRT konsistent in den meisten Szenarien, insbesondere in den LIBERO-90 Study und Living Room Szenen.
• Vergleich zum Multi-Task Learning (MTL):
  • RoboSSM und ICRT übertrafen die Standard-MTL-Baselines (die Sprachinstruktionen benötigen und kein In-Context-Prompting nutzen) signifikant, was die Wirksamkeit des ICIL-Paradigmas bestätigt.
• Latent Space Analyse:
  • Visualisierungen der latenten Repräsentationen zeigten, dass RoboSSM Trajektorien generiert, die nahe an den Clustern der Prompt-Demonstrationen bleiben, was auf eine stabile Extrapolation hindeutet. Im Gegensatz dazu driften die ICRT-Trajektorien mit zunehmender Kontextlänge signifikant aus dem Prompt-Manifold heraus.
• Laufzeit:
  • RoboSSM zeigte eine nahezu lineare Skalierung mit der Prompt-Länge, während die Laufzeit von ICRT quadratisch anstieg, was eine sich weitende Leistungslücke zugunsten von RoboSSM bei langen Kontexten erzeugte.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass RoboSSM einen Paradigmenwechsel für das In-Context Imitation Learning darstellt, indem es State-Space Models als effizientes und skalierbares Backbone etabliert. Die primäre Bedeutung liegt in der Überwindung der rechnerischen und Extrapolations-Engpässe von Transformern, was es Robotern ermöglicht:

• Substanziell längere Kontextfenster zur Testzeit zu handhaben, ohne neu zu trainieren.
• Sich an ungesehene, langfristige und zeitlich variable Aufgaben mittels Few-Shot-Prompts anzupassen.
• Kontinuierliche Anpassung für lebenslanges Lernen zu unterstützen, indem einfach neue Demonstration-Prompts aufgenommen werden, ohne aufgabenspezifische Parameter-Updates vorzunehmen.

Die Autoren schließen, dass während RoboSSM das Potenzial von SSMs für das Long-Context ICIL demonstriert, zukünftige Arbeiten möglicherweise breitere und vielfältigere Trainingskorpora benötigen, um die Generalisierung auf neue Aufgaben weiter zu verbessern. Sie merken zudem an, dass die Validierung an echten Robotern nicht der Kernbeitrag dieser Arbeit ist, dies jedoch ein potenzieller Weg für zukünftige Demonstrationen der Generalisierung bleibt.