Source-linked AI summary
Unsupervised Predictive Memory in a Goal-Directed Agent
Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska-Barwinska, Jack Rae, Piotr Mirowski, Joel Z. Leibo, Adam Santoro, Mevlana Gemici, Malcolm Reynolds, Tim Harley, Josh Abramson, Shakir Mohamed, Danilo Rezende, David Saxton, Adam Cain, Chloe Hillier, David Silver, Koray Kavukcuoglu, Matt Botvinick, Demis Hassabis, Timothy Lillicrap
TL;DR
Partially observed tasks can defeat contemporary reinforcement-learning agents, and memory alone is insufficient when the wrong information is stored. MERLIN combines predictive modeling with memory and reinforcement learning, solving severely partially observed tasks, including long-delay behavioural tasks.
Problem
Contemporary reinforcement-learning agents struggle with simple tasks when critical information is concealed from their sensors, especially tasks involving long delays.
Method
MERLIN is an integrated agent architecture that uses predictive modeling to guide memory formation and compress observations into state representations for storage.
Results
MERLIN found an optimal strategy where RL-LSTM and MEM were unable to solve the task, and reached higher performance than comparison agents and professional human testers.
Takeaways & Limitations
A single MERLIN learning architecture can solve canonical behavioural tasks despite severe partial observability and long memory durations.
Takeaways & Limitations
Performance can fail when the amount of information that must be recalled is sufficiently large or when perceptual representations are trained for memory storage solely by end-to-end reinforcement learning.
Abstract
from arXiv · showhide
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.
Introduction
The paper argues that partial observability makes memory-dependent tasks difficult for contemporary reinforcement-learning agents, and that memory must store task-relevant information in suitable representations. MERLIN uses unsupervised predictive modeling to form and retrieve memories, solving diverse partially observed tasks in 3D environments and behavioral paradigms.
- Problem: Contemporary reinforcement-learning agents can fail on simple partially observed tasks when the information to recall is sufficiently large.RL-LSTM can perform well in partially observed environments but may fail to train when the required memory load is high.
- Approach: MERLIN is an integrated agent architecture that combines reinforcement learning, external memory, state estimation, and inference for partially observed environments.Its memory formation is guided by unsupervised prediction rather than policy learning alone.
- Approach: MERLIN’s predictive memory produced compressed state representations that retained information critical to task-related computations.The paper describes predictive modeling as a way to build compressed representations while preserving task-relevant information.
- Results: On the Memory Game, RL-LSTM and RL-MEM were unable to solve the task, whereas MERLIN found an optimal strategy by retrieving a previously observed matching card before acting.The retrieved memory corresponded to the matching card one time step before MERLIN flipped it.
- Results: MERLIN outperformed comparison agents and professional human testers on navigation tasks, while forming allocentric goal representations and improving return predictions from memory.Goal position decoding from MERLIN’s state and memory reached 1.67 units in Manhattan distance, and return predictions rose even when the goal was out of view.
Correspondence
The model combines a memory-based predictor with a policy, using encoded state variables, recurrent processing, external memory, and multiple decoders. Its memory interface reads content by learned keys and writes interface-generated data for later use.
- MERLIN’s memory-based predictor uses encoders and decoders to represent state variables and predict observations, rewards, actions, text, velocity, and returns.The encoder processes distinct input elements into a concatenated representation, while decoder networks reconstruct or predict corresponding outputs.
- Encoders: The image encoder downsamples 64 × 64 × 3 observations through six ResNet blocks to 8 × 8 × 64 tensors before reducing them to 500 dimensions.The resulting representation is flattened and passed through a linear layer with a tanh nonlinearity.
- Return prediction decoder and value function: The return decoder combines a policy value function with a state-action advantage function to produce a return prediction.The value function receives the latent variable and policy logits, while the advantage function receives the latent variable and action.
- The external memory is a blank-initialised matrix with 2 × |z| columns, storing a history of latent state variables for content-based access.Memory rows are read using keys generated from the recurrent network and written using an interface-generated data vector.
- The MBP and policy each use a two-layer deep LSTM, with the MBP receiving state variables, actions, and memory reads while the policy receives state variables.The MBP output is transformed into a memory-interface vector containing read keys and other interface components.
3 Derivation of the Variational Lower Bound
The derivation uses a variational lower bound on the log marginal likelihood for temporal generative models. It extends the bound to partition target variables and to combine incrementally observable information with delayed future-reward information.
- The log marginal likelihood of a probabilistic generative model is lower-bounded by an expectation under an approximating posterior.This lower bound follows from Jensen’s inequality.
- The temporal model factorises joint observations and latent variables across time, using shorthand for latent-variable sequences.The notation defines z0:t as the sequence of latent states and z0:−1 as the empty set.
- Partitioning target variables into x and y permits an approximate posterior that does not condition on one of the variables.The temporal derivation applies the corresponding conditional structure over time.
- The resulting loss combines prediction of incrementally observable information with delayed information such as the sum of future rewards.The prior can additionally be conditioned on variables such as actions.
4 Cost Functions for the MBP and Policy
The MBP and policy are optimised with separate objectives: a variational prediction loss for the MBP and policy-gradient learning for action selection. The MBP loss combines decoder likelihoods, return prediction, and KL regularisation, while policy updates use advantage estimates and entropy regularisation.
- The policy and memory-based predictor have independent parameters, separated by a gradient stop between the policy and state variable zt.The policy is therefore not updated by gradients from the MBP’s loss functions.
- MBP cost function: The MBP loss uses a variational lower bound whose conditional log-likelihood is factorised across decoder-specific prediction losses.These losses cover action, velocity, reward, return, image reconstruction, and text prediction.
- Return targets: Truncation windows divide long episodes into segments so return targets can use bootstrapping within a window or accumulated rewards at episode end.The return target uses a value estimate when the window ends before the episode and discounted rewards when it reaches the terminal point.
- MBP cost function: The MBP objective includes an analytical KL divergence between diagonal-Gaussian posterior and prior distributions.Each time step contributes to the loss, which is divided by the number of pixel-channels to reduce gradient magnitude.
- Policy gradient: The policy gradient uses Generalised Advantage Estimation with discount and bootstrapping parameters γ and λ.An additional loss term increases the entropy of the policy’s action distribution.
5 Comparison Models
The comparison models replace MERLIN’s predictive-memory design with recurrent or externally controlled memory mechanisms, and agents are trained across navigation and memory-dependent tasks.
- Comparison models: RL-LSTM uses a deep LSTM with separate value-function and action-distribution heads, trained with the policy objective in Eq. 8.
- Comparison models: RL-MEM replaces the policy LSTM with a Differentiable Neural Computer, introducing quadratic computational complexity as memory size scales.
- Comparison models: RL-MEM directly stores an LSTM-produced write vector and omits MERLIN’s retroactive memory updating.
- Training: All models were trained with truncated backpropagation through time using 192 parallel worker threads and asynchronous gradient updates.
- Tasks: The task battery included partially observed navigation, arbitrary visuomotor mapping, object-value learning, episodic water mazes, instruction following, and latent learning in a T-maze.
8 Lesion Comparison Experiments
Lesion experiments removed memory components, predictive objectives, or retroactive updates, while analyses examined representations, memory reads, and predictive accuracy.
- Lesion conditions: The lesions separately removed external memory, observation reconstruction, return prediction, or retroactive memory updating.
- Evaluation: Time-to-goal comparisons used 200 common test mazes and retained episodes in which all three models reached the goal repeatedly.
- Representation analyses: Goal representations were decoded from recurrent or encoder features without backpropagating decoder error into the agent networks.
- Prediction analyses: MERLIN’s value accuracy was assessed by the absolute difference between its value estimate and empirical discounted return.
- Memory-read analyses: One memory read head focused near the agent, while two others read positions closer to the goal; one consistently read slightly ahead toward the goal.
- Reward-memory analyses: Retroactive memory was evaluated for whether it preserved the positive or negative reward value of previously consumed objects.
10 Human Testing
Human testing compared agent performance with human testers and examined transfer, predictive visual representations, gradient pathways, and additional model comparisons.
- Human testing: Two professional game testers practiced each navigation task for 15 minutes and then completed 15 consecutive 90-second runs.
- Human testing: The Arbitrary Visuomotor Mapping task was performed by its creator for six 90-second blocks on naturalistic and synthetic transfer images.
- Representation comparison: MERLIN’s memory-based predictor representations supported more accurate decoding of egocentric goal position than RL-model or encoder representations.
- Performance comparison: MERLIN performed at or above human level on goal-finding tasks after joint training across tasks.
- Gradient analyses: Removing or introducing several gradient pathways did not obviously change performance on the Goal-Finding task.
- Additional comparison: UNREAL performed worse than MERLIN and the comparison models in the additional comparison.
- Credit assignment: Increasing temporal credit-assignment windows improved RL-LSTM and RL-MEM performance, but their performance remained below MERLIN’s.