[Paper Review] Unsupervised Predictive Memory in a Goal-Directed Agent

Greg Wayne, Chia-Chun Hung+22 more|arXiv (Cornell University)|Mar 28, 2018

Reinforcement Learning in Robotics14 references148 citations

TL;DR

MERLIN is an AI agent that uses a memory-based predictor trained with unsupervised predictive modeling to solve highly partial-observability tasks, outperforming standard memory RL agents on psychology/neuroscience benchmarks. It builds compact state variables, stores them in memory, and uses return prediction to shape representations and memory use.

ABSTRACT

Animals execute goal-directed behaviours despite the limited range and scope of their sensors. To cope, they explore environments and store memories maintaining estimates of important information that is not presently available. Recently, progress has been made with artificial intelligence (AI) agents that learn to perform tasks from sensory input, even at a human level, by merging reinforcement learning (RL) algorithms with deep neural networks, and the excitement surrounding these results has led to the pursuit of related ideas as explanations of non-human animal learning. However, we demonstrate that contemporary RL algorithms struggle to solve simple tasks when enough information is concealed from the sensors of the agent, a property called "partial observability". An obvious requirement for handling partially observed tasks is access to extensive memory, but we show memory is not enough; it is critical that the right information be stored in the right format. We develop a model, the Memory, RL, and Inference Network (MERLIN), in which memory formation is guided by a process of predictive modeling. MERLIN facilitates the solution of tasks in 3D virtual reality environments for which partial observability is severe and memories must be maintained over long durations. Our model demonstrates a single learning agent architecture that can solve canonical behavioural tasks in psychology and neurobiology without strong simplifying assumptions about the dimensionality of sensory input or the duration of experiences.

Motivation & Objective

• Motivate memory-enabled agents to operate under partial observability where sensors miss important information.
• Develop MERLIN, a memory-based predictor that compresses observations into state variables and stores them for predictive use.
• Demonstrate that unsupervised predictive modeling guides memory formation and improves performance on psychology/neuroscience-inspired tasks.

Proposed method

• Introduce MERLIN, an agent architecture combining a memory-based predictor (MBP) and a policy with read/write memory mechanisms.
• MBP encodes multi-modal observations into low-dimensional state variables z via a variational autoencoder–like framework and stores them in memory.
• Use a prior p(z_t|z_{1:t-1},a_{1:t-1}) and a posterior q(z_t|z_{1:t-1},a_{1:t-1},o_t) to sample z_t and update memory.
• Train the MBP with a variational lower bound (VLB) consisting of reconstruction losses for multiple modalities and a KL term between p and q, plus a return-prediction decoder that guides z_t toward reward-relevant information.
• Decouple the MBP optimization from the policy to ensure representational learning is driven by predictive modeling rather than reward alone.
• In addition to the MBP, employ retroactive memory updating to attach future information to past memories and explore how return prediction shapes representations.

Experimental results

Research questions

• RQ1 Can unsupervised predictive memory enable memory-based agents to solve tasks with long delays between observations and decisions?
• RQ2 Does compressing sensory inputs into state variables via predictive modeling improve memory formation and retrieval compared to end-to-end memory RL systems?
• RQ3 Do memory reads specialize to information at different temporal distances from goals, enabling hierarchical goal-directed behavior?
• RQ4 Can MERLIN solve one-shot navigation and other psychology/neuroscience-inspired tasks from raw sensory data without strong simplifying assumptions?

Key findings

• MERLIN solves memory-demanding tasks (e.g., Memory game, large-environment navigation) where RL-LSTM and RL-MEM struggle or fail.
• The MBP compresses high-dimensional sensory input to order 10^2 state variables, retaining task-relevant information via predictive modeling.
• Memory reads from MBP specialize to memories formed at different distances from the goal, supporting hierarchical, goal-directed strategies.
• MERLIN demonstrates rapid allocentric goal localization and robust return prediction that guides memory use and planning.
• On a battery of tasks, including arbitrary visuomotor mapping and rapid reward valuation, MERLIN outperforms end-to-end memory baselines and even reaches or exceeds human performance in some cases.
• Latent learning and retroactive memory updating enable MERLIN to recall and utilize information acquired earlier when needed, even beyond conventional backpropagation-through-time windows.