Research map / Architectures and efficiency / RNNs and normalization
Think Before You Act: Decision Transformers with Working Memory
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 DT-Mem, a Decision Transformer with an attention-based working memory matrix (write/erase/read via content addressing) plus LoRA-based memory-only fine-tuning, reporting better generalization and training efficiency on Atari and Meta-World.
Problem
Decision Transformer agents need massive data and compute to generalize across tasks, and training on new tasks can cause forgetting of earlier skills stored implicitly in parameters.
Approach
Adds an explicit working-memory matrix to a Decision Transformer, updated and read via attention-based content addressing (erase/add vectors), and adapts to new tasks by fine-tuning only low-rank (LoRA) adapters on the memory module.
Key results
- On four held-out Atari games, DT-Mem reaches 51.0%, 69.3%, 53.6%, 62.2% DQN-normalized scores versus MDT's 3.8%, 13.2%, 8.6%, 2.3% and RMDT's 22.3%, 22.9%, 17.6%, 27.7%.
- On Meta-World ML45, DT-Mem achieves test (FT) 0.95 versus HDT 0.92 and PDT 0.09, with 147K adaptation parameters (0.7%).
- DT-Mem reduces training time to 50 hours versus 200 hours for MDT-13M and 1600 hours for MDT-200M, per the paper's table.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add the DT-Mem working-memory module rather than using MDT | Decision Transformer · Alien, MsPacman, SpaceInvaders, and StarGunner held-out Atari games | DQN-normalized score: MDT: 3.8%, 13.2%, 8.6%, 2.3%; DT-Mem: 51.0% (± 32.2%), 69.3% (± 19.3%), 53.6% (± 29.0%), 62.2% (± 19.1%) | Helped |
| Add the DT-Mem working-memory module rather than using RMDT | Decision Transformer · Alien, MsPacman, SpaceInvaders, and StarGunner held-out Atari games | DQN-normalized score: RMDT: 22.3% (± 10.7%), 22.9% (± 8.9%), 17.6% (± 9.2%), 27.7% (± 11.5%); DT-Mem: 51.0% (± 32.2%), 69.3% (± 19.3%), 53.6% | Helped |
| Use DT-Mem instead of HDT | Meta-World ML45 · Meta-World ML45 train | score: 0.89±0.00 -> 0.92±0.00 | Helped |
| Use DT-Mem instead of HDT | Meta-World ML45 · Meta-World ML45 test (no-FT) | score: 0.12±0.01 -> 0.20±0.01 | Helped |
| Use DT-Mem instead of HDT | Meta-World ML45 · Meta-World ML45 test (FT) | score: 0.92±0.10 -> 0.95±0.10 | Helped |
| Replace LoRA adaptation with hyper-networks (ablation) | DT-Mem · Meta-World ML45 test (no-FT) | score: 0.20±0.01 -> 0.23±0.10 | Helped |
| Replace LoRA adaptation with hyper-networks (ablation) | DT-Mem · Meta-World ML45 test (FT) | score: 0.95±0.10 -> 0.81±0.15 | Hurt |
| Organize input by grouping states, actions, and rewards by type instead of Decision Transformer ordering (ablation) | DT-Mem · Atari raw scores across Alien, MsPacman, SpaceInvaders, and StarGunner | raw score: Alien: 211.9 vs 239.6; MsPacman: 637.1 vs 713.4; SpaceInvaders: 165.7 vs 171.2; StarGunner: 620.7 vs 709.3 | No clear effect |
| Use LoRA (PEFT) instead of full-parameter fine-tuning on all games (ablation) | DT-Mem · Alien, MsPacman, SpaceInvaders, and StarGunner | DQN-normalized score: PEFT: 127.4%, 130.8%, 100.8%, 158.3%; FFT-All: 113.9%, 77.1%, 73.4%, 40.6% | Helped |
| Add Gaussian noise with standard deviation 1 to Atari input frames (ablation) | DT-Mem · Alien, MsPacman, SpaceInvaders, and StarGunner | DQN-normalized score: 51.0%, 69.3%, 53.6%, 62.2% -> 35.6%, 56.1%, 40.0%, 34.6% | Hurt |
| Add Gaussian noise with standard deviation 2 to Atari input frames (ablation) | DT-Mem · Alien, MsPacman, SpaceInvaders, and StarGunner | DQN-normalized score: 51.0%, 69.3%, 53.6%, 62.2% -> 25.9%, 35.6%, 30.5%, 21.1% | Hurt |
| Fine-tune with LoRA rather than full-parameter fine-tuning on a single game (ablation) | DT-Mem · Alien, MsPacman, SpaceInvaders, and StarGunner | DQN-normalized score: PEFT: 127.4%, 130.8%, 100.8%, 158.3%; FFT-Single: 116.8%, 122.8, 86.8%, 55.7% | Mixed |
Surprising findings
- Fine-tuning DT-Mem on 50k steps of an unseen game outperforms training DT-20M on 200k-step data in the paper's comparison table.
- Full-parameter fine-tuning of memory was worse than LoRA on several games (e.g., StarGunner 40.6-55.7% vs 158.3% for PEFT).
Limitations
- Fine-tuning still needs a sizeable 10% of data, so sample efficiency of memory fine-tuning is limited.
- No theoretical grounding for why memory helps large models.
- Comparison with HDT on Meta-World relies on numbers reported in the HDT paper because code was unavailable.
Experimental setup
Datasets: Atari multi-game DQN trajectories (MDT dataset) Meta-World ML45
Benchmarks: Atari (4 held-out games: Alien, MsPacman, SpaceInvaders, StarGunner) Atari unseen games fine-tuning Meta-World ML45
Metrics: DQN-normalized score human-normalized IQM Meta-World train/test success
Baselines: MDT RMDT HDT PDT
Builds on: Decision Transformer Neural Turing Machine LoRA Multi-game Decision Transformer
Closest papers by meaning
- AUTOMEM: Automated Learning of Memory as a Cognitive Skill (NeurIPS 2026 · similarity 0.88)
- MemGen: Weaving Generative Latent Memory for Self-Evolving Agents (ICLR 2026 · similarity 0.86)
- MemSkill: Learning and Evolving Memory Skills for Self-Evolving Agents (NeurIPS 2026 · similarity 0.86)
- Remember the Decision, Not the Description: A Rate-Distortion Framework for Agent Memory (NeurIPS 2026 · similarity 0.86)
- Meta-Cognitive Memory Policy Optimization for Long-Horizon LLM Agents (NeurIPS 2026 · similarity 0.85)
- Benchmarking Agent Memory in Interdependent Multi-Session Agentic Tasks (ICML 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.