Attention or memory? Neurointerpretable agents in space and time

TL;DR

This paper demonstrates that self-attention mechanisms in deep reinforcement learning agents improve robustness to noise, enable working-memory functions, and enhance interpretability of decision features in dynamic environments.

Contribution

It introduces a self-attention based model in deep RL that improves noise robustness, supports working-memory, and offers interpretability of feature selection and composition.

Findings

Self-attention increases robustness to irrelevant features.

The model can implement a transient working-memory.

Semantic observations reveal feature selection and composition.

Abstract

In neuroscience, attention has been shown to bidirectionally interact with reinforcement learning (RL) processes. This interaction is thought to support dimensionality reduction of task representations, restricting computations to relevant features. However, it remains unclear whether these properties can translate into real algorithmic advantages for artificial agents, especially in dynamic environments. We design a model incorporating a self-attention mechanism that implements task-state representations in semantic feature-space, and test it on a battery of Atari games. To evaluate the agent's selective properties, we add a large volume of task-irrelevant features to observations. In line with neuroscience predictions, self-attention leads to increased robustness to noise compared to benchmark models. Strikingly, this self-attention mechanism is general enough, such that it can be naturally extended to implement a transient working-memory, able to solve a partially observable maze task. Lastly, we highlight the predictive quality of attended stimuli. Because we use semantic observations, we can uncover not only which features the agent elects to base decisions on, but also how it chooses to compile more complex, relational features from simpler ones. These results formally illustrate the benefits of attention in deep RL and provide evidence for the interpretability of self-attention mechanisms.