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

## Abstract

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 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.

## 1 Introduction

Reinforcement Learning (RL) has been responsible for impressive success stories such as
game-playing [Silver et al.,, ; Vinyals et al.,, ; Berner et al.,, ; Patil et al.,, ],
plasma control for fusion [Degrave et al.,, ],
or navigation of stratospheric balloons [Bellemare et al.,, ].
While these successes were based on classical RL approaches, in which agents have been trained online with RL objectives, recently there has been a trend towards offline RL settings [Levine et al.,, ; Schweighofer et al.,, ] and sequence models trained via behavior cloning [Chen et al.,, ; Janner et al.,, ].
Such approaches, in which agents are trained on large-scale
offline datasets with causal sequence modeling objectives,
have been driven by the proliferation of Transformer-based
architectures and gave rise to what we refer to as Large Action Models (LAMs) to highlight their similarity to large language models (LLMs) [Radford et al.,, ].
LAM approaches can also be used in multi-task settings to develop generalist agents such as Gato [Reed et al.,, ].

Existing LAMs are primarily based on the Transformer [Vaswani et al.,, ] architecture.
Because of their powerful predictive performance, robotics has become an emergent application area for large models [; ; Octo Model Team et al.,, ; Gu et al.,, ; Wang et al.,, ] and a number of large multi-task datasets were collected [Jia et al.,, ; Embodiment Collaboration et al.,, ; Jiang et al.,, ; Mandlekar et al.,, ].
This development bears the potential to produce robotics agents that learn to master complex tasks in a wide range of environments and even different embodiments.
For example, recently it has been demonstrated, albeit in restricted settings, that sequence models trained on multi-episodic contexts can perform in-context learning (ICL) [Laskin et al.,, ; Lee et al.,, ].
One potential application of ICL can be to learn new related tasks in robotics without the need for re-training or fine-tuning.

One of the key reasons for the success of Transformer-based models is
their ability to scale to large datasets through their efficient parallelization during training.
However, despite numerous success stories in RL, language modeling [Brown et al.,, ]
or computer vision [Dosovitskiy et al.,, ; He et al.,, ],
a persistent drawback of Transformer-based architectures
is their high inference cost in terms of both speed and memory [Kim et al.,, ].
Consequently, deploying Transformer-based models in resource-constrained scenarios, such as on devices with limited hardware
capacity and/or real-time constraints, e.g., robots or
smartphones, is difficult because of the required
fast inference times
[Firoozi et al.,, ; Hu et al.,, ].
A basic principle of control theory is that the controller sample rate should be at least in the order of magnitude of the sample rate of the sensors [Franklin et al.,, , Ch. 11].
To illustrate this, for typical robots such as drones or industrial robot arms, rates of 100Hz-1000Hz are necessary to keep the system stable [Salzmann et al.,, ; El-Hussieny,, ; Hu et al.,, ; Chignoli et al.,, ].
This implies inference times of less than 10ms.
At 1000Hz, a 15-second movement of the agent corresponds to a sequence of 15K steps [El-Hussieny,, ] resulting in long context lengths even without ICL.
While there exists a range of techniques to make large models faster,
such as quantization [Frantar et al.,, ], distillation [Hinton et al.,, ],
or pruning [LeCun et al.,, ], the quadratic-time complexity of self attention
still remains.

Recently, *modern recurrent architectures* have been proposed, which
exhibit similar parallelization properties during training as the
Transformer architecture while offering linear-time inference complexity.
These modern recurrent architectures include xLSTM [Beck et al.,, ]
and state-space models (SSMs), such as Mamba [Gu and Dao,, ; Dao and Gu,, ] and Griffin/Hawk [De et al.,, ], and have challenged the dominance of the Transformer in language modeling but also in other domains such as computer vision [Alkin et al.,, ; Zhu et al.,, ], and biomedicine [Schmidinger et al.,, ].
More importantly, their linear-time inference makes them suitable for deployment
in scenarios with limited compute, large context sizes, and real-time requirements, such as robotics.

Figure: Figure 1: Illustration of our Large Recurrent Action Model (LRAM) with an xLSTM [Beck et al.,, ] at its core.

In this work, we assess the aptitude of modern recurrent architectures,
such as xLSTM and Mamba, as large action models.
To this end, we introduce a Large Recurrent Action Model (LRAM) with an xLSTM at its core (see Figure 1).
We train our agents on 432 tasks from 6 domains using a supervised learning setting similar
to that of the Decision Transformer [Chen et al.,, , DT].
We use data collected during online-RL training of single-task specialist agents and compile these trajectories alongside other expert demonstrations into a large-scale multi-domain dataset comprising 894M transitions.
Due to their parallelization properties, the modern recurrent architectures
considered in this work can process this large-scale training set as efficiently
as the Transformer while being faster at inference.
Experiments across 4 models sizes with our multi-task models indicate that xLSTM compares favorably to Transformers in terms of both performance and speed.
In addition, we study the effect of modern recurrent architectures on fine-tuning performance and in-context learning abilities, and find that they exhibit strong performance in both dimensions.

The main purpose of this paper is to test the hypothesis
that modern recurrent model architectures are better suited for building LAMs than Transformers.
Hereby, we make the following contributions.

