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

1. Problem Statement

In-context Imitation Learning (ICIL) enables robots to adapt to novel tasks using few-shot demonstrations without parameter updates at deployment. While recent ICIL methods have successfully adopted Transformer-based architectures to formulate imitation learning as a sequence modeling problem, they face significant scalability limitations. Transformers exhibit quadratic time complexity ($O(L^2)$) relative to sequence length and struggle to extrapolate beyond the prompt lengths seen during training. Consequently, existing ICIL methods degrade sharply when handling longer test-time prompts (e.g., more demonstrations or long-horizon tasks) than those used in training, limiting their applicability to complex, real-world scenarios requiring extensive context.

2. Methodology: RoboSSM

The authors propose RoboSSM, a scalable ICIL framework that replaces Transformers with State-Space Models (SSMs), specifically utilizing Longhorn as the backbone architecture.

Architecture and Encoding

• Input Processing: RoboSSM encodes multimodal observations (visual data via CNNs and proprioceptive data via MLPs) into embeddings. Notably, the input representation excludes action tokens and task language instructions to force the model to infer task intent solely from the in-context demonstrations.
• Longhorn Backbone: The sequence of observation embeddings is processed by Longhorn, which maintains a memory state matrix $s_t$ updated via a linear recurrence. Unlike the self-attention mechanism in Transformers, Longhorn uses a recurrent update rule derived from online convex programming:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  This formulation balances retaining prior state information with incorporating new inputs.
• $\beta$-Scaling: To specifically encourage attention to demonstration prompts in the ICIL setting, the authors introduce a test-time scaling mechanism where the weighting vector $\beta_t$ is scaled by a factor $\gamma \in (0, 1]$. This modulates the trade-off between state retention and new information integration.

Training and Inference

• Training: The model is trained on a multi-task dataset where the input consists of a prompt (a set of $N_{train}$ demonstration trajectories) and a query trajectory. The model predicts actions for the query trajectory conditioned on the prompt.
• Inference: At test time, the policy adapts to unseen tasks using a prompt of $N_{test}$ demonstrations. The model iteratively predicts actions based on the initial observation and the context provided by the prompt, without any parameter updates.

3. Key Contributions

1. Scalable ICIL Framework: The introduction of RoboSSM, the first approach to leverage SSMs for in-context imitation learning, enabling linear-time inference ($O(L)$) and strong extrapolation capabilities.
2. Long-Context Extrapolation: Demonstrating that RoboSSM can effectively process test-time prompts up to 16 times longer than those seen during training, a capability where Transformer-based methods fail.
3. Efficient Inference: Achieving significantly lower inference runtimes compared to Transformer baselines as prompt length increases, due to the linear complexity of SSMs versus the quadratic complexity of Transformers.
4. Generalization to Long-Horizon Tasks: Showing that the model can generalize from short-horizon training distributions to long-horizon, multi-stage tasks and time-dilated scenarios.

4. Experimental Results

Experiments were conducted on the LIBERO benchmark, covering diverse manipulation tasks (Object, Goal, Spatial, Long, and 90 suites).

• Prompt Length Generalization:
  • When trained with only 2 demonstrations ($N_{train}=2$), RoboSSM maintained or improved success rates as the number of test demonstrations ($N_{test}$) increased up to 32.
  • In contrast, the Transformer-based baseline (ICRT) experienced a sharp performance drop and collapsed to 0% success when $N_{test} > N_{train}$.
• Long-Horizon and Time-Dilation:
  • RoboSSM successfully solved long-horizon tasks (e.g., "Put both Object A and Object B in Object C") and tasks with time-dilated demonstrations (up to 16x temporal stretching), whereas ICRT failed completely.
• In-Distribution Performance:
  • Even when $N_{train} = N_{test}$, RoboSSM consistently outperformed ICRT across most scenarios, particularly in the LIBERO-90 Study and Living Room scenes.
• Comparison to Multi-Task Learning (MTL):
  • RoboSSM and ICRT significantly outperformed standard MTL baselines (which require language instructions and no in-context prompting), confirming the efficacy of the ICIL paradigm.
• Latent Space Analysis:
  • Visualizations of latent representations showed that RoboSSM generates trajectories that remain close to the prompt demonstration clusters, indicating stable extrapolation. Conversely, ICRT trajectories drifted significantly outside the prompt manifold as context length increased.
• Runtime:
  • RoboSSM demonstrated nearly linear scaling with prompt length, while ICRT's runtime increased quadratically, creating a widening performance gap in favor of RoboSSM for long contexts.

5. Significance and Claims

The paper claims that RoboSSM represents a paradigm shift for in-context imitation learning by establishing State-Space Models as an efficient and scalable backbone. The primary significance lies in overcoming the computational and extrapolation bottlenecks of Transformers, allowing robots to:

• Handle substantially longer context windows at test time without retraining.
• Adapt to unseen, long-horizon, and temporally variable tasks using few-shot prompts.
• Support continual adaptation for lifelong learning by simply ingesting new demonstration prompts without task-specific parameter updates.

The authors conclude that while RoboSSM demonstrates the potential of SSMs for long-context ICIL, future work may require broader and more diverse training corpora to further enhance generalization to novel tasks. They also note that while real-robot validation is not the core contribution of this work, it remains a potential avenue for future demonstration of generalization.