Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?
Authors: Yang Dai
Oubo Ma
Longfei Zhang
Xingxing Liang
Shengchao Hu
Mengzhu Wang
Shouling Ji
Jincai Huang
Li Shen
Venue: NeurIPS 2024 poster (2024)
Keywords: Offline RL; Trajectory Optimization; Mamba
Primary area: reinforcement_learning
TL;DR: An investigation of Mamba in offline RL.

Abstract: Transformer-based trajectory optimization methods have demonstrated exceptional performance in offline Reinforcement Learning (offline RL). Yet, it poses challenges due to substantial parameter size and limited scalability, which is particularly critical in sequential decision-making scenarios where resources are constrained such as in robots and drones with limited computational power. Mamba, a promising new linear-time sequence model, offers performance on par with transformers while delivering substantially fewer parameters on long sequences. As it remains unclear whether Mamba is compatible with trajectory optimization, this work aims to conduct comprehensive experiments to explore the potential of Decision Mamba (dubbed DeMa) in offline RL from the aspect of data structures and essential components with the following insights: (1) Long sequences impose a significant computational burden without contributing to performance improvements since DeMa's focus on sequences diminishes approximately exponentially. Consequently, we introduce a Transformer-like DeMa as opposed to an RNN-like DeMa. (2) For the components of DeMa, we identify the hidden attention mechanism as a critical factor in its success, which can also work well with other residual structures and does not require position embedding. Extensive evaluations demonstrate that our specially designed DeMa is compatible with trajectory optimization and surpasses previous methods, outperforming Decision Transformer (DT) with higher performance while using 30\% fewer parameters in Atari, and exceeding DT with only a quarter of the parameters in MuJoCo.

Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?
Yang Dai
Oubo Ma²
Longfei Zhang' Shouling Ji
Xingxing Liang
Shengchao Hu
Mengzhu Wang
Jincai Huang'
Li Shen
"Laooraoory for Big Data and Decision,
National University Defense Technology
7hejiang University
*Shanghai Jiao Tong University
4 Hebei University of Technology *Shenzhen Campus of Sun Yat-sen University
(daiyang2000, zhanglongfei wwwannxinggxinnne munnggg (ucai))@uude edu.
{mob, sji]ezju. edu. en; charles juues j tu. edu. en; {dreamkil ly mathshenl 1) @gmail com
Abstract
Transformer-based trajectory optimization methods have demonstrated exceptional performance in offline Reinforcement Learning (offline RL). Yet. it poses chal- lenges due to substantial parameter size and limited scalability, which is particularly critical in sequential recision-making scenarios where resources are constrained such as in robots and drones with limited computational power. Mamba, a promis- ing new linear-time sequence model, offers performance par with transformers while delivering substantially fewer parameters on long sequences. As it remains unclear whether Mamba is compatible with trajectory optimization, this work aims to conduct comprehensive experiments to explore the potential of Decision Mamba (dubbed DeMa) in offline RL from the aspect of data structures and essential components with the following insights: (1) Long sequences impose a significant computational burden without contributing to performance improvements since DeMa's 's focus on sequences diminishes approximately exponentially. Consequently, we introduce a Transformer-lik DeMa as opposed to RNN-like DeMa. (2) For the components of DeMa, we identify the hidden attention mechanism critical factor in its success, which can also work well with other residual structures and does not require position embedding. Extensive evaluations demonstrate that our specially designed DeMa is compatible with trajectory optimization and surpasses previous methods, outperforming Decision Transformer (DT) with higher perfor- mance while using 30% fewer parameters in Atari, and exceeding DT with only a quarter of the parameters in MuloCo.
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 [:: (1) model-based algorithm [6-8), (2)
model-free algorithm [9HI]), (3) trajectory optimization 12-166) 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 [[2]. Most trajectory optimization methods rely transformers, which perform credit assignment directly through the
"ooresponding authors: Li Shen and Xingxing Liang.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

attention mechanism. By leveraging the powerful modeling capabilities of transformers, these methods outperform other offline RL algorithms [8 20].
The transformer attention mechanism [1], 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 scala- bility (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 [3] 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 more efficient mechanism with fewer parameters and greater scalability exists for offline RL. Recently, 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 selective hidden attention mechanism [0) for content-hand 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, direct substitution of the transformer [31333 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-requine long-sequence tasks, whereas trajectory optimization methods typically utilize short segments during training and inference, as most RI. tasks are modeled as Markov Decision Processes (MDPss, 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:
Whether Mamba is compatible with trajectory optimization?
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. (I) We explore the data structures with 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 exponential decay pattern. As 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- (p-projection residual structure 3 enabling it to replace the attention layer directly and eliminating the need for position embedding. Extensive evaluations show that with higher average score and nearly 30% fewer parameters, DeMa significantly outperforms DT in eight Atari games. Furthermore, in nine MuJoCo tasks, DeMa's '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 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 ofth-oofthheart performance on both MuJoCo and Atari, our Transformer-I 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
Offline RI. Offline RI. is a data-driven RI, paradigm in which the agent learns solely from a pre-
collected dataset rather than through interaction with the environment [16). Distribution shifts (I]

---

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 dynamics model to
generate additional training data (model-based algorithm) 66 39) (2) learning a policy through
model-free approach by constraining unseen actions or incorporating pessimism into the value functiin (model-free algorithm) [10] 44]. and (3) trajectory optimization [52 The method of trajectory optimization is usually based on causal transformer model and converts an RL problem to a sequence modeling problem [T3]. It performs credit assignment directly through the attention
[
mechanism in contrast to Bellman backups, thus modeling wide distribution of behaviors, enabling
better generalization and transfer [1].
Sequence Modeling in Offline RL. Following DT 12 and Trajectory Transformer (TT) L5
there has been an increasing trend in employing advanced sequence-to-See to-seuuence model to solve RL Unfortunately, these improvements are usually transformer-base and hence suffer from the common dilemma of the attention mechanism, i.e. over-paramettrriatton 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 MIP 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 minimal impact 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 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 further demonstrated by 51), which points to their speed and effectiveness in in-context RI. tasks.
tasks [1 41 46
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-lik DeMa outperforms RNN- like DeMa as the long sequences impose a significant computational burden on Mamba without contributing 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 = {s ae, $1+1, Tr)s), where presents the timestep of transition in the dataset. The states and actions are generated by the behavior policy (st, az) ~ d (-). while the next states and rewards are determined by the unknown transition dynamics p(s', r|s, a). The goal of offline RL is to find an approximate policy r(al-) that maximizes expected return T rt]. 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 E. [(R.s.))<T1
[1 Σi= CMSE//C at), where à = (-1st-- kir, RG Rr-K+itts at- II), and Rú = Σt= Te' is the return-to-go (RTG). At test time, a target RTG Ro is manually set to represent the desired performance. We input the trajectories from the last K timesteps into policy π,, 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 [2 [54] However, these models typically possess a large number of parameters
For detailed insights, one may refer to the relevant comprehensive reviews [6 47
3

---

and struggle with handling long sequences effectively. Fortunately, this issue can be addressed by
using SSMs 28 50 29
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 v(t) to an N-D latent state h(tt) before projecting to a I-D output signal y(t) [55],
(t())= Ah(t) + Bu(t), y(t) Ch(t) + Du(t),
(I)
where A E RNNN B R C' E RIxN and D E R are trainable matrices. As u(t) is typically discretized as f=1......... SSM can be discretized by a step size A. Moreover, recurrent SSM can be written discrete convolution. Let ha 0 and D O, we have
1/1 = CAA Bug + CA' Buz + CABu;-1 + CBus, == u K,
(2)
where A, B is the approximation discrete of A, B, and K is called the SSM convolution kernel and can be represented by filter
K= CCB CAC.C.AB..,,CA'... C.
(3)
S4 and other time- invarianv 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 A 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 as RNN-like Mamba, and the other is the SSM utilizing parallel scamning, 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 (ri-1, az-1, 8,) in conjunction with the hidden state ht. Transformer-Iie consistently truncate the input sequences fixed length of K before their introduction into the
Mamba is a direct replacement for the transformer, where we
model throughout the training and inference phases 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
QK
Self-Attention(z) = aV(x), a a= softmax
VaTe
where Q, K, V represent queries, keys, and values respectively, i.e. input sequences after three linear transformations. d₃ is the dimension of the keys. Similarly, current research suggests that the S6 layer in Mamba can viewed as the hidden attention mechanism with unique data-control linear operator [30). Assuming the initial condition ho O, we can obtain formula similar to Equation
y1 = C. Σ (mi-jt) Byz,, h4 Σ (Dim5t1AN) Byzj.
(5)
1=1
where A, = exp (A)), B, = AB.)). and A, softplus(SAaz))) B, Sp(zz), C. (z,) with Sg. Sc and SA are linear projection layers. Softplus is an elementwise function that is a smooth approximation of ReLU.
Since A, is diagonal matrix, 30 simplifies the hidden matrices and gets the attention mechanism of Mamba:
Hidden-AAtention((s) āx, ăi Q.Hke
i := Sc(a), K2 = ReLU(Sa(2))S Hes == exp
Σ
Sa (xx) Sa(e)))A A.
(6)
k= Sa (a

---

Therefore, we can visualize the hidden attention matrices in DeMa, thus gaining deeper understand- ing 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 [] this section presents an analysis from the perspectives of data structures and essential components. Section 4. 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 the Appendix. Our code is available at https 77 thub. com/ htt///nndssSeeeee /DMMa
6-1
RNN-ilkke DeMa
RNN-IIke DeMa
Mixer
Transformer- like DeMa
Miser
Post up-projection
Pre up uppprojection
Lincar projection
Transformer like DeMa
transforaation
(activation. neemaliration, muliippicaaioo)
(III)
Variant design of the DeMa in trajectory optimization. In left portion, (l) represents the RNN-like DeMa (B3LD), which requires hidden state inputs at each decision step; (II) indicates the transformer-like DeMa (B3LD); and (III) refers to the transformer-like DeMa (BI.3D). The right portion illustrates that both types of these DeMa incorporate two distinct residual structures, ie. the post up-projection residual block and the pre up-projection residual block.
Figure 1:
4.1 Input Data Structures
First, compare the RNN-like DeMa (B3LD) with the Transformer-like DeMa (B3LD The average results are shown in Table (with detailed results in Appendix ED. 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-parameted for the sequence length, we assume that variations in sequence length are likely the primary cause of the observed disparities in results. Therefore, explore the effect of sequence length the Transformer-lie like DeMa in subsequent sections.
identical for both types of DeMa except
BBL3D and B3LD represent different concatenation types. Section 4.1 gives a comprehensive explanation.
Unless otherwise specified, all references to DeMa in this context refer the B3LD type.
5

---

Table The average result of DI. RIN--NNNIE like DeMa and Transformer-lik DeMa in Atari (38) and MuJoCo 59 The results are reported with the normalization following [60] Detailed results can be seen in Appendix
Env
DT
RNN-like DeMa
Transformer-like DeMa
Atari
62.2
67.3
111.8
MuloCo
63.4
61.1
66.0
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-lik 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.
The Sorff mance of Fineged Tipe
The Perfumaaci Traring Time
The Becfemance (EGEL Memory Usage
iak bete 19.00 miviie Be DuMtt 3L307
KMM int Denkodt
20000
RARLEE Detiit S.II
BAR-4t DaPalS 201
11660
2.0
10000
1006 0200
2000 5000 30393
Lenglh
32303 34030
232
600
Lenggh
200
1000
200
seguenca Langtt
3300
Figure 2: The impact of sequence length on single-step forward computation time, single-step training time, and GPU memory usage. The sequence length of RNN-like DeMa is 1000.
Finding 1: Transformer-lits DeMa is not only faster but also more memory-efficien than RNN-
like DeMa for short sequence length. The latter only becomes competitive when processing exceptionally long sequences.
While the computational cost of Transformer- like DeMa increases linearly with the expansion of the sequence length, crucial recognize that the increased computational cost may ensure corresponding enhancement in the model's performance. Transformer-like DeMa's Performance may plateau or even decline as the input sequence length exceeds certain threshold. As illustrated in Figure Transformer-like DeMa's performance reaches a plateau in MuloCo [6] when the input sequence surpasses a specific length; while significantly deteriorates with excessively long input sequences in Atari.
How does sequence length affect the performance of DeMa?
Atari
Mujoce
Noraalzzed Scores
Obert
Hopper-Mede Medium
Breakout
Walker2d-Medium
Sequence Length
Sequence Length
Comparison of Transformer-like DcMa's Performance on Atari and MuloCo Tasks. We report mean values averaged over 3 seeds, shaded areas represent deviations.
Figure 3:
Transformer like DeMa performs well with a short sequence length. Extending the sequence length beyond optimal threshold does not yield further improvements and may adversely affect the model's performance.
Finding 2:
Why does DeMa require merely short input sequences? We calculate the hidden attention scores in DeMa via Eq. (5). which reflect the importance of historical information to DeMa. Figure shows the hidden attention scores 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 pecision-making moment, which aligns with the forgetting property of Markov chain [3]. What's more. the hidden attention across different decision steps exhibits periodic pattern towards the current token, suggesting that the model may have learned kinematic features, as agents these environments engage periodic movements
68 5come
500
3006
Inpett Sequatice
oo 5.3 Input Sequerica
Injit Sequince
Input Supuence
+ 5.3.
Decision on Timestop
de un aso
22 en yp Timestap
18 en wit Decision Timestep
yo Lat Decision Timestop
(a) Layer
(b) Layer 2
(c) Layer
(d) Fused Layer
Figure Hidden attention scores DeMa from the 300th to the 600th timestep in Hopper- ceeper--edium--e medium-rrpplayy The X-axis represents timesteps from 300 to 600, the Y-axis represents the past K tokens, and the Z-axis indicates the attention scores given to the K tokens the time of the current decision. More can be seen in Appendix
The reason that Transformer-like DeMa requires only short input sequences is its hidden attention mechanism primarily focusing on the current token. As a result, Longer sequences can lead to difficulties in training without providing benefits.
Finding 3:
Which type of concatenation is suitable for DeMa? Models like the Transformer and Mamba typically pro- cess inputs token by token. However, given MDP, there are three elements s, a, r to consider. Therefore a significant design consideration is the method of concate- nating these three elements into a suitable token format for the model. We experiment to investigate the suitable design for DeMa. By Table 2 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.
Table 2: Input concatenation types compar- ison: "BL3D" refers to the concatenation of input tokens across the embedding dimension, while "B3LD" indicates concatenation across the temporal dimension, depicted Fig- ure Outcomes are averaged across three random seeds.
Game
BL3D
B3LD
Breakout Qbert
72.8+10.6
314.7=10.7
32.2+14.1
54.426.8
Pong
101.926.9
98.2+12.0
Seaquest
1.3+0.0
2.7+0.002
Asterix
3.9+0.3
7.8+0.4
Frostbite
26.3+220.9
31.1+0.01
Assault
127.947.1
169.4133.1
Finding 4: Concatenating state, action, and rtg along embedding dimension a significant negative impact on the results.
Gopher
190.3+60.1
215.8+29.2
Average
69.6
111.8
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 DT's do not heavily rely on attention mechanism 13 49
is the hidden attention mechanism crucial for DeMa? (2) As the Mamba block is integration of the hidden attention mechanism with pre up-projection residual blocks [38], what impact will have on the performance when integrating 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 G
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 investigate whether similar phenomenon exists in hidden attention.
IIt worth noting that what we want to know the attention scores to the previous K tokens at each decision-making step, which is a bit different from the attention scores between output y: and input æy. which is explained in detail in Appendix
7

---

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 D the performance exhibits 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 certain level of performance. This zeroing of parameters completely removes the hidden attention, ceasing to process historical information and relying solely residual connections to transmit information. This suggests that the residual connections are functional and the role hidden attention crucial for DeMa.
Walker2d
Halfchestah
60
origin
40
Walkerzd
20
0-
Figure 5: Normalized return after swapping the hidden attention of single layer from another DeMa time. The black dashed line represents the evaluation results of the original model. "1", "2", and "3" represent the index of swap layers respectively, and "all" represents the result after swapping all parameters of the hidden attention. It can be seen that swapping the hidden attention has significant impact on the results.
Finding 5: Replacing the hidden attention mechanism would lead to reduction in performance, unlike the attention mechanisms used in transformers. Therefore, the hidden attention mechanism plays crucial role in DeMa.
Table Performance comparison between DT, hidden attention with post up-projection residual block in transformer (DeMa with post.) and hidden attention with pre up up-projection residual block (DeMa) in Atari.
Game
DT
DeMa with post.
DeMa
Breakout
242.4+31.8
296.2+216.3
314.7=10.7
Qbert
28.8+10.3
56.9+10.4
54.426.8
Pong
105.6:2.9
104.6+11.9
98.2+12.0
Seaquest
2.720.7
2.6+0.001
2.720.002
Asterix
5.2+1.2
6.5+1.8
7,8+0,4
Frostbite
25.6+2.1
31.8+4.8
31.120.01
Assault
$2.1+36.2
146.4+16.1
169.4+33.1
Gopher
34.8+10.0
228.9+81.5
215.8:29.2
Average 62.2
109.2
111.8
Table 4: Performance comparison between DT, hidden attention with post up-projection residual block in the transformer (DeMa with post.) and hidden attention with pre up-projection residual block (DeMa) in MuloCo.
Dataset
Eny-Gym
DT
DeMa with post.
DeMa
M
HalfCheetah
42,6
42.7=0.02
43+0.01
M
Hopper
68.4
68.4+2.1
74.5+2.9
Walker
75.5
77.5+2.2
76.6+0.2
M-R
HalfCheetah
37.0
40.8+0.18
40.7+0.03
M-R
Hopper 85.6
86.1+26.9
90.7+6.1
M-R
Walker
71.2
74.43+2.1
70.5+0.1
HalfCheetah
83.8+18
93.2+0.01
M-E
Hopper
109.6
109 (09.8+0.2
11140.03
M-E
Walker 109.3
109.6*0.3
106+11.7
Average-Gym
76.4
77.0
78.5
What occurs when combining hidden attention with post up-projection residual blocks? Mamba represents the integration the hidden attention mechanism with pre up-projection residual blocks discussed in (38.. To determine the contributing factor 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 3 and Table [ although the overall average results of DeMa are slightly better than those of DeMa with post., 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 measurable performance improvement over the DT, it is reasonable to conclude that the hidden attention mechanism within DeMa plays a pivotal role.
The results obtained using both post up-projection and pre po-projection types of residual block structures are similar while they both perform better than DT. Therefore, the hidden attention mechanism is key to its success.
Finding 6:
Evaluations on Offine RL Benchmarks
Table Results for % DQN-replay datasets. We evaluate the performance DeMa on eight Atari games.
Game
CQL
BC
DT
DC
DChybrid
DeMa(Ours)
Breakout
211.1
142.7
242.4+31.8
352.7+44.7
416.0 100..4
314.7+10.7
Qbert
104.2
28.8+10.3
67.0=14.7
62.6
54.4+6.8
Pong
111.9
76,9
105.642.9
106.5+2.0
111.111
98.2+12.0
Scaquest
1.7
2.2
2.7+0.7
2.6+0.3
2.7 #0.04
2.7+0.002
Asterix
4.6
4.7
5.2+1.2
6.5+1.0
613
7.8+0.4
Frostbite
9.4
16.1
25.6+2.1
27.8+3.7
28.0 21.8
31.1+0.01
Assault
73.2
52.1436.2
73.8+20.3
79.0
169.4433.1
Gopher 2.8
33.8
34.8+10.0
52.5+9.3
$1.6
215.8+29.2
Average
64.9
44.9
62.2
86.2
94.7
111.8
Table 6: Results for MuloCo. dataset names are abbreviated as follows: "medium" as "M", "medium-replay" as "M-R" and "medium-expert" as "M-E". The results are reported with the expert-normalized following
Dataset
Enviromment
CQL
DS4
RvS
DT
GDT
DeMa(Ours)
M
HalfCheetah
44.0
42.5
41.6
42.6
42.9
43+0.01
M
Hopper
58.5
54.2
60.2
68.4
65.8
74.5=2.9
Walker
72.5
78.0
71.7
75.5
77.8
76.6+0.2
M-R
HalfChoetah
45.5
15.2
37.0
39,9
40.7=0.03
M-R
Hopper
95.0
49.6
73.5
85.6
81.6
90.7#6,1
M-R
Walker
77.2
69.0
60.6
71.2
74.8
70.5+0.1
M-E
HalfCheetah
91.6
92.7
92.2
88.8
92.4
93.220.01
M-E
Hopper
105.4
110.8
101.7
109.6
110.9
111+0.03
M-E
Walker
108.8
105.7
106.0
109.3
109.3
106+11.7
Average
77.6
68.6
71.7
76.4
76.8
78.5
Table 7: The resource usage for training DT. DC and DeMa on Atari and MuJoCo.
Complexity
DT
DC
DeMa(Ours)
Training time per step(ms)
55
43
50
Atari
GPU memory uSage(GiB)
MACs All params #
12.1G/46.5C
11.1G/40.6G
8.8G/36.3G
2.35M
1.94M
1.7M
Training time per step(ms)
56/58
53.6/53.9
57.6/58.8
Gym
GPU memory UGB))
0.65/0.8
0.55/0.6
1.0/1.0
MACs
2.5G/9.5G
1.6G/6.1G
0.7G/2.1G
All params #
726.2K/2.6M
536K/1.9M
175.5K/500.0K
In this section, we delve into a comparative analysis DeMa's performance against various DT's, Our investigation primarily centers 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 K 20 for DeMa. The detailed eeper-parameter on DeMa are available in Appendix D Our main results are shown in Table and Table 6 DeMa achieves significantly higher average and
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 [ Moreover, DeMa has better scalability for input length which can be seen in Figure 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 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 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 Mar-Maakovian 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.
Acknowledgments and Disclosure of Funding
This work is supported by STI 2030-Majoo Projects (No. 20212ID0201005), National Natural Science Foundation of China (No. 72301289), and the Zhejiang Province Science Foundation under Grants LD24F020002. We thank zigzagcai for his PR support variable-length sequences for mamba block. We thank Liang Zhang and Yang Ma for their valuable suggestions and collaboration. We sincerely appreciate the time and effort invested by the anonymous reviewers in evaluating our work and are grateful for their valuable and insightful feedback.
References
[1] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu, Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0166. 2020.
[2] B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEFE Transactions on Intelligent Transportation Systems, 2909) 2021.
[3] David Brandfonhrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline without off-policy evaluation. Advances in neural information processing systems, 34:4933-4946, 2021.
[4] M Mehdi Afsar, Trafford Crump, and Behrouz Far. Reinforcement learning based recommender systems:
A survey, ACM Computing Surveys, 55(7).1-38, 2022.
10

---

[5] Rafael Figueiredo Prudencio, Marcos ROA Maximo. and Esther Luna Colombini. A survey offline reinforcement learning: Taxonomy, review, and open problems. IEEE Trunsactions Neural Networks and Learning Systems, 2023.
[6] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-hased offline reinforcement learning. Advances neural information processing systems, 33:21810-21823, 2020.
[7] Cong Lu, Philip Ball, Jack Parker-Holder, Michael Osborne, and Stephen Roberts. Revisiting design choices offfine model-basee reinforcement learning. arXiv preprint arXiv:2110.04135. 2021.
[8] Haoyang He. A survey offline model-based reinforcement learning. arXiv preprint arXiv:2305.03336) 2023.
[9] Phillip Swazinna, Steffen Udluft, Daniel Hein, and Thomas Runkler. Comparing model- free and model- based algorithms for offiine reinforeement learning. IFAC--apersOnnLinee, 55(15):19-25 2221999-26666 2022.
[10] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[11] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without explo- ration. In International conference machine learning. pages 2052-2062. PMLR, 2019.
[12] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 01155088-5957, 2021.
[13] Jeonghye Kim, Suyoung Lee, Woojun Kim, and Youngchul Sung. Decision convformer: Local filtering in metaformer is sufficient for decision making. arXiv preprint arXiv:2310.03002. 2023.
[14] Shengehao Hu, Li Shen, Zhang, and Dacheng Tao. Graph decision transformer. arXiv preprint 0XX1:2230.03377,, 2023.
[15] Michael Janner, Qiyang Li, and Sergey Levine. Offine reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34: 1286-1266,, 2021.
[16] Shengehao Hu, Li Shen, Ya Zhang, Yixin Chen, and Dacheng Tao. On transforming reinforcement learning with transformers: The development trajectory. IEEE Transactions on Pattern Analysis and Muchine Intelligence, 2024.
[17] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018
[18] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[19] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)9, 2019.
[20] Tom Brown, Benjamin Mann. Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[21] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[22] Sncha Chaudhari, Varun Mithal, Gungor Polatkan, and Rohan Ramanath. An attentive survey of attention models, ACM Transactions on Intelligent Systems and Technology (TIST), 12(5).1 32, 2021.
[23] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF winter conference applications of computer vision, pages 3531-3539, 2021.
[24] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
[25] Tri Dao. Flashattention-2. 2: Faster attention with better parallelism and work partitioning, arXiv preprint mmXiv2237.08691, 2023,

---

[26] Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng, Synthesizer: Rethinking self-attention for transformer models. In International conference on machine learning, pages 10183- 10192. PMLR, 2021.
[27] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang. Jiashi Feng. and Shuicheng Yan. Metaformer is actually what you need for vision. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10819-10829, 2022.
[28] James D Hamilton. State-space models. Handbook of econometrics, 4:3039-3080, 1994.
[29] Albert Gu and Tri Dao. Mamba: Linear-tima sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
[30] Ameen Ali, Itaaar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590, 2024.
[31] Kunchang Li. Xinhao Li. Yi Wang. Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Videomamba: State space model for efficient video understanding, arXiv preprint arXiv:2403.005 066977 2024.
[32] Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang. Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401. 2024.
[33] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang. Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint urXiv:2401.094177 2024.
[34] Jun Ma, Feifci Li, and Bo Wang. U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722, 2024.
[35] Mohammad Reza Samsami, Artem Zholus, Janarthaman Rajendran, and Sarath Chandar. Mastering memory tasks with world models. arXiv preprint 24X 2024.
[36] Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers, and s4. Advances in Neural Information Processing Systems, 36, 2024.
[37] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahami. Structured state space models for in-context reinforcement lcarning. Advances in Neural Information Processing Systems, 36, 2024.
[38] Maximilian Beek, Korbinian Pöppel. Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint 2405 2024.
[39] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offiine model-based policy optimization. Advances in neural information processing systems, 34:28954 28967, 2021.
[40] Tenglong Liu, Yang Li, Yixing Lan, Hao Gao, Wei Pan, and Xin Xu. Adaptive eerantage-guided policy regularization for offline reinforcement learning. In Forty-first International Conference on Machine Learning, 2024.
[41] Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt- tuninttttnnnggg decision transformer with preference ranking, arXiv preprint arXiv:2305.088 09648, 2023.
[42] Shengehao Hu, Ziqing Fan, Chaoqin Huang, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Q- value regularized transformer for offline reinforcement learning. In International Conference on Machine Learning, 2024.
[43] Shengehao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Harmodt: Harmony multi-task decision transformer for offline reinforcement learning. In International Conference on Machine Learning, 2024.
[44] Shengchao Li Shen, Ya Zhang, and Dacheng Tao. Learning multi-agent communication from graph modeling perspective. In The Twelfth International Conference on Learning Representations, 2024.
[45] Jifeng Hu, Yanchao Sun, Sili Huang. SiYuan Guo, Hechang Chen, Li Shen, Lichao Sun, Yi Chang, and Dacheng Tao. Instructed diffuser with temporal condition guidance for offline reinforcement learning. arXiv preprint 2306. 04875, 2023.
12

---

[46] Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision transformer: Reinforcement learning via hierarchical chain-of-thought. arXiv preprint arXiv:2405.2069) 20692, 2024.
[47] Wenzhe Li. Hao Luo, Zichuan Lin. Chongjie Zhang. Zongqing Lu, and Deheng Ye. A survey on
transformers in reinforcement learning. arXiv preprint 230v:2330.00304 2023.
[48] Scott Emmons, Benjamin Eysenbach, llya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint 2rXi::21 2021.
[49] Daniel Lawson and Ahmed H Qureshi. Merging decision transformers: Weight averaging for forming multi-task policies. arXiv preprint arXiv:2303.0773 0755,, 2023.
[50] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequenees with structured state spaces. arXiv preprint arXiv:21 1eXxx:22111000000 1.00366, 2021.
[51] Sili Huang, Jifeng Hu, Zhejian Yang. Liwei Yang, Tao Luo, Hechang Chen, Lichao Sun, and Bo Yang. Decision mamba: Reinforvement learning via hybrid selective sequence modeling. arXiv preprint arXiv:2406. 00079, 2024.
[52] Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision s4: Efficient sequence- based rl via state spaces layers. In The Eleventh International Conference on Learning Representations, 2022.
[53] Toshihiro Ota. Decision mamba: Reinforcement learning via sequence modeling with selective state spaces. arXiv preprint arXiv:2403. 19925, 2024.
[54] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205 2022.
[55] Sidd Karamcheti Sasha Rush. The annotated s4. https://srusa www.ggi.......nngsiieeeeee =4/ 2023.
[56] Yijun Yang, Zhaohu Xing. and Lei Zhu. Vivim: video vision mamba for medical video object segmenta- tion. arXiv preprint arXiv:240I. J4168, 2024.
[57] Zeyu Zhang, Akide Liu, Ian Reid, Richard Hartley. Bohan Zhuang, and Hao Tang. Motion mamba: Efficient and long sequence motion generation with hierarchical and bidirectional selective ssm. arXiv preprint arXiv:2403.07487, 2024.
[58] Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253 279, 2013.
[59] Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[60] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances neural information processing systems, 34:25476-25488, 2021.
[61] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
[62] Weihao Yu and Xinchao Wang. Mambaout: Do we really need mamba for vision? arXiv preprint vrXiv:2455.07992. 2024.
[63] Justin Fu, Aviral Kumar. Olir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint ar 2004:2000. 2020.
[64] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pages 104-114. PMLR, 2020.
[65] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, ct. al. Human-level control through deep reinforcement learning. nature, 518(7540):529 533,, 2015.
[66] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 13:1199-1191. 2020.
[67] Xiao Zhou, Yujie Zhong, Zhen Cheng, Fan Liang, and Lin Ma. Adaptive sparse pairwise loss for object re-identification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pages 19691-19700 2023.
13

---

DC and Supplementary Material for
Is Mama Compatible with Trajectory Optimization in Offline Reinforcement Learning?
Contents
Environment and Dataset
14
Baselines
14
The Procedure of Training and Inference
14
D Implementation details of DeMa
15
Detailed results
17
Tasks Requires Long Horizon Planning Skills
17
Further Ablation Study
19
H Integrating DeMa with Other Methods
19
MuJoCo and Atari Tasks Scores
19
J Types of Hidde Attention Scores
19
A Environment and Dataset
The MuJoCo domain [59] evaluates the performance of RL algorithms in continuous control tasks. In keeping with previous studies, select three games from the standard locomotion environments (59] Gym 611. namely HalfCheetah, Hopper, and Walker, and three different dataset settings, namely medium, medium-replay, and medium-expert
MuJoCo.
66
Atari. Atari 58 is an ideal platform for evaluating an agent's ability in long-term credit assign- ments. We conduct experiments in eight different games: Breakout, Qbert, Pong, Seaquest, Asterix, Frostbite, Assault, and Gopher.
use 1% DQN Replay Dataset (64) as our training dataset, which encompasses total of 500,000 timesteps worth of samples generated throughout the training process of DQN agent (65). It's worth noting that the version of "atari-py" and "gym" we use 0.2.5 and 0.19.0 respectively, which is noted by the official code in httpss//ghhbbbbooo///////rrr/r/e
rl
B Baselines
Baselines for MuJoCo. To evaluate DeMa's performance in the MuloCo, compare DeMa with one value-base 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 for RvS from and for GDT from 14
Baselines for Atari. In the Atari domain, compare DeMa with CQL 666
DT
12
DCkybrid 13 The results of baselines are directly borrowed from 13
C The Procedure of Training and Inference
Training resources We use one NVIDIA GeForce RTX 4090 train each model in MuJoCo and one NVIDIA GeForce RTX 3090 to train each model in Atari. Training each model typically takes
14

---

3-8 hours and 5-14 hours in MuJoCo and Atari respectively. However, since each environment needs be trained three times with different seeds, the total training time is usually multiplied by three.
The procedure of Transformer-like DeMa The Training and evaluation for Transformer-In like DeMa are similar to variant DTs. Given a dataset of offline trajectories, we randomly select starting point and truncate it into sequence of length K. After forming 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 sequence of length K in the evaluation phase.
The procedure of RNN-like DeMa For RNN-like DeMa, the input during training 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
D Implementation details of DeMa
We implement DeMa based on the official code of DT and the Mamba We have also adopted the code from HiddennammbaAtt/ to calculate the attention scores of DeMa 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 compute the MACs of DeMa.
Tables T0 provide a comprehensive list of hyper-parameter for our proposed transformer-Ile like DeMa and RNN-like DeMa applied to MuJoCo and Atari environments. To ensure fair comparison, we adopt similar
settings to DT [2] and DC [1]
D.1 Hyper-parameters in MuJoCo
Table Hyper-parameters of DeMa for MuloCo.
Hyper-parameter
Value
Layers
3
Embedding dimension
128
Nonlinearity function
|
Batch size
64
Context length
20
Dropout
0.0
Learning rate
10-4
Grad norm clip
0.25
Weight decay Learning rate decay d_model d __state expand
10-4
Linear warmup for first 10 training steps
64
64
2
For our training on MuJoCo, the majority of the hyper-paramente parameters in Table are adapted from For the learning rate, we use a learning rate of 10 for training in medium,,,, hopper- 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-mediun--eplas 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. Table we use ReLU as per convention. For DeMa's parameers,,,
13
while use
d model of 128 in all expert datasets, while use
15

---

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.
D.2 Hyper-parameters in Atari
Transformer-like De.a. For the Atari game we mostly follow those in Table from 12 The only adjustment made is the context length 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.
Table 9: parammeer of Transformer like DeMa for Atari.
Slyper-paramees
Value
Layers
Embedding dimension
256
Nonlinearity function
(state encoder) 128
Batch size
Context length K
30
Return-to-go conditioning
90 Breakout, 12000 Qbert
20 Pong. 1 1750 Seaquest
700 Asterix, 1450 Frostbite
1200 Assault, 6500 Gopher 0.1 6 10-1
Dropout
Learning rate
Grad norm clip
Weight decay
0.1
Learning rate decay
Linear warmup and cosine decay (see code for details)
Max epochs
10
Adam betas
(0.9, 0.95)
Warmup tokens
512x 20
Final tokens
6x 500000 K
d model
128
conv
4
d state
64
expand
2
Table 10: Hyper-parameters of RNN-like DeMa for Atari. The other hyper-parameters are kept consistent with those in Table 9
Hyper-parameter
Value
Context length
all trajectory
Batch size
8
Learning rate
10-4
inner_it
200
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 1 of the DQN-replay dataset utilized. If the prior parameter settings were be used, the training would done after only few hundred upgrades, thereby resulting in an unsatisfactory performance. Therefore, we consider multiple updates for single sample, while simultaneously lowering the learning rate as shown in Table TO) What's more, due the limitation of GPU memory, we can only set batch size of for Atari. Specifically, For the Frostbite, we set a batch size of an epoch of 50. The other hyper- hyper--araman parameters are kept consistent with those in Table
16

---

E
Detailed results
Table and Table 12 show detailed results between RNN-like DeMa Transformer-lik like DeMa.
can be observed that the performance of the RNN-like DeMa is not good as that of the Transformer- like DeMa, and Figure also shows that the RNN-like DeMa requires more computational overhead. Hence, using the RNN model in trajectory optimization seems to be unnecessary. section [4.] finds that past historical information does not provide much assistance to current decision-making.
However, in tasks that require memory capability are model-based, the RNN-like DeMa could be better choice. This could be a direction for deeper future research based on 35 37
Table 11: The Comparison of DT, RNN-like DeMa, and Transformer-lik DeMa in Atari Games.
Env
DT
RNN-like DeMa
Transformer-lik DeMa
Breakout
242.4+31.8
166.0
314.7+10.7
Qbert
10.8+10.3
13.6
54.4+6.8
Pong
105.6*2.9
109.6
98.2+12.0
Seaquest
2.7+0.7
1.7
2.7=0.002
Asterix
5.2+1.2
4.7
7.8+0.4
Frostbite
25.6+2.1
31.1+0.01
Assault
52. 12.1236.2
117.8
169.4433.1
Gopher
34,8+10.0
116.7
215.8+29.2
Average
62.2
67.3
111.8
Table 12: The comparison between DT, RNN-like DeMa, and Transformer-lid DeMa in MuJoCo.
Dataset
Environment
DT
RNN-like DeMa Transformer-lik DeMa
M
HalfChectah
42.6
42.6+0
43=0.01
M
Hopper
68.4
61.7:4.9
74.5=2.9
M
Walker
75.5
76.7=0.2
76.620.2
M-R
HalfCheetah
37.0
36.9=0.3
40.7+0.03
M-R
Hopper
85.6
80.5=25.8
90.726.1
M-R
Walker
71.2
68.128.1
70.5+0.1
Average
63.4
61.1
66.0
F
Tasks Requires Long Horizon Planning Skills
Three experiments involving delayed rewards(MuloCo with delayed rewards) and maze naviga- tion(maze2d, antmaze) are conducted to investigate how would DeMa perform on tasks that require long horizon planning skills.
F.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, shown in Table T3] 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.
train DeMa using the same hyper-parameter settings, and the results
F.2
Maze Navigation
There environments in maze navigation. Maze2d: This environment aims reaching goals with sparse rewards, which suitable for assessing the model's capability to efficiently integrate data and execute long-range planning. The objective of this domain is guide agent through maze to reach a designated goal. Antmaze: This environment is similar to maze2d, while the agent is an ant
'Due the high experimental costs, only the RNN-lik like DeMa in Atari once.
17

---

Table 13: Results for D4RL datasets with delayed (sparse) reward. The "Origin Average" the table represents the normalized scores of evaluations across six datasets under the original dense reward setting.
Dataset
Env-Gym
CQL
DS4
DT
GDT
DeMa(Ours)
M
HalfCheetah
1.0 1.0
42.7+0
42.2 0.2
43+0
42.9=0.01
M
Hopper
1.0
58.2+0.7
2.4
58.2+2,4
69.1#6.5
M
Walker
0.0
75.7+0.5
2.0
78.9+0.1
77.6+1.5
M-R
HalfCheetah
6.9
15,5+0
$3.0
41+0.1
41.1=0.115
M-R
Hopper
5.9
77.5+0.4
50.8 14.3
79.8+15.9
83.8+6.9
M-R
Walker
3.7
69.143.2
51.6 24.6
70.4+8.7
71.7=5
Average-Gym
7.2
56.45
50.8
61.9
64.4
with degrees of freedom. For our training on Maze, the majority of the Syper-parameters in Table 14 and Table 15] For maze2d-medium, we use K 8 and embedding dim=256. For maze2d-umaze
we use hyper-parameters in Table 8
Table 14: Syper--arameterss of DeMa for antmaze.
Hyper-parameter
Value
Layers
3
Embedding dimension
128
Nonlincarity function
\
Batch size
32
Context length K
5
Dropout
0.1
Learning rate
-5
Grad norm clip
0.25
Weight decay
10-4
Learning rate decay d d_model
Linear warmup for first 10 training steps
128
state
64
num_eval_episoder
50
max_iters
50
num_steps_per_iter
2000
Table 15: Hyper-parameters of DeMa for maze2d.
Hyper-parameter
Value
Layers
Embedding dimension
128
Nonlinearity function
\
Batch size
32
Context length K Dropout
20
0.1
Learning rate
2e -5
Grad norm clip
0.25
Weight decay Learning rate decay d model d state num eealeeealleeeeeeeeeeeee _eal_eessoees
10 10-4
Linear warmup for first 10y training steps
64
64
50
max max_iters num_steps_per_itee
50
2000
We compare DeMa with DT 12 GDT 14 and DC 13
The results DT and GDT are directly borrowed from 42). Results in Table [6 show that DeMa performs better compared to DT in the maze
18

---

navigation task. The visualization analysis of the hidden attention mechanism in these environments can be found in Figure I0 and Figure [T]
Table 16: Results for maze2d and antmaze.
Dataset
Env-Gym
DT
GDT
DC
DeMa(Ours)
umaze
31.0
50.4
36.3+3
54.319.4
medium
maze2d
8.2
7.8
2.1+1.02
10.3+3.1
large
2.3
0.7
0.940
2.8+2.2
umaze umaze-diverse
antmaze
59.2
76
85.00
8240
53
69
78.5
80.7+6.2
G Further Ablation Study
Table 17: The affection of position embedding.
Dataset
Env
DeMa with pos. embed.
DeMa without pos. embed.
M
HalfCheetah
42.8+0 71.2
43:0.01
M
Hopper
74.5+2.9
M
Walker
77.2+0.1
76.6+0.2
M-R
HalfCheetah
40.240.1
40.7=0.03
M-R
Hopper
77.2+35
90.726,1
M-R
Walker
69.1#10.2
70.5+0.1
Average
63.0
66.0
All params
431.5K
175.5K
DeMa does not need the position embedding. Position embedding generally used in trans- formers to help the model understand the sequential nature of the data. It's way of encoding the position of tokens the sequence, and it can be crucial in tasks where the order of the data matters. Although
DeMa similar using 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 remember sequential information. As shown in Table [7] the addition of position embedding not only failed to enhance the performance of the model but also led significant decrease in performance embedding significantly increased the model's 's parameter count, therehy adding to its computational burden. This finding highlights the advantage the DeMa in terms of lightweight design, indicating its suitability for tasks with limited resources.
certain tasks. Additionally, the introduction of position
H 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 [8 Q-DeMa achieves performance comparable to state-of-the- art models while utilizing less than one-seventh the parameter size QT. This finding underscores the significant potential of applying Mamba to RL.
I MuJoCo and Atari Tasks Scores
Table 19 shows the normalized scores used in MuJoCo and Atari tasks, followed by 663 and 60
Types of Hidde Attention Scores
In the previous articles triangular matrix. However, this lower-trianglla matrix reflects the attention scores of each generated
67 the visualization of attention in DT was in the form of lower
19

---

Table 18: Q-DeMa's Results for D4RL datasets.
Dataset
Env-Gym
DT
DeMa(Ours)
QT
Q-DeMa(Ours)
M
HalfCheetah
42.6
4340.01
51.4 0,4
51.2+0.04
M
Hopper
68.4
74.5+2.9
96.9 + 3.1
88. 88.149.61
M
Walker
75.5
76.6+0.2
88.8 0.5
89.1+0.2
M-R
HalfCheetah
37.0
40.720.03
48.9 0.3
48.6+0.3
M-R
Hopper
85.6
90.7+6.1
102.0 = 0.2
101.5#0.1
M-R
Walker
71.2
70.5+0.1
98.5 1.1
99.8+1
Average-Gym
63.4
66.0
81.0
79.7
All params #
726.2K/2.6M
175.5K/500.0K
3.7M
500K
Table 19: MuJoCo and Atari baseline scores used for normalization
Env/Game
Random
Expert/Gamer
Hopper
-20.3
3234.3
Gym
Halfcheetah
-280.2
12135
Walker2d
4592.3
Breakout
1.7
30.5
Qbert
163.9
13455
Pong
-20.7
14.6
Atari
Seaquest
68.4
42054.7
Asterix
210
8503
Frostbite
65
4335
Assault
222
742
Gopher 258
2412
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 seen, the element in Tigure/oar the i-th row and j-th column represents presents the output yı attention score to input zz-- In the training phase, 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 L1, D), and the model also outputs L,) At this time, we only use the last of the model's output, which is y/: :l. corresponding to the last row in the matrix. Therefore, not quite appropriate to judge the context information the model focuses on each want understand the
matrix in Figure model's srscisionmmaiing behavior at each step, thus leads to the creation of Figure 4 Figure and Figure It also demonstrates strong forgetting characteristies. This aligns with the properties of Markov chains as described in [13]. where the sequence of states precisely forms 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 11 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 significantly higher focus on states compared to rewards or actions.
decision-making step based the power-triaguule
20

