Technische Samenvatting: RoboSSM – Schaalbare In-context Imitatieleer via State-Space Modellen

1. Probleemstelling

In-context Imitatieleer (ICIL) stelt robots in staat om zich aan te passen aan nieuwe taken met behulp van weinig-schot demonstraties (few-shot demonstrations) zonder parameterupdates tijdens de inzet. Hoewel recente ICIL-methoden succesvol Transformer-gebaseerde architecturen hebben geadoperd om imitatieleer te formuleren als een sequentiemodelleringsprobleem, kampen zij met significante schaalbaarheidsbeperkingen. Transformers vertonen een kwadratische tijdscomplexiteit ($O(L^2)$) ten opzichte van de sequentielengte en hebben moeite met extrapolatie buiten de prompt-lengtes die tijdens de training zijn gezien. Gevolgelijk presteren bestaande ICIL-methoden sterk achteruit wanneer ze te maken krijgen met langere test-time prompts (bijv. meer demonstraties of taken met een lange horizon) dan die gebruikt tijdens de training, wat hun toepasbaarheid in complexe, real-world scenario's die uitgebreide context vereisen, beperkt.

2. Methodologie: RoboSSM

De auteurs stellen RoboSSM voor, een schaalbaar ICIL-framework dat Transformers vervangt door State-Space Modellen (SSM's), waarbij specifiek gebruik wordt gemaakt van de Longhorn backbone-architectuur.

Architectuur en Codering

• Inputverwerking: RoboSSM codeert multimodale observaties (visuele data via CNN's en proprioceptieve data via MLP's) naar embeddings. Opvallend is dat de inputrepresentatie actie-tokens en taalinstructies uitsluit om het model te dwingen de taakintentie uitsluitend af te leiden uit de in-context demonstraties.
• Longhorn Backbone: De sequentie van observatie-embeddings wordt verwerkt door Longhorn, die een geheugen-toestandmatrix $s_t$ bijhoudt die via een lineaire recursie wordt bijgewerkt. In tegenstelling tot het self-attention mechanisme in Transformers, gebruikt Longhorn een recursieve update-regel afgeleid van online convexe programmering:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Deze formulering balanceert het behoud van eerdere toestandsinformatie met het integreren van nieuwe inputs.
• $\beta$-schaling: Om specifiek de aandacht voor demonstratie-prompts in de ICIL-setting te stimuleren, introduceren de auteurs een test-time schaleringsmechanisme waarbij de wectorgewicht $\beta_t$ wordt geschaald met een factor $\gamma \in (0, 1]$. Dit moduleert de afweging tussen toestandbehoud en de integratie van nieuwe informatie.

Training en Inferentie

• Training: Het model wordt getraind op een multi-task dataset waarbij de input bestaat uit een prompt (een set van $N_{train}$ demonstratie-trajecten) en een query-traject. Het model voorspelt acties voor het query-traject, geconditioneerd op de prompt.
• Inferentie: Tijdens de test-fase past het beleid (policy) zich aan aan onbekende taken met behulp van een prompt van $N_{test}$ demonstraties. Het model voorspelt iteratief acties op basis van de initiële observatie en de context die door de prompt wordt geleverd, zonder enige parameterupdates.

3. Belangrijkste Bijdragen

1. Schaalbaar ICIL-framework: De introductie van RoboSSM, de eerste aanpak die SSM's inzet voor in-context imitatieleer, wat lineaire-tijd inferentie ($O(L)$) en sterke extrapolatiecapaciteiten mogelijk maakt.
2. Long-Context Extrapolatie: Het aantonen dat RoboSSM effectief prompts kan verwerken die tot 16 keer langer zijn dan die tijdens de training zijn gezien, een capaciteit waar Transformer-gebaseerde methoden falen.
3. Efficiënte Inferentie: Het bereiken van aanzienlijk lagere inferentie-runtimes vergeleken met Transformer-baselines naarmate de promptlengte toeneemt, dankzij de lineaire complexiteit van SSM's versus de kwadratische complexiteit van Transformers.
4. Generalisatie naar Long-Horizon Taken: Het aantonen dat het model kan generaliseren van korte-horizon trainingsdistributies naar lange-horizon, multi-stage taken en tijd-gedilateerde scenario's.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op de LIBERO benchmark, die diverse manipulatie-taken dekt (Object, Goal, Spatial, Long, en 90 suites).

• Promptlengte Generalisatie:
  • Wanneer getraind met slechts 2 demonstraties ($N_{train}=2$), behield of verbeterde RoboSSM de succesratio's naarmate het aantal test-demonstraties ($N_{test}$) toenam tot 32.
  • In tegenstelling hiertoe ervoer de Transformer-gebaseerde baseline (ICRT) een scherpe prestatiedaling en stortte deze in naar een succesratio van 0% wanneer $N_{test} > N_{train}$.
• Long-Horizon en Tijd-Dilatatie:
  • RoboSSM loste succesvol lange-horizon taken op (bijv. "Zet zowel Object A als Object B in Object C") en taken met tijd-gedilateerde demonstraties (tot 16x temporele rekking), terwijl ICRT volledig faalde.
• In-Distribution Prestaties:
  • Zelfs wanneer $N_{train} = N_{test}$, presteerde RoboSSM consistent beter dan ICRT over de meeste scenario's, met name in de LIBERO-90 Study en Living Room scènes.
• Vergelijking met Multi-Task Learning (MTL):
  • RoboSSM en ICRT presteerden significant beter dan standaard MTL-baselines (die taalinstructies vereisen en geen in-context prompting gebruiken), wat de effectiviteit van het ICIL-paradigma bevestigt.
• Latente Ruimte Analyse:
  • Visualisaties van latente representaties toonden aan dat RoboSSM trajecten genereert die dicht bij de prompt-demonstratieclusters blijven, wat wijst op stabiele extrapolatie. Daarentegen driften de ICRT-trajecten significant buiten de prompt-manifold wanneer de contextlengte toeneemt.
• Runtime:
  • RoboSSM vertoonde bijna lineaire schaling met de promptlengte, terwijl de runtime van ICRT kwadratisch toenam, wat een groeiend prestatie-gat in het voordeel van RoboSSM creëert voor lange contexten.

5. Betekenis en Claims

Het artikel beweert dat RoboSSM een paradigmaverschuiving vertegenwoordigt voor in-context imitatieleer door State-Space Modellen als een efficiënt en schaalbaar fundament te vestigen. De primaire betekenis ligt in het overwinnen van de computationele en extrapolatie-bottlenecks van Transformers, waardoor robots in staat zijn om:

• Substantieel langere contextvensters te verwerken tijdens de test-fase zonder hertraining.
• Zich aan te passen aan onbekende, lange-horizon en temporeel variabele taken met behulp van weinig-schot prompts.
• Continue adaptatie te ondersteunen voor lifelong learning door simpelweg nieuwe demonstratie-prompts te absorberen zonder taakspecifieke parameterupdates.

De auteurs concluderen dat hoewel RoboSSM het potentieel van SSM's voor lange-context ICIL aantoont, toekomstig werk mogelijk bredere en diversere trainingscorpora vereist om de generalisatie naar nieuwe taken verder te verbeteren. Zij merken ook op dat hoewel validatie op echte robots geen kernbijdrage van dit werk is, dit een potentiële avenue blijft voor toekomstige demonstraties van generalisatie.