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.
Paper
Similar papers
Peer review
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.
Strengths
1. The 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. 2. The 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. 3. HM 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. 4. The experimental evaluation, meticulously designed to include several baselines and diverse tasks, demonstrates the algorithm's strengths.
Weaknesses
1. 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. 2. 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.
Questions
1. 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? 2. 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?
Rating
10
Confidence
5
Soundness
4
Presentation
4
Contribution
4
Limitations
The author discusses limitations and potential negative societal impacts in Section 6.
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.
Strengths
1. The paper is well-written and clear to read. 2. HM significantly accelerates testing speed, achieving up to 28 times faster results than baseline methods. 3. 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. 2. 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. 3. 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? 4. 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. 5. 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. 6. 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? 7. (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? *Reference:* [1] Dao T, Fu D, Ermon S, et al. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS 2022. [2] Dao T. Flashattention-2: Faster attention with better parallelism and work partitioning. ICLR 2024. [3] Beck M, Pöppel K, Spanring M, et al. xLSTM: Extended Long Short-Term Memory. arXiv 2024. [4] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv 2023.
Questions
Please see weakness. If I have misunderstood some parts of the paper, I welcome corrections and further discussion.
Rating
4
Confidence
4
Soundness
2
Presentation
3
Contribution
2
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.
Rebuttal of the remaining questions
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.
Looking forward to your reply
Dear Reviewer YUg4, We value your positive feedbacks and constructive suggestions for our paper and sincerely appreciate your effort in reviewing it. We hope we have effectively addressed all the concerns raised. As the end of the discussion is approaching, we are wondering if there are any additional potential clarifications or suggestions that you think would help us improve this manuscript. Thank you again for your dedicated review and invaluable insights. Kind regards, Paper5823 Authors
Response to the authors
Thank you for your rebuttal. However, I have some reservations about the conclusion that the long-horizon experiments can validate HM's long-term memory capabilities. Additionally, considering the feedback from other reviewers, there are several incomplete or unclear parts in this paper. Therefore, I decided to maintain my original score at this stage.
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. [1] Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023). [2] Lu, Chris, et al. "Structured state space models for in-context reinforcement learning." Advances in Neural Information Processing Systems 36 (2024).
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. - 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. - The paper provides extensive empirical analyses. It shares experimental results on multiple benchmarks, including ablation studies and performance changes with varying hyperparameter values. [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. - 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. - 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.
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?
Rating
7
Confidence
5
Soundness
4
Presentation
4
Contribution
3
Limitations
The authors properly addressed their limitations.
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.
Strengths
**Relevance** The paper aims at deploying the Mamba architecture for in-context RL, which is very relevant given the quadratic complexity of the Transformer architecture. This results in clear benefits in terms of time complexity. **Experimental results** Empirical results on simple gridworld environments and D4RL seem convincing and their method exhibits significant gains compared to Transformers.
Weaknesses
**Presentation** The methodology raises some questions and should be improved, in particular: - What is the reasoning behind sampling the sub-goal from a multi-variate Gaussian? - How does this compare to using a fixed representation? (e.g., similar to CLS token) - Why is the done-flag in Hybrid Mamba necessary? Do other methods (e.g., AD [1]) use this as well? - What does “Extra high-value states” mean? - What is the intuition behind removing actions from the Mamba context? - What effect would dropping actions have in other methods? Furthermore, the construction of “valuable sub-goals” is unclear. One way to improve clarity would be to shorten the section on preliminaries and instead add more details to the Method section. Figure 2 and Table 2 are missing the performance curves/scores for HM without valuable subgoals. Finally, Figure 1 can be improved to enhance clarity. **Significance of results** 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 [2]. **Evaluation** The authors change their evaluation methodology from improvement curves on gridworlds (Figure 2) to average performance scores on D4RL (Table2). 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. Can the authors clarify, why no in-context improvement curves are shown for D4RL? **Ablation studies** Some ablation studies are missing and would add more depth to understanding the proposed method, in particular: - What is the impact on performance of including the done-flag in Mamba? - What effect does it have on other methods? - What is the impact on performance of removing the action condition in HM? - What effect does the same intervention have on other methods? [1] Laskin et al., In-context Reinforcement Learning with Algorithm Distillation, ICLR 2023 [2] Raparthy et al., Generalization to New Sequential Decision Making Tasks with In-Context Learning, ICML 2024
Questions
- Did the authors consider techniques such as key-value caching for improving inference speed of Transformers for results reported in Table 4? - Why is Mamba worse in effectiveness (Table 1)? What is a particular (theoretical) reason for this? Why does Mamba shorten the training time? - How does performance generally change, when making the models bigger? Do bigger models help on these tasks? How large are the considered models? - How well does the construction of valuable sub-goals generalize to other environments (e.g., with sparse rewards)? - How do in-context improvement curves look like on D4RL?
Rating
3
Confidence
4
Soundness
2
Presentation
2
Contribution
2
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.
Rebuttal of the remaining questions
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.
Looking forward to your reply
Dear Reviewer yEub, We value your positive feedbacks and constructive suggestions for our paper and sincerely appreciate your effort in reviewing it. We hope we have effectively addressed all the concerns raised. As the end of the discussion is approaching, we are wondering if there are any additional potential clarifications or suggestions that you think would help us improve this manuscript. Thank you again for your dedicated review and invaluable insights. Kind regards, Paper5823 Authors
I would like to thank the authors for the detailed response and the additional results. My concerns about a more detailed presentation have been addressed, however some major issues still persist. **Significance of results** The empirical results are still limited to very simple RL tasks, other works in the field of in-context RL consider more complex tasks such as, e.g. Procgen [1], which also provide the possibility to evaluate on unseen seeds of the same task or entirely unseen tasks. [1] Raparthy et al., Generalization to New Sequential Decision Making Tasks with In-Context Learning, ICML 2024 **Evaluation** One of my major concerns was regarding the lack of in-context evaluation. The whole premise of in-context learning is to learn **new** tasks only by providing them in the context of the model. The authors only show resuls for in-context improvement on the gridworld tasks as for D4RL they evaluate on the training task. The results on those tasks seem promising. However, since the man gist of the paper is in-context learning, I would recommend to add Procgen to their experiments and evaluate on unseen seeds and unseen tasks as in [1] to demonstrate how well HM is able to learn in-context. In order to progress the field, it is vital to move beyond mere gridworld based environments. Based on these reasons I have decided to keep my score.
Reply to the rebuttal
Thank you the authors for your rebuttal. The authors properly addressed my concerns and I keep my score.
Appreciation to Reviewer wVw1
Dear Reviewer wVw1, We are grateful for your constructive suggestions and believe that incorporating the corresponding revision into the manuscript will significantly improve this paper. Thank you again for reviewing our paper and giving valuable feedback! Kind regards, Paper5823 Authors
Looking forward to your reply
Dear Reviewer UASH, We value your positive feedbacks and constructive suggestions for our paper and sincerely appreciate your effort in reviewing it. We hope we have effectively addressed all the concerns raised. As the end of the discussion is approaching, we are wondering if there are any additional potential clarifications or suggestions that you think would help us improve this manuscript. Thank you again for your dedicated review and invaluable insights. Kind regards, Paper5823 Authors
Thanks for the author's clarification. I decided to keep my initial score. I rated this work highly for several reasons. First, Mamba is an emerging foundational model with significant potential, yet its performance in decision-making tasks has been underexplored. This paper fills that gap by systematically comparing Mamba's effectiveness and efficiency against the transformer model on classic RL tasks. Second, the authors propose a hybrid framework that leverages the strengths of both Mamba and transformer models in decision-making tasks. Overall, I believe this work will generate considerable interest within the NeurIPS community and advance the application of foundational models in RL.
Decision
Accept (poster)