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

arXiv: https://arxiv.org/abs/2509.19658 · HTML: https://mlsanity.com/paper/2509.19658/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 32 of its 37 references that are in the map.
- First posted on arXiv 24 Sep 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023)
- [Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success](https://mlsanity.com/paper/2502.19645/) (2025)
- [Open X-Embodiment: Robotic Learning Datasets and RT-X Models](https://mlsanity.com/paper/2310.08864/) (2023)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Behavioral Cloning from Observation](https://mlsanity.com/paper/1805.01954/) (2018)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [Deeply AggreVaTeD: Differentiable Imitation Learning for Sequential Prediction](https://mlsanity.com/paper/1703.01030/) (2017)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)
- [Keypoint Action Tokens Enable In-Context Imitation Learning in Robotics](https://mlsanity.com/paper/2403.19578/) (2024)
- [In-Context Imitation Learning via Next-Token Prediction](https://mlsanity.com/paper/2408.15980/) (2024)
- [Longhorn: State Space Models are Amortized Online Learners](https://mlsanity.com/paper/2407.14207/) (2024)
- [FLaRe: Achieving Masterful and Adaptive Robot Policies with Large-Scale Reinforcement Learning Fine-Tuning](https://mlsanity.com/paper/2409.16578/) (2024)
- [Towards More Generalizable One-shot Visual Imitation Learning](https://mlsanity.com/paper/2110.13423/) (2021)
- [Imitation Learning for Human Pose Prediction](https://mlsanity.com/paper/1909.03449/) (2019)
- [Demonstrate Once, Imitate Immediately (DOME): Learning Visual Servoing for One-Shot Imitation Learning](https://mlsanity.com/paper/2204.02863/) (2022)

## Built on by (most cited first, 5 of 5)

- [From Abstraction to Instantiation: Learning Behavioral Representation for Vision-Language-Action Model](https://mlsanity.com/paper/2605.22671/) (2026)
- [AnoleVLA: Lightweight Vision-Language-Action Model with Deep State Space Models for Mobile Manipulation](https://mlsanity.com/paper/2603.15046/) (2026)
- [RA-VLA: Retrieval-Augmented VLA for Test-Time Adaptation](https://mlsanity.com/paper/2608.25585/) (2026)
- [Non-Commutative State Tracking with Input-Dependent Low-Rank Updates in Mamba-3](https://mlsanity.com/paper/2609.28273/) (2026)
- [ReCAT: Remember, Count, and Time: Structured Recurrent Memory for Robot Manipulation](https://mlsanity.com/paper/2609.35200/) (2026)

## Abstract (arXiv)

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. Through diverse experiments on the LIBERO benchmark, we demonstrate the effectiveness of applying SSMs to ICIL, achieving improved generalization to both unseen and long-horizon tasks than Transformer-based ICIL methods by handling longer contexts at test-time. These results show for the first time that SSMs are an efficient and scalable backbone for ICIL. Our code is available at https://github.com/youngjuY/RoboSSM.