1 Introduction↩︎

Offline Reinforcement Learning (Offline RL) [1] has gained significant attention due to its ability to learn strategies without interacting with the environment, which is particularly beneficial in situations where real-time interaction is expensive or risky [2]–[4]. With a static dataset, offline RL can be implemented through three distinct learning methods [5]: (1) model-based algorithm [6]–[8], (2) model-free algorithm [9]–[11], (3) trajectory optimization[12]–[16]. The first two methods require long-term credit assignment through the Bellman equation, leading to the "deadly triad" problem known to destabilize RL [17]. In contrast, trajectory optimization methods treat RL problems as sequence modeling problems to get better performance and generalization [12]. Most trajectory optimization methods rely on transformers, which perform credit assignment directly through the attention mechanism. By leveraging the powerful modeling capabilities of transformers, these methods outperform other offline RL algorithms [18]–[20].

The transformer attention mechanism [21], which allows the model to focus on the important part of the input sequence [22], has several downsides. The computational demands of the attention mechanism escalate quadratically with the input length, posing a significant constraint on its scalability [23]–[25]. Moreover, some studies [26], [27] suggest that the attention mechanism may not be the primary factor contributing to the effectiveness of transformers. This notion is also supported in offline RL, where [13] discovers that the attention mechanism of Decision Transformer (DT) does not capture local associations effectively, rendering it unsuitable for RL. Given these limitations, we are led to ponder if a more efficient mechanism with fewer parameters and greater scalability exists for offline RL. Recently, a series of state space models (SSMs) [28], particularly Mamba [29], have been proposed as potential solutions with the ability to scale linearly concerning the sequence length. In particular, Mamba introduces a selective hidden attention mechanism [30] for content-based reasoning and employs parallel scan to enhance computational efficiency, resulting in two approaches to employing Mamba in offline RL. The first is the Transformer-like Mamba, a direct substitution of the transformer [31]–[33] while the other is the RNN-like Mamba [34], achieving an inference speed with constant time complexity.

Few studies have explored the application of SSMs in offline RL, though they perform well in model-based algorithms [35], [36] and in-context RL learning [37]. Mamba is tailored for memory-required long-sequence tasks, whereas trajectory optimization methods typically utilize short segments during training and inference, as most RL tasks are modeled as Markov Decision Processes (MDPs), i.e. past information may not influence current decisions. Furthermore, due to the lack of a comprehensive investigation of the key component of Mamba, a question has arisen:

In this work, we aim to undertake a thorough investigation and in-depth analysis to explore this question. Specifically, we focus on the data structures and the essential components in trajectory optimization. The extensive experiments provide strong support for the following key findings. (1) We explore the data structures with an analysis of sequence length and concatenating type. The former reveals that long input sequences present computational challenges without enhancing performance due to the hidden attention scores of DeMa evincing an exponential decay pattern. As a result, we opt for the Transformer-like DeMa as opposed to the RNN-like DeMa for efficiency and effectiveness. The latter finds concatenating in the temporal dimension is better for the Transformer-like DeMa. (2) The hidden attention mechanism plays a pivotal role in DeMa’s effectiveness and is compatible with the transformer’s post up-projection residual structure [38], enabling it to replace the attention layer directly and eliminating the need for position embedding. Extensive evaluations show that with a higher average score and nearly 30% fewer parameters, DeMa significantly outperforms DT in eight Atari games. Furthermore, in nine MuJoCo tasks, DeMa’s performance not only exceeds that of DT but does so with only one-fourth of the parameters, highlighting remarkable improvements in both performance efficiency and model compactness.

In the end, our main contributions can be summarized as follows:

1. We find the Transformer-like DeMa surpasses the RNN-like DeMa in both efficiency and effectiveness for trajectory optimization. Extensive experiments on sequence length and concatenating type show the impact of the input data, which guides the design of DeMa.

2. Through various ablation experiments, we discover that the hidden attention mechanism is the core component in DeMa and does not require position embedding. This finding enhances the effectiveness and efficiency of our Transformer-like DeMa.

3. With state-of-the-art performance on both MuJoCo and Atari, our Transformer-like DeMa significantly addresses the challenges posed by transformer-based trajectory optimization methods, particularly the issues of large parameter sizes and limited scalability.

2 Related Work↩︎

2.0.0.1 Offline RL.

