用状态空间模型提升机器人少样本任务学习的长上下文处理能力
RoboSSM: Scalable In-context Imitation Learning via State-Space Models
- 用长时序状态空间模型替代Transformer，实现线性推理
- 在LIBERO上对长序列任务泛化性能超越基于Transformer的方法
- 适合需要快速适应新任务的机器人系统开发者
上下文内模仿学习（ICIL）使机器人仅需少量示范即可学习新任务，无需部署时参数更新，支持少样本任务适应。然而现有ICIL方法依赖Transformer，存在计算瓶颈，且在测试时处理比训练更长的提示时表现下降。本文提出RoboSSM，基于状态空间模型（SSM）的可扩展模仿学习方案。具体采用Longhorn——一种先进的状态空间模型，具备线性推理速度与强外推能力，适用于长上下文提示。在LIBERO基准上的多组实验表明，该方法能有效处理测试时更长的上下文，在未见任务和长时序任务上均优于基于Transformer的ICIL方法。首次证明状态空间模型是高效、可扩展的ICIL骨干架构。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。