Research map / Architectures and efficiency / Attention and transformers
Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Decision Mamba-Hybrid (DM-H), an in-context RL method where a Mamba model generates sub-goal vectors from long across-episodic contexts and a transformer predicts actions conditioned on them, plus high-value offline sub-goals.
Problem
In-context RL with transformers suffers quadratic attention cost when across-episodic contexts make sequences long; pure Mamba replacement is efficient but less effective on RL tasks.
Approach
Hybrid architecture: Mamba processes long multi-episode context and emits a sub-goal every c steps, which conditions a transformer that predicts c actions; high-value states selected by distance-weighted accumulated reward serve as extra sub-goals to tie the transformer to the sub-goals.
Key results
- On D4RL, DM-H averages 87.53 total normalized score versus 82.84 for AD (Transformer), 77.69 for DT, and 84.83 for TD3+BC.
- DM-H online testing on large Grid World is about 12x faster than AD (Transformer) per the appendix table (e.g., 5.12 vs 45.08 minutes for Large Darkroom); on D4RL the paper claims about 28x speedup.
- DM-H handles tasks with twice the horizon of DT and DM on Tmaze before running out of GPU memory (20k vs 10k context).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the Transformer backbone with Mamba | Algorithm Distillation · D4RL | effectiveness | Mixed |
| Replace the Transformer backbone with Mamba | Algorithm Distillation · D4RL | training time: 37.10 -> 28.56; 22.23 -> 18.15; 33.77 -> 26.25 | Helped |
| Use DM-H instead of TD3+BC | D4RL · Med-Expert HalfCheetah | normalized score: 96.59 -> 96.21 | Hurt |
| Use DM-H instead of TD3+BC | D4RL · Med HalfCheetah | normalized score: 48.93 -> 45.45 | Hurt |
| Use DM-H instead of TD3+BC | D4RL · Med-Replay HalfCheetah | normalized score: 45.84 -> 45.26 | Hurt |
| Use DM-H instead of TD3+BC | D4RL · Total Average | normalized score: 84.83 -> 87.53 | Helped |
| Add valuable sub-goals (ablation) | DM-H · Grid World | mean episode return | Helped |
| Add valuable sub-goals (ablation) | DM-H · D4RL | mean episode return | Helped |
| Set the transformer context length to c in the appropriate range (ablation) | DM-H · Grid World and other evaluated tasks | performance | Mixed |
Surprising findings
- Replacing the transformer with Mamba in AD did not significantly improve effectiveness and was slightly worse on most D4RL datasets, though it was more efficient.
Limitations
- Performance depends on the hyperparameter c (sub-goal interval), which must be hand-tuned; the authors propose adaptive c as future work.
- Evaluated mainly on Grid World, Tmaze, and D4RL locomotion; no evaluation on broader or real-world tasks.
Experimental setup
Datasets: Grid World (Darkroom, Darkroom Hard, Dark Key-to-Door and large variants) Tmaze D4RL HalfCheetah, Hopper, Walker2d (medium, medium-replay, medium-expert)
Benchmarks: Grid World Tmaze D4RL
Metrics: normalized episode return online testing time offline training time
Baselines: AMAGO DPT Algorithm Distillation (Transformer) AD (Mamba) TD3+BC TD3 BC-10% Decision Transformer DM (Decision Mamba)
Builds on: Decision Transformer Algorithm Distillation Mamba S4
Closest papers by meaning
- Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL (NeurIPS 2024 · similarity 0.92)
- Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning? (NeurIPS 2024 · similarity 0.91)
- Rethinking Decision Transformer via Hierarchical Reinforcement Learning (ICML 2024 · similarity 0.90)
- In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.89)
- MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning (ICLR 2024 · similarity 0.89)
- Decision ConvFormer: Local Filtering in MetaFormer is Sufficient for Decision Making (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.