- •
We propose a Large Recurrent Action Model (LRAM) with an xLSTM at its core that enables efficient inference.
- •
We assess the aptitude of modern recurrent architectures as backbones for large-action models with respect to their efficiency at inference time and overall performance in multi-task, fine-tuning, and in-context learning settings.
- •
To foster further research on large action models, we release our data preparation pipeline and generated datasets(^1^11GitHub: [https://github.com/ml-jku/LRAM](https://github.com/ml-jku/LRAM)).

## 2 Related work

Sequence Models in RL.
LSTM [Hochreiter and Schmidhuber,, ] is the dominant backbone architecture for partially observable online RL problems and has been behind achievements such as mastering Starcraft II [Vinyals et al.,, ], Dota 2 [Berner et al.,, ], and Atari [Espeholt et al.,, ; Kapturowski et al.,, ].
After the success of the Transformer in NLP [Devlin et al.,, ; Radford et al.,, ; Brown et al.,, ], computer vision [Dosovitskiy et al.,, ; He et al.,, ; Radford et al.,, ; Fürst et al.,, ] and speech recognition [Radford et al.,, ; Baevski et al.,, ], the architecture has found its way into RL.
Chen et al., [] proposed the Decision Transformer (DT) a GPT-style model [Radford et al.,, ], that learns to predict actions from offline trajectories via behavior cloning.
Trajectory Transformer [Janner et al.,, ] predicts action tokens along with states and rewards, which allows for dynamics modeling. A number of follow-up works build on the DT-architecture [Zheng et al.,, ; Wang et al.,, ; Shang et al.,, ; Meng et al.,, ; Siebenborn et al.,, ; ].
Furthermore, sequence models trained on offline data were found to exhibit in-context learning if conditioned on previous trajectories [Laskin et al.,, ; Lee et al.,, ; Kirsch et al.,, ], albeit only in limited scenarios.

Large Action Models (LAMs).
LAMs, such as the Decision Transformer,
are well suited for multi-task settings.
Lee et al., [] found that a multi-game DT can learn to play 46 Atari games.
Reed et al., [] introduced a generalist agent trained on over 600 tasks from different domains, ranging from Atari to manipulation of a robot arm.
Jiang et al., [] a Transformer for robot manipulation based on multi-modal prompts, that allow to steer the model to perform new tasks.
Recently, Raad et al., [] introduced an agent instructable via language to play a variety of commercial video games.
Since then, robotics has become an emergent area for developing LAMs [; ; Octo Model Team et al.,, ; Gu et al.,, ; Wang et al.,, ; Di Palo et al.,, ; Kim et al.,, ], also due to the availability of large-scale robotics datasets [Jia et al.,, ; Embodiment Collaboration et al.,, ; Jiang et al.,, ; Mandlekar et al.,, ].

Next-generation Sequence Modeling Architectures.
Linear recurrent models, such as state-space models (SSM, Gu et al.,, ; ; Smith et al.,, ; Orvieto et al.,, ) have challenged the dominance of the Transformer [Vaswani et al.,, ] architecture on long-range tasks [Tay et al.,, ].
The key insight of those linear RNNs was to diagonalize the recurrent state matrix and enforce stable training via an exponential parameterization [; Orvieto et al.,, ].
Since then, there have been efforts to include features such as gating from RNNs [Elman,, ; Jordan,, ; Hochreiter and Schmidhuber,, ; Cho et al.,, ]. Non-linear gates are believed to have higher expressivity, but are harder to train.
Griffin [De et al.,, ] mixes gated linear recurrences with local attention to achieve more training data efficiency than Llama-2 [Touvron et al.,, ] and better sequence extrapolation.
Mamba [Gu and Dao,, ] introduces a selection mechanism similar to gating into SSMs, which makes its state and input matrix time dependent.
This is similar to the gating mechanism of RNNs but also bears resemblance to approaches like fast weights [Schmidhuber,, ] and Linear Attention [Katharopoulos et al.,, ].
Mamba-2 [Dao and Gu,, ] highlight the connection between SSMs with input dependent state and input matrices and (Gated) Linear attention variants.
Most recently, the xLSTM [Beck et al.,, ] was proposed as an improvement over the classic LSTM [Hochreiter and Schmidhuber,, ] that combines gating, linear recurrences and recurrent weights into a single architecture for language modeling.
First, xLSTM leverages exponential gating with stabilization to RNNs for stronger emphasis on important inputs.
Second, xLSTM is composed of two variants, the mLSTM variant with an emphasis on memory that proves important in language modeling and the sLSTM variant that keeps the non-diagonalized recurrent matrix to enable state-tracking [Merrill et al.,, ].
State tracking is important in logic tasks and cannot be modeled fundamentally by linearized recurrent or state-space models like Mamba, Griffin or Transformers.

## 3 Large Recurrent Action Models

### 3.1 Background

Reinforcement Learning.
We assume the standard RL formulation via a Markov Decision Process (MDP) represented by a tuple of $(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R})$, where $\mathcal{S}$ and $\mathcal{A}$ denote state and action spaces, respectively.
At every timestep $t$ the agent observes state $s_{t}\in\mathcal{S}$, predicts action $a_{t}\in\mathcal{A}$, and receives a scalar reward $r_{t}$.
The reward is determined by the reward function $\mathcal{R}(r_{t}\mid s_{t},a_{t})$.
$\mathcal{P}(s_{t+1}\mid s_{t},a_{t})$ defines the transition dynamics and constitutes a probability distribution over next states $s_{t+1}$ when executing action $a_{t}$ in state $s_{t}$.
The goal of RL is to learn a policy $\pi(a_{t}\mid s_{t})$ that predicts an action $a_{t}$ in state $s_{t}$ that maximizes $r_{t}$.

Decision Transformer [Chen et al.,, ] casts the RL problem setting as next action prediction task via causal sequence modeling.
At training time, DT aims to learn a policy $\pi_{\theta}$ that maps future rewards to actions, which is often referred to as upside-down RL [Schmidhuber,, ].
At inference time, the DT is conditioned via a target return to emit high-reward actions.
Consequently, we assume access to a dataset $\mathcal{D}=\{\tau_{i}\}_{i=1}^{N}$ containing $N$ trajectories $\tau_{i}$ consisting of quadruplets $\tau_{i}=(s_{1},\hat{R}_{1},a_{1},r_{1},\dots,s_{T},\hat{R}_{T},a_{T},r_{T})$ of state $s_{t}$, return-to-go (RTG) $\hat{R}_{t}=\sum_{t^{\prime}=t}^{T}r_{t^{\prime}}$, action $a_{t}$, and reward $r_{t}$.
Here, $T$ refers to the length of the trajectory.
The DT $\pi_{\theta}$ is trained to predict the ground-truth action $a_{t}$ conditioned on sub-trajectories from the dataset:

$$ $\hat{a}_{t}\sim\pi_{\theta}(\hat{a}_{t}\mid s_{t-C},\hat{R}_{t-C},a_{t-C},r_{t -C},\dots,s_{t-1},\hat{R}_{t-1},a_{t-1},r_{t-1},s_{t},\hat{R}_{t}),$ (1) $$

where $C\leq T$ is the size of the context window.
In fact, Equation 1 describes the setting of the multi-game DT [Lee et al.,, ], which also includes rewards in the sequence representation.

### 3.2 Large Recurrent Action Models (LRAMs)

Our LRAM has a modern recurrent architecture at its core (see Figure 1), which comes with a parallel training and a recurrent inference mode.
We instantiate LRAM with three different variants, two different xLSTM configurations and Mamba.
Furthermore, we use a training protocol similar to that of Lee et al., [] and Reed et al., [] with some differences.

Multi-modal sequence representation.
To encode input from different environments with varying state and action spaces, we use separate encoders per modality that are shared across tasks and domains.
For encoding images we use a CNN similar to Espeholt et al., [],
whereas for low-dimensional inputs we use a fully connected network.
We refrain from patchifying images and tokenizing continuous states to avoid unnecessarily long sequences.
Similarly, we use linear layers to encode rewards and RTGs.
We omit actions in our sequence formulation, as we found that this can be detrimental to performance, in particular for continuous control tasks (see Section 4.3).
Consequently, our trajectories have the form $\tau_{i}=(s_{1},\hat{R}_{1},r_{1},\dots,s_{T},\hat{R}_{T},r_{T})$ and we train our policy $\pi_{\rho}$ to predict the ground-truth action $a_{t}$ as:

$$ $\hat{a}_{t}\sim\pi_{\rho}(\hat{a}_{t}\mid s_{t-C},\hat{R}_{t-C},r_{t-C},\dots, s_{t-1},\hat{R}_{t-1},r_{t-1},s_{t},\hat{R}_{t}).$ (2) $$

Shared action head.
Action spaces in RL typically vary across environments.
For example, in the environments we consider, there are 18 discrete actions and a maximum of 8 continuous dimensions for continuous control environments.
Therefore, we employ discretization of continuous action dimensions into 256 uniformly-spaced bins, similar to Reed et al., [] and .
Unlike prior work, we leverage a shared action head to predict all discrete actions or continuous action dimensions jointly.
We found this setup significantly reduces inference time compared to using autoregressive action prediction of continuous actions.

Recurrent inference mode.
At inference time, we leverage the recurrent backbone and maintain the hidden states of the last timestep.
This enables fast inference with linear-time complexity along the sequence length.
In addition, the recurrent-style inference is well suited for online fine-tuning via RL objectives, similar to LSTM-based policies in online RL.
To further speed-up inference, we leverage custom kernels for the xLSTM backbone (see Appendix 22).

Our unified discrete action representation enables consistent training of our agents via the cross-entropy loss as training objective across all tasks and domains, similar to Reed et al., [].
We use separate reward scales per domain and target returns per task.
Furthermore, we do not make use of timestep encodings as used by Chen et al., [], which are detrimental when episode lengths vary.
We provide additional implementation details in Appendix B.

## 4 Experiments

We study the aptitude of modern recurrent architectures as LAMs on 432 tasks from 6 domains: Atari [Bellemare et al.,, ], Composuite [Mendez et al.,, ], DMControl [Tassa et al.,, ], Meta-World [], Mimicgen [Mandlekar et al.,, ], and Procgen [].
To this end, we compile a large-scale dataset containing 894 million transitions (see Section 4.1).

Across all experiments, we compare four backbone variants: xLSTM [7:1], xLSTM [1:0] [Beck et al.,, ], Mamba [Gu and Dao,, ], and the GPT-2 style Transformer employed in the DT [Chen et al.,, ].
Following [Beck et al.,, ], we use the bracket notation for xLSTM, which indicates the ratio of mLSTM to sLSTM blocks.
For example, xLSTM [1:0] contains only mLSTM blocks.

In Section 4.2, we conduct a scaling comparison for four model sizes ranging from 16M to 208M parameters that shows that modern recurrent architectures achieve performance comparable or favorable to the Transformer baseline across different model sizes.
In Section 4.3, we study the impact of the recurrent backbones on fine-tuning performance and ICL abilities, and further analyze our trained recurrent backbones.
Finally, in Section 4.4, we empirically examine the differences at inference time in terms of latency and throughput between xLSTM-based and Transformer-based agents, which indicate a clear advantage for the recurrent backbone.

### 4.1 Datasets & Environments

**Table 1: Dataset statistics for all 432 training tasks.**
| Dataset | Tasks | Trajectories | Mean Trj. Length | Total Transitions | Repetitions |
| --- | --- | --- | --- | --- | --- |
| Atari | 41 | 136K | 2733 | 205M | 1.03$\times$ |
| Composuite | 240 | 480K | 500 | 240M | 0.87$\times$ |
| DMControl | 11 | 110K | 1000 | 110M | 1.92$\times$ |
| Meta-World | 45 | 450K | 200 | 90M | 2.34$\times$ |
| Mimicgen | 83 | 83K | 300 | 25M | 8.5$\times$ |
| Procgen | 12 | 2185K | 144 | 224M | 0.94$\times$ |
| Total | 432 | 3.4M | - | 894M | - |

Datasets. We compile a large-scale dataset comprising 432 tasks from six domains.
We leverage datasets from prior works.
For Atari, we extract 5M transitions per task from the DQN-Replay dataset released by Agarwal et al., [].
For Composuite, we leverage the datasets released by [Hussing et al.,, ].
For Meta-World, we use 2M transitions per task released by [].
For DMControl, we generate 10M transitions per task using task-specific RL agents.
For Mimicgen, we use the datasets for the 21 tasks released by [Mandlekar et al.,, ] and generate trajectories for the remaining 62 tasks.
Finally, for Procgen, we extract 20M transitions from the datasets released by [].
Our final dataset contains 3.4M trajectories and in total 894M transitions (see Table 1).
We reserve an additional 37 tasks from the same domains for zero-shot evaluation.
To foster future research, we release our data-preparation pipeline and generated data.

Environments.
Atari and Procgen come with image observations and discrete actions.
In contrast, the remaining four domains exhibit state-based observations and continuous actions.
Consequently, our experiments involve a mixture of state and action spaces as well as varying episode lengths (see Table 1).
Periodically evaluating the trained agents on all 432 tasks sequentially is time-consuming and we, therefore, distributed the evaluation across GPUs and parallel processes (see Appendix B).

Additional details on our datasets, environments are available in Appendix A.

Figure: (a) Sequence prediction
Refer to caption: extracted/5960624/figures/scaling_laws/legend.png

### 4.2 Scaling comparison

To conduct our main comparisons, we train our four backbone variants on the full training task mixture of 432 tasks.
For each architecture backbone, we report performance scores for four model sizes: 16M, 48M, 108M, and 206M parameters.
We train all models for 200K updates with a batch size of 128 and context length of 50 timesteps.
All domains are represented with approximately equal proportion, resulting in 33K updates per domain.
Additional implementation details and hyperparameters for every backbone variant and model size are available in Appendix B.

Sequence prediction performance.
In Figure 2a, we report the validation set perplexity for all backbones and model sizes averaged over the individual scores from all domains.
To achieve this, we maintain a hold-out set of trajectories for each training task (2.5%) and compute the perplexities after every 50K steps.
Both recurrent backbones outperform the Transformer baseline considerably, especially as the model sizes increase.
We provide the perplexities on the training set in Figure 13.

Figure: Figure 3: Normalized scores per domain for model size 206M. For Meta-World, DMControl, Mimicgen, Composuite and Procgen we report data-normalized scores, for Atari we report human-normalized scores.
Refer to caption: extracted/5960624/figures/per_domain/legend.png

Evaluation performance.
During training, we evaluate our agents after every 50K step in all 432 training environments.
In Figure 2b, we report the resulting normalized performances averaged across all six domains.
The recurrent backbones outperform the Transformer one across model sizes.
While xLSTM and Mamba performs similarly at smaller scales, xLSTM tends to outperform Mamba at larger scales (206M).
This is an important advantage of xLSTM, as LRAM agents can strongly benefit from more data and consequently larger models.
Note, that Mamba has a significantly higher number of parameters than competitors. For the zero-shot evaluation performances on the 37 hold-out tasks, we refer to Figure 15 in Appendix C.2.

Performance per domain.
In Figure 3, we report the normalized scores for the 206M parameter models attained on all six domains.
For Meta-World, DMControl, Mimicgen, Composuite, and Procgen we use data-normalized scores, as suggested by [Levine et al.,, ].
For Atari, we report human-normalized scores.
Overall, we observe that the xLSTM backbone outperforms competitors on three of the six domains, while all methods perform similarly on the remaining 3 domains.

These experiments suggest that modern recurrent backbones can be attractive alternatives to the Transformer architecture for building LAMs in terms of final performance.

### 4.3 Analyses & Ablations

Fine-tuning.
To assess the effect of the recurrent backbones on fine-tuning performance, we fine-tune our models on 37 held-out environments from all 6 domains.
We evaluate the fine-tuning performance of the xLSTM architecture for both the 16M parameter pretrained models and compared it against an xLSTM trained from scratch.
The pretrained LRAM outperforms the randomly initialized xLSTM model in most domains.
For detailed results, see Appendix C.3.
This suggests that fine-tuning performance is not negatively affected by switching the backbone.

Figure: Figure 4: ICL with modern recurrent architectures on Dark-Room $10\times 10$.
Refer to caption: extracted/5960624/figures/ablations/icl/legend.png

In-context Learning.
Next, we study the ICL abilities of our recurrent backbones on the Dark-Room environment considered in prior work on in-context RL [Laskin et al.,, ; Lee et al.,, ; ].
To study ICL in isolation, we train models from scratch with a multi-episodic context, which results in a large context length (we refer to Appendix C.4 for details on the experiment setup).
In particular, we adopt the Algorithm Distillation (AD, Laskin et al.,, ) framework and exchange the Transformer backbone architecture with modern recurrent architectures.
In Figure 17, we report the ICL performance on (a) 80 train and (b) 20 hold-out tasks.
We find that xLSTM [7:1] attains the highest overall scores both on training and hold-out tasks, which we attribute to the state-tracking abilities [Merrill et al.,, ] of sLSTM blocks.

Embedding space analysis.
In Figure 5, we analyze the representations learned by our model.
To this end, we sample 32 sub-trajectories from every task, extract the sequence representation at the last layer, cluster them using UMAP [McInnes et al.,, ], and color every point by its domain.
Appendix E describes the setup in greater detail.
We find that tasks from the same domain cluster together.
Furthermore, xLSTM exhibits a more refined domain separation compared to DT, which may contribute to the better down-stream performance.

Figure: (a) DT
Refer to caption: extracted/5960624/figures/embedding_space/legend_noframe.png

Removing Actions & Effect of Context Length.
We found that removing actions from the context results in better performance across backbones.
While context lengths beyond 1 hurt performance on Meta-World when training with actions, the reverse is true when training without actions (see Figure 23).
This is in contrast to recent works, which did not benefit from longer contexts [Octo Model Team et al.,, ].
While removing actions improves performance on the robotics domains, it does not affect performance on discrete control.
For robotics, we observed that the models become overly confident (high action logits), which is problematic if poor initial actions are produced.
We assume this is because in robotics actions change smoothly and by observing previous actions the agent learns shortcuts.
Removing actions from the input prevents the agent from using shortcuts.
Importantly, the evaluation performance improves as the sequence length increases, which indicates that the history helps to predict the next action (e.g., by observing mistakes made in the recent past).

We present additional ablations on the effect of reducing the number of layers in xLSTM and disabling Dropout on DT in Appendix D.3 and D.2, respectively.

### 4.4 Inference Time Comparison

Finally, we empirically examine the difference between xLSTM-based and Transformer-based agents at inference time.
Similar to De et al., [], we report both latency and throughput.
We focus our analysis on latency, as it is the more important dimension for real-time applications.

Figure: (a) $B=1$
Refer to caption: extracted/5960624/figures/inference_time/latency/legend.png

Setup.
We conduct all inference time tests on A100 GPUs with 40GB of RAM using 206M parameter models.
For the Transformer, we use KV-caching and FlashAttention [Dao,, ] as supported by PyTorch [Paszke et al.,, ].
For xLSTM, we use recurrent-style inference using custom kernels to accelerate the computations (see Figure 22 for the impact of kernel acceleration).
For both backbones, we use torch.compile.
The Transformer with KV-caching has a linear time complexity per step and quadratic in the sequence length.
In contrast, the xLSTM has a constant time complexity per step and linear in the sequence length.
Therefore, we expect speed-ups especially for longer sequences and larger batch sizes, as observed by De et al., [].
To ensure a fair comparison, we compare DT and xLSTM with the same number of layer blocks and increase the hidden size of xLSTM to match the number of parameters of DT (see Appendix D.3 for evaluation performance of these models).
We provide further details on our inference time tests in Appendix C.5.

Environment.
We conduct all inference time tests on the environment that exhibited the longest average episode lengths in our experiments, the Atari game Freeway.
Every episode in Freeway lasts for 8192 steps, which is equivalent to 24576 tokens (s/rtg/r).
We evaluate all models for 5 episodes and preserve the KV-cache/hidden state across episode boundaries.
The reported latencies and throughputs are averaged across all evaluation episodes, except for the first episode, which we discard to exclude compilation times and prefilling.
We opted for measuring the inference times during environment interaction, i.e., including simulator latency, rather than mere token generation.

Figure: Figure 7: Memory consumption during Latency comparison on A100 (% of GPU memory) for varying context lengths and $B=1$.
Refer to caption: extracted/5960624/figures/inference_time/latency/legend.png

Latency.
Similar to De et al., [], we measure latency by the average time (in seconds) taken to perform a single inference step with a fixed batch size $B$ (lower is better).
In Figure, 6, we report the latencies for varying context lengths, $C\in[50,25600]$ and two batch sizes $B\in\{1,16\}$.
Note that $C$ is in time steps and every time step contains 3 tokens (state, reward-to-go, reward).
Hence, the effective sequence length for the largest $C$ is 76800.
As expected, we find that the recurrent backbone attains lower inference latencies than the Transformer one.
As the sequence length increases, DT runs out of memory due to the increasing size of the KV cache (see Figure 7).
In contrast, the inference speeds for xLSTM are independent of the context length, and therefore enable significantly longer context lengths.
This property is particularly interesting for in-context RL, which requires keeping multiple episodes in the context [Laskin et al.,, ].
Nevertheless, our experiments highlight that the materialization of the complexity advantage (quadratic vs. linear) depends on the device, model size, batch size and the context length, which is similar to findings by De et al., [].

Throughput.
Throughput is measured by the total amount of inference steps performed per second for a model with a fixed context length.
In Figure, 8, we report the throughputs for varying batch sizes, $B\in[1,128]$ for a fixed context length of $C=1600$.
Here, the batch size can be interpreted as the number of parallel environments the agent interacts with.
As expected, we find that xLSTM attains considerably higher throughputs than the DT.
The benefit of xLSTM increases with larger batch sizes.
While the DT with quadratic complexity in the sequence length goes OOM for batch sizes above 64, the xLSTM with linear complexity can easily handle larger batch sizes.
In both experiments, the recurrent xLSTM performs favorably over the Transformer backbone.

Figure: Figure 8: Throughput comparison on A100 for varying batch sizes with $C=1600$ timesteps on the Atari Freeway environment. Missing bars for DT indicate OOM.
Refer to caption: extracted/5960624/figures/inference_time/latency/legend.png

## 5 Conclusion

In this work, we study the aptitude of modern recurrent architectures as alternatives to Transformers for building LAMs.
We found that our LRAM with an xLSTM or Mamba at its core compare favorably to the Transformer in terms of evaluation performance across different model scales.
Moreover, we demonstrated that xLSTM-based LRAMs exhibit higher inference speeds, especially at large context sizes.
Thus, the empirical evidence suggests, that recurrent backbones such as the xLSTM can be attractive alternatives for LAMs.
Notably, the linear-time inference complexity of xLSTM may enable applications that require long context lengths, such as in-context RL, and facilitate the application of large-scale agents for real-time applications, such as robotics.

Limitations.
The primary target application of LAMs is robotics.
While the majority of our experiments involve robotic simulations, we do not yet provide empirical evidence for real robots.
We do, however, believe that our findings translate to real-world scenarios and aim to provide further evidence in future work.
Moreover, the fine-tuning experiments in this work are limited to offline RL.
We envision that an agent pre-trained by behavioral cloning on large-scale offline RL datasets may be successfully fine-tuned in an online RL setting to explore new strategies that do not appear in the training data.
Modern recurrent architectures offer both parallel and recurrent training mode, which might be the key to success for such applications.
While we provide initial evidence of improved ICL abilities of modern recurrent architectures, we only consider a limited grid-world setting.
Consequently, we aim to further investigate the in-context RL abilities of recurrent backbones on more complex environments in future work.

## 6 Ethics Statement

While we conduct all our experiments in simulated environments, the primary target application of our method is robotics.
We believe that our work can positively impact applications in the near future, which require efficient inference, on-device processing, or have real-time constraints.
However, robotics applications in the real world are not without risks.
In particular, in areas where humans are involved, such as factory settings, special care is required.
LAMs are trained via next-action prediction similar to LLMs.
Consequently, LAMs may also suffer from hallucinations in unknown scenarios.
We therefore strongly discourage users from blindly following the predictions made by real-world LAMs without appropriate safeguards regarding safety and robustness.
It is essential to ensure responsible deployment of such future technologies, and we believe that more research on the robustness of LAMs is necessary.

## 7 Reproducibility

Our code-base used for our experiments and the datasets we generated are publicly available at: [https://github.com/ml-jku/LRAM](https://github.com/ml-jku/LRAM).
We describe the environments we use for our experiments and provide dataset statistics in Appendix A.
Furthermore, in Appendix B, we provide implementation details for all methods and a list of hyperparameters used for our experiments.
In Appendix C, we present additional figures that accompany our results in the main text (e.g., all model sizes).
Finally, in Appendices D and E, we provide further details on the conducted ablation studies and the embedding space analysis, respectively.

## Acknowledgments and Disclosure of Funding

We acknowledge EuroHPC Joint Undertaking for awarding us access to Karolina at IT4Innovations, Czech Republic, MeluXina at LuxProvide, Luxembourg, and Leonardo at CINECA, Italy.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-871302), EPILEPSIA (FFG-892171), AIRI FG 9-N (FWF-36284, FWF-36235), AI4GreenHeatingGrids (FFG- 899943), INTEGRATE (FFG-892418), ELISE (H2020-ICT-2019-3 ID: 951847), Stars4Waters (HORIZON-CL6-2021-CLIMATE-01-01). We thank NXAI GmbH, Audi.JKU Deep Learning Center, TGW LOGISTICS GROUP GMBH, Silicon Austria Labs (SAL), FILL Gesellschaft mbH, Anyline GmbH, Google, ZF Friedrichshafen AG, Robert Bosch GmbH, UCB Biopharma SRL, Merck Healthcare KGaA, Verbund AG, GLS (Univ. Waterloo), Software Competence Center Hagenberg GmbH, Borealis AG, TÜV Austria, Frauscher Sensonic, TRUMPF and the NVIDIA Corporation.

## Appendix

## Appendix A Environments & Datasets

### A.1 General

We compile a large-scale dataset comprising 432 tasks from six domains, 3.4M trajectories, and 894M transitions in total (see Table 1).
To enable fast and targeted data-loading, every trajectory is stored in a separate hdf5 file.
We trade off some data-loading speed for disk space efficiency, by compressing trajectories that contain image-based observations.

### A.2 Atari

The Arcade Learning Environment (ALE) [Bellemare et al.,, ] is the standard benchmark for evaluating RL agents and consists of 57 Atari games.
Input observations in Atari are RGB images, but as is standard practice we gray-scale and crop frames ($|\mathcal{S}|=1\times 64\times 64$).
There are 18 discrete action across all 57 Atari games ($|\mathcal{A}|=18$), but individual games may use only use a subset of these actions.
Furthermore, we adopt the standard Atari recipe as used in prior works, including a frame skip of 4, maximum number of no-ops of 30, resetting on life loss, and reward clipping to $[-1,1]$ [Mnih et al.,, ; Hessel et al.,, ].

Tasks.
Similar to Lee et al., [], we assign 41 games to the training set, and 5 additional tasks to the hold-out set.
The 41 training tasks include:

amidar, assault, asterix, atlantis, bank-heist, battle-zone,
beam-rider, boxing, breakout, carnival, centipede, chopper-command, crazy-climber,
demon-attack, double-dunk, enduro, fishing-derby, freeway, frostbite, gopher,
gravitar, hero, ice-hockey, jamesbond, kangaroo, krull, kung-fu-master,
name-this-game, phoenix, pooyan, qbert, riverraid, road-runner, robotank,
seaquest, time-pilot, up-n-down, video-pinball,
wizard-of-wor, yars-revenge, zaxxon

The 5 hold-out tasks include: alien, pong, ms-pacman, space-invaders, star-gunner

**Table 2: Atari Dataset Statistics.**
| Task | # of Trajectories | Mean Length | Mean Return |
| --- | --- | --- | --- |
| amidar | 1813 | 2753 | 145 |
| pooyan | 2773 | 1800 | 176 |
| frostbite | 5218 | 766 | 18 |
| video-pinball | 1023 | 3902 | 266 |
| wizard-of-wor | 3059 | 1314 | 15 |
| chopper-command | 5452 | 738 | 18 |
| breakout | 3780 | 1300 | 39 |
| phoenix | 3307 | 1509 | 49 |
| asterix | 5250 | 951 | 55 |
| enduro | 571 | 8720 | 636 |
| kung-fu-master | 1775 | 2812 | 131 |
| hero | 3022 | 1345 | 168 |
| assault | 3782 | 1170 | 77 |
| demon-attack | 1649 | 2431 | 116 |
| qbert | 3939 | 1138 | 155 |
| jamesbond | 2841 | 1758 | 11 |
| bank-heist | 4146 | 1204 | 62 |
| up-n-down | 3246 | 1538 | 99 |
| centipede | 6879 | 582 | 81 |
| boxing | 4796 | 1041 | 63 |
| battle-zone | 1933 | 2134 | 15 |
| name-this-game | 988 | 5049 | 389 |
| zaxxon | 2561 | 1950 | 12 |
| beam-rider | 1232 | 3248 | 77 |
| time-pilot | 3886 | 1029 | 11 |
| ice-hockey | 1465 | 3407 | -6 |
| riverraid | 2645 | 1512 | 143 |
| krull | 3032 | 1319 | 528 |
| gopher | 1817 | 2338 | 185 |
| freeway | 2438 | 2048 | 33 |
| seaquest | 2807 | 1779 | 150 |
| double-dunk | 1774 | 2815 | 0 |
| road-runner | 3308 | 1217 | 135 |
| atlantis | 186 | 26349 | 1394 |
| gravitar | 6187 | 646 | 1 |
| yars-revenge | 4094 | 1036 | 96 |
| crazy-climber | 1105 | 3954 | 572 |
| kangaroo | 1787 | 2792 | 50 |
| fishing-derby | 2737 | 1825 | 0 |
| carnival | 21131 | 194 | 37 |
| robotank | 747 | 6652 | 56 |
| Average | 3321 | 2734 | 153 |

Dataset.
For Atari, we leverage the DQN-Replay dataset released by Agarwal et al., [].
The dataset contains the trajectories seen over the entire training of the DQN agent (50M frames), We extract a subset of the last 5M transitions for every task, amounting to 205M transitions in total for the 41 training tasks.
The number of episodes, the episodes lengths and total achieved rewards vary across tasks, as shown in Table 2.

### A.3 Meta-World

The Meta-World benchmark [] consists of 50 manipulations tasks using a Sawyer robotic arm, ranging from opening or closing windows, to pressing buttons.
Meta-World is based on the MuJoCo physics engine [].
Observations in Meta-World are 39-dimensional continuous vectors ($|\mathcal{S}|=1\times 64\times 39$), and actions are represented by 6 continuous dimensions ($|\mathcal{A}|=18$) in range $[-1,1]$.
All tasks share a common action and state space.
Following Wolczyk et al., [] and , we limit the episode lengths to 200 interactions.

Tasks.
We follow and split the 50 Meta-World tasks into 45 training tasks (MT45) and 5 evaluation tasks (MT5).

The 45 training tasks are:

reach, push, pick-place, door-open, drawer-open, drawer-close, button-press-topdown, peg-insert-side, window-open, window-close, door-close, reach-wall, pick-place-wall, push-wall, button-press, button-press-topdown-wall, button-press-wall, peg-unplug-side, disassemble, hammer, plate-slide, plate-slide-side, plate-slide-back, plate-slide-back-side, handle-press, handle-pull, handle-press-side, handle-pull-side, stick-push, stick-pull, basketball,soccer, faucet-open, faucet-close, coffee-push, coffee-pull, coffee-button, sweep, sweep-into, pick-out-of-hole, assembly, shelf-place, push-back, lever-pull, dial-turn

The 5 evaluation tasks are: bin-picking, box-close, door-lock, door-unlock, hand-insert

Dataset.
For Meta-World, we use the datasets released by [], which contain 2M transitions per tasks and consequently 90M transitions in total for the training set.
All episodes last for 200 environment interaction steps, and consequently there are 10K episodes for every task.
For detailed dataset statistics per task, we refer to their publication.

Figure: (a) IIWA
Refer to caption: extracted/5960624/figures/envs/composuite/iiwa.png

### A.4 DMControl

The DMControl benchmark [Tassa et al.,, ] consists of 30 different robotic tasks.
Unlike Meta-World, the benchmark contains robots with different morphologies instead of a single common Sawyer arm.
Due to the different robot morphologies, the state, and action spaces vary across tasks ($3\leq|\mathcal{S}|\leq 24$, $1\leq|\mathcal{A}|\leq 6$), with all actions in range $[-1,1]$.

Tasks.
We do not use all 30 tasks contained in the DMControl benchmark, but select 16 of the 30 tasks that have been used in prior works [Hafner et al.,, ; ; ], which we refer to as DMC11 and DMC5 respectively.

The 11 training tasks are:

finger-turn_easy, fish-upright, hopper-stand,
point_mass-easy, walker-stand, walker-run, ball_in_cup-catch,
cartpole-swingup, cheetah-run, finger-spin, reacher-easy

The 5 evaluation tasks are:

cartpole-balance, finger-turn_hard, pendulum-swingup, reacher-hard, walker-walk

Dataset.
For DMControl, we generate 10M transitions per task by training task-specific SAC [Haarnoja et al.,, ] agents, using the same setup as .
Episodes in all DMControl tasks last for 1000 environment steps and per time-step a maximum reward of +1 can be achieved, which results in a maximum reward of 1000 per episode.
Consequently, our training set contains 10K episodes per tasks, amounting to 110K episodes and 110M transitions in total across all tasks.
We list the dataset statistics for all 11 tasks in Table 3.

**Table 3: DMControl Data statistics.**
| Task | # of Trajectories | Mean Length | Mean Return |
| --- | --- | --- | --- |
| point_mass_easy | 10K | 1K | 851 |
| cheetah_run | 10K | 1K | 385 |
| walker_run | 10K | 1K | 230 |
| ball_in_cup_catch | 10K | 1K | 969 |
| hopper_stand | 10K | 1K | 460 |
| walker_stand | 10K | 1K | 939 |
| finger_turn_easy | 10K | 1K | 954 |
| reacher_easy | 10K | 1K | 938 |
| cartpole_swingup | 10K | 1K | 817 |
| fish_upright | 10K | 1K | 815 |
| finger_spin | 10K | 1K | 966 |
| Average | 19628 | 152 | 8.2 |

### A.5 Composuite

The Composuite benchmark [Mendez et al.,, ], is a robotics benchmark for grasping and object manipulation.
The benchmark is implemented on top of robotsuite [Zhu et al.,, ], which in turn leverages the MuJoCo simulator under the hood [].
Composuite contains a mix of 4 simulated robot arms: IIWA, Jaco, Gen3, and Panda (see Figure 9).
All arms share a common state and action space containing 93 continuous state dimensions and 8 continuous action dimensions, respectively ($|\mathcal{S}|=93$, $|\mathcal{A}|=8$).

Tasks.
CompoSuite is designed as a compositional multi-task benchmark for RL, in which a particular robot manipulates a particular object given an objective, while avoiding obstacles.
Overall, there are 4 robots arms, 4 objects, 4 obstacles, and 4 task objectives.
This results in 256 possible robot/object/objective/obstacles combinations.
For our experiments, we assign 240 tasks to the training set and use the remaining 16 tasks as hold-out set (Panda and Object_Wall) combinations.
For a list of all 256 tasks, we refer to Mendez et al., [].

Dataset.
For Composuite, we leverage the datasets released by Hussing et al., [].
For every task, we select 2000 episodes, which last on average for 500 steps.
This amounts to 1M transitions per task, and 240M transitions across all 240 training tasks.
For dataset statistics, we refer to Hussing et al., [].

### A.6 Mimicgen

Similar to Composuite, Mimicgen [Mandlekar et al.,, ] is based on robosuite and the MuJoCo simulator.
Mimicgen is designed for automatically synthesizing large-scale datasets from only a handful of human demonstrations.
Observations in Mimicgen can be represented as images (from multiple cameras) or low dimensional continuous states.
For our experiments, we opt for the low-dimensional state representation to simplify learning.
Therefore, observations and actions are represented by 37-dimensional and 7-dimensional continuous vectors, respectively ($|\mathcal{S}|=37$, $|\mathcal{A}|=7$).
Similar to Composuite, Mimicgen supports 4 different robot arms: Panda, IIWA, Sawyer, and UR5e (see Figure 10).

Figure: (a) IIWA
Refer to caption: extracted/5960624/figures/envs/mimicgen/iiwa.png

Tasks.
Mimicgen consists of 24 diverse tasks, including stacking blocks, re-assembling objects, and even long-horizon tasks like coffee preparation.
These 24 tasks can be performed with the four supported robot arms, amounting to 96 tasks in total.

Dataset.
Mandlekar et al., [] released dataset for the 24 tasks using the default robot arm Panda.
To increase the dataset diversity, we additionally generated data for the remaining 3 robot arms.
However, not all data generation runs produce successful trajectories, and we discard with too few successful trajectories.
Our final dataset for Mimicgen contains 83 training and 2 evaluation tasks.
For each task, we collect 1000 successful demonstrations (we do not include unsuccessful trajectories).
Episode lengths vary across tasks, ranging from 260 to 850 environment steps.

### A.7 Procgen

Procgen benchmark consists of 16 procedurally-generated video games [].
Observations in Procgen are RGB images of dimension $3\times 64\times 64$.
However, for training efficiency, we apply gray-scaling to image observations ($|\mathcal{S}|=1\times 64\times 64$).
All 16 environments share a common action space of 15 discrete actions ($|\mathcal{A}|=16$).
Procgen is designed to test the generalization abilities of RL agents.
Consequently, procedural generation is employed to randomize background and colors, while retaining the game dynamics.

Tasks.
Following prior works [Raparthy et al.,, ; ], we assign 12 and 4 tasks to training and hold-out set, respectively.
The 12 training tasks are:

bigfish, bossfight, caveflyer, chaser, coinrun, dodgeball,
fruitbot, heist, leaper, maze, miner, starpilot

The 4 hold-out tasks are: climber, ninja, plunder, jumper

Dataset.
We leverage the datasets released by , which contain 20M transitions per task.
The datasets were generated by recording all transitions observed by training RL agents for 25M steps, followed by uniform subsampling to 20M transitions.
Consequently, the dataset contains mixed quality trajectories ranging from random (beginning of training) to expert (end of training).
We list the dataset statistics for all 16 tasks in Table 4.

**Table 4: Procgen Data statistics.**
| Task | # of Trajectories | Mean Length | Mean Return |
| --- | --- | --- | --- |
| bigfish | 82835 | 230 | 6.251 |
| bossfight | 112459 | 141 | 1.946 |
| caveflyer | 151694 | 105 | 7.745 |
| chaser | 93612 | 212 | 3.248 |
| coinrun | 261117 | 51 | 9.473 |
| dodgeball | 144364 | 137 | 2.884 |
| fruitbot | 73653 | 270 | 16.094 |
| heist | 101361 | 196 | 8.405 |
| leaper | 296084 | 67 | 4.446 |
| maze | 482245 | 41 | 9.432 |
| miner | 288818 | 68 | 11.8 |
| starpilot | 96468 | 206 | 17.3 |
| Average | 182059 | 144 | 8.3 |

## Appendix B Experimental & Implementation Details

### B.1 Training & Evaluation.

In our experiments, we compare two variants of xLSTM, Mamba and DT.
For our main experiments in Section 4.2, we train all models for 200K updates, and evaluate after every 50K update steps.
We report the mean and 95% confidence intervals over three seeds in our experiments, as suggested by Agarwal et al., [].
For every evaluation tasks, we take the average of 3 evaluation seeds.

We train our agents with a batch size of 128 and gradient accumulation across the 6 domains, such that every domain is represented with the same proportion.
Consequently, the effective batch size is 768.
We use a learning rate of $1e^{-4}$, 4000 linear warm-up steps followed by a cosine decay to $1e^{-6}$, and train using the AdamW optimizer [Loshchilov and Hutter,, ].
In addition, we employ gradient clipping of 0.25, weight decay of 0.01 for all models.
We do not employ Dropout, as is standard practice in DTs, as we found that it negatively affects performance (see Section 4.3).
We use separate reward scales of 200, 100 and 20 for Meta-World, DMControl and Atari, respectively.
Furthermore, for all domains, we set the target return to the maximum return achieved for a particular task in the training datasets.
This is particularly useful for domains, where the maximum returns differ heavily across tasks (e.g., Atari).
We list all hyperparameters in Table 5.

**Table 5: Hyperparameters for RA-DT.**
| Parameter | Value |
| --- | --- |
| Gradient steps | 200K |
| Evaluation frequency | 50K |
| Evaluation episodes | 5 |
| Optimizer | AdamW |
| Batch size | 128 |
| Gradient accumulation | 6 |
| Lr schedule | Linear warm-up + Cosine |
| Warm-up steps | 4000 |
| Learning rate | 1e-4 $\rightarrow$ 1e-6 |
| Weight decay | 0.01 |
| Gradient clipping | 0.25 |
| Dropout | 0.2 |
| Context len (timesteps) | 50 |
| Reward scale | per-domain |
| Target return | per-task |

### B.2 Context Lengths.

By default, we train all models with a context length $C=50$ timesteps.
For every timestep there are three tokens (s/rt/r) and consequently, the effective context length is 150.
We found that performance improves for longer context lengths (see Section D.1), but limit our experiments to $C=50$ to reduce the computational cost.

### B.3 Model Architectures.

We train models across 4 models sizes: 16M, 48M, 110M, and 206M.
We follow Lee et al., [] in selecting the number of layers and hidden dimensions.
For xLSTM and Mamba, we use twice the number of layers blocks to match the number of parameters of the Transformer [Beck et al.,, ; Gu et al.,, ] (see Table 6)
For our xLSTM [7:1] variant, which contains sLSTM blocks, we strive to maintain the same ratio as proposed by Beck et al., [].
Not all our model sizes are divisible by 8 and only the 16M and 110M models exhibit the exact 7:1 ratio of mLSTM to sLSTM blocks.
For consistency, however, we maintain the same notation as Beck et al., [].
We place sLSTM blocks at positions [1], [1, 3], [1, 3], and [1, 3, 5] for the 16M, 48M, 110M, 206M, respectively.

Across backbones, we use linear layers to encode continuous states, reward returns-to-go, similar to Chen et al., [].
The maximal state-dimension across continuous control environments is 204 in our experiments.
To use a shared linear embedding layer for continuous states, we pad states that have lower number of dimensions to 204 dimensions using zeros.
To encode image inputs on visual domains, we use the IMPALA-CNN proposed by Espeholt et al., [] and adopted by previous works on Procgen [] and Atari [Schmidt and Schmied,, ; Schwarzer et al.,, ].
Consequently, we do not make use of discretization of continuous states or patchification of images.
This design choice significantly reduces the sequence length to only three tokens per time-step (see Appendix B.2) and consequently results in faster inference.

For continuous actions, we make use of discretization and discretize of every action dimension into 256 uniformly-spaced bins, similar to Reed et al., [] and .
We experimented with lower/higher number of bins, but did not observe a benefit beyond 256 bins.
Consequently, this resolution is sufficient for the environments we consider.
We use a shared action head to predict the action bins of all continuous dimensions jointly.
The maximum number of continuous action dimensions is 8 in our experiments and consequently the number of discrete action classes is 2048.
In addition, there are 18 discrete actions originating from Atari and Procgen.
Therefore, our action head learns to predict the correct action among the 2066 discrete classes.
We opt for the shared action head representation, as this further speeds up inference and does not require autoregressive action prediction.

For the Transformer baseline, we use global positional embeddings similar to Chen et al., [].
For the recurrent backbones, we do not make use of positional encodings.

### B.4 Hardware & Training Times.

We train all our models on a server equipped with 4 A100 GPUs.
We use distributed data parallel to distribute the workload, as supported in PyTorch [Paszke et al.,, ].
Training times range from 5 hours for the smallest DT model to 30 hours for the largest Mamba model.
Throughout all our experiments, we use mixed precision training [Micikevicius et al.,, ] as supported in PyTorch to speed up training time.

**Table 6: Model Sizes.**
| Model | Layers | Hidden Dim | Heads | Parameters |
| --- | --- | --- | --- | --- |
| Transformer | 4 | 512 | 8 | 16M |
| Transformer | 6 | 768 | 12 | 48M |
| Transformer | 8 | 1024 | 16 | 110M |
| Transformer | 10 | 1280 | 20 | 206M |
| Mamba | 8 | 512 | - | 16M |
| Mamba | 12 | 768 | - | 48M |
| Mamba | 16 | 1024 | - | 110M |
| Mamba | 20 | 1280 | - | 206M |
| xLSTM | 8 | 512 | 4 | 16M |
| xLSTM | 12 | 768 | 4 | 48M |
| xLSTM | 16 | 1024 | 4 | 110M |
| xLSTM | 20 | 1280 | 4 | 206M |

We evaluate our models after every 50K steps.
However, periodically evaluating the trained agents on all 432 tasks sequentially is time-consuming.
Therefore, we perform parallel evaluation with 4 processes at a time.
For multi-GPU setups, we distribute the evaluation workload among the available GPUs.
For example, with 4 available GPUs and 4 evaluation processes per GPU, 16 environments are evaluated simultaneously.
Consequently, the total evaluation time for all 432 tasks, ranges from 18 minutes for the smallest DT model to roughly 2 hours for the largest Mamba model.

## Appendix C Additional Results

### C.1 Training Tasks

In Figures 11 and 12, we report the normalized scores obtained per domain and the average learning curves across tasks for all four model sizes.

Figure: (a) 16M
Refer to caption: extracted/5960624/figures/per_domain/legend.png

Figure: (a) 16M
Refer to caption: extracted/5960624/figures/learning_curves/legend.png

In Figure 13, we report the training perplexity on the 432 training tasks over 200K updates.
Here, we observe that the training perplexity behaves similar to the validation perplexity.
This is expected, as our models see most transitions only a single time (see Table 1 for the number of repetitions per domain).

Figure: (a) Training Perplexity
Refer to caption: extracted/5960624/figures/scaling_laws/legend.png

Furthermore, we report the scaling curves with an additional model size of 408M parameters in Figure 14.
Due to the high computational cost of the 408M models, we were currently only able to conduct a single run for this size.
However, we aim to provide further empirical evidence for this model sizes in future work.

Figure: (a) Sequence prediction
Refer to caption: extracted/5960624/figures/scaling_laws/legend.png

### C.2 Hold-out Tasks

In Figure 15, we show the zero-shot evaluation performance on the hold-out tasks 15.
We want to highlight, that the performance declines for all methods and model sizes compared to performance on training tasks.
This is because, hold-out tasks exhibit severe shifts in state-spaces, action-spaces and reward functions.

Figure: Figure 15: Scaling comparison. Zero-shot performance on hold-out tasks at four models sizes, 16M, 48M, 110M, and 206M. Note that performance declines for all methods and model sizes compared to performance on training tasks. This is because, hold-out tasks exhibit severe shifts in state-spaces, action-spaces and reward functions.
Refer to caption: extracted/5960624/figures/scaling_laws/eval_tasks/legend.png

### C.3 Fine-Tuning

In Figure 16, we present the fine-tuning evaluation performance on the held-out tasks.
We compare an xLSTM trained from scratched against an xLSTM initialized with the pre-trained weights.
We use full fine-tuning for our experiments, but highlight that parameter-efficient approaches are attractive alternatives [Hu et al.,, ; Liu et al.,, ; Paischer et al.,, ].
We do observe consistent improvement of the pre-trained models over the models trained from scratch.
This indicates that fine-tuning performance is not negatively impacted by the recurrent backbone.
While we train on a substantial number of environments, the total amount of data used is, however, still only a fraction of that employed in training other large-scale models, such as LLMs.
Consequently, we do not observe comparable few-shot generalization.
We anticipate that few-shot generalization capabilities will emerge as we increase both data volume and model size.

Figure: Figure 16: Fine-tune performance on hold-out tasks. We compare the performance of a pretrained xLSTM against an xLSTM trained from scratch, both with 16 million parameters. We select the top 5% percent of trajectories from our held-out tasks based on performance and used this subset to fine-tune the models. We perform 25K update steps during fine-tuning and show the normalized scores, averaged across held-out tasks from each domain.
Refer to caption: extracted/5960624/figures/finetune/bar_legend_0.png

### C.4 In-context Learning

We assess the ICL abilities of modern recurrent architectures on the Dark-Room environment considered in prior works on in-context RL [Laskin et al.,, ; Lee et al.,, ; Kirsch et al.,, ; Sinii et al.,, ; Huang et al.,, ; ].
In Dark-Room, the agent is located in a dark room.
The task is to navigate to an invisible goal location in that dark room.
The state is partially observable, as the agent only observes its own x-y position on the grid ($|\mathcal{S}|=2$).
The action space consists of 5 discrete actions: move up, move down, move left, move right, stay ($|\mathcal{A}|=5$).
Upon reaching the goal location, the agent receives a reward of +1 for every step in the episode it resides on the goal location.
Consequently, the agent first has to explore the room to find the goal.
Once the goal location is found (as indicated by the positive reward), the agent can exploit this knowledge.
Given a multi-episodic context, the agent should be able to exploit information contains in the previous trials (e.g., exploiting one path vs. avoiding another).

In our experiments, the Dark-Room is a $10\times 10$ grid and episodes last for 100 steps, starting in the top left corner of the grid.
We adopt the same experiment setup as and leverage their datasets.
We train 16M parameter agents on datasets from 80 randomly selected goal locations in the grid.
The datasets contain 100K transitions per task and are obtained by training task-specific PPO [Schulman et al.,, ] agents.
Then, we evaluate the in-context abilities of our agents on 20 hold-out goal locations.
During evaluation, the agent is given 40 episodes to interact with the environment, which we refer to as ICL-trials.
Furthermore, we adopt the AD [Laskin et al.,, ] framework for training our agents with a multi-episodic context.
We use the same sequence representation as used in our main experiments, consisting of states, returns-to-go (target return set to 80 during evaluation), and rewards.
Note that this differs from the sequence representation used by Laskin et al., [].
We set the context length for all agents to the equivalent of two episodes, which amounts to 200 timesteps in total.

In Figure 17, we report the ICL performance over the 40 ICL trials on (a) 80 training locations and (b) 20 hold-out locations for the 4 different backbones considered in this work.
We observe that the recurrent backbones attain considerably higher scores than the Transformer backbone.
Furthermore, we find that xLSTM [7:1] attains the highest overall scores, which we attribute to the state-tracking abilities [Merrill et al.,, ] of sLSTM blocks.
We aim to explore the ICL abilities of modern recurrent backbones on more complex benchmarks such as XLand-minigrid [Nikulin et al.,, , ] in future work.

Figure: (a) 80 training tasks
Refer to caption: extracted/5960624/figures/ablations/icl/legend.png

### C.5 Inference Time Comparisons

We empirically examine the difference in inference speed between of our models.
Similar to De et al., [], we report both latency and throughput.
For real-time applications, latency is the more important dimension, and therefore we focus our analysis on latency.

#### C.5.1 Latency

In Figures 18 and 19, we report the latencies for DT and xLSTM with the same number of layer blocks as DT, and twice the number of layers blocks as DT, respectively.
We conduct our comparison for two different batch sizes and across varying sequence lengths.

Figure: (a) $\text{batch_size}=1$
Refer to caption: extracted/5960624/figures/inference_time/latency/legend.png

Figure: (a) $\text{batch_size}=1$
Refer to caption: extracted/5960624/figures/inference_time/latency/legend.png

#### C.5.2 Throughput

In Figures 20 and 21, we similarly report the attained throughput for DT and xLSTM with the same number of layer blocks as DT, and twice the number of layers blocks as DT, respectively.
We conduct our comparison for two fixed context lengths and varying batch sizes.

Figure: (a) $\text{context_length}=800$
Refer to caption: extracted/5960624/figures/inference_time/throughput/legend.png

Figure: (a) $\text{context_length}=800$
Refer to caption: extracted/5960624/figures/inference_time/throughput/legend.png

#### C.5.3 xLSTM Kernel Comparisons

We leverage custom kernels for xLSTM to conduct our inference-speed comparisons.
In particular, we compare 4 variants: recurrent-style inference with and without kernel acceleration, and chunkwise inference with and without kernel acceleration.
In our experiments, every timestep contains 3 individual tokens.
Consequently, regular recurrent-style inference requires iterating over the token sequence of length 3 in a loop given the hidden state of the previous timestep.
This requires 3 forward passes.
In contrast, the chunkwise implementation operates on chunks of timesteps given a hidden state.
Consequently, this only requires a single forward pass.
In Figure 22, we illustrate the impact of kernel acceleration.
We find that our chunkwise kernels result in considerably lower latencies.
Interestingly, we find that for $B=1$, our chunkwise implementation without kernel acceleration is faster than the recurrent-style inference with kernel acceleration.
However, as the batch size increases, this trend reverses.
This highlights the importance of kernel acceleration for efficient inference.

Figure: (a) $\text{batch_size}=1$
Refer to caption: extracted/5960624/figures/inference_time/kernels/legend.png

## Appendix D Ablations

### D.1 Removing action condition

We found that removing actions from the context results in better performance across backbones.
In Figure 23, we report the learning curves over 200K updates for DT with varying context lengths, both with and without actions in the context.
While context lengths beyond 1 hurt performance on meta-world when training with actions, the reverse is true when training without actions.
This is in contrast to recent works, which did not benefit from longer contexts [Octo Model Team et al.,, ].
However, while removing actions improves performance on the robotics domains, it does not affect performance on discrete control.
For robotics, we observed that the models become overly confident (high action logits), which is problematic if poor initial actions are produced.
We assume this is because in robotics actions change smoothly and by observing previous actions the agent learns shortcuts.
Thus, removing actions from the input prevents the agent from using shortcuts.

Figure: (a) w/ actions
Refer to caption: extracted/5960624/figures/ablations/action_condition/legend.png

### D.2 Effect of Dropout in DT

DTs use by default a Dropout [Srivastava et al.,, ] rate of 0.1.
However, during our experiments, we found that Dropout has detrimental effects on the evaluation performance, particularly on continuous control domains like Composuite.
In Figure 24, we show the validation perplexities and evaluation performance for a DT trained with and without Dropout.
Consequently, we remove Dropout from our DT variant.

Figure: (a) Sequence Prediction Performance
Refer to caption: extracted/5960624/figures/ablations/dt_dropout/legend.png

### D.3 Effect of reducing number of layers in xLSTM

In prior works, xLSTM and Mamba use twice the number of layers blocks as the Transformer baseline, while maintaining the same hidden dimension [Gu and Dao,, ; Beck et al.,, ].
For our inference-time comparisons, we therefore reduce the number of layer blocks in xLSTM by half.
To ensure a fair comparison, we consequently adjust the hidden size of xLSTM to match the number of parameters of the Transformer baseline.
In this section, we investigate the effect of these modifications of the xLSTM architecture on the model performance.

In Figure 25, report the validation perplexities and evaluation performance for the *regular* xLSTM with twice the number of layer blocks as DT, and an xLSTM with *half* the number of blocks.
Reducing the number of layer blocks results in slight decrease in performance on both metrics.
However, xLSTM still outperforms the Transformer baseline (see Figure 2).

Figure: (a) Sequence Prediction Performance
Refer to caption: extracted/5960624/figures/ablations/xlstm_half/legend.png

## Appendix E Embedding Space Analysis

In Figure 5, we analyze the representations learned by our models using UMAP [McInnes et al.,, ].
Here, we explain the clustering procedure in more detail.
For every task, we sample 32 sub-trajectories containing 50 timesteps (150 tokens) and encode them using our sequence models.
Then, we extract the hidden states at the last layer of our model and aggregate them via mean pooling.
We cluster all vectors using default hyperparameters of UMAP into a two-dimensional space.
Finally, we color the resulting points by their domain.
Generally, we find that tasks from the same domain cluster together.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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}$
.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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?
soundness: 2
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.
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.
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

Meta-review: 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.

Decision: Reject