Predictive auxiliary objectives in deep RL mimic learning in the brain

TL;DR

This paper explores how predictive auxiliary objectives in deep reinforcement learning enhance representation learning, stabilize training, and mimic neural activity patterns observed in the brain, especially in the hippocampus, visual cortex, and striatum.

Contribution

It demonstrates that predictive objectives improve learning stability and transfer in resource-limited RL architectures and establish parallels between RL representations and brain regions.

Findings

Predictive objectives stabilize learning in resource-limited models.

Longer predictive horizons support better transfer of representations.

RL representational changes resemble neural activity in hippocampus, visual cortex, and striatum.

Abstract

The ability to predict upcoming events has been hypothesized to comprise a key aspect of natural and machine cognition. This is supported by trends in deep reinforcement learning (RL), where self-supervised auxiliary objectives such as prediction are widely used to support representation learning and improve task performance. Here, we study the effects predictive auxiliary objectives have on representation learning across different modules of an RL system and how these mimic representational changes observed in the brain. We find that predictive objectives improve and stabilize learning particularly in resource-limited architectures, and we identify settings where longer predictive horizons better support representational transfer. Furthermore, we find that representational changes in this RL system bear a striking resemblance to changes in neural activity observed in the brain across various experiments. Specifically, we draw a connection between the auxiliary predictive model of the RL system and hippocampus, an area thought to learn a predictive model to support memory-guided behavior. We also connect the encoder network and the value learning network of the RL system to visual cortex and striatum in the brain, respectively. This work demonstrates how representation learning in deep RL systems can provide an interpretable framework for modeling multi-region interactions in the brain. The deep RL perspective taken here also suggests an additional role of the hippocampus in the brain -- that of an auxiliary learning system that benefits representation learning in other regions.