Offline RL is a data-driven RL paradigm in which the agent learns solely from a pre-collected dataset rather than through interaction with the environment [16]. Distribution shifts [11] can severely impact performance when RL algorithms are deployed directly in offline environments, leading to significant degradation. To mitigate this problem, several methods have been introduced, which the study [5] categorizes into three primary approaches: (1) learning a dynamics model to generate additional training data (model-based algorithm) [6], [39], (2) learning a policy through a model-free approach by constraining unseen actions or incorporating pessimism into the value function (model-free algorithm) [10], [11], [40], and (3) trajectory optimization [12], [15]. The method of trajectory optimization is usually based on a causal transformer model and converts an RL problem to a sequence modeling problem [13]. It performs credit assignment directly through the attention mechanism in contrast to Bellman backups, thus modeling a wide distribution of behaviors, enabling better generalization and transfer [12].

2.0.0.2 Sequence Modeling in Offline RL.

Following DT [12] and Trajectory Transformer (TT) [15], there has been an increasing trend in employing advanced sequence-to-sequence model to solve RL tasks [14], [41]–[46].¹ Unfortunately, these improvements are usually transformer-based and hence suffer from the common dilemma of the attention mechanism, i.e. over-parameterization and inability to scale to long sequence tasks. What’s more, Emmons et al. [48] find that simply maximizing likelihood with a two-layer feedforward MLP is close to the results of substantially more complex methods based on sequence modeling with Transformers. Similarly, Lawson et al. [49] find that replacing the attention parameters with those learned in other environments has a minimal impact on the performance. Besides, Decision ConvFormer (DC) [13] indicates that substituting the attention layers with learnable parameters can lead to improved outcomes. These observations suggest significant redundancy in the Transformer architecture, highlighting the potential to explore lighter and more scalable networks for implementation in offline RL. Building on this, the Structure SSM (S4) [50] has emerged as a promising alternative. Studies [35] and [36] use S4 in model-based RL, outperforming traditional Transformer and RNN approaches. The capabilities of S4 and Mamba are further demonstrated by [37], [51], which points to their speed and effectiveness in in-context RL tasks.

The most related work to ours is Decision S4 (DS4) [52] and Decision Mamba (DMamba) [53], where the former uses an RNN-like S4 for inference, and the latter replaces the attention mechanism with Mamba directly. In contrast, our work finds that Transformer-like DeMa outperforms RNN-like DeMa as the long sequences impose a significant computational burden on Mamba without contributing to performance improvements. What’s more, DMamba simply substitutes Mamba for the attention block rather than the transformer block while our investigation shows the key component is the hidden attention mechanism, which eliminates the need for position embedding and hence achieves better performance with fewer parameters.

3 Preliminaries↩︎

In this section, we present several necessary preliminaries and terminologies of offline RL, trajectory optimization, state space model, and hidden attention in Mamba.

3.1 Offline RL with Trajectory Optimization↩︎

Given a static dataset of transitions \(\tau=\{(s_t,a_t,s_{t+1},r_t)_i\}\), where \(i\) presents the timestep of a transition in the dataset. The states and actions are generated by the behavior policy \((s_t,a_t)\sim d^{\pi_\beta}(\cdot)\), while the next states and rewards are determined by the unknown transition dynamics \(p(s^{\prime},r|s, a)\). The goal of offline RL is to find an approximate policy \(\pi(a|\cdot)\) that maximizes expected return \(\mathbb{E}[\sum_{t=0}^{T}r_{t}]\), where \(T\) represents the time step at which the episode terminates. Due to the lack of interaction with the environment, trajectory optimization methods transform the goal into minimizing reconstruction loss, i.e. minimizing loss \(\mathbb{E}_{(\hat{R},s, a)\sim\tau}[\frac{1}{T}\sum_{t=1}^{T}\mathcal{L}_{\mathrm{MSE/CE}}(\hat{a}_{t}; a_{t})]\), where \(\hat{a}_{t}=\pi(\cdot|s_{t-K+1:t},\hat{R}_{t-K+1:t}, a_{t-K:t-1})\), and \(\hat{R}_{t} = \sum_{t^{\prime}=t}^{T}r_{t^{\prime}}\) is the return-to-go (RTG). At test time, a target RTG \(R_0\) is manually set to represent the desired performance. We input the trajectories from the last \(K\) timesteps into policy \(\pi\), which then generates an action for the current timestep. Subsequently, the next state and reward are received from the environment. These elements are concatenated and also input into the model. The policy is approximated through the sequential model [12], [54]. However, these models typically possess a large number of parameters and struggle with handling long sequences effectively. Fortunately, this issue can be addressed by using SSMs [28], [29], [50].

