Résumé Technique : RoboSSM – Apprentissage par Imitation In-Context Scalable via les Modèles d'Espace d'États

1. Énoncé du Problème

L'apprentissage par imitation in-context (ICIL - In-context Imitation Learning) permet aux robots de s'adapter à de nouvelles tâches à l'aide de démonstrations à faible échantillonnage (few-shot) sans mise à jour des paramètres lors du déploiement. Bien que les méthodes récentes d'ICIL aient adopté avec succès des architectures basées sur les Transformers pour formuler l'apprentissage par imitation comme un problème de modélisation de séquences, elles font face à des limitations de scalabilité significatives. Les Transformers présentent une complexité temporelle quadratique ($O(L^2)$) par rapport à la longueur de la séquence et peinent à extrapoler au-delà des longueurs de prompts observées durant l'entraînement. Par conséquent, les méthodes d'ICIL existantes se dégradent brutalement lorsqu'elles gèrent des prompts de test plus longs (par exemple, plus de démonstrations ou des tâches à long horizon) que ceux utilisés pendant l'entraînement, ce qui limite leur applicabilité dans des scénarios réels complexes nécessitant un contexte étendu.

2. Méthodologie : RoboSSM

Les auteurs proposent RoboSSM, un cadre d'ICIL scalable qui remplace les Transformers par des Modèles d'Espace d'États (SSM - State-Space Models), en utilisant spécifiquement l'architecture Longhorn comme colonne vertébrale.

Architecture et Encodage

• Traitement des Entrées : RoboSSM encode les observations multimodales (données visuelles via des CNN et données proprioceptives via des MLP) en embeddings. Notamment, la représentation d'entrée exclut les tokens d'action et les instructions linguistiques de la tâche pour forcer le modèle à inférer l'intention de la tâche uniquement à partir des démonstrations in-context.
• Colonne Vertébrale Longhorn : La séquence d'embeddings d'observation est traitée par Longhorn, qui maintient une matrice d'état de mémoire $s_t$ mise à jour via une récurrence linéaire. Contra�que le mécanisme d'auto-attention des Transformers, Longhorn utilise une règle de mise à jour récurrente dérivée de la programmation convexe en ligne :
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Cette formulation équilibre la rétention de l'information d'état préalable avec l'incorporation de nouvelles entrées.
• Mise à l'échelle $\beta$ ($\beta$-Scaling) : Pour encourager spécifiquement l'attention portée aux prompts de démonstration dans le cadre de l'ICIL, les auteurs introduisent un mécanisme de mise à l'échelle au moment du test où le vecteur de pondération $\beta_t$ est multiplié par un facteur $\gamma \in (0, 1]$. Cela module le compromis entre la rétention de l'état et l'intégration de nouvelles informations.

Entraînement et Inférence

• Entraînement : Le modèle est entraîné sur un jeu de données multi-tâches où l'entrée consiste en un prompt (un ensemble de $N_{train}$ trajectoires de démonstration) et une trajectoire de requête. Le modèle prédit les actions pour la trajectoire de requête conditionnée par le prompt.
• Inférence : Au moment du test, la politique s'adapte à des tâches non vues en utilisant un prompt de $N_{test}$ démonstrations. Le modèle prédit de manière itérative les actions basées sur l'observation initiale et le contexte fourni par le prompt, sans aucune mise à jour de paramètres.

3. Contributions Clés

1. Cadre ICIL Scalable : L'introduction de RoboSSM, la première approche exploitant les SSM pour l'apprentissage par imitation in-context, permettant une inférence en temps linéaire ($O(L)$) et de fortes capacités d'extrapolation.
2. Extrapolation de Long Contexte : Démontrer que RoboSSM peut traiter efficacement des prompts de test jusqu'à 16 fois plus longs que ceux vus durant l'entraînement, une capacité où les méthodes basées sur les Transformers échouent.
3. Inférence Efficace : Atteindre des temps d'exécution d'inférence nettement inférieurs aux baselines de Transformers à mesure que la longueur du prompt augmente, grâce à la complexité linéaire des SSM par rapport à la complexité quadratique des Transformers.
4. Généralisation aux Tâches à Long Horizon : Montrer que le modèle peut généraliser de distributions d'entraînement à court horizon vers des tâches à long horizon, multi-étapes et des scénarios de dilatation temporelle.

4. Résultats Expérimentaux

Les expériences ont été menées sur le benchmark LIBERO, couvrant diverses tâches de manipulation (Object, Goal, Spatial, Long, et 90 suites).

• Généralisation de la Longueur du Prompt :
  • Lorsqu'il est entraîné avec seulement 2 démonstrations ($N_{train}=2$), RoboSSM maintient ou améliore les taux de réussite à mesure que le nombre de démonstrations de test ($N_{test}$) augmente jusqu'à 32.
  • En revanche, la baseline basée sur les Transformers (ICRT) subit une chute brutale de performance et s'effondre à 0 % de succès lorsque $N_{test} > N_{train}$.
• Long Horizon et Dilatation Temporelle :
  • RoboSSM a résolu avec succès des tâches à long horizon (ex: "Mettre à la fois l'Objet A et l'Objet B dans l'Objet C") et des tâches avec des démonstrations dilatées temporellement (jusqu'à 16x l'étirement temporel), là où ICRT a totalement échoué.
• Performance In-Distribution :
  • Même lorsque $N_{train} = N_{test}$, RoboSSM surpasse systématiquement ICRT dans la plupart des scénarios, particulièrement dans les scènes LIBERO-90 Study et Living Room.
• Comparaison avec l'Apprentissage Multi-Tâche (MTL) :
  • RoboSSM et ICRT surpassent de manière significative les baselines standards de MTL (qui nécessitent des instructions linguistiques et n'utilisent pas de prompting in-context), confirmant l'efficacité du paradigme ICIL.
• Analyse de l'Espace Latent :
  • Les visualisations des représentations latentes ont montré que RoboSSM génère des trajectoires qui restent proches des clusters de démonstration du prompt, indiquant une extrapolation stable. Inversement, les trajectoires d'ICRT dérivent significativement hors de la variété (manifold) du prompt à mesure que la longueur du contexte augmente.
• Temps d'Exécution :
  • RoboSSM a démontré une mise à l'échelle quasi linéaire avec la longueur du prompt, tandis que le temps d'exécution d'ICRT augmentait de manière quadratique, créant un écart de performance croissant en faveur de RoboSSM pour les contextes longs.

5. Signification et Revendications

L'article affirme que RoboSSM représente un changement de paradigme pour l'apprentissage par imitation in-context en établissant les Modèles d'Espace d'États comme une colonne vertébrale efficace et scalable. La principale signification réside dans le dépassement des goulots d'étranglement computationnels et d'extrapolation des Transformers, permettant aux robots de :

• Gérer des fenêtres de contexte substantiellement plus longues au moment du test sans réentraînement.
• S'adapter à des tâches non vues, à long horizon et temporellement variables en utilisant des prompts à faible échantillonnage.
• Soutenir l'adaptation continue pour l'apprentissage à vie (lifelong learning) en ingérant simplement de nouveaux prompts de démonstration sans mises à jour de paramètres spécifiques à la tâche.

Les auteurs concluent que bien que RoboSSM démontre le potentiel des SSM pour l'ICIL à long contexte, les travaux futurs pourraient nécessiter des corpus d'entraînement plus larges et plus diversifiés pour améliorer davantage la généralisation. Ils notent également que bien que la validation sur robot réel ne soit pas la contribution centrale de ce travail, elle reste une avenue potentielle pour une démonstration future de la généralisation.