Research map / Reinforcement learning / Offline RL
Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?
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.
Investigates whether Mamba (Decision Mamba, DeMa) is compatible with transformer-style trajectory optimization for offline RL, finding short-context Transformer-like DeMa with hidden attention beats Decision Transformer on Atari and MuJoCo with far fewer parameters.
Problem
Transformer-based trajectory optimization for offline RL has large parameter counts and poor scalability to long sequences; it is unclear whether Mamba, a linear-time sequence model, is compatible with trajectory optimization, which typically uses short segments.
Approach
Analyzes Mamba's components in the offline RL setting and finds that short-context Transformer-like DeMa works best, that the hidden attention mechanism (not the RNN recurrence or position embeddings) is the key component, and that it works with either pre or post up-projection residual blocks without position embedding.
Key results
- Transformer-like DeMa averages 111.8 on 8 Atari games vs 62.2 for DT, with 1.7M vs 2.35M parameters.
- On MuJoCo gym tasks DeMa averages 78.5 normalized score vs 76.4 for DT, using 175.5K vs 726.2K parameters for medium datasets.
- RNN-like DeMa averages 67.3 on Atari, well below Transformer-like DeMa, and its inference time is independent of sequence length but training is costlier.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace DT with DeMa | Offline trajectory optimization · Atari 1% DQN-replay, average | normalized score: 62.2 -> 111.8 | Helped |
| Replace DT with DeMa | Offline trajectory optimization · MuJoCo, average | expert-normalized score: 76.4 -> 78.5 | Helped |
| Use RNN-like full-trajectory input instead of Transformer-like fixed-length input (ablation) | DeMa · Atari, average | normalized score: 111.8 -> 67.3 | Hurt |
| Use RNN-like full-trajectory input instead of Transformer-like fixed-length input (ablation) | DeMa · MuJoCo, average | expert-normalized score: 66.0 -> 61.1 | Hurt |
| Concatenate state, action, and RTG across embedding dimension (BL3D) rather than temporal dimension (B3LD) (ablation) | DeMa · Atari, average | normalized score: 111.8 -> 69.6 | Hurt |
| Replace pre up-projection residual blocks with post up-projection residual blocks (ablation) | DeMa · Atari, average | normalized score: 111.8 -> 109.2 | No clear effect |
| Replace pre up-projection residual blocks with post up-projection residual blocks (ablation) | DeMa · MuJoCo Gym, average | expert-normalized score: 78.5 -> 77.0 | No clear effect |
| Add position embedding (ablation) | DeMa · MuJoCo, average | expert-normalized score: 66.0 -> 63.0 | Hurt |
| Add position embedding (ablation) | DeMa · MuJoCo, average | parameters: 175.5K -> 431.5K | Hurt |
| Replace hidden attention weights with weights from another environment or randomized weights (ablation) | DeMa · Atari and MuJoCo experiments | normalized return | Hurt |
| Remove hidden attention by zeroing its parameters (ablation) | DeMa · Atari and MuJoCo experiments | normalized return | Hurt |
| Increase input sequence length beyond the optimal threshold (ablation) | Transformer-like DeMa · Atari and MuJoCo | performance | Mixed |
Surprising findings
- Hidden attention scores decay approximately exponentially with distance, so longer context windows bring no gains and can hurt Atari performance.
- Zeroing the hidden attention still leaves a fair level of performance, yet swapping or randomizing it hurts, and removing position embedding improves results.
Limitations
- RNN-like DeMa's memory capacity and performance in POMDPs or long-horizon non-Markovian tasks are not explored.
- Unclear whether DeMa suits multi-task or online RL settings.
- Atari RNN-like results were run only once, and Atari gains are within high variance on some games.
Experimental setup
Datasets: 1% DQN Replay Dataset (Atari) D4RL MuJoCo (HalfCheetah, Hopper, Walker2d) D4RL maze2d D4RL antmaze
Benchmarks: Atari (8 games, 1% DQN-replay) D4RL MuJoCo (gym locomotion) D4RL maze2d D4RL antmaze
Metrics: normalized return average score parameter count MACs training time per step GPU memory usage
Baselines: Decision Transformer Decision ConvFormer CQL BC RvS DS4 GDT QT
Builds on: Decision Transformer Mamba Decision S4 Hidden Attention of Mamba Quantum Transformer
Closest papers by meaning
- Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling (NeurIPS 2024 · similarity 0.91)
- Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL (NeurIPS 2024 · similarity 0.89)
- Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression (NeurIPS 2025 · similarity 0.86)
- When should we prefer Decision Transformers for Offline Reinforcement Learning? (ICLR 2024 · similarity 0.86)
- In-Dataset Trajectory Return Regularization for Offline Preference-based Reinforcement Learning (AAAI 2025 · similarity 0.86)
- Rethinking Decision Transformer via Hierarchical Reinforcement Learning (ICML 2024 · similarity 0.86)
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.