Attention or memory? Neurointerpretable agents in space and time

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.