Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling

• URL: http://arxiv.org/abs/2406.00079v1
• Date: Fri, 31 May 2024 10:41:03 GMT
• Title: 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,
• Abstract summary: We propose a Decision Mamba-Hybrid (DM-H) for in-context reinforcement learning. DM-H generates high-value sub-goals from long-term memory through the Mamba model. Online testing of DM-H in the long-term task is 28$times$ times faster than the transformer-based baselines.
• Score: 13.253878928833688
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• The Mamba in the Llama: Distilling and Accelerating Hybrid Models [76.64055251296548]
  We show that it is feasible to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources. The resulting hybrid model, which incorporates a quarter of the attention layers, achieves performance comparable to the original Transformer in chat benchmarks.
  arXiv  Detail & Related papers  (2024-08-27T17:56:11Z)
• Integrating Multi-Modal Input Token Mixer Into Mamba-Based Decision Models: Decision MetaMamba [0.0]
  Sequence modeling with State Space models (SSMs) has demonstrated performance surpassing that of Transformers in various tasks. However, decision models based on Mamba, a state-of-the-art SSM, failed to achieve superior performance compared to enhanced Decision Transformers. We propose the Decision MetaMamba (DMM), which augments Mamba with a token mixer in its input layer.
  arXiv  Detail & Related papers  (2024-08-20T03:35:28Z)
• DeciMamba: Exploring the Length Extrapolation Potential of Mamba [89.07242846058023]
  We introduce DeciMamba, a context-extension method specifically designed for Mamba. We show that DeciMamba can extrapolate context lengths 25x longer than the ones seen during training, and does so without utilizing additional computational resources.
  arXiv  Detail & Related papers  (2024-06-20T17:40:18Z)
• An Empirical Study of Mamba-based Language Models [69.74383762508805]
  Selective state-space models (SSMs) like Mamba overcome some shortcomings of Transformers. We present a direct comparison between 8B-context Mamba, Mamba-2, and Transformer models trained on the same datasets. We find that the 8B Mamba-2-Hybrid exceeds the 8B Transformer on all 12 standard tasks.
  arXiv  Detail & Related papers  (2024-06-12T05:25:15Z)
• Mamba as Decision Maker: Exploring Multi-scale Sequence Modeling in Offline Reinforcement Learning [16.23977055134524]
  We propose a novel action predictor sequence, named Mamba Decision Maker (MambaDM) MambaDM is expected to be a promising alternative for sequence modeling paradigms, owing to its efficient modeling of multi-scale dependencies. This paper delves into the sequence modeling capabilities of MambaDM in the RL domain, paving the way for future advancements.
  arXiv  Detail & Related papers  (2024-06-04T06:49:18Z)
• Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning? [32.33214392196923]
  Transformer-based trajectory optimization methods have demonstrated exceptional performance in offline Reinforcement Learning (offline RL) Yet, it poses challenges due to substantial parameter size and limited scalability, which is particularly critical in sequential decision-making scenarios. Mamba, a promising new linear-time sequence model, offers performance on par with transformers while delivering substantially fewer parameters on long sequences.
  arXiv  Detail & Related papers  (2024-05-20T15:05:47Z)
• Decision Mamba Architectures [1.4255659581428335]
  Decision Mamba architecture has shown to outperform Transformers across various task domains. We introduce two novel methods, Decision Mamba (DM) and Hierarchical Decision Mamba (HDM) We demonstrate the superiority of Mamba models over their Transformer counterparts in a majority of tasks.
  arXiv  Detail & Related papers  (2024-05-13T17:18:08Z)
• Is Mamba Effective for Time Series Forecasting? [30.85990093479062]
  We propose a Mamba-based model named Simple-Mamba (S-Mamba) for time series forecasting. Specifically, we tokenize the time points of each variate autonomously via a linear layer. Experiments on thirteen public datasets prove that S-Mamba maintains low computational overhead and achieves leading performance.
  arXiv  Detail & Related papers  (2024-03-17T08:50:44Z)
• Is Mamba Capable of In-Context Learning? [63.682741783013306]
  State of the art foundation models such as GPT-4 perform surprisingly well at in-context learning (ICL) This work provides empirical evidence that Mamba, a newly proposed state space model, has similar ICL capabilities.
  arXiv  Detail & Related papers  (2024-02-05T16:39:12Z)
• Emergent Agentic Transformer from Chain of Hindsight Experience [96.56164427726203]
  We show that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches. This is the first time that a simple transformer-based model performs competitively with both temporal-difference and imitation-learning-based approaches.
  arXiv  Detail & Related papers  (2023-05-26T00:43:02Z)
• On Transforming Reinforcement Learning by Transformer: The Development Trajectory [97.79247023389445]
  Transformer, originally devised for natural language processing, has also attested significant success in computer vision. We group existing developments in two categories: architecture enhancement and trajectory optimization. We examine the main applications of TRL in robotic manipulation, text-based games, navigation and autonomous driving.
  arXiv  Detail & Related papers  (2022-12-29T03:15:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.