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

Authors: Thomas Schmied, Thomas Adler, Vihang Prakash Patil, Maximilian Beck, Korbinian Pöppel, Johannes Brandstetter, Günter Klambauer, Razvan Pascanu, Sepp Hochreiter
Primary area: reinforcement learning
Keywords: reinforcement learning, rnn, xlstm, mamba, multi-task, robotics

Abstract:
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 results 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.

Decision: reject
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730203342605

thread_type:
official_review

<thread_object>
creator: reviewer_uztk

<content>
summary:
This paper investigates recurrent architectures for large action models that enable efficient inference. It proposes the Large Recurrent Action Model (LRAM), featuring an xLSTM core, which achieves linear-time inference complexity and supports natural extrapolation to longer sequence lengths. Experiments conducted on 432 tasks across six domains demonstrate that LRAM performs competitively with Transformers, offering advantages in both speed and performance.

strengths:
1. The paper is well-written and easy-to-read.
2. The study compiles 894 million transitions and conducts experiments across 432 tasks in six domains, demonstrating a substantial research effort.

weaknesses:
1. There is a lack of analysis regarding certain experimental observations. Specifically, an explanation for the lower performance of the Decision Transformer (DT) compared to recurrent-based agents would be valuable, especially since DT has more parameters, which, according to scaling laws, would be expected to yield higher performance.
2. Additionally, I am interested in understanding why xLSTM outperforms Mamba, as both architectures are designed for linear-time inference complexity. In Line 300, what is meant by "This is an important advantage of xLSTM, as LRAM agents can strongly benefit from more data and consequently larger models"?
3. It would enhance clarity to make key findings more prominent, as in DeMa [1], which also investigates the benefits of integrating Mamba within sequence modeling frameworks.

[1] Yang Dai, et al. Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning? 2024

questions:
1. Why is only xLSTM mentioned as the core of LRAM in the abstract and introduction? The authors also indicate that both xLSTM and Mamba provide parallelization benefits during training similar to the Transformer architecture while enabling faster inference.
2. How is the shared action head used across all environments, given that each environment may have a different output dimension and that continuous actions require discretization?
3. In Figure 2(b), should the performance of different environments be weighted according to the number of tasks in each environment to prevent any single environment from disproportionately influencing the results? For example, only in ProcGen, xLSTM [1:0] appears to show a distinct advantage.
4. Why did the authors choose to place the results of fine-tuning in Section 4.3 in the appendix, considering that there is still space to present these results in the main paper?
5. I am interested in the training details for multi-task training. With 432 tasks, could the learning curve be affected by conflicting gradients inherent to multi-task learning?
6. In Section 4.4, the authors investigate inference time within a simulation environment. However, during inference, it is common to set $B=1$, where the latency of DT and xLSTM is comparable. A notable difference in inference time only appears when $B=16$. Given this, is the use of xLSTM truly necessary?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730443343842

thread_type:
official_review

<thread_object>
creator: reviewer_c3kt

<content>
summary:
This paper explores the application of a recurrent model (xLSTM) in robotic policy learning, introducing a novel approach called the Large Recurrent Action Model (LRAM), which enhances inference efficiency. In comparison to conventional sequential models used in Offline RL, such as the Decision Transformer, LRAM demonstrates superior performance on tasks within the Procgen and DMControl environments, while achieving competitive results on other benchmarks like Meta-World, Atari, MimicGen, and Composuite. Notably, LRAM benefits from significantly higher throughput, presenting a distinct advantage. Comprehensive experiments and ablation studies highlight the effectiveness of LRAM, incorporating xLSTM or Mamba as its core, showing favorable evaluation performance across various model scales when compared to Transformer-based approaches.

strengths:
1. The paper is well-written and organized, making it easy to follow.

2. The reviewer appreciates the solid and robust experiments across diverse benchmarks, including thorough ablation studies on context length and action token.

3. Analysis provided in Section 4.3 makes the paper an insightful and valuable study, which may benefit the community.

4. The inclusion of detailed explanations for dataset preparation and training processes is appreciated.

weaknesses:
The reviewer believes the following aspects can be further improved:

1. The authors are encouraged to discuss the key differences between xLSTM, Mamba, and a typical transformer (such as GPT-2) used in the Decision Transformer.

2. It would be better if the authors could present the normalized performance of the offline datasets (ground truth) presented in Figure 2b.

**Minor**

1. The bottom border of Table 5 is missing.

2. In Lines 1282-1283, the scientific notation should be written as $1e-4$ instead of $1e^{-4}$
.

questions:
1. How does the ratio of mLSTM to sLSTM blocks impact performance? Could you provide further observations on specific configurations, such as xLSTM[3:1]?

2. As a follow-up, what would happen if Transformer layers in DT were gradually replaced with sLSTM or mLSTM layers?

3. Could you elaborate on the visual encoder referenced in Lines 186-187? Specifically, is it jointly trained with the downstream policy, and how is it initialized?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730682385960

thread_type:
official_review

<thread_object>
creator: reviewer_vb2e

<content>
summary:
This paper conducts a comprehensive empirical study on architectural choices for sequence modeling in offline reinforcement learning. The results show that xLSTM models scale effectively with training tasks and data, outperforming transformers and mamba in both policy performance and inference speed.

strengths:
* The empirical evaluation is thorough and extensive. THe paper also provides rich experimental details in appendix, providing valuable comparisons across different architectures.
* The release of the data pipeline and dataset can be a useful resource for the community.

