RoboSSM: Scalable In-context Imitation Learning via State-Space Models

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Yoo, Youngju, Hu, Jiaheng, Zhu, Yifeng, Liu, Bo, Liu, Qiang, Martín-Martín, Roberto, Stone, Peter
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866916966924550144
author Yoo, Youngju
Hu, Jiaheng
Zhu, Yifeng
Liu, Bo
Liu, Qiang
Martín-Martín, Roberto
Stone, Peter
author_facet Yoo, Youngju
Hu, Jiaheng
Zhu, Yifeng
Liu, Bo
Liu, Qiang
Martín-Martín, Roberto
Stone, Peter
contents In-context imitation learning (ICIL) enables robots to learn tasks from prompts consisting of just a handful of demonstrations. By eliminating the need for parameter updates at deployment time, this paradigm supports few-shot adaptation to novel tasks. However, recent ICIL methods rely on Transformers, which have computational limitations and tend to underperform when handling longer prompts than those seen during training. In this work, we introduce RoboSSM, a scalable recipe for in-context imitation learning based on state-space models (SSM). Specifically, RoboSSM replaces Transformers with Longhorn -- a state-of-the-art SSM that provides linear-time inference and strong extrapolation capabilities, making it well-suited for long-context prompts. We evaluate our approach on the LIBERO benchmark and compare it against strong Transformer-based ICIL baselines. Experiments show that RoboSSM extrapolates effectively to varying numbers of in-context demonstrations, yields high performance on unseen tasks, and remains robust in long-horizon scenarios. These results highlight the potential of SSMs as an efficient and scalable backbone for ICIL. Our code is available at https://github.com/youngjuY/RoboSSM.
format Preprint
id arxiv_https___arxiv_org_abs_2509_19658
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle RoboSSM: Scalable In-context Imitation Learning via State-Space Models
Yoo, Youngju
Hu, Jiaheng
Zhu, Yifeng
Liu, Bo
Liu, Qiang
Martín-Martín, Roberto
Stone, Peter
Robotics
Artificial Intelligence
In-context imitation learning (ICIL) enables robots to learn tasks from prompts consisting of just a handful of demonstrations. By eliminating the need for parameter updates at deployment time, this paradigm supports few-shot adaptation to novel tasks. However, recent ICIL methods rely on Transformers, which have computational limitations and tend to underperform when handling longer prompts than those seen during training. In this work, we introduce RoboSSM, a scalable recipe for in-context imitation learning based on state-space models (SSM). Specifically, RoboSSM replaces Transformers with Longhorn -- a state-of-the-art SSM that provides linear-time inference and strong extrapolation capabilities, making it well-suited for long-context prompts. We evaluate our approach on the LIBERO benchmark and compare it against strong Transformer-based ICIL baselines. Experiments show that RoboSSM extrapolates effectively to varying numbers of in-context demonstrations, yields high performance on unseen tasks, and remains robust in long-horizon scenarios. These results highlight the potential of SSMs as an efficient and scalable backbone for ICIL. Our code is available at https://github.com/youngjuY/RoboSSM.
title RoboSSM: Scalable In-context Imitation Learning via State-Space Models
topic Robotics
Artificial Intelligence
url https://arxiv.org/abs/2509.19658