Current browse context:
cs.LG
Change to browse by:
References & Citations
Computer Science > Machine Learning
Title: Working Memory Graphs
(Submitted on 17 Nov 2019 (v1), last revised 18 Aug 2020 (this version, v4))
Abstract: 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.
Submission history
From: Ricky Loynd [view email][v1] Sun, 17 Nov 2019 03:14:02 GMT (1311kb,D)
[v2] Thu, 20 Feb 2020 02:09:20 GMT (1665kb,D)
[v3] Wed, 8 Jul 2020 16:51:20 GMT (2789kb,D)
[v4] Tue, 18 Aug 2020 15:56:25 GMT (2760kb,D)
Link back to: arXiv, form interface, contact.