# Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL

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

## Key facts

- Cited by 15 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 52 of its 63 references that are in the map.
- First posted on arXiv 8 Jun 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 52)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [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)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [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)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [On the Parameterization and Initialization of Diagonal State Space Models](https://mlsanity.com/paper/2206.11893/) (2022)
- [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022)
- [VideoMamba: State Space Model for Efficient Video Understanding](https://mlsanity.com/paper/2403.06977/) (2024)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [RvS: What is Essential for Offline RL via Supervised Learning?](https://mlsanity.com/paper/2112.10751/) (2021)

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

- [Optimus-1: Hybrid Multimodal Memory Empowered Agents Excel in Long-Horizon Tasks](https://mlsanity.com/paper/2408.03615/) (2024)
- [Spatial-Temporal Graph Diffusion Policy with Kinematic Modeling for Bimanual Robotic Manipulation](https://mlsanity.com/paper/2503.10743/) (2025)
- [Mamba-FSCIL: Dynamic Adaptation with Selective State Space Model for Few-Shot Class-Incremental Learning](https://mlsanity.com/paper/2407.06136/) (2024)
- [Overflow Prevention Enhances Long-Context Recurrent LLMs](https://mlsanity.com/paper/2505.07793/) (2025)
- [Flow Actor-Critic for Offline Reinforcement Learning](https://mlsanity.com/paper/2602.18015/) (2026)
- [Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient](https://mlsanity.com/paper/2410.08893/) (2024)
- [On the Expressivity of Selective State-Space Layers: A Multivariate Polynomial Approach](https://mlsanity.com/paper/2502.02209/) (2025)
- [Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning](https://mlsanity.com/paper/2412.14834/) (2024)
- [RLBenchNet: The Right Network for the Right Reinforcement Learning Task](https://mlsanity.com/paper/2505.15040/) (2025)
- [Mamba-FETrack V2: Revisiting State Space Model for Frame-Event based Visual Object Tracking](https://mlsanity.com/paper/2506.23783/) (2025)
- [HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning](https://mlsanity.com/paper/2605.05951/) (2026)
- [Decoupling Return-to-Go for Efficient Decision Transformer](https://mlsanity.com/paper/2601.15953/) (2026)
- [Decision MetaMamba: Enhancing Selective SSM in Offline RL with Heterogeneous Sequence Mixing](https://mlsanity.com/paper/2408.10517/) (2024)
- [Decision MetaMamba: Enhancing Selective SSM in Offline RL with Heterogeneous Sequence Mixing](https://mlsanity.com/paper/2602.19805/) (2026)
- [Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer](https://mlsanity.com/paper/2605.06104/) (2026)

## Abstract (arXiv)

While the conditional sequence modeling with the transformer architecture has demonstrated its effectiveness in dealing with offline reinforcement learning (RL) tasks, it is struggle to handle out-of-distribution states and actions. Existing work attempts to address this issue by data augmentation with the learned policy or adding extra constraints with the value-based RL algorithm. However, these studies still fail to overcome the following challenges: (1) insufficiently utilizing the historical temporal information among inter-steps, (2) overlooking the local intrastep relationships among return-to-gos (RTGs), states, and actions, (3) overfitting suboptimal trajectories with noisy labels. To address these challenges, we propose Decision Mamba (DM), a novel multi-grained state space model (SSM) with a self-evolving policy learning strategy. DM explicitly models the historical hidden state to extract the temporal information by using the mamba architecture. To capture the relationship among RTG-state-action triplets, a fine-grained SSM module is designed and integrated into the original coarse-grained SSM in mamba, resulting in a novel mamba architecture tailored for offline RL. Finally, to mitigate the overfitting issue on noisy trajectories, a self-evolving policy is proposed by using progressive regularization. The policy evolves by using its own past knowledge to refine the suboptimal actions, thus enhancing its robustness on noisy demonstrations. Extensive experiments on various tasks show that DM outperforms other baselines substantially.