weaknesses:
* The paper is built upon the decision transformer approach to offline RL. However, this approach is not widely used in the recent recent developments of large-scale embodied or robotic foundation models. To show that xLSTM is a competative architecture, it would be important to validate its advantages for other learning approaches, e.g., generative behavior cloning.
* The ablation study examines the influence of context length, both including and excluding action; however, it looks that even without actions, all choices of context length converge to similar performance. This result goes against the necessity of using xLSTM or other forms of recurrent models. It would be good to more clearly justify when xLSTM could benefit from extended history. 
* In terms of specific architecture designs, the paper considers xLSTM [7:1] and [1:0]. However, the rationale behind these choices is unclear. Moreover, the results look puzzling: the former works better in atari, while the latter works better in procgen. Providing deeper insight into these choices would strengthen the paper.

questions:
In line 380, the authors note `While removing actions improves performance on the robotics domains, it does not affect performance on discrete control. We assume this is because ... by observing previous actions, the agent learns shortcut`.

To verify this explanation, could you provide correlation metrics between sequeantial actions in robotic vs discrete tasks, similar to Fig 3 in [1]?

[1] Fighting Copycat Agents in Behavioral Cloning from Observation Histories, NeurIPS 2020
</content>
</thread_object>

</thread>
<thread>
creation date:
1730809007839

thread_type:
official_review

<thread_object>
creator: reviewer_oyy1

<content>
summary:
The paper introduces a Large Recurrent Action Model (LRAM) that utilizes xLSTM to address the limitations of Transformer-based models, specifically their slow inference times, which make them impractical for real-time applications like robotics. The authors claim that LRAM offers linear-time inference complexity and natural sequence length extrapolation capabilities. They support these claims with experiments across various domains and 432 tasks, demonstrating favorable performance and speed compared to Transformers.

strengths:
1. The paper provides a thorough experimental evaluation, covering a wide range of domains and tasks, including common RL benchmarks (Atari, DMControl, etc), robotics simulation environments, and generalization benchmarks.
2. The empirical studies on these benchmarks, especially the comparison between Transformer-based models and LSTM-based models, have brought some useful practices to the community.

weaknesses:
1. From current status,  the contribution of this paper may be limited. I appreciate the comprehensive experiments in this paper, while the use of xLSTM for faster inference is not surprising, as LSTM-based models are typically faster than Transformer-based models.
2. I think it would be better if the paper discuss when and why to use LSTM-based models versus Transformer-based models. If the primary goal is to improve the inference speed of Transformers, there are already several engineering solutions, such as FlashAttention, that can significantly accelerate Transformer inference.
3. From Figure 6, it appears that the inference speed of Decision Transformers (DT) is acceptable when the batch size (B) is 1, which suggests that the need for an alternative model might be less urgent in some scenarios. As far as I known, in real-world applications, most robotics models perform inference at batch size 1. I suggest the authors do some investigation into scenarios where LSTM-based models are most beneficial.
4.I did some comparative experiments between Transformer and LSTMs architectures. My own visualization experiments suggest that Transformers and LSTMs have different focus points. Transformers tend to have longer memory, while LSTMs often focus on recent states. For more challenging tasks, such as those with sparse rewards, Transformer-based models may perform better due to their ability to handle long-term dependencies, which is consistent to the original DT paper. Maybe the authors could do some similar visualization experiments to investigate the difference among vanilla LSTM, xLSTM and Transformers on sparse and dense reward tasks, respectively.

questions:
I suggest the authors provide a more detailed comparison with other recurrent architectures and explain why xLSTM is a superior choice for this specific application. Additionally, the authors may offer more insights into the trade-offs between LSTM-based and Transformer-based models, and under what conditions each is more suitable.
This may not be a question, but I think it will be valuable to the community if the paper can further investigate the above point.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732178205179

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers, 

We thank you for your helpful comments and constructive feedback! We addressed your questions and comments in our individual responses. Your comments convinced us to perform additional experiments, which we report in the revised manuscript.
We believe that these changes considerably improved our paper - thank you! We appeal to the reviewers to re-check their assessment in light of our responses and the new version of our paper.

Kind Regards
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work conducts an empirical study of using xLSTM as a large-scale recurrent neural networks for sequence modeling in offline RL. The experimental result shows that the xLSTM models scale better when training with more tasks and tasks, outperforming transformers and mamba in performance and inference speed. 

After rebuttal and discussion, this submission receives review scores of 5,6,6,6. So it is a borderline paper. 

Since it is an experiment paper, reviewers mainly raised concerns about more comparisons and more experiments, which the authors did a good job addressing. 

The reviewer who gave a 5 score argued that the contribution of this paper may be limited, which is echoed by the AC. 

AC has checked the submission, the reviews, and the discussions. AC appreciated the comprehensive experiments done by the authors in this paper. However, AC couldn't find much depth and intellectual merits in this submission. This work is indeed a wonderful experimental report, but AC expected an ICLR paper also needs to go beyond an experimental report and exhibit some intellectual novelty or methodology contributions.   For example, after investigating the architecture comparison, the authors can offer more insights into the trade-offs, strengths, and weaknesses of LSTM and Transformer models under certain conditions and develop a new hybrid method to take the best of the two worlds. Furthermore, no real-world robotics tasks were evaluated in the submission. It remains unknown how much this evaluated architecture can generalize to the real world. 

Given the above reasons, AC recommends a rejection. 

AC has communicated this decision with the Senior AC, and the Senior AC has gone through the submission and discussion and confirmed the AC's rejection decision. AC and senior AC believe the paper would be a better fit for a robotics conference, where the robotics community will be excited about applying existing but new architecture to robotic tasks.