# Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?

arXiv: https://arxiv.org/abs/2405.12094 · HTML: https://mlsanity.com/paper/2405.12094/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Video, motion and generative Mamba models, one of 10 directions in State-space sequence models (S4, Mamba).
- Builds on 51 of its 61 references that are in the map.
- First posted on arXiv 20 May 2024.

## Placed in

- State-space sequence models (S4, Mamba) / [Video, motion and generative Mamba models](https://mlsanity.com/ssm/video-motion-and-generative-mamba-models/)
- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

## Builds on (most cited first, 30 of 51)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model](https://mlsanity.com/paper/2401.09417/) (2024)
- [VMamba: Visual State Space Model](https://mlsanity.com/paper/2401.10166/) (2024)
- [Deep Reinforcement Learning for Autonomous Driving: A Survey](https://mlsanity.com/paper/2002.00444/) (2020)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [MetaFormer Is Actually What You Need for Vision](https://mlsanity.com/paper/2111.11418/) (2021)
- [U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation](https://mlsanity.com/paper/2401.04722/) (2024)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [VideoMamba: State Space Model for Efficient Video Understanding](https://mlsanity.com/paper/2403.06977/) (2024)
- [xLSTM: Extended Long Short-Term Memory](https://mlsanity.com/paper/2405.04517/) (2024)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [A Survey on Offline Reinforcement Learning: Taxonomy, Review, and Open Problems](https://mlsanity.com/paper/2203.01387/) (2022)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [Synthesizer: Rethinking Self-Attention in Transformer Models](https://mlsanity.com/paper/2005.00743/) (2020)
- [Online Decision Transformer](https://mlsanity.com/paper/2202.05607/) (2022)
- [Reinforcement learning based recommender systems: A survey](https://mlsanity.com/paper/2101.06286/) (2021)

## Built on by (most cited first, 7 of 7)

- [Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling](https://mlsanity.com/paper/2406.00079/) (2024)
- [A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks](https://mlsanity.com/paper/2410.22391/) (2024)
- [Skill Expansion and Composition in Parameter Space](https://mlsanity.com/paper/2502.05932/) (2025)
- [UNIDOOR: A Universal Framework for Action-Level Backdoor Attacks in Deep Reinforcement Learning](https://mlsanity.com/paper/2501.15529/) (2025)
- [TrojanTO: Action-Level Backdoor Attacks against Trajectory Optimization Models](https://mlsanity.com/paper/2506.12815/) (2025)
- [Analytic Energy-Guided Policy Optimization for Offline Reinforcement Learning](https://mlsanity.com/paper/2505.01822/) (2025)
- [Continual Task Learning through Adaptive Policy Self-Composition](https://mlsanity.com/paper/2411.11364/) (2024)

## Abstract (arXiv)

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