A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks

📄 arXiv: 2410.22391v3 📥 PDF

作者: Thomas Schmied, Thomas Adler, Vihang Patil, Maximilian Beck, Korbinian Pöppel, Johannes Brandstetter, Günter Klambauer, Razvan Pascanu, Sepp Hochreiter

分类: cs.LG, cs.AI

发布日期: 2024-10-29 (更新: 2025-06-04)

💡 一句话要点

提出基于xLSTM的大型循环动作模型LRAM，加速机器人任务推理。

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

关键词: 机器人控制 强化学习 动作模型 循环神经网络 xLSTM 序列建模 实时推理

📋 核心要点

  1. Transformer动作模型推理速度慢，难以应用于机器人等实时性要求高的场景。
  2. 提出LRAM，核心是xLSTM循环架构，兼顾训练并行性和快速推理能力。
  3. 实验表明，LRAM在多个机器人任务上性能与Transformer相当，但速度更快。

📝 摘要（中文）

近年来，强化学习领域出现了一种趋势，即通过序列建模在大型数据集上离线训练大型动作模型。现有的模型主要基于Transformer架构，产生了强大的智能体。然而，由于推理速度慢，基于Transformer的方法对于机器人等实时应用来说是不切实际的。最近，诸如xLSTM和Mamba等现代循环架构被提出，它们在训练期间表现出与Transformer架构类似的并行化优势，同时提供快速推理。在这项工作中，我们研究了这些现代循环架构在大型动作模型中的适用性。因此，我们提出了一个以xLSTM为核心的大型循环动作模型（LRAM），它具有线性时间推理复杂度和自然的序列长度外推能力。在来自6个领域的432个任务上的实验表明，LRAM在性能和速度方面与Transformer相比具有优势。

🔬 方法详解

问题定义：现有基于Transformer的动作模型在机器人等实时性要求高的场景中，由于推理速度慢而难以应用。Transformer的自注意力机制虽然在训练时可以并行化，但在推理时需要逐层计算，导致计算复杂度较高。

核心思路：利用现代循环神经网络架构xLSTM的优势，xLSTM在训练时可以并行化，同时在推理时具有线性时间复杂度，从而实现快速推理。LRAM旨在替代Transformer，成为一种更适合实时机器人任务的大型动作模型。

技术框架：LRAM（Large Recurrent Action Model）的核心是一个xLSTM层。整体架构上，LRAM接收环境状态作为输入，经过xLSTM层处理后，输出动作序列。模型训练采用离线方式，在大规模数据集上进行序列建模。

关键创新：LRAM的关键创新在于使用xLSTM作为其核心循环层。与传统的LSTM相比，xLSTM具有更好的长期依赖建模能力，同时在推理速度上优于Transformer。此外，LRAM还具有自然的序列长度外推能力，可以处理不同长度的动作序列。

关键设计：LRAM的具体参数设置（如xLSTM的层数、隐藏层大小等）需要根据具体的任务进行调整。损失函数通常采用序列建模中常用的交叉熵损失或均方误差损失。网络结构方面，除了xLSTM层外，还可以包含一些额外的线性层或非线性激活函数，以提高模型的表达能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在包含432个任务的6个领域上的实验表明，LRAM在性能上与Transformer相当，但在推理速度上明显优于Transformer。具体来说，LRAM在某些任务上甚至超过了Transformer的性能，同时保持了线性时间复杂度的推理速度。这些结果表明，LRAM是一种很有前途的机器人动作模型。

🎯 应用场景

LRAM适用于各种需要实时控制的机器人任务，例如自动驾驶、机械臂操作、无人机控制等。其快速推理能力使得机器人能够更快地响应环境变化，从而提高任务的完成效率和安全性。此外，LRAM的序列长度外推能力也使得其能够处理更复杂的任务，例如需要规划较长时间的动作序列的任务。

📄 摘要（原文）

In recent years, there has been a trend in the field of Reinforcement Learning (RL) towards large action models trained offline on large-scale datasets via sequence modeling. Existing models are primarily based on the Transformer architecture, which result in powerful agents. However, due to slow inference times, Transformer-based approaches are impractical for real-time applications, such as robotics. Recently, modern recurrent architectures, such as xLSTM and Mamba, have been proposed that exhibit parallelization benefits during training similar to the Transformer architecture while offering fast inference. In this work, we study the aptitude of these modern recurrent architectures for large action models. Consequently, we propose a Large Recurrent Action Model (LRAM) with an xLSTM at its core that comes with linear-time inference complexity and natural sequence length extrapolation abilities. Experiments on 432 tasks from 6 domains show that LRAM compares favorably to Transformers in terms of performance and speed.