Riepilogo Tecnico: RoboSSM – In-context Imitation Learning Scalabile tramite State-Space Models

1. Definizione del Problema

L'In-context Imitation Learning (ICIL) consente ai robot di adattarsi a nuovi compiti utilizzando poche dimostrazioni (few-shot) senza aggiornamenti dei parametri durante l'impiego. Sebbene i recenti metodi ICIL abbiano adottato con successo architetture basate su Transformer per formulare l'apprendimento per imitazione come un problema di modellazione di sequenze, essi affrontano significativi limiti di scalabilità. I Transformer esibiscono una complessità temporale quadratica ($O(L^2)$) rispetto alla lunghezza della sequenza e faticano a estrapolare oltre le lunghezze dei prompt utilizzate durante l'addestramento. Di conseguenza, i metodi ICIL esistenti degradano drasticamente quando devono gestire prompt di test più lunghi (ad esempio, più dimostrazioni o compiti a lungo orizzonte) rispetto a quelli usati durante l'addestramento, limitandone l'applicabilità in scenari reali complessi che richiedono un contesto esteso.

2. Metodologia: RoboSSM

Gli autori propongono RoboSSM, un framework ICIL scalabile che sostituisce i Transformer con gli State-Space Models (SSM), utilizzando specificamente Longhorn come architettura di base (backbone).

Architettura e Codifica

• Elaborazione dell'Input: RoboSSM codifica le osservazioni multimodali (dati visivi tramite CNN e dati propriocettivi tramite MLP) in embedding. Notevolmente, la rappresentazione dell'input esclude i token delle azioni e le istruzioni testuali del compito per costringere il modello a inferire l'intento del compito esclusivamente dalle dimostrazioni in-context.
• Backbone Longhorn: La sequenza di embedding delle osservazioni è elaborata da Longhorn, che mantiene una matrice di stato di memoria $s_t$ aggiornata tramite una ricorrenza lineare. A differenza del meccanismo di self-attention dei Transformer, Longhorn utilizza una regola di aggiornamento ricorrente derivata dalla programmazione convessa online:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Questa formulazione bilancia la ritenzione delle informazioni dello stato precedente con l'integrazione di nuovi input.
• $\beta$-Scaling: Per incoraggiare specificamente l'attenzione ai prompt di dimostrazione nel contesto ICIL, gli autori introducono un meccanismo di scaling al tempo di test in cui il vettore di ponderazione $\beta_t$ è scalato da un fattore $\gamma \in (0, 1]$. Questo modula il compromesso tra la ritenzione dello stato e l'integrazione di nuove informazioni.

Addestramento e Inferenza

• Addestramento: Il modello è addestrato su un dataset multi-task in cui l'input consiste in un prompt (un set di $N_{train}$ traiettorie di dimostrazione) e una traiettoria di query. Il modello predice le azioni per la traiettoria di query condizionata dal prompt.
• Inferenza: Al tempo di test, la policy si adatta a compiti non visti utilizzando un prompt di $N_{test}$ dimostrazioni. Il modello predice iterativamente le azioni basandosi sull'osservazione iniziale e sul contesto fornito dal prompt, senza alcun aggiornamento dei parametri.

3. Contributi Chiave

1. Framework ICIL Scalabile: L'introduzione di RoboSSM, il primo approccio che sfrutta gli SSM per l'in-context imitation learning, consentendo un'inferenza a tempo lineare ($O(L)$) e forti capacità di estrapolazione.
2. Estrapolazione del Lungo Contesto: Dimostrare che RoboSSM può elaborare efficacemente prompt di test fino a 16 volte più lunghi di quelli visti durante l'addestramento, una capacità in cui i metodi basati su Transformer falliscono.
3. Inferenza Efficiente: Raggiungere tempi di esecuzione dell'inferenza significativamente inferiori rispetto ai baseline Transformer all'aumentare della lunghezza del prompt, grazie alla complessità lineare degli SSM rispetto alla complessità quadratica dei Transformer.
4. Generalizzazione a Compiti a Lungo Orizzonte: Mostrare che il modello può generalizzare da distribuzioni di addestramento a breve orizzonte a compiti a lungo orizzonte e multi-stadio, e a scenari con dilatazione temporale.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul benchmark LIBERO, coprendo diversi compiti di manipolazione (Object, Goal, Spatial, Long e 9-suite).

• Generalizzazione della Lunghezza del Prompt:
  • Quando addestrato con sole 2 dimostrazioni ($N_{train}=2$), RoboSSM ha mantenuto o migliorato i tassi di successo all'aumentare del numero di dimostrazioni di test ($N_{test}$) fino a 32.
  • Al contrario, il baseline basato su Transformer (ICRT) ha subito un drastico calo delle prestazioni ed è collassato allo 0% di successo quando $N_{test} > N_{train}$.
• Lungo Orizzonte e Dilatazione Temporale:
  • RoboSSM ha risolto con successo compiti a lungo orizzonte (ad esempio, "Metti sia l'Oggetto A che l'Oggetto B nell'Oggetto C") e compiti con dimostrazioni con dilatazione temporale (fino a 16 volte l'allungamento temporale), mentre ICRT è fallito completamente.
• Prestazioni In-Distribution:
  • Anche quando $N_{train} = N_{test}$, RoboSSM ha costantemente superato ICRT in quasi tutti gli scenari, in particolare nelle scene LIBERO-90 Study e Living Room.
• Confronto con l'Apprendimento Multi-Task (MTL):
  • RoboSSM e ICRT hanno superato significativamente i classici baseline MTL (che richiedono istruzioni linguistiche e non utilizzano il prompting in-context), confermando l'efficacia del paradigma ICIL.
• Analisi dello Spazio Latente:
  • Le visualizzazioni delle rappresentazioni latenti hanno mostrato che RoboSSM genera traiettorie che rimangono vicine ai cluster delle dimostrazioni del prompt, indicando un'estrapolazione stabile. Al contrario, le traiettorie di ICRT si sono allontanate significativamente dal manifold del prompt all'aumentare della lunghezza del contesto.
• Runtime:
  • RoboSSM ha dimostrato una scalabilità quasi lineare con la lunghezza del prompt, mentre il runtime di ICRT aumentava quadraticamente, creando un divario di prestazioni crescente a favore di RoboSSM per contesti lunghi.

5. Significato e Rivendicazioni

Il paper sostiene che RoboSSM rappresenti un cambio di paradigma per l'in-context imitation learning, stabilendo gli State-Space Models come un backbone efficiente e scalabile. Il significato primario risiede nel superamento dei colli di bottiglia computazionali e di estrapolazione dei Transformer, permettendo ai robot di:

• Gestire finestre di contesto sostanzialmente più lunghe al tempo di test senza ri-addestramento.
• Adattarsi a compiti non visti, a lungo orizzonte e temporalmente variabili utilizzando prompt few-shot.
• Supportare l'adattamento continuo per l'apprendimento a vita (lifelong learning) semplicemente ingerendo nuovi prompt di dimostrazione senza aggiornamenti dei parametri specifici del compito.

Gli autori concludono che, sebbene RoboSSM dimostri il potenziale degli SSM per l'ICIL a lungo contesto, il lavoro futuro potrebbe richiedere corpora di addestramento più ampi e diversificati per migliorare ulteriormente la generalizzazione. Notano inoltre che, sebbene la validazione su robot reali non sia il contributo principale di questo lavoro, essa rimane una possibile via per future dimostrazioni di generalizzazione.