Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling
Authors: Sili Huang
Jifeng Hu
Zhejian Yang
Liwei Yang
Tao Luo
Hechang Chen
Lichao Sun
Bo Yang
Venue: NeurIPS 2024 poster (2024)
Keywords: In-context Reinforcement Learning
Mamba
Primary area: reinforcement_learning

Abstract: Recent works have shown the remarkable superiority of transformer models in reinforcement learning (RL), where the decision-making problem is formulated as sequential generation. Transformer-based agents could emerge with self-improvement in online environments by providing task contexts, such as multiple trajectories, called in-context RL. However, due to the quadratic computation complexity of attention in transformers, current in-context RL methods suffer from huge computational costs as the task horizon increases. In contrast, the Mamba model is renowned for its efficient ability to process long-term dependencies, which provides an opportunity for in-context RL to solve tasks that require long-term memory. To this end, we first implement Decision Mamba (DM) by replacing the backbone of Decision Transformer (DT). Then, we propose a Decision Mamba-Hybrid (DM-H) with the merits of transformers and Mamba in high-quality prediction and long-term memory. Specifically, DM-H first generates high-value sub-goals from long-term memory through the Mamba model. Then, we use sub-goals to prompt the transformer, establishing high-quality predictions. Experimental results demonstrate that DM-H achieves state-of-the-art in long and short-term tasks, such as D4RL, Grid World, and Tmaze benchmarks. Regarding efficiency, the online testing of DM-H in the long-term task is 28$\times$ times faster than the transformer-based baselines.

Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling
Sili Huang"
Jifeng Hu?
Zhejian Yang*
Liwei Yang" Bo Yang
Tao Luo
Hechang Chen 6+
Lichao Sun
1,8 Key Laboratory of Symbolic Computation and Knowledge Engineering of Ministry of Education 5,2,3, of Artificial Intelligence, Jilin University, China Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore Lehigg University, Bethlehem, Pennsylvania, USA
Abstract
Recent works have shown the remarkable superiority of transformer models in reinforcement learning (RL), where the decision-making problem is formu- lated as sequential generation. Transformer-basse agents could emerge with self- improvement in online environments by providing task contexts, such as multiple trajectories, called in-context RL. However, due to the quadratic computation complexity of attention in transformers, current in-context RL methods suffer from huge computational costs as the task horizon increases. In contrast, the Mamba model is renowned for its efficient ability to process long-term dependen- cies, which provides an opportunity for in-context RL to solve tasks that require long-term memory. To this end, first implement Decision Mamba (DM) by replacing the backbone of Decision Transformer (DT). Then, we propose a De- cision Mamba-Hybrid (DM-H) with the merits of transformers and Mamba in high-quality prediction and long-term memory. Specifically, DM-H first generates high-value sub-goals from long-term memory through the Mamba model. Then, we use sub-goals to prompt the transformer, establishing high-quality predictions. Experimental results demonstrate that DM-H achieves state-of-the-re -o-tthe-rrtt long and short-term tasks, such as D4RL, Grid World, and Tmaze benchmarks. Regarding efficiency, the online testing DM-H the long-tew tag-term task is 28x times faster than the transformer-bast baselines.
1 Introduction
Large transformer models 43 have achieved notable successes across a variety of domains, including text [4], image [9]. and audio 1 In the field of reinforcement learning (RL), large transformer models can treat RL tasks as a type of sequential prediction problem and have shown impressive results with offline training [31 (3] However, these methods lack seff-improvement when used in online environments, where online environments could differ from offline training. To overcome this, in-context RL has been proposed, which enables continued policy improvement [29].
Recent works demonstrated that in-context RL methods can automatically improve online perfor- mance by providing prompt conditions called across-episodis contexts (30]. The construction of the across-episodic context is flexible and easy implement, such as multiple historical trajecto- ries arranged in ascending order of returns [20]. Although no gradient updates are required for self-improvemmnt,, current in-context RL still suffers from high computational costs on long-term tasks [29] This arises from (1) the quadratic complexity of the self-attention mechanism and (2) the multiplicative growth of the long-term sequence caused by across-episodic contexts.
"Corresponding Author. Bo Yang and Hechang Chen.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Facing the challenge of handling long-term sequences. a novel foundation model called Mamba has attracted widespread attention for its ability capture long-term dependencies with linear computational costs 17 6 Mamba is a state space model-based framework with good potential in natural language processing and vision tasks [47]. Motivated by the success of Mamba in language and vision modeling, it is appealing that we can also transfer this success to RI. tasks. Therefore, natural question arises:
"Can Mamba boost in-context RL with both effectiveness and efficiency on long-term tasks?"
Taking the famous in-context RL method Decision Transformer (DT) [29 as an example, we replace its transformer backbone with a Mamba backbone. Regarding efficiency, there is no doubt that Mamba is superior to transformers as the task length increases because Mamba uses a data-dependent selection mechanism that computes independently on each input of sequences [7]. Compared with the attention mechanism that computes input pairs, the data-dependent selection mechanism makes it easier for Mamba to handle long sequences but also introduces additional independence assumptions. In RL tasks, since there is a sequential relationship between states rather than independence, the attention mechanism may be intuitively more suitable for capturing the information between states, thereby outperforming Mamba in effectiveness.
To this end, instead of implement Decision Mamba (DM) that simply replaces the backbone of DT, we propose the Decision Mamba-Hybrid (DM-H) with the merits of transformers and Mamba in high-quality prediction and long-term memory. Specifically, the Mamba model first generates sub-goals, represented by vector, based on long-term contexts. Then, we combine the sub-goal and short-term contexts as a prompt condition to the transformer. The setting of sub- sub-goals enables DM-H to leverage Mamba's 's ability to efficiently recall long-term contexts while using the transformer to predict high-quality actions. However, remains to be seen whether the transformer will benefit from sub-goals or ignore them and predict actions focused on short-term contexts. Therefore, we select high-value states from offline trajectories to form extra sub-goals and serve them as input conditions for the transformer. Then, the predicted actions will associate with these selected sub-goals and, in turn, encourage Mamba to generate high-value sub-goals. Our contributions are as follows:
We investigate Mamba model compared to the transformer model in traditional RL tasks, D4RL, and find that Mamba model is more efficient but slightly inferior to the transformer model in terms of effectiveness.
We propose DM-H, an in-context RI. method that connects Mamba and the transformer with high-value sub-goals. DM-H inherits the merits of Mamba and transformers, achieving both high effectiveness and efficiency in long-tern tasks.
Our extensive experiments across Grid World, D4RL, and Tmaze reveal the superiority of DM-H over other baselines. In the online testing of long-term tasks, DM-H can be 28x times faster than baselines and more than double the effectiveness.
2 Related Work
Mamba for Long Sequence Modeling, The structured State-Space Sequence (S4) model is a novel alternative to CNNs or transformers model the long-term dependency [8] The promising property of linearly scaling in sequence length attracts further exploration. Based on the $4, Smith et al. [40] propose new S5 layer by introducing MIMO SSM and an efficient parallel scan into the S4 layer. Fu et al. [12 design a new SSM layer, H3, that nearly fills the performance gap between SSMs and transformers in language modeling. Mehta et al. (32 build the Gated State Space layer on $4 by introducing more gating units to improve the expressivity. Recently, Gu and Dao [17] propose data-dependent SSM layer and build generic language model backbone, Mamba, which outperforms transformers at various sizes large-scale real data and enjoys linear scaling in sequence length. Later, Vision Mamba adds position coding and bidirectional scanning to extend to visual tasks [47). In this work, explore transferring Mamba's 's success to RL, i.e., achieving high effectiveness and efficiency for long-term memory.
Transformer for Decision-Making damental online paradigm [4]] The nature of online learning comes with some limitations when meeting the applications for which it is impossible to gather online data and learn simultaneously. such as autonomous driving. To this end, offline RL proposed that the agent can learn from fixed
In general, reinforcement learning was proposed as a fun-
offline dataset without gathering new data during learning [15] 28 45 27
211
23
19 24]. In the

---

context of offline RL, recent works explored using transformer-basee policy by treating RL tasks as a type of sequential prediction problem [22]. Among them, decision transformer [S] was pro-
posed to model trajectories as sequences and autoregressively predict action conditioning on desired return-to-go, past states, and actions. Trajectory transformer [26) demonstrated that transformer could learn single-task policies from offline data. Subsequently, the multi-game decision transformer [31 and Gato [39] further showed that transformer-based policies could address multi-tasks in the same domain and cross-domain tasks. However, these works focused on distilling expert policies from offline data and failed to enable self-imprveement like DM-H. When the offline data are sub-optimal, or the agent is required to adapt to new tasks, the multi- game decision transformers need to finetune the model parameters, while Gato is required to get prompted with expert demonstrations.
Meta RL. DM-H falls into the category of methods of learning to learn, which is also known as meta-learning. More precisely, recent in-context RI. methods can categorized as in-context meta- RL methods. The general idea of learning seff-imprrvement has a long history in RL but is limited to nyper-parameters in the early stages (25). In-context meta-RL. methods 44 T0 are commonly trained in the online setting by maximizing multi-episodic value functions with memory- architectures through environment interactions. Another online meta-RL attempts to find good network parameter
initializations and then quickly adapt through additional gradient updates [1 336. More recently, meta-RL has seen substantial breakthroughs, from performance gains on popular benchmarks to offline settings, such as Bayesian RL [8] and optimization-base based meta-RL [33]. Considering the
difficulty of a completely offline setting. recent work has explored hybrid offline-online settings DM-H is similar to the hybrid offline-online setting but saves more computing resources
46 37.
because the online phase does not involve gradient updates.
In-Context RL. In-context RL is the one that addresses tasks by providing prompts or demonstrations
26] By training agents at a large scale, transformer-based policies usually have the ability to learn in context [31 [39] The learning process is performed entirely in context and does not involve parameter updates of neural networks. In this work, we consider incremental in-context RL, which involves learning from one's own behaviors in trial-and-error manner. Laskin et al. [29] proposed Algorithm Distillation (AD). which automatically improved its performance by providing multiple historical trajectories. Subsequently, Lee et al. [30] proposed a Decision-Pretrainee Transformer, which trained the agent to find optimal behaviors faster by only predicting the optimal trajectory. More recently, Hao Liu (20) further demonstrated that across-episodic contexts encourage large transformer models' emerging self-in eef-iippoveeent behaviors. However, these methods suffer from huge computational costs as across-episodic contexts induce too-long sequences. In contrast, DM-H leverages Mamba's 's ability to efficiently process long-term dependencies while using the transformer to establish high-quality predictions.
3 Preliminaries
Partially Observable Markov Decision Process. We consider learning problems in the context of Partially Observable Markov Decision Processes (POMDP) represented by tuple The POMDP tuple consists of states s E S, observations o E O, ac- tions a E A. rewards r E R. and transition probability function P(s (,, where is an integer denoting the timestep. At each timestep f, the agent receives the observation or, selects an action a, T(-(0t) based on its policy, and then receives the next observation 0++1- For con- venience, we uniformly use to denote the observations or states received from the environment. trajectory sequence that consists of observations, actions, and rewards and is denoted by (So, 4g. .o. aT, PT). In addition, a completion token d, binary identifier, is used to indicate whether a trajectory ends at time f.
M = AS
Transformers. The Transformer 43 architecture consists of multiple layers of self-attention operation and MLP. The self-attention begins by projecting input data X with three separate matrices onto D-dimensional vectors called queries Q, keys K, and values V. These vectors are then passed through the attention function:
Attention(Q K, V) = softmax(QK") IVDV
The QKT term computes an inner product between two projections of the input data X. The inner product then normalized and projected back to a D-dimensional vector with the scaling term V. Transformers utilize self-attention core part of the architecture to process sequential data 77.77
3

---

Table I: Mamba Transformer D4RL datasets.
Davironment
HallCheetah
Hoopper
Walker2d
Exper.
Medium
Replay
Expert
Medium
Replay
Med Expert
Replay
Effectiveness
Mamba
Transforer
94
92
41
42.28: 14
41.28
110.8: ude
73.65 123
111.6 91.
108.31
78. 78.26 cus
70.92 121
Eificiecy
Mamba
Transfommrr
38.77a
In this work, we use GPT 38 architecture that modifies the transformer with a causal self-attention mask to focus on the previous tokens in the sequence (j E [1, enabling us to do autoregressive generation at test time.
S4 and Mamba. S4 18] and Mamba 17 are inspired by the continuous system, which maps a
1-D function sequence xtt) IR y(t) E R through hidden state h(t) E R.V. The mapping
process can be represented as the following linear ordinary differential equation:
N((1) AA((() + Bx((), (t) -Ch(t)
where A E RNN denotes the evolution parameter, B E R and C E I denote the projection parameters. For application to a discrete input sequence instead of a continuous function, S4 uses zero-order hold to transform the continuous parameters A, B discrete parameters A, B. Then, the Equation 2 can be rewritten as:
h4 = Ah + Brg
yt CCh,,
where A = exp(A)A, (8A)- (AB), and A is timescale parameter. Based on the S4 framework, Mamba introduces a data-dependent selection mechanism while leveraging a hardware-awee aware parallel algorithm in recurrent mode. Compared with the Transformer, the combined architecture of Mamba empowers it to capture contexts effectively and maintains computational efficiency, particularly for long sequences.
4 Method
In this section, we first compare Mamba and transformer models in the D4RL dataset, and inves- tigate the potential of Mamba RI. tasks. Then, We present DM-H, which can handle long-term dependencies from contexts with high effectiveness and efficiency, as shown in Figure
4.1 Mamba vs. Transformer in RL tasks
We first consider the Algorithm Distillation (AD) as the baseline, which a classic in-context method using transformer as the backbone 229. AD can predict high-quality actions by recalling the historical trajectories from the context, the same settings, we replace the transformer in the AD algorithm with Mamba to compare their effectiveness and efficiency.
also incurs higher computational costs. Under
As shown in Table I the simple backbone replacement did not significantly improve effectiveness. Regarding efficiency, Mamba brings predictable improvements, thus saving training time under the same settings. Compared with the attention mechanism acting on state pairs, Mamba uses data-dey dete-deeendent selection mechanism acting on each state independently, which brings a more efficient method of recalling long-term memory. However, since states in RL tasks commonly exhibit sequential relationships,
attention mechanism more suitable for capturing the information between states, thereby outperforming Mamba in terms of effectiveness. To this end, we aim at new in-context RL approach that leverages Mamba's 's strengths processing long-term memory while preserving high-quality predictions from transformers.
4.2 Decision Mamba-Hybrid
In-context RI. can automatically improve its performance through trial-and-error when across- episodic contexts serve as prompt conditions. Specifically, across-episodic context consisting of

---

Qt-c-1
Generating actions by transformer
linear layer
valuable sub-gon sub-goals
Sz
at
Zt
s
attc-1
rete-1
offline training
Generating sub-goals from long-term experience by mamba
offline training and online testing
St
d
St
de
SE+t
(+t
d
previous episodes
current episode
The architecture of DM-H. During offline training, Mamba module generates sub-goals from long-term experience, where the long-term experience consists multiple historical trajectories arranged in ascending order of the total rewards. Based the generated sub-goals, the transformer is required to predict better actions by supervising the expert behaviors. Meanwhile, the linear layer feeds the valuable sub-goals into the transformer module and associates them with the generated actions. During online testing, DM-H can automatically improve its performance in a trial-and-error manner without requiring gradient updates.
Figure 1:
trajectories represented
(T',72, ,...,,"" r"), where
+¹ (sh, ag. ri dy..... ... s/, ar,, rT,, dr-).
(4) 1
The trajectories are sorted according their total rewards, i.e.. i.e. Et - With autoregressive training and generation, the transformer can uncover meaningful patterns from multiple trajectories and improve itself conditioned on experience. However, the quadratic complexity of the attention mechanism suffers from huge computational costs along with the growth in task horizon. Inspired by Mamba's success with long sequences, we propose that Mamba handles across-episodis contexts and preserves local short-term sequences for the transformer. For Mamba, we reconstruct the long-term sequences (r),, Tm. Tm) from across-episodic contexts. Each m is denoted as
m = (s Pú. dy, s:. Pi d kke Ple. dke),
(5)
where represents the local sequence length for the transformer, T S ke S T, and ft = ri the sum of steps rewards. Mamba module will generate sub-goals to prompt the transformer, where the sub-goal is represented by vector z sampled from multivariate Gaussian distribution. Then, local short-term sequence is represented as:
(z),
#+1 (411
......,ke where T2 starts from the generation step j E {0, ...., ke} and completes steps actions based on the generated sub-goal z).
4.3 Decision Hamba--Hbbrid with Valuable Sub-goals
DM-H links Mamba and the transformer through the sub-goal z to efficiently recall long-ten contexts while ensuring high-quality predictions. As sub-goal are commonly hidden in the offline data, Mamba must infer them along with the transformer training. However, remains be seen whether the transformer will benefit from z or ignore it and only imitate expert behaviors based on the local context. Therefore, we select extra high-value states from the offline data and transform them into sub-goals to align actions generated by the transformer.
5

---

Valuable sub-goals. Intuitively, a sub-goal should be highly valuable for the agent to reach and have a high probability of uppearing at subsequent time steps of the current state in the trajectory. When reached sequentially, these states should mark milestones in the trajectory, making it highly probable that the agent successfully performs the task. this point, the valuable sub-goals guide the agent through the task, meaning they have the same purpose the returns-to-go in the Decision Transformer [5]. Therefore, we can model this behavior by finding states with high accumulated reward values in the trajectory. Specifically, for state 8; at timestep i, the value of state sj at timestep
j is Ei=++1 Tkk where: 1 S S T. However, this may prioritize the last state of trajectories when the environment only provides positive rewards. To encourage selecting states that are close to the current state 8i, we divide the accumulated rewards by the distance between their timesteps Ei= rw/(j - i). With the punishment of the distance, the weighted average of accumulated rewards can identify the short-term and important future states.
Based on the selected sub-goals, we reconstruct the local short-term sequence (Equation (6) by replacing z generated from Mamba. The reconstructed local short-term sequence is represented as:
T (f(sg) s),a rz, f(s)), s'+1 a' 11+1: "":1:" f(s')) Sj (e 2)t--1
(7)
where s is the most valuable sub- for state sÿ and f is a linear layer that maps sy to the same dimension of zi. The reconstructed local short-term sequence aligns the actions generated by the transformer with subsequent high-valued states. Since Equation (6) is consistent with the reconstructed sequence except for zj. this encourages Mamba module generate high-value sub- goals from the long-term context to ensure that the transformer module predicts similar actions.
4.4 Implementation of DM-H
Architecture. We feed trajectories into Mamba module, which results in 3 x c/c tokens, with one token for each of the three modalities: state, reward, and completion. In the transformer module, we feed 4 x tokens, with one token for each of the four modalities: sub-goal, state, action, and reward. To create the token embeddings, we train a linear layer for each modality, which transforms the raw inputs into the desired embedding dimension, followed by layer normalization (2). Finally, we freeze a linear layer that maps the high-value sub-goals& $y in Equation (7) to the same dimensions as the sub-goals generated by Mamba.
Offiine Training and Online Testing. During offline training, we are given a dataset of offline trajectories, where the trajectories can be suboptimal. In each iteration, we sample minibatches of trajectories from the dataset. Then, Mamba module first predicts the sub-goals z, every e steps, given the input token and past trajectories. Then, the transformer module autoregressively predicts steps of actions {at..... at+c-1} } given Z, and {st,.............. .., St+c-1]. Meanwhile, we use the weighted average of accumulated rewards to select valuable sub-goals from the offline data and feed them to the transformer model to predict the same c steps actions. The predicted actions are evaluated with discrete or
either cross-entropy loss mean-squared error, depending continuous. The losses from each step are averaged and updated in all modules end-to and-to-end. At online testing, we roll out the DM-H with multiple trajectories and report the return of the last trajectory. Following the configuration from related works 20 29
whether the actions
we set a context size across n = 4 episodes. The pseudocode for DM-H is summarized in Appendix A. Source code and more hyperparameters
are described Appendix B
Experiments
In this section, we will introduce datasets and baselines in Section 5.1 Then, in Section 5.2 Section 5.3] Section 5.4 [53 we report the comparison results, ablation study, and parameters sensitivity analysis. In Appendix C we report additional results about offline training time, online testing time, and ablation study.
5.1 Environmental Settings
Dataset: Grid World. We first consider the discrete control environments from the Grid World 31. which a commonly used benchmark for recent in-context RL methods. The environments support many tasks that cannot be solved through zero-shot generalization after pre-training because

---

large Darkroom
Large Darkroom Hard
Large Dark Key-to-Don
15
1.0
1.6
Return
10
Return
Retur
DM-H (ours)
0.5
0.8
Long- terg-teemm horizon
AD (mamba)
episodes
20
0
episodes
20
0
episodes
20
AD
Darkroom
Darkroom Hard
Key-to-Dor to-ooor
(transformer
15
1.0
2.0
AMAGO
Retur
10
Return
Return
0.5
1.0
DPT
episodes
20
episodes
20
0
episodes
20
sparse reward
Figure 2: Results for Grid World. An agent is expected solve a new task by interacting with the environments for 20 episodes without online model updates. Our DM-H significantly outperforms baselines on long-term tasks with sparse rewards because it inherits the merits of transformers and Mamba in high-quality prediction and long-term memory.
these tasks cannot be inferred easily from the observation. Specifically, test method three environments: Darkroom, Darkroom Hard, and Dark Key-to-Door. In addition, we create long-term variant of Large Darkroom, Large Darkroom Hard, and Large Darkroom Key-to-Door, where the coordinate space of each environment is expanded to 20 times, and the episode length expanded 10 times. The dataset is collected from learning histories that generated by training gradient- based RL algorithms, such as Deep Q-Network [34]. For each environment, we randomly create 60 tasks from the coordinate space and collect data for I million steps.
Tmaze We also evaluate our method on the Tmaze 35 a benchmark for testing the recall ability of in-context agents. In Tmaze, any policy that achieves the maximum return must be able to recall information from the first step at the final step. Since the task horizon can be set arbitrarily, it is often used to test the limits of the model's processing of long-term memory.
Dataset:
commonly used offline RL benchmark, including continuous control tasks. The dataset is collected from Mujoco environments, including HalfCheetah, Hopper, and Walker. The episode length in D4RL 1000, which is far more than that of Grid World. Therefore, current in-context RI methods require huge computational costs in D4RL, even though it a commonly used benchmark for conventional RI. algorithms.
Dataset: D4RL. D4RL 13
Baselines. We investigate the effectiveness and efficiency of DM-H relative to in-context RI,, dedicated offline RL, and imitation learning algorithms. Our baselines can be categorized follows:
In-context RI: These methods use the transformer to model trajectory sequences and predict actions autoregressively. We compare with recent methods, AMAGO (16), Decision Pretrained Transformer (DPT) [30], and Algorithm Distillation (AD) 229 which achieve impressive results based on the setting of across-episodic contexts.
learning: Most -ifforel-difference (TD) learning methods use an action space constraint or value pessimism and will serve as faithful comparisons DM-H, representing standard RL methods. Following recent work (20), we consider state-of-the-ie TD3+BC [4] that is demonstrated be effective on D4RL.
Imitation learning: Imitation learning methods similarly utilize supervised losses for training, such Behavior Cloning (BC) 42 and Decision Transformer (DT) [5]. We compare with BC-10%
which is shown to be competitive with state-of-the-an on D4RL. DT also uses a transformer to predict actions autoregressively but limited single episode context.
For all comparison methods, we adhere closely to the original hyper-parameter settings. To evaluate DM-H and other in-context RI. algorithms, we roll out episodes D4RL. 20 episodes Grid World and Tmaze. For each result, report mean and standard error across 10 random seeds.
7

---

DM-H (Ours)
DM
DT
X
X out of memory
80
1.0
1.0.
1,0
1,0
1.0
1.0
1.0
1.0
0.96
0.97
DM-H (Ours)
0.96
minute
Return
40
0.95
DM
0.92
0.91
0.9
Memory length (a) Performance
5
10
20
XX
0-
Memory length (b) learning time
20
Figure 3: Results for (a) performance and (b) online testing times on Tmaze tasks. We train each method address Tmaze tasks that have different horizons until we run out of GPU memory at context length to achieve 10k (DT, DM) or 20k (our DM-H). We report the online testing time for 20 episodes of Tmaze tasks.
Table 2: Results for D4RL datasets. DM-H H outperforms both in-context RL (DT, AD) and supervised learning (BC) and performs competitively with conventional RL algorithms (TD3+BC and TD3) on almost all tasks.
Dataset
Environment
BC-10%
TD3+BC
T133
DT
AD (Mamba)
AD (Transtormer)
DM-H
Med-Expert
HaltChestah
94.11
96.59
87.60
93.40
=20
94.21
96.21
Mcd-Expert
Hopper
113.13
113.22
98.41
111.18
110.82= 0.5
108.32 #.0
117.19- 645
Med-Expert
Walker2d
109.90
112.21
100.52
108.71
108.3 108.31=om
111.36
118.21
Med
HallCheetah
43.90
48.93
34.60
42.73
41.92= 0.11
42.28 + 1.8
45.45= 0as
Med
Hopper
73.84
70.44
56.98
69.42
73.65=
72.58 #:05
83.15= 663
Med
Walker2d
82.05
86,91
70.95
74.70
78.26
85,96 4:046
88.29 6.76
Med Med-Replay
HaltChectah
42.27
45.84
38.81
40.31
39.68_
41,28 1:021
45.26 8.43
Replay
Hopper
90.57
78.90
82.65
91.32 1.06
98.36 6.31
Med-Replay
Walker2d
76.09
91.17
65.94
71.22
70.92 1:21
89.21 4:1
95.66 1.36
Total Average
80.65
84.83s
70.28
77.69
76.97=
82.84
87.53 0.59
5.2 Grid World Results
To evaluate DM-H's celf-improvement capabilities in unseen tasks, we compared recent in-context RL methods in the Grid World environments. The agent is required to solve an unseen task by interacting with the environments for 20 episodes without online model updates. In addition, added a variant of the representative in-context RL AD, replacing the transformer with Mamba as the backbone.
As shown in Figure 2 DM-H achieves of-te-of-the-art performance in wide range of tasks. DM-H achieves 28x times faster than baselines and more than double the effectiveness as the task horizon increases and rewards become sparse. In long-term tasks, DM-H leverages Mamba's ability to obtain long-term memories while retaining high-quality predictions established by the transformer in decision-making. On the contrary, AD (transformer) and AD (Mamba) only retain one aspect of the advantages of DM-H. In terms of sparse rewards, DM-H is similar to the structure of hicrarchical RL, in which Mamba performs a decision every e steps and is, therefore, better at handling reward sparse scenarios. Overall, DM-H demonstrated that the hybrid method is not only feasible but combines more advantages synergetically.
5.3 Tmaze Results
In this section, we test the limits of DM-H's recall capabilities in the Tmaze task. Solving this task requires accurate recall the first step at the last step. Therefore, we set the context length equal to the task horizon and compare it with DT and DM, where DM replaces the DT's transformer backbone with a Mamba backbone. We train each method to recover the optimal policy until we run out of GPU memory
context size equal to 20k (DM-H) or 10k (DT and DM).
As shown in Figure 3 DM-H achieves the maximum reward with minimum online testing costs at any task horizon, demonstrating that Mamba's 's recalled memory positively prompts the transformer. Regarding efficiency, DM-H comparable to DM using only Mamba during online testing and outperforms DM during offline training (Figure 5 in Appendix C). This is because (1) the context of
8

---

valuable sub-goals
2.0
DM-H without valuable sub-goals
120
1.5
Large Dark Key to Door
Retu
Retur
100
Med Med-Expert
Retuurn
1.0
Hopper Med-Roplay
Large Darkrbom
80
Hopper
0.5
Hard
0.0
Dar barkroom
Dark Key-to-Door
Large Darkroom Large Dark
15
10
15
20
Hard
Hard Key to Door
Ablation study valuable sub-goals
(b) Parameter sensitivity analysis "c"
Figure 4: (a) The ablation study on DM-H with or without valuable sub-goals. (b) The parameter sensitivity analysis of "c."
the transformer in DM-H is fixed to the hyperparameter c, which does not change with the task length. (2) Mamba model generates a sub-goal every steps, thereby shortens the sequence it processes by ex times. Benefiting from these merits, DM-H handles tasks twice the length and has fewer memory resource requirements than DT and DM.
5.4 D4RL Results
In addition to navigation tasks, we also test DM-H the control tasks from the D4RL dataset, which commonly used in conventional offline RL methods. Based on previous work [13], the results on D4RL are normalized so that 100 denotes expert policy. Baseline numbers
reported by
the Agentic Transformer (20] and from the D4RL paper. As shown in Table DM-H outperforms baselines in majority of the tasks and is competitive with the state-of-the-an art the remaining tasks. In the TD learning and imitation learning categories, TD3+BC is generally the most remarkable algorithm. Compared with them, the superior performance demonstrates the self-improvement of
DM-H on suboptimal data.
5.5 Ablation Study and Parameter Sensitivity Analysis
Ablation Study on Valuable Sub-goals. To validate the effectiveness of valuable sub-goals conduct an ablation study of DM-H in Grid World tasks. Figure (a) presents the ablation results, which report the mean episode return across 10 seeds. We can observe that sub-goals can significantly improve DM-H's performance, proving that the sub-goal strengthens the transformer's dependency on long-term contexts from Mamba. In Appendix [ also carry out ablation studies in the D4RL dataset. Please refer to Figure
Sensitivity of Key Hyperparameters. In this experiment, we introduce important hyper- parameter c. A large e enables Mamba model to sean multiple episodes with smaller context sizes, significantly improving the computational efficiency. In addition, also controls the short-term context length in the transformer, which affects the quality of generated actions. As shown in Figure (b), DM-H performs well within the appropriate range of i... 10 20.
6 Conclusion, Limitations, and Broader Impacts
In this work, we propose an in-context RL method DM-H that achieves both high effectiveness and efficiency in long-term tasks. The core idea DM-H is to use sub-goal settings to leverage Mamba's ability efficiently recall long sequence contexts while using the transformer to perform high-quality predictions. Unlike current in-context RL methods limited to short-term tasks, DM-H is also good at standard RL benchmarks, which typically have long-term sequences under the across-episodic context setting. On the Grid World, D4RL, and Tmaze benchmarks, we demonstrate that DM-H can outperform baselines in both efficiency and effectiveness.
Regarding limitations, our method has an important hyperparameter, which is the context length of the transformer "c." As "c" increases, Mamba model can sean multiple episodes with smaller context sizes, significantly improving the computational efficiency. On contrary, as "c" decreases, Mamba can generate high-quality sub-goals from the context that is closer to the original episode. However, this reduces the short-term context available to the transformer when making predictions. Therefore, meaningful future direction is for "c" to adapt autonomously different tasks. In terms of the

---

potential broader impact, we do not anticipate any negative ethical and societal impacts of our work while using our method in practice.
Acknowledgments
This work was supported by the National Key R&D Program of China under Grant No. 20217D01 12500; the National Natural Science Foundation China under Grant Nos. U22A2098, U2341229, 62172185, 61976102, 62206105, 62202200 and 624761 10; the Key R&D Project of Jilin Province under Grant Nos. 2024021 20003XX and F0240304200SF; the International Cooperation Project of Jilin Province under Grant No. 20220402009GH.
References
[1] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35: 23716-23736, 2022.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization, arXiv preprint 1rXiv:16070064500 2016.
[3] 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.
[4] 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 few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[5] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lec, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34: 34:15084-15097 2021.
[6] Yang Dai, Oubo Ma, Longfei Zhang, Xingxing Liang, Shengchao Hu, Mengzhu Wang. Shouling Ji, Jincai Huang, and Li Shen. Is mamba compatible with trajectory optimization in offline reinforcement learning? arXiv preprint arXiv:2405. /2094, 2024.
[7] Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1800v:1810.04805, 2018.
[8] Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
[9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale, arXiv preprint arXiv:2010.1 11999, 2020.
[10] Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 161 2016.
(11) Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. International conference on machine learning, pages 1126-1133 PMLR, 2017.
[12] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, arXiv preprint arX1v:2212.1405, 2022.
10

---

[13] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[14] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. Advances in neural information processing systems, 01120132-20145, 2021.
[15] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019.
[16] Jake Grigsby, Linxi Fan, and Yuke Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. 2024.
[17] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. 2023.
arXiv preprint
[18] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.0030 2021.
[19] Siyuan Guo, Yanchao Sun, Jifeng Hu, Sili Huang, Hechang Chen, Haiyin Piao, Lichao Sun, and Yi Chang. simple unified uncertainty-guided framework for offline-to-online reinforcement learning. arXiv preprint arXiv:2306.00 07541, 2023.
[20] Pieter Abbeel Hao Liu. Emergent agentic transformer from chain of hindsight experience. In Proceedings of the 20th International Conference on Machine Learning (ICML 2023), 2023.
[21] Jifeng Hu. Li Shen, Sili Huang. Zhejian Yang, Hechang Chen, Lichao Sun, Yi Chang, and Dacheng Tao. Continual diffuser (cod): Mastering continual offline reinforcement learning with experience rehearsal, arXiv preprint XXiv:2409.025/2 2024.
[22] Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision trans- former: Reinforcement learning via hierarchical chain-of-thoughh In Forty-first International Conference on Machine Learning.
[23] Sili Huang, Hechang Chen, Haiyin Piao, Zhixiao Sun, Yi Chang, Lichao Sun, and Bo Yang. Boosting weak-to toak-to--t agents in multiagent reinforcement learning via balanced ppo. IEEE Transactions on Neural Networks and Learning Systems, 2024.
[24] Sili Huang, Yanchao Sun, Jifeng Hu, Siyuan Guo, Hechang Chen, Yi Chang, Lichao Sun, and Bo Yang. Learning generalizable agents via saliency-guided features decorrelation. Advances in Neural Information Processing Systems, 36, 2024.
[25] Shin Ishii, Wako Yoshida, and Junichiro Yoshimoto. Control of exploitation meta-parameter in reinforcement learning. Neural networks, 15(4-6)6665-68) 2002.
[26] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[27] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy lelearinn via bootstrapping reduction. Advances in Neural Information Processing Systems, 32, 2019.
[28] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020.
[29] Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. In The Eleventh International Conference on Learning Representations, 2023.
[30] Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306. 14892, 2023.

---

[31] Kuang- Suunn Lee, Ofrr Nachum. Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang. Henryk Michalewski, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35:27921-27936, 2022.
[32] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur, Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 13947, 2022.
[33] Eric Mitchell, Rafael Rafailov, Xue Bin Peng. Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021.
[34] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1125v213 2013.
[35] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rI? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36, 2024.
[36] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv: 1808.0039999 2018.
[37] Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reinforcement learning with online self-supervision. In International Conference on Machine Learning, pages 17811-17829. PMLR, 2022.
[38] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[39] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint 2rXiv:2205.0611 2022.
[40] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933 2022.
[41] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
[42] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint arXiv: 1805.01944 2018
[43] 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.
[44] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning reinforcement learn. arXiv preprini arXiv: /6/ 1.05763, 2016.
[45] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances neural information processing systems, 34:28954 28967, 2021.
[46] Tom Zahavy, Zhongwen Xu, Vivek Veeriah, Matteo Hessel, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. A self-tuning actor-critic algorithm. Advances in neural information processing systems, 33:20913-20922. 2020.
[47] 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 240. 2024,
12

---

Appendix of paper "Decision Mamba: Sequence Modeling"
Reinforcement Learning via Hybrid Selective
A Pseudocode of Decision Mamba-Hybrid
Algorithm 1: Decision Mamba-Hybrid.
Input: dataset Trajectories, Max Iterations as training phase, Max episodes m at testing phase, A number of trajectories in across-episodic contexts used in Mamba model, A number of steps of actions c for one sub-goals
Output: The generated actions //raining
for i to MM do
Randomly sample episodes from dataset - ("..,,,)
Sort episodes ascending according their returns Σ rt S =1 Ti
ET
E
Concatenate n episodes as a across-episodic context 8/1 (F T ......:.. m based on
Equation (5)
Mamba module predicts the next sub-goals tokens from the across-episodic context Build a local short-term sequence every steps based Equation (6)
Select the high-value states from the offline data based on the weighted average of
steps based on Equation ++1 T// i) Build similar local short-term sequence every c
accumulated rewards Σ
The transformer module predicts the next steps action tokens for each predicted sub-goals token from Mamba and selected high-value sub-goals
10
Train Mamba and transformer models based on the loss of predicted actions end-to-end
11
end
12
//Testing
13
for i to m do
14
Start a new episode and reset the timestep while /. 7 do
15
16
Mamba model generates next sub-goal token zł based on the historical trajectories
(T - T 1-1 "t,=,,..,,,))))) where Tm expressed as Equation (5)
for k 00
do
18
The transformer model generates next action d'te based on the local context (z), st, ot,, ri.....
........ z s1t1)
19
end
20
Compute the sum of e steps rewards ff Lkm Tk
21
Receive the next observation or state sitc.
22
Update the across-episodic context (Ty), Tm T .....
8 8
23
Update time step
24
end
end
In Algorithm [ we introduce the training and testing process DM-H. each iteration, we first construct long-term sequence consisting of multiple trajectories, as described in lines 3-5. Mamba model predicts sub-goals based on the long-term sequence (line 6). Then, each sub-goal will correspond to short sequence of steps actions, described in line 7. To encourage the transformer to rely on Mamba's 's predictions, we select the high-value states from the offline data and transform them into sub-goals to reconstruct another short sequence of steps actions (line 8). Based on the short sequence and reconstructed sequences, the transformer model predicts e steps actions (line 9). Finally, the predicted actions are evaluated with either cross-entropy loss mean-squared error, depending on whether the actions are discrete or continuous. The losses from each time step are averaged and updated in all models end-to-end, described in line 10.
During online testing, DM-H needs to generate actions autoregressively and interact with the environ- ment in episodes. At step f of episode (line 16), Mamba module first generates sub-goal token z
13

---

Table 3: Hyperparameters of DM-H.
Hyperparameterss
Value
Number of layers
2
Embedding dimension
128
Mamba
Expand factor Convolution size
2
4
Number of layers
Number of atlention heads Embedding dimension
3
Transforme
Activation function
ReLU
c steps controlled by one sub-goal
20 D4RL and Large Grid World
5 Grid World and Tmaze
Batch size
128
Dropout
0.1
Learning
lc-
Training
Learning rate decay
Linear warmup for le5 steps
Grad norm clip
0.25
Weight decay
le-4
Number of trajectories to form across across-episodes episodic contexts
(Large) Dark Key-to-
10 other tasks Grid World
4 D4RL.
Targel return for Tmaze Target return for HallCheetah
12000
Target return for Hopper
3600
Target return for Walker
5000
Target return for Darkroom
20
Target return for Darkroom Hard
1
Targel return for Durkroom Key-to-Duor
2
Testing
Target return for Large Darkroom
15
Target return for Large Darkroom Hard
1
Target return for Large Darkroom Key-to-Door
Number of trajectories form across-episodic contexts
(1 (aage)) Dark Key-to-Door
10 other tasks Grid World
D4RL
...... st), s.)..,(),, where TH is expressed as Equation (5) Then, the transformer will generate the following c steps actions (ac..... a(+c-1) autoregressively. described in lines 17-19. Finally, we update the across-episodic context for generating the next sub-goal zeter as described lines 20-23.
conditioned on the historical context (T)) Tm T
B Experimental Details
Dataset: Grid World. The evaluation environments of Grid World provide 2D discrete POMDP where an agent spawns in a room and must find a goal location. The agent only observes its own (x. y) coordinates but does not know the goal location, which is required to deduce from the rewards received. The room dimensions are x with the agent's possible actions, including moving one step either left, right, up, down, staying idle. In Darkroom, an episode lasts 20 steps, and the agent obtain a reward (r 1) each time the goal is achieved. The Darkroom Hard is a variant of Darkroom. In the Darkroom Hard, agents only obtain reward when the goal is achieved first. In the Dark Key-to-Door, the length of an episode is 50, where the agent is required to locate an invisible key to receive a one-time reward first and then identify an invisible door to obtain another one-time reward. In addition, we create a long-term variant of Large Darkroom, Large Darkroom Hard, and Large Darkroom Key-to-Door, where the coordinate space of each environment is expanded 40 40, and the episode length is expanded 10 times.
Dateset: D4RL. D4RL 3 is commonly used offline RL benchmark, including continuous
control tasks. The different dataset settings are described below.
Medium: 1 million timesteps generated by "medium" policy that performs approximately one- third well as expert policy.
14

---

Offline training time (hour)
20
DM-H (Ours)
DT
DM
2
10 Memory length (k)
20
Figure 5: Results for offline training times on Tmaze tasks. We train each method to address Tmaze tasks that have different horizons until we run out of GPU memory at context length to achieve 10k (DT, DM) 20k (DM-H) H). We report the training times for 10k gradient updates on Tmaze tasks.
Table Results for offline training and online testing times. We report the offline training time per 10k gradient updates, the online testing time for 20 episodes over Grid World, and 10 episodes over D4RL. As the task length increases, the context length is forced to grow exponentially, resulting in square increase in computational costs. In contrast, DM-H completes trial-and-eree on Mamba model sizes smaller than other baselines, significantly reducing computational costs.
Context
Tisks
Oflline training (hour)
Online testing (minute)
AD'
DMLH
Darknoom
0.21
0.23
0.18
0,20
0.61
0.21
Hard
0.19
Dynamic
Dark Key-te Key-te-Door
0.65
1.01
0.41
0.45
1.50
0.46
Large
evo
2000
Large Durlroom Hard
Large Durknoom Dymaiic
4.26
2.71
6.69
5.84
2.78
2.63
5.21
5,61
42.12 exo
18.95 (s))
5.36
$.59
Large Dark Key-to- Door
6.87
18.23
3.16
5.88
76.79
6.08
HalfCheetah
28.56
31.10
20.96
6.16
173.11 (800
6.21
4000
Walker2d
26.25
33.77
19.96
6.01
172.34
6.11
Hopper
22.23
Medium- million timesteps collected from the replay buffer of an agent trained the performance of "medium" policy.
Redium-Expeer It consists of 1 million timesteps generated by the "medium" policy and another 1 million timesteps generated by the expert policy.
Dataset: Tmaze. The Tmaze task provides T-shaped maze where the agent is rewarded only by reaching the end point at the last step. However, the endpoint's location unknown; it is only told once, at oracle early in the maze. Therefore, any policy that achieves the maximum return must be able to recall information from the first step at the final step.
Compute. Experiments are carried out on NVIDIA GeForce RTX 3090 GPUs and NVIDIA A10 GPUs. Besides, the CPU type Intel(R) Xeon(R) Gold 6230 CPU 2.10GHz. In particular, our method has lower memory requirements because it naturally shortens the across-episodic contexts.
Hyperparameters. The default length of across-episodic four trajectories unless mentioned otherwise. In D4RL, Tmaze, and Large Grid World, the transformer model generates c 20 steps actions while Mamba model generates one sub-gou aob-goal. conventional Grid World, we set c because the task is too short. In summary, Table shows the hyperparameters used in our DM-H model.
C Additional Experimental Results
Additional Tmaze Results. As described in the experiments section, we use Tmaze tasks to test the limits of DM-H's recall capabilities. We train each method to recover the optimal policy until
15

---

DM-H with valuable sub sub-goals
DM-H without valuable sub-goals
120
Retur
75
30 Halfcheetah Halfcheetah Halfcheetah m m-r
Hopper m-e
Hopper
Hopper m-r
Walk2d
Walk2d
Walk2d m-r
m-ee
m-e
m
Figure 6: The ablation study on DM-H with or without valuable sub-goals.
we run out of GPU memory at a context size equal to 20k (DM-H) 10k (DT and DM). As shown Figure 5 DM-H achieves the maximum reward with minimum offline training costs at any task horizon. Regarding efficiency, DM-H is even faster than DM using only Mamba. This is because (1) context the transformer in DM-H fixed to the hyperparameter c, which does not change with the task length. (2) Mamba model generates sub-goal every steps, which shortens the sequence it processes by cx times.
Additional Evaluation of Computing Costs. An important property of in-context RL is that it can improve itself without expensive gradient updates during online testing. However, the computational costs of forward propagation are hidden in short-horizon tasks. Therefore, we reported the offline training time per 10k gradient updates, the online time for 20 episodes over Grid World, and 10 episodes over D4RL. As shown in Table 4 our DM-H has efficient training and significantly reduces the online testing time compared to the baselines, approximately 28x times faster D4RL and 12x times faster in large Grid World. As the task length increases, the online testing time of AD grows quadratically. This because the across-episodic contexts multiply the sequence length, leading to intolerable computational costs in the self- mechanism. In contrast, DM-H leverage Mamba to process long-term terg-termm sequence, where the computational complexity increases linearly with length.
Additional Ablation Study Valuable Sub-goals. To validate the effectiveness of valuable sub- sub-goals, we also conduct an ablation study of DM-H in D4RL tasks. [igure/preeentss the ablation results, which report the mean episode return across 10 seeds. We can observe that sub- goals can significantly improve DM-H's performance, proving that the sub-goal strengthens transformer's dependency on long-term contexts from Mamba.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [TODO]
Guidelines:
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 include aspirational goals as motivation as long is 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:
We discuss the limitations in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in 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 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 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]
17

---

Justification: [TODO]
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 the statement of any theorems.
The proofs can either appear 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 of the paper to the extent that it 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: The experiments and implementation details are shown Section and Appendix A
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 are provided or not.
If 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 is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of 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) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a 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 or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In 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 to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
18

---

Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
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 NeurIPS code and data submission guidelines (https
/ /pipp/.cc/pullc//uuuii wwwtc//gguie///eeeeceoocc CoodeSubeiiieeee ssi wswonnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how 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. If 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 possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: 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?
Answer: [Yes]
Justification: Please refer to Appendix B for more details.
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, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: [TODO]
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, at least for the experiments that support 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
19

---

It 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 enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or 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 in the text how they were calculated and reference the corresponding figures or 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:
We report the computer resources of our experiments in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or 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 well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htppprs////uripsss http://uurrpsccpppbbpi/ea httppp///ccppp...tttcccccccceeea
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: In Section 6 discuss the broader impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or 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.
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 improvement in the quality of generative models could be used to
20

---

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 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
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: [TODO]
Guidelines:
The answer NA means that the paper poses no 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 to 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.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, wwppppswpptseeee whw.ree.wthhhe....c.. 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 has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
21

---

Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coddemoodll as part 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 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 as details about compensation (if any)?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, 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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor 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 expect authors to adhere the NeurIPS 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.
22


Reviews:
[{"review_id": "g4MnMSbKBh", "reviewer": "Reviewer_yEub", "summary": "The paper proposes Hybrid Mamba (HM) for in-context RL. Existing in-context RL methods are predominantly based on the Transformer architecture. Transformers come with quadratic complexity of self-attention and are computationally costly. Consequently, the authors propose a hybrid architecture that uses Mamba to compute sub-goals from long-context, which are fed into a low-level Transformer policy. The authors conduct experiments on grid-worlds and D4RL to evaluate their method.", "questions": "- Did the authors consider techniques such as key-value caching for improving inference speed of Transformers for results reported in Table 4?\n- Why is Mamba worse in effectiveness (Table 1)? What is a particular (theoretical) reason for this? Why does Mamba shorten the training time?\n- How does performance generally change, when making the models bigger? Do bigger models help on these tasks? How large are the considered models?\n- How well does the construction of valuable sub-goals generalize to other environments (e.g., with sparse rewards)?\n- How do in-context improvement curves look like on D4RL?", "limitations": "The authors highlight that setting the context length c to a fixed value as a current limitation of their method. However, a notable limitations are missing, namely that their evaluation is limited to simple environments, while it is unclear how well HM performs on more complex or new tasks.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "**Relevance**\n\nThe paper aims at deploying the Mamba architecture for in-context RL, which is very relevant given the quadratic complexity of the Transformer architecture.\nThis results in clear benefits in terms of time complexity.\n\n**Experimental results**\n\nEmpirical results on simple gridworld environments and D4RL seem convincing and their method exhibits significant gains compared to Transformers.", "weaknesses": "**Presentation**\n\nThe methodology raises some questions and should be improved, in particular:\n - What is the reasoning behind sampling the sub-goal from a multi-variate Gaussian?\n - How does this compare to using a fixed representation? (e.g., similar to CLS token)\n - Why is the done-flag in Hybrid Mamba necessary? Do other methods (e.g., AD [1]) use this as well?\n - What does “Extra high-value states” mean?\n - What is the intuition behind removing actions from the Mamba context?\n - What effect would dropping actions have in other methods?\n\nFurthermore, the construction of “valuable sub-goals” is unclear.\nOne way to improve clarity would be to shorten the section on preliminaries and instead add more details to the Method section.\nFigure 2 and Table 2 are missing the performance curves/scores for HM without valuable subgoals.\nFinally, Figure 1 can be improved to enhance clarity.\n\n**Significance of results**\n\nThe authors evaluate primarily on simple grid-world environments and rather simple robotics tasks. However, it is unclear how well HM generalizes to more complex tasks as used in other works [2].\n\n**Evaluation**\n\nThe authors change their evaluation methodology from improvement curves on gridworlds (Figure 2) to average performance scores on D4RL (Table2).\nOn D4RL, HM seems to clearly outperform other methods.\nHowever, the authors do not show in-context improvemenst which raises the question whether HM actually learns to improve in-context. Can the authors clarify, why no in-context improvement curves are shown for D4RL?\n\n**Ablation studies**\n\nSome ablation studies are missing and would add more depth to understanding the proposed method, in particular:\n- What is the impact on performance of including the done-flag in Mamba?\n- What effect does it have on other methods?\n- What is the impact on performance of removing the action condition in HM?\n- What effect does the same intervention have on other methods?\n\n\n [1] Laskin et al., In-context Reinforcement Learning with Algorithm Distillation, ICLR 2023\n [2] Raparthy et al., Generalization to New Sequential Decision Making Tasks with\nIn-Context Learning, ICML 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the reviewer's positive appraisal, insightful comment, and criticism of our paper.\n\n### [I].Reply to the Weakness\n>**[1/6]W1.What is the reasoning behind sampling the sub-goal from a multi-variate Gaussian? How does this compare to using a fixed representation?**\n\nIt is possible to predict the sub-goal directly by using a fixed representation. The multivariate Gaussian distribution is inspired by the method of predicting continuous variables in RL tasks, such as SAC[1]. In fact, sampling from a multivariate Gaussian distribution introduces variability and flexibility in representing information. This approach can generate diverse sub-goals or latent variables, allowing for exploration and capturing more nuanced aspects of the input data.\n\n>**[2/6]W2. Why is the done-flag in Hybrid Mamba necessary? Do other methods (e.g., AD) use this as well? What does “Extra high-value states” mean?**\n\nOur use of the \"done-flag\" is inspired by the Amago [2] and AT [3], which investigated the contexts consisting of different tokens and different numbers of episodes in D4RL. The ablation experiment found that the \"done-flag\" is beneficial for in-context RL to improve its performance from multiple historical trajectories.\n\nThe \"Extra high-value state\" denotes the state with high-value weighted average of accumulated rewards in the future steps. During the training phase, we introduce extra high-value states to associate the actions generated by the transformer and assist the Mamba to generate similar sub-goals that align with these states. On the other hand, although the high-value sub-goal usually appears in the future time step of the current trajectory, it may also appear in the historical trajectory, encouraging the Mamba to extract information from the long sequence context.\n\n>**[3/6]W3. What is the intuition behind removing actions from the Mamba context?**\n\nWe clarify that adding action tokens to the context of Mamba is possible. There are two reasons why we do not introduce action tokens: First, unlike transformers, the Mamba model does not generate actions, but rather representations of subgoals. Due to such subgoals do not appear explicitly in the training data, Mamba cannot use them as input for autoregressive generation. Therefore, we do not add action tokens to Mamba's input to distinguish this from the autoregressive generation of actions in the transformer. Second, our Mamba model predicts a sub-goal every $c$ steps, resulting in the context not being composed on consecutive time steps. Therefore, one-step actions can hardly help Mamba predict future states after multiple time steps, as they are used to assist in predicting the next adjacent state. In particular, we are not saying that actions are unimportant. In the context of consecutive time steps, actions are essential tokens, also included in our transformer and other baselines.\n\n>**[4/6]W4.The authors evaluate primarily on simple grid-world environments and rather simple robotics tasks. However, it is unclear how well HM generalizes to more complex tasks as used in other works.**\n\nThe grid world is one of the most commonly used benchmarks for in-context RL, such as AMAGO[2], AD[4], and DPT[5]. This is because the grid world provides many tasks that are difficult to achieve zero-shot transfer, which is very suitable for testing the ability of methods to learn from context. In addition, we also tested on large variants of grid world, where the difficulty of the tasks is greatly increased. Figure 2 in our manuscript shows that the baselines all showed significant performance degradation in large variant tasks. We are also encouraged by the new work pointed out by the reviewers, published in ICML 2024, which exploits the potential of in-context RL in a completely different setting. Due to time constraints for the rebuttal, we will add their proposed new benchmark in the revised manuscript.\n\n>**[5/6]W5.On D4RL, HM seems to clearly outperform other methods. However, the authors do not show in-context improvemenst which raises the question whether HM actually learns to improve in-context.**\n\nIn the current manuscript, we followed the prior in-context RL work by showing different styles of results for d4rl [3] and grid-world [4]. We supplement the D4RL plots in Figure 1 in the global pdf to avoid reviewer misunderstandings.\n\n\n>**[6/6]W6.What is the impact on performance of including the done-flag in Mamba? What is the impact on performance of removing the action condition in HM?**\n\nAccording to the reviewer's suggestions, we report the ablation studies on the done-flag and action condition. As we discussed in W2 and W3, the results in Table 1 show the corresponding conclusions. In addition, we show the improvement curve in Figure 1 in the global PDF.\n\n**Table 1. The ablation study on the done-flag and action condition.**\n|Token-flag|HalfCheetah-Med-Expert|HalfCheetah-Med|HalfCheetah-Med-Replay|Hopper-Med-Expert|Hopper-Med|Hopper-Med-Replay|Walker2d-Med-Expert|Walker2d-Med|Walker2d-Med-Replay|\n|-|-|-|-|-|-|-|-|-|-|\n| HM | 96.12 $\\pm$ 0.28 | 45.45 $\\pm$ 0.35 | 45.26 $\\pm$ 0.43 | 117.19 $\\pm$ 0.65 | 83.15 $\\pm$ 0.63 | 98.36 $\\pm$ 0.51 | 118.21 $\\pm$ 0.56 | 88.29 $\\pm$ 0.76 | 95.66 $\\pm$ 1.16 |\n| HM without done-flag | 93.95 $\\pm$ 0.26 | 42.21 $\\pm$ 0.38 | 41.92 $\\pm$ 0.36 | 112.21 $\\pm$ 0.62 | 78.89 $\\pm$ 0.58 | 93.21 $\\pm$ 0.32 | 113.57 $\\pm$ 0.75 | 86.24 $\\pm$ 0.65 | 90.52 $\\pm$ 0.66 |\n| HM with actions token | 96.03$\\pm$ 0.12 | 45.51 $\\pm$ 0.28 | 44.98 $\\pm$ 0.18 | 117.22 $\\pm$ 0.60 | 83.14 $\\pm$ 0.65 | 98.25 $\\pm$ 0.46 | 118.31 $\\pm$ 0.58 | 88.15 $\\pm$ 0.67 | 95.61 $\\pm$ 1.05 |\n\nDue to word limitations, we will answer the remaining questions in the next comment.", "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": "RXOiXhJg9a", "reviewer": "Reviewer_wVw1", "summary": "This paper investigates to utilize the Mamba [1] architecture for In-Context RL task. Addressing this task with Transformer architecture is effective while it is very inefficient due to the quadratic computation overhead of Transformer. The Mamba can reduce this overhead dramatically while sustain the performance somewhat. The application of State-Space Models (SSMs) to In-Context RL task is studied in [2], but different from [2], they combinationally utilize Mamba and Transformer as high-level memory and low-level (short-term) memory. Additionally, as Mamba predicts the sub-goal for the Transformer short-term memory, they improved the performance. Through this modeling, they can achieve better performance than previous works while improving the efficiency.\n\n[1] Gu, Albert, and Tri Dao. \"Mamba: Linear-time sequence modeling with selective state spaces.\" arXiv preprint arXiv:2312.00752 (2023).\n\n[2] Lu, Chris, et al. \"Structured state space models for in-context reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2024).", "questions": "- I am curious why AD (transformer) shows worse performance than HM. I thought AD (transformer) performance would be the upper bound of HM while HM is more efficient. However, AD (transformer) performance is generally worse than HM in your tests, especially for Grid World in Figure 2. Why is AD (transformer) performance poor in Grid World? Did you use a smaller context size for the Grid World test? If not (using the same context size), what could be the reasons for the significant performance gap?", "limitations": "The authors properly addressed their limitations.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Appropriate modeling is applied in this study. While the effectiveness of hybrid modeling of SSMs and local Attention has been previously explored in [1], the authors effectively implement this concept for the In-Context RL task with new functionalities, such as predicting high-value sub-goals.\n- The introduction and methodology sections are well written. The motivation is clearly articulated, and the logical flow of their method proposal is coherent. The empirical analysis comparing Mamba and Transformer in RL tasks convincingly demonstrates the need for more advanced modeling.\n- The paper provides extensive empirical analyses. It shares experimental results on multiple benchmarks, including ablation studies and performance changes with varying hyperparameter values.\n\n[1] De, Soham, et al. \"Griffin: Mixing gated linear recurrences with local attention for efficient language models.\" arXiv preprint arXiv:2402.19427 (2024).", "weaknesses": "- The high-level encoding is done by encoding the intervalled trajectories (e.g., every $c$ -th trajectory), which might miss important information in the middle of the interval.\n- The section on Hybrid Mamba with Valuable Sub-goals is initially confusing, especially regarding the relationship between Mamba’s sub-goal prediction and the collected valuable sub-goals. Discussing this relationship at the beginning of the Valuable Sub-goal section could help readers understand the content more easily.\n- One of the experimental results differs from my expectations, but the paper does not provide an analysis for this. I will address this in the Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly encouraged that Reviewer wVw1 finds our method effective.\n\n### [I].Reply to the Weakness\n>**[1/2]W1. The high-level encoding is done by encoding the intervalled trajectories (e.g., every $c$-th trajectory), which might miss important information in the middle of the interval.**\n\nIt is possible to provide complete trajectories to the high-level encoding. In fact, we tested non-intervalled trajectories during the paper's preparation and found that this setting did not produce the expected results. We design an HM variant that constructs a context sequence containing the complete trajectories but still generates subgoals for every $c$ steps.\n\n**Table 1. The performance of intervalled trajectories and non-intervalled trajectories.**\n| Trajectories | HalfCheetah-Med-Expert | HalfCheetah-Med| HalfCheetah-Med-Replay | Hopper-Med-Expert | Hopper-Med| Hopper-Med-Replay | Walker2d-Med-Expert | Walker2d-Med| Walker2d-Med-Replay |\n|-|-|-|-|-|-|-|-|-|-|\n| Intervalled trajectories | 96.12 $\\pm$ 0.28 | 45.45 $\\pm$ 0.35 | 45.26 $\\pm$ 0.43 | 117.19 $\\pm$ 0.65 | 83.15 $\\pm$ 0.63 | 98.36 $\\pm$ 0.51 | 118.21 $\\pm$ 0.56 | 88.29 $\\pm$ 0.76 | 95.66 $\\pm$ 1.16 |\n| Non-intervalled trajectories | 95.98 $\\pm$ 0.21 | 45.22 $\\pm$ 0.18 | 44.97 $\\pm$ 0.28 | 117.21 $\\pm$ 0.61 | 82.98 $\\pm$ 0.65 | 98.12 $\\pm$ 0.38 | 117.97 $\\pm$ 0.62 | 88.25 $\\pm$ 0.68 | 95.62 $\\pm$ 0.96 |\n\nTable 1 shows no significant performance gap between the two variants of HM. This is because (1) the transformer retains the complete short-term trajectory for action prediction. (2) The cross-episodic context in Mamba is long enough so that the uniformly sampled trajectory preserves sufficient information for predicting the sub-goals.\n\n>**[2/2]W2. The section on Hybrid Mamba with Valuable Sub-goals is initially confusing, especially regarding the relationship between Mamba’s sub-goal prediction and the collected valuable sub-goals. Discussing this relationship at the beginning of the Valuable Sub-goal section could help readers understand the content more easily.**\n\nThanks for the reviewer's suggestion. We will explain their relationship in the revised manuscript. In summary, the valuable sub-goals collected can be regarded as additional signals, and we hope that the Mamba model can generate similar sub-goals to prompt the transformer's prediction.\n\n\n### [II].Reply to the Questions\n>**[1/1]Q1. I am curious why AD (transformer) shows worse performance than HM. I thought AD (transformer) performance would be the upper bound of HM while HM is more efficient. However, AD (transformer) performance is generally worse than HM in your tests, especially for Grid World in Figure 2. Why is AD (transformer) performance poor in Grid World? Did you use a smaller context size for the Grid World test? If not (using the same context size), what could be the reasons for the significant performance gap?**\n\nAs far as we know, AD was the first to introduce in-context learning into transformer-based RL. It uses the simplest (state, action, reward) tuples to form the context sequence, so it is not outstanding in performance and usually serves as a baseline for subsequent advanced methods, such as Amago [1] and DPT[2]. In the Grid World, both AD and our method HM set the same context that covers the same number of trajectories. This performance gap can be attributed to two aspects: token modality and model architecture.\nOur approach adds two different tokens compared to AD: the c-step cumulative reward and the done flag. Since the environment provided by Grid World is mostly sparsely rewarded, cumulative reward tokens are more advantageous. On the other hand,  recent methods, such as AT [3], demonstrated that done flag tokens are beneficial for in-context RL. In the model architecture, the hierarchical architecture of HM also brings about performance improvements, which is analogous to the fact that hierarchical RL learning is better at handling long-term and sparse reward tasks.\n\nAccording to the reviewer's suggestion, we will incorporate the above analysis into the revised manuscript.\n\n\n[1]Amago: Scalable in-context reinforcement learning for adaptive agents. ICLR 2024.  \n[2]Supervised Pretraining Can Learn In-Context Reinforcement Learning. NeurIPS 2023.  \n[3]Emergent Agentic Transformer from Chain of Hindsight Experience. ICML 2023.", "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": "0U27B892i7", "reviewer": "Reviewer_YUg4", "summary": "This paper presents Hybrid Mamba (HM), a method that combines the Mamba model and Transformer to enhance reinforcement learning (RL) performance. HM leverages Mamba to generate high-value sub-goals, which then condition the transformer, leading to significant improvements in online testing efficiency and task-solving capabilities.", "questions": "Please see weakness. If I have misunderstood some parts of the paper, I welcome corrections and further discussion.", "limitations": "The authors raise some limitaions, for example, how to control the setting of hyperparameter $c$, which is not addressed in this paper but is claimed to be solved in the future.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and clear to read.\n2. HM significantly accelerates testing speed, achieving up to 28 times faster results than baseline methods.\n3. HM demonstrates superior performance across various benchmarks, such as D4RL, Grid World, and Tmaze.", "weaknesses": "1. This paper claims to present a in-context RL approach. The motivation of this paper is concerned with the problems encountered with the no-gradient updates in-context approach (line 28), where the policy network does not require parameter updates. However, this paper uses a global update approach, which is closer to gradient-based and conditional-based offline RL. It seems to contradict the original intention of this paper. \n\n2. HM benefits from using a powerful subgoal encoder (Mamba in this case) and conditioning the policy network with subgoals. The performance improvement is expected and unsurprising due to the advantages inherent in conditional-based RL algorithms. Hence, it is necessary to further explain the unique contributions of combining Mamba and causal transformer in this paper.\n\n3. If the sub-goal encoder are replaced with other advanced and efficient sequence encoders (e.g., flash-attention1/2 [1,2], x-lstm [3]), would it also yield better or more efficient performance? \n\n4. The experiments demonstrating HW's efficacy in capturing long-term dependencies are unconvincing. Achieving good results in tasks with an arbitrarily horizon (e.g., Tmaze) does not necessarily prove effective long-term memory embedding. It is crucial to test the stability and performance of HM with varying horizon lengths or other length-based settings. For example, Mamba’s original paper [4] demonstrated the ability to capture long-term dependencies through the scaling laws.\n\n5. Could the authors clarify in which specific aspects HM's training time is faster than DT's? Since HM appears to be a combination of Mamba and DT.\n\n6. There are parts of the paper that are not clearly explained. For instance, in lines 228-233, it is mentioned that the transformer predicts a c-step action sequence (named $a_1$ here) through the sub-goal $z_t$ and another c-step action sequence (named $a_2$) through valuable sub-goals from offline data. How are $a_1$ and $a_2$ subsequently updated or processed? \n\n7. (minor) The paper contains some typos and inconsistencies in tense usage. For example, in the related work section, the section on Mamba uses the present tense, while the section on in-context RL uses the past tense. These should be corrected for consistency. In addition, what's the meaning of the different gaussian distribution figures in Figure 1?\n\n*Reference:*\n\n[1] Dao T, Fu D, Ermon S, et al. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS 2022.\n\n[2] Dao T. Flashattention-2: Faster attention with better parallelism and work partitioning. ICLR 2024.\n\n[3] Beck M, Pöppel K, Spanring M, et al. xLSTM: Extended Long Short-Term Memory. arXiv 2024.\n\n[4] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the reviewer's positive appraisal, insightful comment, and criticism of our paper.\n\n### [I].Reply to the Weakness\n>**[1/7]W1. This paper claims to present a in-context RL approach. The motivation of this paper is concerned with the problems encountered with the no-gradient updates in-context approach (line 28), where the policy network does not require parameter updates. However, this paper uses a global update approach, which is closer to gradient-based and conditional-based offline RL. It seems to contradict the original intention of this paper.**\n\nWe want to clarify the setting of transformer-based in-context RL, which is divided into two phases: training and testing. In the training phase, in-context RL designs the architecture of policy models and trains the model by predicting the action tokens in carefully designed sequences, such as multiple trajectories in AD [1]. This process involves parameter updates and can be regarded as a kind of pre-training. In the testing phase, we deploy the model in a new task and can only make inferences in the context obtained by interacting with the environment without involving parameter updates.\n\nIn the method section, we introduced the architecture of HM and how to construct sequences for training. Then, at the end of the section, we summarized its testing process. We roll out the HM with multiple trajectories and report the return of the last trajectory. In-context learning is reflected by the HM, which can predict better actions by recalling historical trajectories from the context. In Figure 2 in the manuscript, the ascending reward curves are achieved by learning in the contexts without requiring parameter updates.\n\n\n>**[2/7]W2. HM benefits from using a powerful subgoal encoder (Mamba in this case) and conditioning the policy network with subgoals. The performance improvement is expected and unsurprising due to the advantages inherent in conditional-based RL algorithms. Hence, it is necessary to further explain the unique contributions of combining Mamba and causal transformer in this paper.**\n\nThanks for the reviewer's suggestion. Mamba is well known for its competitive effectiveness in transformer-based models while achieving linear complexity in long sequences. Its outstanding performance in NLP and visual tasks has attracted more and more attention. Therefore, our first contribution is to explore the potential of Mamba in decision-making tasks. In traditional RL tasks, we demonstrate that the Mamba model is more efficient but slightly inferior to the transformer model in terms of effectiveness. Therefore, we propose the HM method that inherits the merits of Mamba and transformers, achieving both high effectiveness and efficiency in long-term RL tasks. HM is inspired by the idea of conditional RL, proving that this hybrid structure is a promising way to leverage the strengths of Mamba and Transformer to complement their weaknesses.\n\n>**[3/7]W3. If the sub-goal encoder are replaced with other advanced and efficient sequence encoders (e.g., flash-attention1/2, x-lstm), would it also yield better or more efficient performance?**\n\nIt is possible to use other sequence encoders to generate sub-goals, such as the reviewer's suggested flash-attention and x-lstm. The x-lstm has been proposed recently, and we did not find its open source code. Therefore, we tested HM by replacing the Mamba module with the flash-attention. Due to the limited time and device, we report the results of flash-attention1 across three random seeds.\n\n**Table 1. The performance of different sub-goal encoder.**\n|Sub-goal encoder|HalfCheetah-Med-Expert|HalfCheetah-Med|HalfCheetah-Med-Replay|Hopper-Med-Expert|Hopper-Med|Hopper-Med-Replay|Walker2d-Med-Expert|Walker2d-Med|Walker2d-Med-Replay|\n|-|-|-|-|-|-|-|-|-|-|\n|Mamba|96.12 $\\pm$ 0.28|45.45 $\\pm$ 0.35|45.26 $\\pm$ 0.4|117.19 $\\pm$ 0.65|83.15 $\\pm$ 0.63|98.36 $\\pm$ 0.51|118.21 $\\pm$ 0.56|88.29 $\\pm$ 0.76|95.66 $\\pm$ 1.16|\n|flash-attention|94.91 $\\pm$ 0.18|44.22 $\\pm$ 0.16|43.98 $\\pm$ 0.35|116.09 $\\pm$ 0.68|81.58 $\\pm$ 0.33|96.85 $\\pm$ 0.64|117.19 $\\pm$ 0.49|86.26 $\\pm$ 0.61|94.12 $\\pm$ 0.88|\n\nThe results show that flash attention is slightly inferior to Mamba on d4rl tasks. This is because the process of predicting sub-goals is non-autoregressive. As we clarified in the manuscript, the sub-goal is represented by a vector that is not explicitly present in the training data. This non-autoregressive process may not fully exploit the power of flash-attention. On the contrary, although Mamba was also first proposed for NLP tasks, its structure is closer to RNN, a commonly used architecture in RL. Similarly, X-LSTM should yield efficient performance because it is also an RNN-like model.\n\n>**[4/7]W4.The experiments demonstrating HW's efficacy in capturing long-term dependencies are unconvincing. Achieving good results in tasks with an arbitrarily horizon (e.g., Tmaze) does not necessarily prove effective long-term memory embedding. It is crucial to test the stability and performance of HM with varying horizon lengths or other length-based settings.**\n\nThere is a misunderstanding about our testing of HM recalling long-term memory. We followed the setting of the previous in-context RL method AMAGO [2] to extend the horizon of tasks until we ran out of GPU memory. The horizon of HM also varies with the tasks from short-horizon to long-horizon. Although unlike indicators such as perplexity or accuracy in NLP tasks, the cumulative reward (return) in RL tasks has a similar meaning and is used to evaluate the model's performance. Figure 3(a) results in the manuscript show that HM with varying horizon lengths are stable and high-performance. In summary, we want to show that HM can effectively handle tasks that require recalling long-term memory, as tested by previous methods.\n\nDue to word limitations, we will answer the remaining questions in the next comment.", "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": "fa59IP0Nn9", "reviewer": "Reviewer_UASH", "summary": "This paper investigates an emerging foundation model, Mamba, in Reinforcement Learning (RL) scenarios and compares it with Transformer in terms of effectiveness and efficiency. The authors find that in-context RL methods with Mamba as the backbone are generally more efficient than Transformer, but there is no significant improvement in effectiveness. Then, this paper proposes a Hybrid Mamba (HM) with the merits of transformers and Mamba in high-quality prediction and long-term memory. Finally, this paper conducts experiments on three benchmarks to exhibit its improved effectiveness and efficiency.", "questions": "1.\tIn Table 1 and Table 2, the author used AD (Mamba) as the primary baseline. However, in Figure 3, the author used the DM baseline instead. What is the main difference between AD (Mamba) and DM?\n2.\tThe experiments demonstrated that the online testing of HM in the long-term task is 28 times faster than the transformer-based method. However, can this hybrid structure also inherit Mamba's high efficiency in terms of training cost?", "limitations": "The author discusses limitations and potential negative societal impacts in Section 6.", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1.\tThe paper is commendably well-written and coherent, effectively explaining complex ideas in an accessible manner. The authors explored the potential of the widely discussed model Mamba in the context of RL and compared it with Transformer in terms of effectiveness and efficiency.\n2.\tThe authors proposed a novel hybrid model that inherits the merits of both Transformer and Mamba in a goal-conditional manner. The main advantage of using the hybrid structure is that when the time horizon is very long, as in the D4RL tasks, several episodes/trials are required for good in-context learning, as in the larger Dark Room and Tmaze environments.\n3.\tHM improves training and inference speed by reducing the horizon of the transformer model. This can be particularly important in applications such as robotics, which require high-frequency control.\n4.\tThe experimental evaluation, meticulously designed to include several baselines and diverse tasks, demonstrates the algorithm's strengths.", "weaknesses": "1.\tThe baseline AD (Mamba) in Figure 2 and the baseline DM in Figure 3, which appear to be AD (Transformer) and DT variants, are crucial for the readers' understanding of how Mamba replaces the Transformer architecture. However, the lack of explanation of these two baselines in the experimental setup section might confuse readers.\n2.\tSome experimental settings are not explained clearly. In Section 5.3, the authors do not explain what GPU device they used. Although the device is introduced in Appendix A, it is recommended that it be explained clearly in the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly encouraged that Reviewer UASH finds our method effective.\n\n### [I].Reply to the weakness\n>**[1/2]W1. The baseline AD (Mamba) in Figure 2 and the baseline DM in Figure 3, which appear to be AD (Transformer) and DT variants, are crucial for the readers' understanding of how Mamba replaces the Transformer architecture. However, the lack of explanation of these two baselines in the experimental setup section might confuse readers.**\n\nThanks for the reviewer's suggestions. The difference between AD (Mamba) and DM is whether the context covers multiple trajectories. We highlight the advantages and disadvantages of Mamba by comparing it with two transformer-based methods, AD (Transformer) and DT. In the revised manuscript, we will add more detailed instructions in Section 5.1 of the experiment.\n\n>**[2/2]W2. Some experimental settings are not explained clearly. In Section 5.3, the authors do not explain what GPU device they used. Although the device is introduced in Appendix A, it is recommended that it be explained clearly in the main text.**\n\nExperiments are carried out on NVIDIA GeForce RTX 3090 GPUs and NVIDIA A10 GPUs. Besides, the CPU type is Intel(R) Xeon(R) Gold 6230 CPU @ 2.10GHz. As suggested, we will emphasize the device in the revised manuscript.\n\n### [II].Reply to the questions\n>**[1/2]Q1. In Table 1 and Table 2, the author used AD (Mamba) as the primary baseline. However, in Figure 3, the author used the DM baseline instead. What is the main difference between AD (Mamba) and DM?**\n\nThe main difference between AD (Mamba) and DM is whether the context covers multiple trajectories. Since the Tmaze requires the policy to recall the first step at the current trajectory, we show the DT and DM comparison results in Figure 3. In contrast, the Grid world and d4rl tasks require the policy to improve itself based on the historical trajectories. Thus, we use AD (Mamba) as the primary baseline. \n\n>**[2/2]Q2. The experiments demonstrated that the online testing of HM in the long-term task is 28 times faster than the transformer-based method. However, can this hybrid structure also inherit Mamba's high efficiency in terms of training cost?**\n\nYes, we are glad that the reviewer found our method also highly efficient in terms of training cost. Assume the sequence length is $n$. The computational complexity of the transformer is $O(n^2)$. In HM, the sequence is divided into a hierarchical structure. In the Mamba level, the sequence length is $\\frac{n}{c}$, and the complexity is $O(\\frac{n}{c})$, where the hyperparameter $c$ denotes the timestep interval at which each sub-goal is set to guide the transformer. In the transformer level, the sequence is divided into $\\frac{n}{c}$ subsequences of length $c$, and the complexity becomes $O(\\frac{n}{c}\\cdot c^2)$. As the sequence length $n$ increases, the computational complexity of HM will be significantly lower than that of the transformer-based method, and thus, the training will be faster.", "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:
Dear Reviewers,

We are very grateful to the reviewers for their valuable suggestions, which further improved our work. We provide the learning curve of our HM and ablation studies in d4rl tasks with a submitted 1-page pdf.

Thank you again for your careful review and helpful comments.

Kind regards,

Paper5823 Authors

Meta-review: The authors propose a model called HM that combines transformers and the Mamba model to efficiently handle long memory in in-context reinforcement learning. HM generates high-value sub-goals through the Mamba model, which are then used to prompt the transformer for high-quality predictions.

The paper is well-written and easy to follow. The experiments are thorough and demonstrate strong empirical performance. While additional experiments on Procgen would be interesting, I believe the paper has already shown significant potential. Therefore, I recommend acceptance.
Decision: Accept (poster)