---

(a) channel 30 to channel 36, layer 1
(b) channel to channel 36, layer
(c) channel 30 channel 36, layer
(d) fused channel, layer
(e) fused channel, layer 2
fused channel, layer (g) fused channel and layer
Figure 6: Hidden Attention Score Matrix of each channel and layer of DeMa, trained on the Hopper- medium dataset. The element Aij present the attention score between output yi and input æ;
21

---

Score
Sc00
6 b & & & Input Sequence
+ 4 to t Input Sequence
300 150 .00 450 400 559
300 350 400 as0 500 Decision Timestep
Decision Timestep
(a) Layer 1
(b) Layer 2
Score
6 % lo % Input Sequence
0 2 % 4
Input Sequence
350 450 you 550 600 Decision Timestep
Decision Timestep
aon 450 seo 550 snc
(c) Layer
(d) Fused Layer
Figure 7: Hidden attention scores of DeMa from the 300th the 600th timestep. trained the Walker2d- medium dataset.
Score
lo +b b to to Input Sequence
fo db & to Input Sequence
300 409 450 ,00 Decision Timestep
300 350 200 a50 500 Decision Timestep
(a) Layer I
(b) Layer 2
25
25
Score
2.0
2.0
155
1.0
0.
1.0
05
Input Sequence
to is 5 f
le d 15-5 fo Input Sequence
300 yo0 sse 600 Decision Timestep
,50 500 sso soo Decision Timestep
(c) Layer
(d) Fused Layer
Figure 8: Hidden attention scores of DeMa from the 300th to the 600th timestep. trained on the Halfcheetah- medium-ee expert dataset.
22

