Timezone: »

Working Memory Graphs
Ricky Loynd · Roland Fernandez · Asli Celikyilmaz · Adith Swaminathan · Matthew Hausknecht

Tue Jul 14 08:00 AM -- 08:45 AM & Tue Jul 14 09:00 PM -- 09:45 PM (PDT) @ None #None

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.

Author Information

Ricky Loynd (Microsoft Research)
Roland Fernandez (Microsoft Research)

Research interested include learning and reasoning with symbolic representations and lifelong learning.

Asli Celikyilmaz (Microsoft Research)
Adith Swaminathan (Microsoft Research)
Matthew Hausknecht (Microsoft Research)

More from the Same Authors