Skip to content


Switch branches/tags

Name already in use

A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?

Latest commit


Git stats


Failed to load latest commit information.
Latest commit message
Commit time

Working Memory Graphs

This repository accompanies the ICML 2020 paper Working Memory Graphs by Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan and Matthew Hausknecht. (Video)

WMG is a Transformer-based RL agent that attends to a dynamic set of vectors representing observed and recurrent state.


(The code has been tested with Python 3.6, on Windows and Linux.)

  • Create a new virtual environment.
  • Activate the virtual environment.
  • Install PyTorch 1.3.1.
  • Install BabyAI (from any directory):
  • Clone the wmg_agent repository.
  • cd wmg_agent
  • pip install -r requirements.txt

Running experiments

Execute all run commands from the wmg_agent directory, using this format:

python <runspec>

Execute the runspecs in the specs directory to reproduce results from the paper. They can also be used as examples for new experiments.


If using this code in your work, please cite as follows:

  author = {Loynd, Ricky and Fernandez, Roland and Celikyilmaz, Asli and Swaminathan, Adith and Hausknecht, Matthew},
  booktitle = {Proceedings of Machine Learning and Systems 2020},
  pages = {2928--2937},
  title = {Working Memory Graphs},
  year = {2020}

This repo's implementation of Sokoban was derived in part from that of If using the Sokoban environment in your work, please honor that source's license, and cite as follows:

  author = {Schrader, Max-Philipp B.},
  title = {gym-sokoban},
  year = {2018},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{}},
  commit = {#CommitId}

For convenience, we provide some of the predefined Boxoban levels from If using this data in your work, please honor that source's license, and cite as follows:

  author = {Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sebastien Racaniere, Theophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, Greg Wayne, David Silver, Timothy Lillicrap, Victor Valdes},
  title = {An investigation of Model-free planning: boxoban levels},
  howpublished= {},
  year = "2018",


This project welcomes contributions and suggestions. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit

When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact with any additional questions or comments.