---

(a) Layer
(b) Layer
(c) Layer
(d) Fused Layer
Figure 9: Hidden attention scores DeMa from 540th 600th timestep, trained on the Walker2d- medium dataset with delayed rewards. Hidden attention mechanism highlighting more focus on historical observations. Top: 2D Representation, Bottom: 3D Representation.
(a) Layer
(b) Layer
(c) Layer 3
(d) Fused Layer
Figure 10: Hidden attention scores DeMa from 80th to 140th timestep in maze2d-umaze. Representation, Bottom: 3D Representation.
Top:
Mad Ne
bark
(a) Layer
(b) Layer
(c) Layer
(d) Fused Layer
Figure 11: Hidden attention scores of DeMa from the 80th the 140th timestep in antmaze divze-
Top: 2D Representation, Bottom: 3D Representation.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's s contributions and scope?
Answer: [Yes]
Justification:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
In the last paragraph in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: There is theoretical result, this study is empirical research.

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: In Appendix D
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c)
If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
25

---

Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
Stt////ooSSSesssiiieieeee wwi.iiosiiiiiycyooo. for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Mostly following previous studies, some parameters have been modified according to our investigation can be seen in Section 4.1 and Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The results are averaged across three random seeds, with deviation for Fig and variance in others.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
26

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
In Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://aaaraa...cpppaaaa//aaaaaa
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
27

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
28

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "IHkKHwLq2y", "reviewer": "Reviewer_BojN", "summary": "The work introduces Decision Mamba (DeMa) to address the challenges in offline RL posed by the large parameter size and limited scalability of Transformer-based methods. DeMa aims to achieve similar performance to Transformers with significantly fewer parameters.  DeMa surpasses the DT with significantly fewer parameters in the benchmarks.", "questions": "1. What is the definition of $L_{MSE/CE}$ and $_{-K:t}$?\n\n2. Is DeMa applicable to all trajectory optimization methods?", "limitations": "The suggestions have been claimed in \"Weaknesses\" and \"Questions\".", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Extensive evaluations demonstrate the effectiveness of DeMa, highlighting its superior performance and efficiency compared to existing methods.\n\n2. DeMa provides a novel solution to the parameter size and scalability issues in trajectory optimization.", "weaknesses": "1. Some symbols are not defined before use.\n\n2. This paper seems to have little relation to RL and appears more like a method applicable to all trajectory optimization.\n\n3. There is too little discussion on the relationship to RL in sections 3.2 and 3.3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your insightful comments, which have catalyzed numerous enhancements and refinements to the paper. In the following, we reply to the questions one by one for the convenience of checking.\n\n---\n**Weakness 1 & Question 1**: *Some symbols are not defined before use.*\n\n**Response**: Thanks for your comment. We have thoroughly reviewed the manuscript and provided definitions for all symbols. Additionally, we have removed some redundant symbols to enhance the article's readability. Major modifications include:\n1. $L_{MSE/CE}$ in line 123. It represents the loss function of DeMa, when the output corresponds to continuous actions (as in MuJoCo), the loss function is Mean Squared Error (MSE), while when the output pertains to discrete actions, the loss function is the cross-entropy loss.\"\n2. $s_{-K:t}$ in line 124. Subscript denotes a range from time step t-K+1 to time step t. We have corrected it to $_{t-K+1:t}$ to accurately represent the intended range from time step $t-K+1$ to time step $t$.\n3. $A,B,C,D$ in line 133. $A\\\\in\\\\mathbb{R}^{N\\\\times N} , B\\\\in\\\\mathbb{R}^{N\\\\times 1}, C\\\\in\\\\mathbb{R}^{1\\\\times N}, D\\\\in\\\\mathbb{R}$ are parameter matrices in State Space Model.\n4. $\\\\bar{K}$ and $\\\\bar{C}$ in line 136-137. We have removed the redundant symbol L and eliminated the symbol $\\\\bar{C}$, which was used incorrectly. Now it becomes\n$$\ny_{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},\\\\\\\\\n\\\\bar{K}=(C\\\\bar{B},C\\\\bar{A}\\\\bar{B},\\\\ldots,C\\\\bar{A}^{k}\\\\bar{B},\\\\ldots).\n$$\n5. Line 158-159. To provide a clearer formulation of hidden attention, we have revised it as follows\n$$\ny_i=C_i\\sum_{j=1}\\^i \\\\big(\\\\Pi_{k=j+1}\\^i \\\\bar{A}\\_k \\\\big) \\\\bar{B}\\_j x_j,\\\\\\\\\nh_i=\\sum_{j=1}\\^i \\\\big(\\\\Pi_{k=j+1}\\^i \\\\bar{A}\\_k \\\\big) \\\\bar{B}\\_j x_j,\n$$\nwhere $B_i=S_B (\\\\hat{x}\\_i)$, $C_i=S_C(\\\\hat{x}\\_i)$, $\\\\Delta_i=\\\\text{softplus}(S_{\\\\Delta}(\\\\hat{x}\\_i))$ with $S_B$, $S_C$ and $S_{\\\\Delta}$ are linear projection layers, and SoftPlus is an elementwise function that is a smooth approximation of ReLU. $\\\\bar{A},\\\\bar{B}$ is the discretization of $A, B$, that is, $\\\\bar{A}\\_{i}=\\\\exp(\\\\Delta\\_{i}(A))$ and $\\\\bar{B}\\_{i}=\\\\Delta_{i}(B_i)$.\n\n---\n**Weakness 2**: *This paper seems to have little relation to RL and appears more like a method applicable to all trajectory optimization.*\n\n**Response**: Thank you for your comments. \n\n1.Trajectory optimization methods treat RL problems as sequence modeling problems to get better performance and generalization[1]. Decision Mamba (DeMa) proposed in this paper is designed to address the challenges posed by transformer-based trajectory optimization methods, particularly the issues of large parameter sizes and limited scalability, which are long-standing issues that have been largely overlooked. Exploring the potential of DeMa in RL tasks will contribute to the development of scalable and efficient decision models, which will significantly enhance the practical applications of RL. \n\nWe revise Sections 1-2 to emphasize the significance of DeMa in the context of RL.\n\n2.Our paper primarily focuses on the analysis of DeMa in RL to provide valuable insights for the community. Indeed, DeMa can be combined with other trajectory optimization methods to achieve even better performance. \n\nWe conduct additional experiments, the results of which are summarized in the table below. By integrating DeMa with QT [2], we develop Q-DeMa, which achieves performance comparable to state-of-the-art (SOTA) models while utilizing less than one-seventh of the parameter size of QT. This finding underscores the significant potential of applying Mamba to RL and emphasizes the critical importance of the research presented in this paper.\n\nTable A: Results for D4RL datasets. The dataset names are abbreviated as follows: \"medium\" as \"M\", \"medium-replay\" as \"M-R\". We report the mean across 3 seeds.\n|Dataset|Env-Gym|DT|DeMa (Ours)|QT|Q-DeMa (Ours)|\n|-|-|-|-|-|-|\n|M|HalfCheetah|42.6|43|51.4|51.2|\n|M|Hopper|68.4|74.5|96.9|88.1|\n|M|Walker|75.5|76.6|88.8|89.1|\n|M-R|HalfCheetah|37.0|40.7|48.9|48.6|\n|M-R|Hopper|85.6|90.7|102.0|101.5|\n|M-R|Walker|71.2|70.5|98.5|99.8|\n|Average|-|63.4|66.0|81.0|79.7|\n|All params #|-|0.7M/2.6M|0.2M/0.5M|3.7M|0.5M|\n\n[1] On Transforming Reinforcement Learning with Transformers: The Development Trajectory\n\n[2] Q-value Regularized Transformer for Offline Reinforcement Learning.\n\n---\n\n**Weakness 3**: *There is too little discussion on the relationship to RL in Sections 3.2 and 3.3.*\n\n**Response**: Thanks for your comments. We revise Sections 3.2 and 3.3 to enhance the discussion regarding their relation to RL. Specifically, Section 3.2 provides a concise introduction to the two types of Mamba, allowing for the utilization of both types of DeMa in RL. To expand our analysis, we introduce hidden attention in Section 3.3. This enables us to visualize the hidden attention matrices within DeMa, thereby gaining a deeper understanding of the model's internal behaviors.\n\n---\n**Question 2**: *Is DeMa applicable to all trajectory optimization methods?*\n\n**Response**: Thanks for your comments. DeMa is applicable to most transformer-based trajectory optimization methods, as it is designed to address the challenges posed by transformers, particularly the issues of large parameter sizes and limited scalability. By integrating DeMa with QT [2], we develop Q-DeMa, which achieves performance comparable to state-of-the-art (SOTA) models while utilizing only one-seventh of the parameter size of QT. This finding underscores the significant potential of applying Mamba to RL and emphasizes the critical importance of the research presented in this paper.\n\n---\nWe hope that our answers can address your concerns satisfactorily and improve the clarity of our contribution. We would be grateful if you could re-evaluate our paper. We look forward to receiving your further feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OLd6YCKqkx", "reviewer": "Reviewer_qC8w", "summary": "This paper investigates how Mamba perform in trajectory optimization in offline RL with ablation analysis on mamba's data input structures and architectural structures and shows Mamba DT can achieve SOTA performance with less parameters.", "questions": "please see weakness.", "limitations": "please see weakness", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper writing is good, the visualizations look good.\n2. The input concatenation experiments provides useful practical insight also for other sequence-based decision-making models \n3. The paper provides a detailed analysis of how various components of Mamba, such as the hidden attention mechanism and different residual structures, influence performance.", "weaknesses": "1. Finding 3 is not very surprising on the tested MDP environment, since they by definition should focus only on recent states. It will be interesting to explore how this mechanism might perform in environments with long-term dependencies where the Markov property does not hold strictly.\n2. Only tested on standard Atari and mujoco tasks. How would mamba perform on tasks that requires long horizon planning skills? such as maze navigation or tasks with delayed rewards.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your insightful comments, which have catalyzed numerous enhancements and refinements to the paper. In the following, we reply to the questions one by one for the convenience of checking.\n\n---\n\n**Weakness 1**: *Finding 3 is not very surprising on the tested MDP environment. It will be interesting to explore how this mechanism might perform in environments with long-term dependencies where the Markov property does not hold strictly.*\n\n**Response**: Thanks for your comments. We conduct additional explorations in environments involving maze navigation(maze2d, antmaze) and delayed rewards(MuJoCo with delayed rewards). The performance of the hidden attention mechanism is illustrated in Figures B-C in the one-page .pdf file, which shows that 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. \n\nThere remains a significant gap in the exploration of the analysis of Markovian property and attention mechanism. Most studies focus on the attention scores of outputs $y_j$ relative to inputs $x_i$ in the training phase. However, it does not effectively evaluate the context attention of the model during decision-making. Our analysis reveals that the hidden attention mechanism predominantly emphasizes the current token at each decision-making step, even when the Markov property of the environment is relaxed. Thus, our finding (Finding 3) offers valuable insights and guidance for the community.\n\n---\n**Weakness 2**: *How would mamba perform on tasks that requires long horizon planning skills? such as maze navigation or tasks with delayed rewards.*\n\n**Response**: Thanks for your comment. We conduct three new experiments involving maze navigation(maze2d, antmaze) and delayed rewards(MuJoCo with delayed rewards). \n\n1. 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.\n\n2. antmaze: This environment is similar to maze2d, while the agent becomes an ant with 8 degrees of freedom.\n\nIn these environments, we adjust the hyper-parameters, which will be provided in the Appendix of our revised manuscript. Results in Table B show that DeMa performs better compared to DT in the maze navigation task.\n\nTable B: Results for maze2d and antmaze. We report the mean across 3 seeds.\n|Dataset|Env|DT|GDT|DC|DeMa(Ours)|\n|-|-|-|-|-|-|\n|umaze|maze2d|31|50.4|36.3|54.5|\n|medium|maze2d|8.2|7.8|2.1|16.7|\n|large|maze2d|2.3|0.7|0.9|6.6|\n|umaze|antmaze|59.2|76|85|96.7|\n|umaze-diverse|antmaze|53|69|78.5|96.7|\n\n3. MuJoCo with delayed rewards: This is a delayed return version of the D4RL benchmarks where the agent does not receive any rewards along the trajectory, and instead receives the cumulative reward of the trajectory in the final timestep.\n    \nIn this environment, we train DeMa using the same hyper-parameters settings, and the results are shown in Table C. Results show that CQL is the most affected, while DT also experiences a certain degree of influence. In contrast, DeMa and GDT are relatively less impacted. The results indicate that DeMa demonstrates effective performance in tasks with delayed rewards.\n\nTable C: Results for D4RL datasets with delayed (sparse) reward. The \"Origin Average\" in the table represents the normalized scores of evaluations across six datasets under the original dense reward setting. We report the mean across 3 seeds. The dataset names are abbreviated as follows: \"medium\" as \"M\", \"medium-replay\" as \"M-R\".\n|Dataset|Env|CQL|DT|GDT|DeMa(Ours)|\n|-|-|-|-|-|-|\n|M-Delayed|HalfCheetah|1.0|42.2|43|42.9|\n|M-Delayed|Hopper|23.3|57.3|58.2|69.1|\n|M-Delayed|Walker|0.0|69.9|78.9|77.6|\n|M-R-Delayed|HalfCheetah|7.8|33.0|41|41.1|\n|M-R-Delayed|Hopper|7.7|50.8|79.8|83.8|\n|M-R-Delayed|Walker|3.2|51.6|70.4|71.7|\n|Average|-|7.2|50.8|61.9|64.4|\n|Origin Average|-|65.5|63.4|63.8|66|\n\nOverall, DeMa achieves better performance than DT with fewer parameters in tasks that require long-horizon planning skills.\n\n---\n\nWe hope that the above answers can address your concerns satisfactorily. We would be grateful if you could re-evaluate our work based on the above responses. We look forward to receiving your further feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m4uv2OLFQ0", "reviewer": "Reviewer_85kG", "summary": "This paper comprehensively investigates the possibility of leveraging Mamba for trajectory learning. The authors take Decision Mamba as a playground and analyse the performance of this model over trajectory learning scenarios (gym/mujoco) from several aspects. A group of conclusions are attained through rigorous experiments, which is solid and potentially valuable for further researches realted with Mamba.", "questions": "see Weaknesses.", "limitations": "The authors provide a brief limitation summary in Conclusion. I would appreciate the authors if they could refine this part since the \"limitations\" listed there do not seem like limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Novelty: given that Mamba is still at its exploratory stage, this paper positively probes Mamba's potential for tranjectory learning, with surprising results that with some specific pre-conditions, Mamba is more suited than Transformer.", "weaknesses": "1. Most discoveries in this paper have been implicitly discussed for several months within the community, while it is firstly presented officially in this paper. Besides, these dicoveries lean to be emparical  evidence, which is relatively shallow. This would make this paper's technical contribution weak. I would appreciate the authors if they could provide more in-depth explanation over these discoveries, in particular: 1) Transfomer-like model favors short sequence. 2) The significant role of the hidden attention. \n\n2. Although the experimental results are solid, I found that this paper is more suitable for Benchmark Track, since the technical novelty revolves around benchmarking Decision Mamba. \n\n3. Figure 1 (the title and the pic) should be improved. For now, it confuses me, especially the corresponding relationship between the text content (title) and the illustration.\n\n3. Minor: line 295: may more suitable -> may be more suitable", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your insightful comments, which have catalyzed numerous enhancements and refinements to the paper. In the following, we reply to the questions one by one for the convenience of checking.\n\n---\n**Weakness 1**: *Most discoveries in this paper have been implicitly discussed for several months within the community, while it is firstly presented officially in this paper. Besides, these dicoveries lean to be emparical evidence, which is relatively shallow. This would make this paper's technical contribution weak. I would appreciate the authors if they could provide more in-depth explanation over these discoveries, in particular: 1. Transfomer-like model favors short sequence. 2. The significant role of the hidden attention.*\n\n**Response**: Thanks for your comment. Past research on trajectory optimization usually approaches RL problems as sequence modeling problems in various ways. However, these studies are often constrained by issues related to large parameter counts and poor scalability, and there has been relatively little investigation into attention mechanisms. This study on DeMa contributes to the development of scalable and efficient decision models. Moreover, our empirical study lays a foundation for the further application of Mamba in RL. \n\nMore explanations of our discoveries are as follows:\n\n1. Most environments in trajectory optimization are Markovian. In these environments, a model inputs the current state and returns current action can also work(e.g. CQL), which means that excessively long historical information provides limited assistance for current decision-making. As shown in Figures B-C in the one-page PDF file, this result still holds in environments where the Markov property is relaxed. The results from the maze navigation task demonstrate that while DeMa's attention to past information increases, the hidden attention mechanism still primarily focuses on current information. This explains why the Transformer-like model favors short sequences.\n\n2. The significant role of hidden attention lies in its ability to aggregate contextual information efficiently. Hidden attention is the third view of Mamba, which shows that such models can be viewed as attention-driven models[1]. It allows the model to selectively propagate or forget information along the sequence length dimension depending on the current token. This nature significantly reduces the number of parameters and aligns with the Markov property of the RL environment. As a result, DeMa can achieve better performance than DT with fewer parameters.\n\n[1] The hidden attention of mamba models.\n\n---\n\n**Weakness 2**: *Although the experimental results are solid, I found that this paper is more suitable for Benchmark Track, since the technical novelty revolves around benchmarking Decision Mamba.*\n\n**Response**: Thanks for your comment. DeMa proposed in this paper can significantly address the challenges of large parameter counts and limited scalability in transformer-based trajectory optimization methods while maintaining good performance. Furthermore, our analysis of DeMa and the hidden attention mechanism provides valuable practical insights that apply to other sequence-based decision-making models. \n\n---\n**Weakness 3**: *Figure 1 (the title and the pic) should be improved. For now, it confuses me, especially the corresponding relationship between the text content (title) and the illustration.*\n\n**Response**: Thanks for your comment, We have redrawn Figure 1 to improve clarity and readability. We provide it as Figure A in the one-page .pdf file. Overall, we have added numbering to each subplot and revised the descriptions in the title. We have arranged them side by side according to two structures: architecture and block.\n\n---\n**Weakness 4**: *Minor: line 295: may more suitable -> may be more suitable*\n\n**Response**: Thanks for your comment, we have addressed the aforementioned errors. Additionally, we conducted a thorough review of the entire manuscript and corrected similar typos throughout.\n\n---\n**Limitations**: *The authors provide a brief limitation summary in Conclusion. I would appreciate the authors if they could refine this part since the \"limitations\" listed there do not seem like limitations.* \n\n**Response**: Thanks for your comment, we have refined the limitations at the end of our revised manuscript.\n\nLimitations: 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. The potential of RNN-like DeMa warrants further exploration, particularly in POMDP environments and tasks that require long-term decision-making and memory. 2. Although we examine the importance of the hidden attention mechanism in Section 4.2, our exploration is still in its nascent stages. 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.\n\n---\nWe hope that the above answers can address your concerns satisfactorily. We would be grateful if you could re-evaluate our work based on the above responses. We look forward to receiving your further feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
**We want to thank all the reviewers for their thoughtful suggestions on our submission**, and we appreciate that the reviewers have multiple positive opinions of our work, including:

