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

📄 arXiv: 2509.19658v1 📥 PDF

作者: Youngju Yoo, Jiaheng Hu, Yifeng Zhu, Bo Liu, Qiang Liu, Roberto Martín-Martín, Peter Stone

分类: cs.RO, cs.AI

发布日期: 2025-09-24

备注: 8 pages, 11 figures

🔗 代码/项目: GITHUB

💡 一句话要点

RoboSSM：基于状态空间模型实现可扩展的上下文模仿学习

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 上下文模仿学习 状态空间模型 机器人学习 长序列建模 少样本学习

📋 核心要点

  1. 现有基于Transformer的上下文模仿学习方法在计算上存在瓶颈，且难以泛化到更长的上下文序列。
  2. RoboSSM采用Longhorn状态空间模型替代Transformer，利用其线性推理时间和强大的外推能力处理长上下文提示。
  3. 实验表明，RoboSSM在LIBERO基准测试中，能够有效处理不同长度的演示，并在新任务和长时程场景中表现出色。

📝 摘要（中文）

本文提出RoboSSM，一种基于状态空间模型（SSM）的可扩展上下文模仿学习（ICIL）方法。ICIL允许机器人仅从少量演示中学习任务，无需部署时进行参数更新，从而支持对新任务的少样本适应。然而，现有的ICIL方法依赖于Transformer，存在计算限制，并且在处理比训练期间更长的提示时表现不佳。RoboSSM用Longhorn（一种最先进的SSM）取代Transformer，Longhorn提供线性时间推理和强大的外推能力，非常适合长上下文提示。在LIBERO基准上评估了该方法，并与基于Transformer的ICIL基线进行了比较。实验表明，RoboSSM能有效地外推到不同数量的上下文演示，在新任务上产生高性能，并在长时程场景中保持鲁棒性。这些结果突出了SSM作为ICIL高效且可扩展骨干网络的潜力。

🔬 方法详解

问题定义：现有的上下文模仿学习（ICIL）方法，特别是基于Transformer的方法，在处理长序列的上下文信息时面临计算效率和泛化能力的挑战。Transformer的计算复杂度随序列长度呈平方增长，限制了其在长时程任务中的应用。此外，当部署时遇到的上下文长度超过训练时长度时，Transformer的性能会显著下降。

核心思路：RoboSSM的核心思路是利用状态空间模型（SSM），特别是Longhorn SSM，来替代Transformer作为ICIL的骨干网络。SSM具有线性时间复杂度的推理能力，使其能够高效地处理长序列。Longhorn SSM进一步增强了SSM的外推能力，使其能够更好地泛化到不同长度的上下文序列。

技术框架：RoboSSM的整体框架包括以下几个主要步骤：1) 将上下文演示（包括状态和动作序列）输入到Longhorn SSM编码器中，生成上下文嵌入。2) 将当前状态输入到另一个Longhorn SSM解码器中。3) 解码器利用上下文嵌入来预测下一步的动作。整个模型通过模仿学习的方式进行训练，目标是最小化预测动作与真实动作之间的差异。

关键创新：RoboSSM最重要的技术创新点在于将Longhorn SSM引入到上下文模仿学习中。与Transformer相比，Longhorn SSM具有线性时间复杂度的推理能力和强大的外推能力，使其能够更有效地处理长上下文序列，并更好地泛化到不同长度的上下文。

关键设计：RoboSSM的关键设计包括：1) 使用Longhorn SSM作为编码器和解码器，以实现高效的长序列处理。2) 使用模仿学习损失函数，最小化预测动作与真实动作之间的差异。3) 针对LIBERO基准测试，对模型结构和超参数进行了优化，例如，调整了Longhorn SSM的隐藏层维度和学习率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RoboSSM在LIBERO基准测试中取得了显著的成果。实验结果表明，RoboSSM能够有效地外推到不同数量的上下文演示，并在新任务上实现了高性能。此外，RoboSSM在长时程场景中表现出很强的鲁棒性。与基于Transformer的ICIL基线相比，RoboSSM在多个指标上都取得了显著的提升，证明了SSM作为ICIL骨干网络的有效性。

🎯 应用场景

RoboSSM具有广泛的应用前景，例如机器人操作、自动驾驶、游戏AI等领域。它可以使机器人能够仅通过少量演示快速学习新任务，而无需进行大量的训练或参数调整。这对于需要快速适应新环境或新任务的机器人应用至关重要。此外，RoboSSM的强大外推能力使其能够处理长时程任务，例如复杂的装配或导航任务。

📄 摘要（原文）

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.