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

Thank you for your constructive feedback. We appreciate your positive assessment and are very glad that you believe that our work may mark a potential turning point for real-world agent design.

Real-World Experiments: We agree that experiments in real-world robotics settings would be valuable. While we believe that such experiment are out-of-scope for the current work, evaluating modern recurrent architectures in real-world robotics settings is an exciting direction for future research. Nevertheless, we believe our findings will transfer to real-world scenarios and that the advantages of recurrent architectures may be even more pronounced in real-world settings. For example, we conduct our inference time comparison on a high-end data-center GPU with 40GB of RAM and the Transformer runs OOM for larger context sizes. In contrast, applications on edge devices in real-world scenarios, may have to deal with less powerful accelerators, making the use of modern recurrent architectures attractive. Moreover, their ability to handle long sequences without increasing computational cost/requirements can be particularly beneficial for complex real-world applications, which often exhibit longer-term dependencies.

Memory-intensive Environments: We again agree that exploring the advantages of modern recurrent architectures in long-term memory retention more would be interesting. While we cannot provide a detailed study within the time-frame of the rebuttal, we believe that this direction is interesting for future work. Note that the experiments on Dark-Room, which exhibits sparse rewards and a partially-observable observation space, go in a related direction (Figure 4 and 16). There we find that recurrent backbones compare favorably to Transformers (especially xLSTM [7:1], which enables state-tracking). In the meantime, we refer to [1] for a comparison of vanilla LSTM to the Transformer in memory intensive RL environments. Furthermore, we refer to the Figure 5 in [2], for a comparison outside the field of RL of xLSTM and Transformer on associative recall tasks.

Discretization for Continuous Actions: The reviewer is correct, that we make use of discretization of continuous actions similar to prior works. The 432 task we consider exhibit both discrete/continuous action inputs and image/vector-based state representations. The use of discretization in our large action models is motivated by the need to handle both discrete and continuous action spaces. To this end, every action dimension is discretized into 256 uniformly spaced bins. The shared action head is used to predict the action bins of all continuous action dimensions jointly, which allows handling envs with different action spaces. Furthermore, because of the shared action head, we do not (have to) rely on autoregressive prediction of continuous action dimensions, which further speeds up inference for all backbones. At inference time, the number of dimensions of the current environment is known, and we extract the respective dimensions from the joint predictions. We describe this procedure in detail in Appendix B.3.

Alternative Designs: In this work, our goal is to better understand whether modern recurrent backbones can be alternatives for large action models (LAMs) and focus our analyses on the DT setup as correctly identified by the reviewer. DTs rely on return-conditioning, but existing large-scale robotic models instead rely on behavior cloning (e.g., [2]). Therefore, we validate that our findings also transfer to a behavior cloning setting at the 206M scale (Section 4.3 and Appendix E.2). Exploring alternative architectures beyond the DT and BC settings is an exciting direction for future work. We believe that exploring MoEs [4,5] for multi-task settings, which may require specialization, can be a promising approach. Moreover, we would like to point the reviewer to an interesting concurrent work that studies the design space of imitation learning policies [6].

Once again, thank you for your helpful comments and positive assessment of our work.

[1] “When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment”, NeurIPS, 2023
[2] “xLSTM: Extended Long Short-Term Memory”, NeurIPS, 2024
[3] “RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control”, ArXiv, 2023
[4] “Mixture of Experts in a Mixture of RL settings”, RLC, 2024
[5] “Mixtures of Experts Unlock Parameter Scaling for Deep RL”, ICML, 2024
[6] “X-IL: Exploring the Design Space of Imitation Learning Policies”, ArXiv, 2025

Thank you for your helpful feedback and positive assessment of our work.

Imitation Learning: We agree with the reviewer that studying modern recurrent architectures in settings other than the DT setting is important. In this work, our goal is to better understand whether modern recurrent backbones can be alternatives for large action models (LAMs) and focus our analyses on the DT setup. However, we want to highlight that we already provide experiments in an imitation learning / behavior cloning (BC) setting at the 206M scale, as suggested by the reviewer (Section 4.3 and Appendix E.2). Notably, performance trends across backbones mirror the results in the DT setting (Figures 27, 28), which indicates that our findings generalize beyond the DT framework. Moreover, we agree with the reviewer that studying modern recurrent architectures in RL is interesting. While we believe that such experiments are out-of-scope for this work, we listed online RL with modern recurrent architectures in our future work section (Lines 431-435).