* novelty (BojN, 85kG)
* good writing, good visualizations (qC8w)
* the detailed analysis provides useful practical insight (BojN, qC8w)

We provide a summary of our responses, and we will carefully revise our manuscript by adding suggested evaluations, providing more detailed explanations, and fixing the typos.

**Introduction, Related Works, and Preliminaries**:

* We strengthen the relationship between DeMa, trajectory optimization, and RL. (for Reviewer BojN)
* We refine the notations for improved clarity. (for Reviewer BojN)
* We enhance the discussion regarding their relation to RL. (for Reviewer BojN)

**Experiments**:

* We redraw Figure 1 to improve clarity and readability. (for Reviewer 85kG)
* We explore the applicability of combining DeMa with another method. (for Reviewer BojN)
* We analyze the hidden attention in tasks that require long horizon skills and tasks with delayed rewards. (for Reviewer qC8w)
* We explore the potential of DeMa in tasks that require long horizon skills and tasks with delayed rewards. (for Reviewer qC8w)
* We provide more explanations for our discoveries. (for Reviewer 85kG)

**Conclusion**：

* We provide more details about our limitations. (for Reviewer 85kG)

**We appreciate all the reviewers' time and effort again**. All these comments and suggestions are very insightful and beneficial for us to improve the quality of this work.