3.2 State Space Model and Mamba↩︎

There are two approaches to utilizing Mamba in RL, which are both closely related to the modeling methods of SSM. SSM is defined by the following first-order differential equation, which maps a 1-D input signal \(u(t)\) to an \(N\)-D latent state \(h(t)\) before projecting to a 1-D output signal \(y(t)\) [55], \[\begin{align} h^{\prime}(t)=Ah(t)+Bu(t),\quad y(t)=Ch(t)+Du(t), \end{align}\] where \(A\in\mathbb{R}^{N\times N} , B\in\mathbb{R}^{N\times 1}, C\in\mathbb{R}^{1\times N}\) and \(D\in\mathbb{R}\) are trainable matrices. As \(u(t)\) is typically discretized as \(\{u_i\}_{i=1,2,...}\), SSM can be discretized by a step size \(\Delta\). Moreover, recurrent SSM can be written as a discrete convolution. Let \(h_{0}=0\) and \(D=0\), we have \[\begin{align} y_{i}=C\bar{A}^i\bar{B}u_1+C\bar{A}^{i-1}{\bar{B}}u_2+\cdots+C\bar{A}\bar{B}u_{i-1}+C\bar{B}u_i,\quad y=u*\bar{K}, \end{align}\label{eq:2}\tag{1}\] where \(\bar{A},\bar{B}\) is the approximation discrete of \(A, B\), and \(\bar{K}\) is called the SSM convolution kernel and can be represented by filter \[\bar{K}=(C\bar{B},C\bar{A}\bar{B},\ldots,C\bar{A}^{i}\bar{B},\ldots).\]

S4 and other time-invariant models cannot select the previous tokens to invoke from their history records. To solve this problem, Mamba merges the sequence length and batch size of the inputs, allowing the matrices \(B, C\) and the step size \(\Delta\) to depend on the inputs. Therefore, it is a time-varying system and cannot use the convolution view. To ensure efficient training and inference with Mamba, techniques such as parallel scanning, kernel fusion, and recomputation are employed, resulting in two types of Mamba. One type is the SSM using the recursive view, referred to as RNN-like Mamba, and the other is the SSM utilizing parallel scanning, known as Transformer-like Mamba. RNN-like Mamba is akin to DS4 [52], wherein the complete trajectory is taken as a sample and fully inputted into the model for training. Utilizing this approach, which capitalizes on the ability to capture long-term dependencies, the inference speed can be significantly increased. During the inference process, it is sufficient to input only the current tuple (\(r_{t-1}\), \(a_{t-1}\), \(s_t\)) in conjunction with the hidden state \(h_t\). Transformer-like Mamba is a direct replacement for the transformer, where we consistently truncate the input sequences to a fixed length of \(K\) before their introduction into the model throughout the training and inference phases [34], [53], [56], [57].

3.3 Hidden Attention in Mamba↩︎

Although the role of the self-attention mechanism in offline RL remains uncertain, it is known that this mechanism allows the model to dynamically focus on different parts of the input sequences, following the Equation 2 . \[\text{Self-Attention}(x)=\alpha V(x),\quad\alpha=\mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right),\label{eq:1}\tag{2}\] where \(Q, K, V\) represent queries, keys, and values respectively, i.e. input sequences after three linear transformations. \(d_k\) is the dimension of the keys. Similarly, current research suggests that the S6 layer in Mamba can be viewed as the hidden attention mechanism with a unique data-control linear operator [30]. Assuming the initial condition \(h_0=0\), we can obtain a formula similar to Equation 1 \[y_i=C_i\sum_{j=1}^i\big(\Pi_{k=j+1}^i\bar{A}_k\big)\bar{B}_jx_j,\;\; h_i=\sum_{j=1}^i\big(\Pi_{k=j+1}^i\bar{A}_k\big)\bar{B}_jx_j, \label{eq:33}\tag{3}\] where \(\bar{A}_{i}=\exp(\Delta_{i}(A))\), \(\bar{B}_{i}=\Delta_{i}(B_i)\), and \(\Delta_i=\text{softplus}(S_{\Delta}(x_i))\). \(B_i=S_B(x_i)\), \(C_i=S_C(x_i)\), with \(S_B\), \(S_C\) and \(S_{\Delta}\) are linear projection layers. Softplus is an elementwise function that is a smooth approximation of ReLU.