Differences between Backbones: While we discuss when and why to use which backbone (Section 5), we agree with the reviewer that an extended discussion on the reasons for performance differences between backbones is useful:

• We observed that xLSTM outperforms Transformers/Mamba in terms of sequence prediction and env performance (Figure 2). One benefit of xLSTM is that it enables state tracking [1] via sLSTM blocks, which Transformers/Mamba cannot. This property can be useful for partially observable environments and may be a reason for the enhanced ICL performance (Figure 4). Another reason may be the enhanced domain-separation in the embedding space (Figure 5), which could facilitate task identification at inference time. Moreover, our results reflect performance improvements in language settings [2].
• Nevertheless, the differences may depend on the environment. For example, Transformers may have advantages for tasks where exact recall is required. For such tasks, self-attention is typically superior (Figure 5 in [2]) and can be important for decision-making tasks [3]. Therefore, the choice for the right backbone may depend on the task at hand.

Long Context:

• Generally, performance improves when training with longer sequences both for DT and xLSTM (Section 4.3). For DT, the avg. normalized score increases from 0.3 with C=1 to 0.7 with C=50 (Figure 23). Similarly, for xLSTM it increases from 0.4 to 0.8 (Figure 25). Furthermore, domains with longer episode lengths, like DMC or Atari, benefit more from increasing context (Figures 24, 26). This is because the history helps to predict the next action (e.g., by observing past mistakes). This highlights that LAMs can strongly benefit from increased context length, even on the simulated environments we consider.
• We agree with the reviewer that for our environments, increasing the context to several thousand timesteps may not result in better performance. However, we believe that the ability to handle longer sequences can be beneficial for complex real-world applications that exhibit longer-term dependencies. Similarly, longer context can be beneficial for ICL applications, which benefits from keeping multiple episodes in the context, as indicated in Figure 4.

We added a short discussion on long sequences in real-world tasks to our manuscript.

Experiment Setup: Yes, we train a single model on datasets comprising 432 tasks. The environments contain both discrete/continuous action inputs and image/vector-based state representations. To handle both discrete/continuous actions, we make use of discretization. Every action dimension is discretized into 256 uniformly spaced bins. The shared action head is used to predict the action bins of all continuous dimensions jointly, which allows handling envs with different action spaces. At inference time, the number of dimensions of the current environment is known, and we extract the respective dimensions from the joint predictions. We describe this procedure in detail in Appendix B.3.

Inference Speed: The reviewer is correct that the benefits of modern recurrent architectures become more important at longer sequence lengths. Importantly, the inference speed and memory consumption of xLSTM/Mamba does not change with increasing sequence lengths, which is particularly apparent in memory consumption (Figure 6c). It is true that in our experiments there are some inference time speed-ups at shorter context lengths, which may result from the different kernel implementations. Note that the sequence lengths are measured in timesteps with each timestep containing 3 tokens (largest C is 76K tokens).

We revised our manuscript, and hope to have clarified your questions.

[1] The Illusion of State in State-Space Models, ICML 2024
[2] xLSTM: Extended Long Short-Term Memory, NeurIPS, 2024
[3] When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment, NeurIPS, 2023

Thank you for your helpful feedback. We are glad that you consider our experiments well-founded and the paper well written.

Embedding space:

• Regarding your question on the embedding space, we want to clarify that Figure 5 is constructed using the aggregated hidden states (averaged across the sequence) from the final layer of the respective agents and visualized via UMAP (see Appendix F for details).
• The purpose of this visualization is to examine how the models organize their representations of different environments. In general, tasks within the same domain tend to share similar input characteristics - such as visual inputs (e.g., image frames) possible actions to perform, and reward structures - and are therefore more likely to be “grouped” together in the embedding space. Consequently, when embeddings of Atari games are closer to each other than to Procgen games, it indicates that Atari games share more similar underlying dynamics or inputs structures compared to Procgen. For reference, we now also include the embedding space plots for Mamba in our updated manuscript, as suggested by Reviewer 8NAr14.
• We observe that xLSTM exhibits a slightly more refined and better-separated embedding space, which may be a reason for its better final performance. In contrast, DT produces embeddings with less clear separation between domains. This suggests that for the environments we consider in this work, it may be beneficial for the model to learn more “separate” representations, potentially because it facilitates task identification at inference time. However, we fully agree with the reviewer that generalization across domains is generally desirable if environments share structural similarities. Therefore, we believe that studying the learned embedding spaces of multi-task agents in environments that share more structure across domains is interesting for future work.