Meta-review: This paper presents a comprehensive investigation of leveraging Mamba architecture for trajectory optimization. The proposed DeMa achieves comparable or better performance as compared with the DT class of method, but with much fewer parameters. A group of findings is attained through comprehensive experiments, which can potentially be valuable for future research based on the Mamba architecture.

**Strengths:**

- The paper provides extensive evaluations and ablation analyses to show how various components of Mamba, such as the hidden attention mechanism, way of concatenation, and different residual structures, influence performance. These could provide useful practical insights for other sequence-based decision-making models.

**Weaknesses:**

- I share similar feelings with Reviewer 85kG that some of the findings are somewhat shallow, it would be better if the authors could provide more in-depth explanations. Nevertheless, the compressive experiments and investigation are meaningful and can potentially benefit the community.
- The paper primarily conducted evaluations on the D4RL datasets. Note that most of the D4RL datasets are generated using RL polices, which by nature exhibit strong Markovian properties. Considering D4RL datasets are somewhat biased if compared with most real-world offline decision-making datasets, solely conducting experiments on D4RL might lead to some misleading insights. The study can benefit from exploring its effectiveness on datasets that are more non-Markovian, such as some human demonstration data in the robotic manipulation domain. This is also partly mentioned by Reviewer 85kG and qC9w.

The paper received unanimously positive feedback after the rebuttal. Although the paper still has some issues, I think it is a timely investigation of Mamba architecture for the community.

 I suggest the authors include the additional Antmaze results in the main paper, as well as conduct additional experiments on some non-Markovian human demonstration data for robotics tasks in their final version. It should be noted that D4RL actually is not a very good benchmark to evaluate the real-world practical performance of offline RL methods, which is heavily biased (strongly Markovian, unrealistic state-action space coverage, unreasonably large amount of data for simple tasks). I suspect in some long-horizon non-Markovian data/tasks, some of the findings might have a different conclusion, we should be careful to not release misleading insights to the community.
Decision: Accept (poster)