Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling

We are particularly encouraged that Reviewer UASH finds our method effective.

[I].Reply to the weakness

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

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

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

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

[II].Reply to the questions

[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?

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

[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?

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

Due to word limitations, we answer the remaining questions in this comment. Please review this response after the rebuttal part.

[I].Reply to the Weakness

[5/7]W5.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.

Thanks for the reviewer's suggestion. Although HM combines Mamba and DT, its sequence length is the same as that of DT. Assume the sequence length is $n$. The computational complexity of DT 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 DT. In the DT 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 DT, and thus, the training will be faster. We will incorporate the above complexity analysis into the revised manuscript.

[6/7]W6.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?

In lines 228-233, we introduce the training process of HM using a sub-goal $z_t$ as an example. Assume a c-step sequence $(s_g,s_t,a_t^*,s_{t+1},a_{t+1}^*,\dots,s_{t+c-1},a_{t+c-1}^*)$ exists in the offline data, where $s_g$ is the valuable sub-goal selected by the weighted average of accumulated rewards. In the training process, the Mamba model in HM generates a sub-goal $z_t$ and guides the transformer to predict the c-step action sequence (named $a_1$ here). Meanwhile, we replace the $z_t$ with $s_g$ and predict the c-step action sequence again (named $a_2$ here). Finally, we update the model parameters so that its predictions for these two action sequences ($a_1$ and $a_2$) are close to $(a_t^*,a_{t+1}^*,\dots,a_{t+c-1}^*)$. In the testing process, the trained HM is deployed in a new environment without parameter updates. This process will not have access to $s_g$, where HM relies on generating better sub-goals $z$ from the context to improve its performance.

[7/7]W7.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?

Thanks for the reviewer's corrections. We will revise the typos and inconsistencies in tense usage. In Figure 1 in the manuscript, the different Gaussian distributions indicate that Mamba generates a different subgoal for the transformer's predictions every c steps. It is possible to predict the sub-goal directly by using a fixed representation. However, sampling from a multi-variate 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.

[1] In-context Reinforcement Learning with Algorithm Distillation. ICLR 2023.
[2] Amago: Scalable in-context reinforcement learning for adaptive agents. ICLR 2024.

Thanks for the reviewer's positive appraisal, insightful comment, and criticism of our paper.

[I].Reply to the Weakness

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

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

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

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

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

[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?

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

Table 1. The performance of different sub-goal encoder.

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

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

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

Due to word limitations, we will answer the remaining questions in the next comment.

We are particularly encouraged that Reviewer wVw1 finds our method effective.

[I].Reply to the Weakness

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

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

Table 1. The performance of intervalled trajectories and non-intervalled trajectories.

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

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

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

[II].Reply to the Questions

[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?

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

According to the reviewer's suggestion, we will incorporate the above analysis into the revised manuscript.

[1]Amago: Scalable in-context reinforcement learning for adaptive agents. ICLR 2024.
[2]Supervised Pretraining Can Learn In-Context Reinforcement Learning. NeurIPS 2023.
[3]Emergent Agentic Transformer from Chain of Hindsight Experience. ICML 2023.

Due to word limitations, we answer the remaining questions in this comment. Please review this part of the response after the rebuttal part.

[II].Reply to the Questions

[1/5]Q1. Did the authors consider techniques such as key-value caching for improving inference speed of Transformers for results reported in Table 4?

The transformer model we counted in Table 4 has applied key-value caching. When calculating the action of a new time step, we do not recalculate the key and value of the previous time step. However, the transformer's computational complexity grows quadratic with the sequence length, which makes it much slower than the Mamba method in the context of long sequences.

[2/5]Q2. Why is Mamba worse in effectiveness (Table 1)? What is a particular (theoretical) reason for this? Why does Mamba shorten the training time?

Mamba models long sequence dependencies through an input-dependent selection mechanism. This mechanism acts independently on each token, so its computational complexity is less than the attention mechanism based on token pairs. Therefore, Mamba's computational complexity grows linearly with the sequence, making it naturally superior to transformers in terms of training speed and inference speed. However, even under the MDP assumption, the states in RL tasks are not independent, which causes Mamba to perform lower than the transformer in terms of effectiveness.

[3/5]Q3. Do bigger models help on these tasks? How large are the considered models?

As shown in the hyperparameters in Table 3 of the manuscript, our Transformer model follows the setting of the AT method [3], and the Mamba model follows the setting of the DM method [6]. According to the conclusions in the Mamba paper, larger models generally improve performance. Due to resource and time constraints, we will add tests of larger models in the revised manuscript. However, it is worth mentioning that HM, with the current model size, can handle these tasks well.

[4/5]Q4. How well does the construction of valuable sub-goals generalize to other environments (e.g., with sparse rewards)?

The construction of valuable sub-goals also generalizes well to environments with sparse rewards. Figures 2 and 3 in the manuscript report the tasks in a sparse rewards setting, showing sub-goals' effectiveness. In particular, we will try our best to test HM on the benchmarks suggested by the reviewers in the revised manuscript.

[5/5]Q5. How do in-context improvement curves look like on D4RL?

As we clarify in W5, the curves of D4RL are shown in Figure 1 in the global PDF file.

[1]Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. ICML 2018.
[2]Amago: Scalable in-context reinforcement learning for adaptive agents. ICLR 2024.
[3]Emergent Agentic Transformer from Chain of Hindsight Experience. ICML 2023.
[4]In-context Reinforcement Learning with Algorithm Distillation. ICLR 2023.
[5]Supervised Pretraining Can Learn In-Context Reinforcement Learning. NeurIPS 2023.
[6]Decision Mamba: Reinforcement learning via sequence modeling with selective state spaces. 2024.

Thanks for the reviewer's positive appraisal, insightful comment, and criticism of our paper.

[I].Reply to the Weakness

[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?

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

[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?

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

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

[3/6]W3. What is the intuition behind removing actions from the Mamba context?

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

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

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

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

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

[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?

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

Table 1. The ablation study on the done-flag and action condition.

Due to word limitations, we will answer the remaining questions in the next comment.