Research map / Reinforcement learning / Visual and representation RL
Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient
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 Drama, a Mamba-2-based model-based RL agent with a 7M-parameter world model and a dynamic frequency-based sampling method, evaluated on Atari100k.
Problem
Model-based RL needs accurate world models, but RNN world models struggle with long-term dependencies and transformer world models have O(n^2) cost; early-training model inaccuracy also leads to suboptimal policies.
Approach
Use a Mamba-2 state space sequence model as the world model's sequence backbone for linear-cost long-context learning, plus dynamic frequency-based sampling that balances how often transitions are used for world-model training versus behaviour-policy imagination.
Key results
- DramaXS (7M world model) reaches a normalised mean of 105% and median of 27% on Atari100k, comparable to IRIS (105%) and TWM (96%), below STORM (127%) and DreamerV3 (125%).
- On Atari100k, DramaXS achieves a 105% normalised mean vs 37% for a 12M DreamerV3XS.
- DFS raises the normalised mean to 105% versus 80% for uniform sampling, with similar median (27% vs 28%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Dynamic frequency-based sampling (DFS) instead of uniform sampling (ablation) | Drama XS (Mamba-2) · Atari100k (26 games) | Normalised Mean (%): 105 vs 80 | Helped |
| Dynamic frequency-based sampling (DFS) instead of uniform sampling (ablation) | Drama XS (Mamba-2) · Atari100k (26 games) | Normalised Median (%): 27 vs 28 | No clear effect |
| Dynamic frequency-based sampling (DFS) instead of uniform sampling (ablation) | Drama XS (Mamba-2) · Atari100k (26 games) | Per-game wins: DFS better in 11 games, worse in 2, tied in 13 | Mixed |
| Mamba-2 sequence model instead of Mamba (ablation) | Drama world model (DFS) · Krull, Boxing, Freeway, Kangaroo | Score (per-game learning curves) | Helped |
| Mamba-2 sequence model instead of Transformer/GRU/Mamba on long grid world (ablation) | Grid world sequence modelling (l=1664) · Grid World l=1664 | Training time per step (ms): 214 ms vs 299 (Mamba), 628 (GRU), OOM (Transformer) | Helped |
| Mamba or Mamba-2 instead of GRU on long grid world (ablation) | Grid world sequence modelling (l=1664) · Grid World l=1664 | Reconstruction error (%): 14.0 (Mamba), 14.2 (Mamba-2) vs 34.7 (GRU) | Helped |
| Drama (Mamba-2 world model, XS) instead of DreamerV3 at similar parameter count | Atari100k · Atari100k (26 games) | Normalised Mean (%): 105 vs 37 | Helped |
| Drama XS vs DreamerV3 (200M) at full scale | Atari100k · Atari100k (26 games) | Normalised Mean (%): 105 vs 125 | Hurt |
| Larger autoencoder and larger SSM hidden state (S model, ~3M more params) (ablation) | Drama · Atari100k subset (Breakout, BankHeist) | Score | Helped |
Surprising findings
- Drama underperforms badly on Breakout (7 vs 30 human-normalised limited) and PrivateEye, attributed to the autoencoder failing to reconstruct the ball.
Limitations
- Does not clearly outperform other world models on Atari100k; longer sequences do not translate into decisive gains.
- Weak performance on sparse-reward and small-object games such as Breakout, Jamesbond and PrivateEye.
- Hyperparameters for the many components were not extensively tuned due to compute cost.
Experimental setup
Datasets: Atari100k
Benchmarks: Atari100k custom grid world long-sequence prediction
Metrics: normalised mean score normalised median score training time memory usage reconstruction error
Baselines: SimPLe SPR TWM IRIS STORM DreamerV3 PPO DreamerV3XS Uniform sampling
Builds on: DreamerV3 STORM Mamba Mamba-2 IRIS
Closest papers by meaning
- Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL (NeurIPS 2024 · similarity 0.86)
- Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling (NeurIPS 2024 · similarity 0.86)
- MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning (ICLR 2024 · similarity 0.86)
- Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.85)
- Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining (ICLR 2025 · similarity 0.85)
- Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL (NeurIPS 2026 · similarity 0.85)
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.