Since \(\bar{A_t}\) is a diagonal matrix, [30] simplifies the hidden matrices and gets the attention mechanism of Mamba: \[\begin{align} \text{Hidden-Attention}&(x)=\tilde{\alpha} x,\quad\tilde{\alpha}_{i,j}\approx\tilde{Q}_i\tilde{H}_{i,j}\tilde{K}_j\\ \tilde{Q}_{i}:=S_{C}(x_{i}), \tilde{K}_{j}:=\mathrm{ReLU}(S_{\Delta}&(x_{j})S_{B}(x_{j}), \tilde{H}_{i,j}:=\exp\Big(\sum\limits_{\substack{k=j+1\\S_{\Delta}(x_{k})>0}}^{i}S_{\Delta}(x_{k})\Big)A. \end{align}\label{eq:44}\tag{4}\] Therefore, we can visualize the hidden attention matrices in DeMa, thus gaining a deeper understanding of the behavior inside the model in the setting of offline RL.

4 The Analysis of DeMa↩︎

Considering most trajectory optimization methods use short segments during both training and inference, the compatibility of Mamba with these methods remains an open question. As shown in Figure 1, this section presents an analysis from the perspectives of data structures and essential components. Section 4.1 discusses the impact of data structure on trajectory optimization. Our study reveals that the RNN-like DeMa does not offer substantial benefits in terms of effectiveness or efficiency. Therefore, we investigate three critical factors: sequence length, the hidden attention mechanism, and the input concatenation types. We find that the balance between performance and efficiency highly depends on the appropriate sequence length selection. Moreover, the input concatenation method significantly influences the results, with temporal concatenation (i.e., B3LD) demonstrating its effectiveness. Section 4.2 conducts ablation studies to identify the hidden attention mechanism as a key component of DeMa, facilitating better utilization and component replacement. Detailed experiments and additional results are in the Appendix. Our code is available at https://github.com/AndssY/DeMa.

4.1 Input Data Structures↩︎

First, we compare the RNN-like DeMa (B3LD) with the Transformer-like DeMa (B3LD)². The average results are shown in Table 1 (with detailed results in Appendix 11), where the performance of the RNN-like DeMa is significantly inferior to that of the Transformer-like DeMa, especially in Atari games. These findings suggest that the recurrent mode may be unnecessary in trajectory optimization methods. Given that the hyper-parameters are identical for both types of DeMa except for the sequence length, we assume that variations in sequence length are likely the primary cause of the observed disparities in results. Therefore, we explore the effect of sequence length on the Transformer-like DeMa in subsequent sections.

4.1.0.1 How does sequence length affect the computational load?

We investigate the impact of sequence length on single-step training time, single-step inference time and GPU memory usage for models including DT, Transformer-like DeMa, and RNN-like DeMa. Figure 2 shows that the Transformer-like DeMa operates faster than the RNN-like DeMa when dealing with short sequence lengths, despite that the inference time of RNN-like DeMa is independent of the sequence length. With conventional sequence lengths (such as 20), Transformer-like DeMa holds an advantage in forward speed, training speed, and GPU memory consumption.

4.1.0.2 How does sequence length affect the performance of DeMa?

While the computational cost of Transformer-like DeMa increases linearly with the expansion of the sequence length, it is crucial to recognize that the increased computational cost may not ensure a corresponding enhancement in the model’s performance. Transformer-like DeMa’s Performance may plateau or even decline as the input sequence length exceeds a certain threshold. As illustrated in Figure 3, Transformer-like DeMa’s performance reaches a plateau in MuJoCo [61] when the input sequence surpasses a specific length; while significantly deteriorates with excessively long input sequences in Atari.

4.1.0.3 Why does DeMa require merely short input sequences?

We calculate the hidden attention scores in DeMa via Eq. 3 4 , which reflect the importance of historical information to DeMa. Figure 4 shows the hidden attention scores of the last \(K\) tokens at each decision-making step (from the 300th to the 600th step). It can be seen that the attention scores exhibit exponential decay as the tokens become increasingly distant from the current decision-making moment, which aligns with the forgetting property of a Markov chain [13]. What’s more, the hidden attention across different decision steps exhibits a periodic pattern towards the current token, suggesting that the model may have learned kinematic features, as agents in these environments engage in periodic movements.³

r0.41

4.1.0.4 Which type of concatenation is suitable for DeMa?

Models like the Transformer and Mamba typically process inputs token by token. However, given an MDP, there are three elements \(s, a, r\) to consider. Therefore a significant design consideration is the method of concatenating these three elements into a suitable token format for the model. We experiment to investigate the suitable design for DeMa. By Table ¿tbl:tab:input32tokens?, concatenating the three elements in the temporal dimension yields better results. This may be due to the significant differences between the three elements of the MDP. As illustrated in [14], states and actions symbolize fundamentally dissimilar notions, concatenating them in the embedding dimension directly may make it more difficult for the model to recognize, leading to poorer results.

4.2 The Essential Components of DeMa↩︎

Aside from the perspective of input data, this section delves into DeMa from the standpoint of network components. We primarily investigate the following questions: (1) Considering that some DTs do not heavily rely on attention mechanism [13], [49], is the hidden attention mechanism crucial for DeMa? (2) As the Mamba block is an integration of the hidden attention mechanism with pre up-projection residual blocks [38], what impact will it have on the performance when integrating it with other residual structures (i.e. the post up-projection residual block in the transformer)? (3) With the inherent recurrent nature of SSM [62], does DeMa need position embedding? (Appendix 13)

4.2.0.1 Is the hidden attention mechanism crucial for DeMa?

[27] shows that the transformer does not heavily rely on attention, and [13] finds the attention mechanism of DT is not suitable for RL. Given these insights, we aim to investigate whether a similar phenomenon exists in hidden attention. In line with [49], we evaluate DeMa by swapping the hidden attention weights trained in different environments, in addition to randomizing and zeroing these weights. As depicted in Figure 5, the performance exhibits a marked decrease regardless of whether the parameters are replaced with those pre-trained in other environments or randomized. Interestingly, when the parameters of hidden attention are set to zero, the model still maintains a certain level of performance. This zeroing of parameters completely removes the hidden attention, ceasing to process historical information and relying solely on residual connections to transmit information. This suggests that the residual connections are functional and the role of hidden attention is crucial for DeMa.

4.2.0.2 What occurs when combining hidden attention with post up-projection residual blocks?

Mamba represents the integration of the hidden attention mechanism with pre up-projection residual blocks as discussed in [38]. To determine the contributing factor to the model’s enhanced performance, we explore the combination of hidden attention with post up-projection residual blocks in transformer. According to the results in Table ¿tbl:tab:residual? and Table [tbl:tab:residual2], although the overall average results of DeMa are slightly better than those of DeMa with post., it is observable that they each have advantages in different environments. Hence, we believe that the performance differences when integrating with the two types of residual blocks are not statistically significant. It suggests that the structure of the residual blocks exerts minimal influence on the outcome. Given that both configurations yield a measurable performance improvement over the DT, it is reasonable to conclude that the hidden attention mechanism within DeMa plays a pivotal role.

5 Evaluations on Offline RL Benchmarks↩︎

In this section, we delve into a comparative analysis of DeMa’s performance against various DTs. Our investigation primarily centers on the influence of disparate network architectures on the experimental outcomes. Consistent with antecedent studies, we assessed both discrete (Atari [58]) and continuous control tasks (MuJoCo [63]), presenting the normalized scores accordingly. Given that the sequence length considerably affects the results, we selected the optimal outcomes from sequence lengths \(K=8\) to \(K=20\) for DeMa. The detailed hyper-parameters on DeMa are available in Appendix 10. Our main results are shown in Table [tbl:tab:atari32result] and Table [tbl:tab:MuJoCo]. DeMa achieves a significantly higher average score compared to DT in Atari games, while the number of parameters and the number of MACs in DeMa are each five times fewer than those in DT, as shown in Table [tbl:tab:param]. Moreover, DeMa has better scalability for input length which can be seen in Figure 2, it maintains a slow linear growth with the input sequence length increases while the computational cost of the Transformer grows quadratically. These results demonstrate that our transformer-like DeMa is well-suited for integration with trajectory optimization methods.

6 Conclusion↩︎

To investigate Mamba’s compatibility with trajectory optimization, this work conducts comprehensive experiments from the aspect of data structures and network architectures. Our findings reveal that (1) DeMa benefits from short sequence lengths due to its exponentially decaying focus on sequences. Consequently, we incorporate a Transformer-like DeMa. (2) The hidden attention mechanism plays a crucial role in DeMa. It can combine with other residual structures and does not require position embedding. Based on the insights gained from the investigation, our DeMa surpasses previous methods, achieving higher performance over the DT while using 30% fewer parameters in eight Atari games. In the MuJoCo, our DeMa outperforms DT with only a quarter of the parameters. In conclusion, our DeMa is compatible with trajectory optimization in offline RL.

Limitations. We investigate the application of Mamba in trajectory optimization and present findings that provide valuable insights for the community. However, there remain several limitations: (1) Trajectory optimization tasks typically involve shorter input sequences, raising questions about how well the RNN-like DeMa performs in terms of memory capacity in RL compared to models such as RNNs and LSTMs. Furthermore, the potential of both types of DeMa warrants further exploration, particularly in some POMDP environments and long-horizon non-Markovian tasks that require long-term decision-making and memory. (2) We examine the importance of the hidden attention mechanism in Section 4.2, future work could leverage interpretability tools to examine further the causal relationship between memory and current decisions in DeMa, ultimately contributing to the development of interpretable decision models. (3) While we have assessed the properties of DeMa and identified improvements in both performance efficiency and model compactness compared to DT, it remains unclear whether DeMa is suitable for multi-task RL and online RL environments.

Supplementary Material for Is Mama Compatible with Trajectory Optimization in Offline Reinforcement Learning?↩︎

7 Environment and Dataset↩︎

7.0.0.1 MuJoCo.

The MuJoCo domain [59] evaluates the performance of RL algorithms in continuous control tasks. In keeping with previous studies, we select three games from the standard locomotion environments [59] in Gym [61], namely HalfCheetah, Hopper, and Walker, and three different dataset settings, namely medium, medium-replay, and medium-expert [63].

7.0.0.2 Atari.

Atari [58] is an ideal platform for evaluating an agent’s ability in long-term credit assignments. We conduct experiments in eight different games: Breakout, Qbert, Pong, Seaquest, Asterix, Frostbite, Assault, and Gopher. We use 1% DQN Replay Dataset [64] as our training dataset, which encompasses a total of 500,000 timesteps worth of samples generated throughout the training process of a DQN agent [65]. It’s worth noting that the version of "atari-py" and "gym" we use is 0.2.5 and 0.19.0 respectively, which is noted by the official code in https://github.com/google-research/batch_rl.

8 Baselines↩︎

8.0.0.1 Baselines for MuJoCo.

To evaluate DeMa’s performance in the MuJoCo, we compare DeMa with one value-based method: CQL [66] and four trajectory optimization methods with different network architectures: DS4 [52], RvS [4], DT [12], GDT [14] and obtain baseline performance scores for CQL and DS4 from [13], for RvS from [4] and for GDT from [14].

8.0.0.2 Baselines for Atari.

In the Atari domain, we compare DeMa with CQL [66], DT [12], DC and DC\(^{hybrid}\) [13]. The results of baselines are directly borrowed from [13].

9 The Procedure of Training and Inference↩︎

9.0.0.1 Training resources

We use one NVIDIA GeForce RTX 4090 to train each model in MuJoCo and one NVIDIA GeForce RTX 3090 to train each model in Atari. Training each model typically takes 3-8 hours and 5-14 hours in MuJoCo and Atari respectively. However, since each environment needs to be trained three times with different seeds, the total training time is usually multiplied by three.

9.0.0.2 The procedure of Transformer-like DeMa

The Training and evaluation for Transformer-like DeMa are similar to variant DTs. Given a dataset of offline trajectories, we randomly select a starting point and truncate it into a sequence of length \(K\). After forming a batch of data, it is input into the model for training. We minimize the reconstruction loss between the predicted action and the actual action, i.e. the cross-entropy loss for discrete actions and Mean Square Error (MSE) for continuous actions. The input data is also a sequence of length \(K\) in the evaluation phase.

9.0.0.3 The procedure of RNN-like DeMa

For RNN-like DeMa, the input during training is a batch of complete trajectories. As different trajectories have different lengths, we pad the trajectories to the same length before inputting them into the model and mask the loss of the padding. However, training with full trajectories rather than truncated sequences may be more inefficient, especially in scenarios where sequence lengths vary widely. In the DQN Replay Dataset in Atari, the lengths of different trajectories varied dramatically. Some trajectories might only be 500 timesteps long, while others could contain a sample with a length of 10,000 timesteps. This causes a lot of computing resources to be wasted on meaningless padding, resulting in inefficiency and ineffectiveness. Some techniques can avoid this issue. One can refer to this PR.

10 Implementation details of DeMa↩︎

We implement DeMa based on the official code of DT and the Mamba. We have also adopted the code from HiddenMambaAttn to calculate the attention scores of DeMa on the current input sequence at each decision step. Given that the official Mamba code utilizes Triton, we also employ Mamba-minimal which is fully based on pytroch to compute the MACs of DeMa.

Tables 2-4 provide a comprehensive list of hyper-parameters for our proposed transformer-like DeMa and RNN-like DeMa applied to MuJoCo and Atari environments. To ensure a fair comparison, we adopt similar hyper-parameter settings to DT [12] and DC [13].

10.1 Hyper-parameters in MuJoCo↩︎

For our training on MuJoCo, the majority of the hyper-parameters in Table 2 are adapted from [13]. For the learning rate, we use a learning rate of \(10^{-4}\) for training in hopper-medium, hopper-medium-replay, and walker2d-medium and use \(10^{-3}\) for other environments. For the embedding dimension, we use an embedding dimension of 256 in hopper-medium and hopper-medium-replay, while use 128 in the other environments. What’s more, as DeMa does not use multilayer perceptron (MLP), so there is no nonlinearity function for DeMa. As for DeMa with post. in Table ¿tbl:tab:residual?, we use ReLU as per convention. For DeMa’s hyper-parameters, we use a d_model of 128 in all expert datasets, while use 64 in the other environments. As for d_state and expand, we set 64 and 2 respectively for all env. We keep the experimental parameters consistent for all types of DeMa in MuJoCo.

10.2 Hyper-parameters in Atari↩︎

10.2.0.1 Transformer-like DeMa.

For the Atari game we mostly follow those in Table 3 from [12]. The only adjustment made is to the context length \(K\) and return-to-go conditioning. As revealed in Figure 3, the sequence length is not always better when it’s longer. Thus for Qbert and Frostbite we use \(K=8\). For other games, we keep \(K=30\). As for the return-to-go conditioning, we find the return obtained by DeMa in some games has already exceeded the initial "return to go" set for DT. Therefore, we increase the "return to go" so that DeMa can fully demonstrate its performance.

10.2.0.2 RNN-like DeMa.

Since the RNN-like DeMa utilizes trajectories for training and the trajectories in Atari are exceptionally lengthy, the available sample size becomes significantly limited when only 1% of the DQN-replay dataset is utilized. If the prior parameter settings were to be used, the training would done after only a few hundred upgrades, thereby resulting in an unsatisfactory performance. Therefore, we consider multiple updates for a single sample, while simultaneously lowering the learning rate as shown in Table 4. What’s more, due to the limitation of GPU memory, we can only set a batch size of 8 for Atari. Specifically, For the Frostbite, we set a batch size of 1, an epoch of 50. The other hyper-parameters are kept consistent with those in Table 3.

11 Detailed results↩︎

Table 5 and Table 6 show detailed results between RNN-like DeMa⁴ and Transformer-like DeMa. It can be observed that the performance of the RNN-like DeMa is not as good as that of the Transformer-like DeMa, and Figure 2 also shows that the RNN-like DeMa requires more computational overhead. Hence, using the RNN model in trajectory optimization seems to be unnecessary, as section 4.1 finds that past historical information does not provide much assistance to current decision-making. However, in tasks that require memory capability or are model-based, the RNN-like DeMa could be a better choice. This could be a direction for deeper future research based on [35]–[37].

12 Tasks Requires Long Horizon Planning Skills↩︎

Three experiments involving delayed rewards(MuJoCo with delayed rewards) and maze navigation(maze2d, antmaze) are conducted to investigate how would DeMa perform on tasks that require long horizon planning skills.

12.1 MuJoCo with Delayed Rewards↩︎

To investigate DeMa’s performance on tasks with delayed rewards, we conduct an experiment on a delayed return version of the D4RL benchmarks [12], in which the agent does not receive any rewards along the trajectory but instead receives the cumulative reward of the trajectory in the final timestep. In this environment, we train DeMa using the same hyper-parameters settings, and the results are shown in Table 7. Results show that CQL is the most affected, while DT also experiences a certain degree of influence. In contrast, DeMa is relatively less impacted. The results indicate that DeMa demonstrates effective performance in tasks with delayed rewards.

12.2 Maze Navigation↩︎

There are two environments in maze navigation. Maze2d: This environment aims at reaching goals with sparse rewards, which is suitable for assessing the model’s capability to efficiently integrate data and execute long-range planning. The objective of this domain is to guide an agent through a maze to reach a designated goal. Antmaze: This environment is similar to maze2d, while the agent is an ant with 8 degrees of freedom. For our training on Maze, the majority of the hyper-parameters in Table ¿tbl:tab:hyper32antmaze? and Table [tbl:tab:hyper32maze2d]. For maze2d-medium, we use \(K=8\) and embedding_dim=256. For maze2d-umaze, we use hyper-parameters in Table 2.

We compare DeMa with DT [12], GDT [14] and DC [13]. The results of DT and GDT are directly borrowed from [42]. Results in Table 8 show that DeMa performs better compared to DT in the maze navigation task. The visualization analysis of the hidden attention mechanism in these environments can be found in Figure 9 (q) and Figure 9.

13 Further Ablation Study↩︎

13.0.0.1 DeMa does not need the position embedding.

Position embedding is generally used in transformers to help the model understand the sequential nature of the data. It’s a way of encoding the position of tokens in the sequence, and it can be crucial in tasks where the order of the data matters. Although we can use DeMa similar to using a transformer, which has input and output dimensions of (B, L, D) during training and inference, it differs in that it does not require position embedding to help the model have the ability to remember sequential information. As shown in Table 9, the addition of position embedding not only failed to enhance the performance of the model but also led to a significant decrease in performance on certain tasks. Additionally, the introduction of position embedding significantly increased the model’s parameter count, thereby adding to its computational burden. This finding highlights the advantage of the DeMa in terms of lightweight design, indicating its suitability for tasks with limited resources.

14 Integrating DeMa with Other Methods↩︎

We conduct additional experiments to show that DeMa can be combined with other trajectory optimization methods to achieve even better performance. By integrating DeMa with QT [42], we develop Q-DeMa. As shown in Table 10, Q-DeMa achieves performance comparable to state-of-the-art models while utilizing less than one-seventh of the parameter size of QT. This finding underscores the significant potential of applying Mamba to RL.

15 MuJoCo and Atari Tasks Scores↩︎

Table 11 shows the normalized scores used in MuJoCo and Atari tasks, followed by [63] and  [60].

16 Types of Hidde Attention Scores↩︎

In the previous articles [13], [67], the visualization of attention in DT was in the form of a lower triangular matrix. However, this lower-triangular matrix reflects the attention scores of each generated token to the input sequence during the training phase, and it cannot accurately illustrate the context information that the model focuses on at each decision-making step. As can be seen, the element in Figure 6 at the i-th row and j-th column represents presents the output \(y_i\)’s attention score to input \(x_j\). In the training phase, all corresponding predicting actions are used to calculate the reconstructed loss with target actions. However, during the evaluation phase, i.e. when interacting with the environment, we input \(x:(1, L, D)\), and the model also outputs \(y:(1, L, D)\). At this time, we only use the last one of the model’s output, which is \(y[:,-1,:]\), corresponding to the last row in the matrix. Therefore, it is not quite appropriate to judge the context information the model focuses on at each decision-making step based on the lower-triangular matrix in Figure 6, as we want to understand the model’s decision-making behavior at each step, thus leads to the creation of Figure 4, Figure 7 and Figure 8. It also demonstrates strong forgetting characteristics. This aligns with the properties of Markov chains as described in  [13], where the sequence of states precisely forms a Markov chain. We also conduct additional explorations in environments involving delayed rewards(MuJoCo with delayed rewards) and maze navigation(maze2d, antmaze). The performance of the hidden attention mechanism is illustrated in Figures 9 (i)-9. Although DeMa’s attention to past information increases, the hidden attention mechanism still prioritizes the current information when the Markov property of the environment is relaxed. Furthermore, among historical information, the hidden attention mechanism demonstrates a significantly higher focus on states compared to rewards or actions.

References↩︎