We hope to have clarified your remaining question. Thank you again for your positive assessment of our work.

Thank you for your helpful feedback on our work and your positive assessment. We are glad that you consider it a clear demonstration, robust experimental validation and valuable dataset pipeline. We address your open points in the following.

Policy formulation: We agree with the reviewer that the differences between our policy formulation that omits actions (Equation 2) and the formulation in Equation 1 can be highlighted more clearly. Many thanks for the suggestion of using color coding. We now introduced color coding and believe it enhances readability.

Data ratios:

• A key motivation behind our dataset compilation was the scarcity of suitable existing datasets that span a large number of simulated tasks. To address this, our primary target was to assemble a collection of datasets that span as many tasks as possible to enable a robust comparison of sequence model architectures. To facilitate usability for future works, we consider standard benchmarks (e.g., Meta-World, Atari, Procgen) that are widely adopted by the community. Therefore, we hope that our data pipeline can serve as a solid basis for future research on multi-task agents.
• Note that during pre-training, every domain is represented with approximately equal proportion in every update step (see Section 4.2 and Appendix C). Because the dataset sizes vary across domains (due to different numbers of tasks), this results in different numbers of total repetitions per dataset (Table 1). We opt for equal proportions, because we aim to study how the different backbones perform across domains, rather than optimizing performance on specific domains.
• We agree with the reviewer that understanding the impact of the data ratios on multitask capabilities would be insightful. Varying the data ratios would, for example, allow studying potential interferences between the 432 tasks. While we believe that a thorough investigation on data ratios is out-of-scope for this work, because of the excessive cost or pre-training with varying ratios, we believe that studying data ratios in large action models represents an interesting direction for future work. We hope that our data pipeline and datasets may facilitate such future studies.

We updated our manuscript to include a short discussion about the rationales behind our data ratios.

Embedding plots for Mamba: Thank you for the suggestion. We added the embedding plots for Mamba to Appendix F in our manuscript (we cannot update the PDF). Mamba exhibits a more refined embedding space than DT, but slightly less refined than xLSTM. However, while a more refined embedding may help with task identification at inference time, it is unclear how strongly it benefits final performance.

Inference speed of Mamba: We agree with the reviewer that the inference speeds for Mamba were missing from our original manuscript and are important to provide the entire picture. Initially, we focused our inference time tests on DT and xLSTM, as xLSTM tends to be slightly faster than Mamba in prior work (see Figure 9 in [1]). Therefore, we now repeated the inference time analyses conducted in Section 4.4 with Mamba and provide the detailed results in our updated manuscript. As expected, we find that Mamba exhibits the same linear scaling properties as xLSTM. Consequently, Mamba exhibits advantages over DT in terms of speed and memory for longer sequence lengths. In our experiments, Mamba runs slightly slower compared to xLSTM (0.002 sec/step vs. 0.0048 for B=16 in Figure 6b), but requires slightly less GPU RAM. We hypothesize that this is because the Mamba kernels are not compatible with torch.compile, which may result in a slowdown. With compatible kernels, that gap might be closed. Moreover, Mamba achieves throughputs similar to xLSTM. When using larger batch sizes with xLSTM, we found that decreasing the head dimension (more heads, same total hidden dim) is important for enabling higher throughput, and added an ablation on this to the Appendix. This is because a higher head dimension incurs more FLOPS. Generally, these results reflect findings in existing works [1,2].

Our manuscript improved considerably by addressing and incorporating your comments: thank you. In particular, believe that the additional results for Mamba strengthen the paper.

[1] “xLSTM: Extended Long Short-Term Memory”, NeurIPS, 2024
[2] “xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference”, Arxiv, 2024