A neural network model that learns to encode and retrieve memories for spatial navigation

Numerous machine learning algorithms incorporate memory models that were inspired by episodic memory. However, almost all of them employ pre-defined and rigid mechanisms for the encoding and retrieval of memories. Usually, a pre-determined type of information is stored at each timestep regardless of whether anything interesting has happened. Here, we investigate a model based on a Memory-Augmented Neural Network, which learns autonomously what and when to store into, and retrieve from, an external memory buffer, while solving a navigation task in a simulated maze. The agent learns to navigate to an unmarked goal whose location changes at fixed intervals. As expected, the agent learns to store information in memory when it reaches the goal, and suppresses storage afterwards to avoid interference by irrelevant information. Surprisingly, even though the model receives only camera images as inputs, the information encoded in memory reflects the 2D spatial structure of the maze. Intriguingly, the outcome of memory retrieval already reflects the information about what action the agent will select next. We develop a geometrical theory which explains how the representations in the memory and the computations during retrieval give rise to the correct solutions to the navigation task. Our modelling results show how a pure memory structure, such as for example the hippocampus, can develop spatial and action representations. We believe that it is the structure of the spatial task that molds the encoding and retrieval strategy that the agent learns, and different tasks would eventually lead to different memory strategies and representations.