# Working Memory Graphs

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

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 17 Nov 2019.

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

- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [MoCoDA: Model-based Counterfactual Data Augmentation](https://mlsanity.com/paper/2210.11287/) (2022)
- [Transformers are Meta-Reinforcement Learners](https://mlsanity.com/paper/2206.06614/) (2022)
- [Relational Attention: Generalizing Transformers for Graph-Structured Tasks](https://mlsanity.com/paper/2210.05062/) (2022)
- [On Transforming Reinforcement Learning by Transformer: The Development Trajectory](https://mlsanity.com/paper/2212.14164/) (2022)
- [Factored Adaptation for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2203.16582/) (2022)
- [Measuring Sample Efficiency and Generalization in Reinforcement Learning Benchmarks: NeurIPS 2020 Procgen Benchmark](https://mlsanity.com/paper/2103.15332/) (2021)
- [Recurrent Action Transformer with Memory](https://mlsanity.com/paper/2306.09459/) (2023)
- [A System for Morphology-Task Generalization via Unified Representation and Behavior Distillation](https://mlsanity.com/paper/2211.14296/) (2022)
- [MemoNav: Working Memory Model for Visual Navigation](https://mlsanity.com/paper/2402.19161/) (2024)
- [Explain My Surprise: Learning Efficient Long-Term Memory by Predicting Uncertain Outcomes](https://mlsanity.com/paper/2207.13649/) (2022)
- [Consistent Dropout for Policy Gradient Reinforcement Learning](https://mlsanity.com/paper/2202.11818/) (2022)
- [PLEX: Making the Most of the Available Data for Robotic Manipulation Pretraining](https://mlsanity.com/paper/2303.08789/) (2023)
- [Snowflake: Scaling GNNs to High-Dimensional Continuous Control via Parameter Freezing](https://mlsanity.com/paper/2103.01009/) (2021)
- [Automated Metaheuristic Algorithm Design with Autoregressive Learning](https://mlsanity.com/paper/2405.03419/) (2024)
- [On learning history based policies for controlling Markov decision processes](https://mlsanity.com/paper/2211.03011/) (2022)
- [Compositional Generalization in Grounded Language Learning via Induced Model Sparsity](https://mlsanity.com/paper/2207.02518/) (2022)
- [GriddlyJS: A Web IDE for Reinforcement Learning](https://mlsanity.com/paper/2207.06105/) (2022)
- [MemoNav: Selecting Informative Memories for Visual Navigation](https://mlsanity.com/paper/2208.09610/) (2022)
- [Inductive Biases for Zero-shot Systematic Generalization in Language-informed Reinforcement Learning](https://mlsanity.com/paper/2501.15270/) (2025)
- [Multi-Task Learning for User Engagement and Adoption in Live Video Streaming Events](https://mlsanity.com/paper/2106.10305/) (2021)
- [AIGenC: An AI generalisation model via creativity](https://mlsanity.com/paper/2205.09738/) (2022)
- [Geometry of naturalistic object representations in recurrent neural network models of working memory](https://mlsanity.com/paper/2411.02685/) (2024)

## Abstract (arXiv)

Transformers have increasingly outperformed gated RNNs in obtaining new state-of-the-art results on supervised tasks involving text sequences. Inspired by this trend, we study the question of how Transformer-based models can improve the performance of sequential decision-making agents. We present the Working Memory Graph (WMG), an agent that employs multi-head self-attention to reason over a dynamic set of vectors representing observed and recurrent state. We evaluate WMG in three environments featuring factored observation spaces: a Pathfinding environment that requires complex reasoning over past observations, BabyAI gridworld levels that involve variable goals, and Sokoban which emphasizes future planning. We find that the combination of WMG's Transformer-based architecture with factored observation spaces leads to significant gains in learning efficiency compared to baseline architectures across all tasks. WMG demonstrates how Transformer-based models can dramatically boost sample efficiency in RL environments for which observations can be factored.