Unsupervised Predictive Memory in a Goal-Directed Agent

Introduction

• End-to-end RL agents can fail on behavioral tasks involving long delays, including navigation to previously visited goals (14–16), rapid reward valuation (17), and latent learning (18,19).These failures arise even though external-memory systems have achieved successes on other tasks (10–13).Source: Introduction · pp. 2–3
• MERLIN is an integrated agent for partially observed virtual-reality environments, developed as a pragmatic approach to predictive modeling and hippocampal memory.The introduction links its inspiration to predictive sensory coding (20–22), Gluck and Myers’ hippocampal representation theory (23, 24), and temporal context and successor representations (25–27).Source: Introduction · pp. 3, 7
• MERLIN combines external memory, reinforcement learning, and state-estimation models, with memory formation guided by unsupervised prediction (20–22).Its memory-based predictor compresses observations into low-dimensional state variables, stores them, and uses them to make predictions from past observations.Source: Introduction · pp. 3, 5
• MERLIN finds solutions to canonical psychology and neuroscience tasks that pose severe challenges to existing AI agents.The authors present it as a step beyond end-to-end RL for studying memory in computational agents.Source: Introduction · p. 3
• External memory alone does not ensure high-fidelity storage: reward-driven end-to-end learning can fail when tasks require perceptual memory or long delays between encoding and action.Policy-gradient learning uses a fixed credit-assignment window τ, and models struggle when the needed storage-and-retrieval timescale exceeds it.Source: Introduction · pp. 4–5
• The proposed thesis is that predictive modeling should build compressed representations of the environment that memory stores directly for the policy to use.In MERLIN, the memory-based predictor models state variables and memory contents, while the policy receives those variables and contents downstream.Source: Introduction · p. 5

Predictive memory model

• MERLIN encodes multimodal observations into latent states, using a history-conditioned prior and an observation-conditioned posterior to update its estimate.The policy reads memory but does not write to it; it is the system component trained conventionally for the task.Source: Introduction · p. 7
• The memory model learns a predictive world model by optimizing a variational lower bound that combines sensory reconstruction with prior–posterior agreement.The KL-divergence term encourages the predictive prior to remain consistent with the posterior after new observations.Source: Introduction · pp. 7–8

Reward-shaped representations

• Return prediction shapes MERLIN’s compressed state to retain task-relevant information while reducing sensory input from order 10^4 dimensions to order 10^2.The return decoder predicts task return from the latent state, encouraging it to preserve information relevant to reward.Source: Introduction · p. 8

Retroactive memory updating

• MERLIN’s retroactive memory updating stores a representation of later events with each state, allowing a landmark memory to be associated with a subsequently experienced goal.The paper illustrates this with navigation, where a waypoint landmark is stored alongside information about a later goal.Source: Introduction · p. 8

Navigation

• MERLIN learned faster and reached higher performance than comparison agents and professional human testers on partially observed navigation tasks.Across repeated visits, it returned to the goal in less time as it learned the environment layout.Source: Introduction · p. 9, Figure 2 · p. 10

Memory game

• On the Memory game, RL-LSTM and RL-MEM could not solve the task, whereas MERLIN found an optimal strategy.Its policy retrieved a memory of a previously seen matching card immediately before flipping that card.Source: Introduction · p. 9, Figure 2 · p. 10

Navigation analysis

• From its memory-model state, MERLIN rapidly formed an allocentric representation of goal location and anticipated the goal in its return predictions even when it was out of view.Its return-prediction error on approach was lower than the analogous value-function prediction errors of RL-LSTM and RL-MEM.Source: Figure 3 · p. 11, Introduction · p. 12
• Performance tolerated a range of return-prediction cost weights, but very low or high weights impaired it by making states insensitive to reward or other sensors, respectively.Higher weights also made return prediction explain more variance in the state, while lower weights produced cleaner reconstructed images with more retained information.Source: Figure 3 · p. 11, Introduction · p. 12

Predictive representations and read behavior

• Return-prediction gradients through the state supported goal-centered receptive fields, while three memory read heads robustly specialized across distances from the goal.The heads read memories linked to locations ahead, the goal, and important sub-goals such as nearby doorways.Source: Introduction · pp. 12–13

Arbitrary Visuomotor Mapping

• In Arbitrary Visuomotor Mapping, MERLIN performed above human level with essentially no errors and generalized immediately to novel synthetic images.It learned to distinguish complex images through unsupervised mechanisms, including previously unseen images with different statistics.Source: Introduction · p. 13, Figure 4 · p. 14

Rapid Reward Valuation

• MERLIN learned to probe and remember object values, focus reads on matching object memories, and predict future values; first-approach memory rows decoded consumption reward with 93% accuracy.The decoding result used N = 25,000 and 5-fold cross-validation.Source: Introduction · pp. 13, 15, Figure 4 · p. 14

Episodic Water Mazes and Transient Instructions

• Across six minutes and three water mazes, MERLIN stored and retrieved maze-specific information to relocate hidden platforms, while its state and memory readouts retained transient instructions throughout episodes.Instruction tokens were decodable on held-out data at any point in an episode, even after the instruction was no longer present.Source: Figure 4 · p. 14, Introduction · p. 15

Latent Learning

• In latent learning, MERLIN used phase-1 observations to find a cued object after a distractor delay longer than its backpropagation-through-time window.Its read-only policy recalled the earlier object observation, and memory readouts decoded the correct goal arm.Source: Figure 5 · p. 16, Introduction · p. 17

Architecture and ablations

• Removing memory, sensory or return prediction, or retroactive memory updating harmed performance or data efficiency; alternative policy-network arrangements performed comparably but obscured module independence.The gradient block between policy and memory was not necessary to the predictive-modeling principle.Source: Introduction · p. 17

Cross-task performance

• Across tasks, MERLIN performed qualitatively better than comparison models, which often failed to store or recall environmental information.Figure 4 summarizes learning curves for the task battery, and Figure 5 shows the latent-learning comparison.Source: Figure 4 · p. 14, Figure 5 · p. 16, Introduction · p. 17

Long-term memory limitation

• For long-lived agents, end-to-end optimization over very long intervals would require keeping an exact record of network states, which the paper describes as practically prohibitive.The discussion illustrates the issue with retrieving an event from 24 hours earlier.Source: Introduction · p. 18

Introduction

• Neural implausibility suggests better algorithmic solutions than end-to-end gradient computation for memory and other aspects of temporal credit assignment.Source: Introduction · p. 18
• A 1.3s window enabled MERLIN to solve tasks requiring memory over longer intervals, just 0.36% of the longest task’s 6-minute duration.The passage emphasizes that tasks required memory over much longer intervals than the model’s window.Source: Introduction · p. 18
• Although sufficiently large end-to-end RL networks could theoretically learn to store relevant information, the actual requirements are prohibitive.The passage compares this need for architectural innovation to object recognition and vision, where convolutional neural networks (52) are critical for practical systems.Source: Introduction · p. 18
• The authors believe that memory combined with predictive modeling will prove essential to future large-scale agent models in AI and neuroscience (53).They note that implementation details will likely change.Source: Introduction · p. 18

1 Environment Software

• Experiments used DeepMind Lab at 60 frames per second, passing one observation per four-frame packet and repeating each selected action across that packet.Rewards across each packet were summed and assigned to its first frame; reported Environment Steps are four times Agent Steps. Most tasks used six actions, while Arbitrary Visuomotor Mapping used four viewing-angle actions.Source: §1 · p. 26

2 Model

• MERLIN pairs a memory-based predictor, which maintains prior and posterior state distributions, with a policy that uses the state variable and external memory.The predictor’s external memory stores a history of state variables.Source: §2 · p. 26

2.1 Encoder

• The encoder independently processes images, velocity, previous action, previous reward, and text, then concatenates their outputs into a single feature vector.DM Lab images pass through six ResNet blocks; velocity is a six-dimensional egocentric signal, actions remain one-hot, rewards remain scalar, and text is encoded with an LSTM.Source: §2.1 · p. 26, §2.1.1 · p. 26, §2.1.2 · p. 27

2.2 Decoder

• The decoder mirrors five encoder networks and adds a return predictor, with separate decoders for images, instructions, reward, velocity, and actions.The return prediction combines a value function of the latent state and policy logits with an advantage function of the latent state and action.Source: §2.2 · p. 27, §2.2.1 · p. 28, §2.2.2 · p. 28

2.3 Memory

• The external memory is a two-dimensional matrix sized by the number of memory rows and twice the latent-state dimensionality, initialized blank at each episode.Source: §2.3 · p. 28

2.3.1 Deep LSTMs; 2.3.2 MBP LSTM; 2.3.3 Reading; 2.3.4 Writing

• A two-layer LSTM in the memory-based predictor uses the latent state, action, and prior memory reads to generate read keys and memory-interface outputs.Memory is read by content similarity before writing; writes append to the current row, retroactive updates use a discount-based weighting, and least-used rows are selected when overwriting is needed.Source: §2.3.1 · p. 29, §2.3.2 · p. 29, §2.3.3 · p. 29

2.4 Prior Distribution; 2.5 Posterior Distribution; 2.6 State Variable Generation

• The predictor’s prior is generated from its previous recurrent output and memory reads, while the posterior combines encoded features, recurrent outputs, and prior parameters before sampling the state.Source: §2.4 · p. 30, §2.5 · pp. 30–31, §2.6 · p. 31

2.7 Policy

• The policy uses a two-layer LSTM that reads memory with one key, then combines its outputs with the latent state to produce action logits and sample an action.Action-space dimensionality varies by environment: four for Arbitrary Visuomotor Mapping, six for other DM Lab tasks, and 16 for Memory.Source: §2.7 · p. 31

• Jensen’s inequality yields a lower bound on the log marginal likelihood using an approximating posterior.Source: §3 · p. 31
• The temporal formulation defines z0:t as the sequence of latent variables through time t and z0:−1 as the empty set.Source: §3 · p. 31
• For partitioned target variables, the stationary and temporal formulations use an approximate posterior that does not condition on one of the variables.Source: §3 · p. 32
• The resulting loss combines prediction of incrementally observable information with delayed information, such as the sum of future rewards, and the prior may additionally condition on actions.Source: §3 · p. 32

4 Cost Functions for the MBP and Policy

• Policy parameters are independent of the memory-based predictor’s and receive no gradients from its losses, enforced by a gradient stop at zt.Source: §4 · p. 33

4.1 Conditional Log-Likelihood

• The predictor’s conditional log-likelihood factorizes across decoders, using action and word cross-entropy, pixel-wise Bernoulli cross-entropy, and mean-squared error for velocity, reward, and return.The likelihood is conditioned on a sample from the approximate posterior, yielding a stochastic approximation to the variational lower-bound expectation; long episodes use truncation windows to define target returns.Source: §4.1 · pp. 33–34

4.2 Kullback-Leibler Divergence

• The predictor’s KL-divergence loss is computed analytically between the diagonal Gaussian posterior and prior distributions.Source: §4.2 · p. 34

4.3 Practical Details

• The total memory-based predictor loss is divided by the number of image pixel channels to reduce gradient magnitude.Source: §4.3 · pp. 34–35

4.4 Policy Gradient

• Policy updates use Generalised Advantage Estimation with discount and bootstrapping parameters γ and λ, plus an entropy-increasing loss for the action distribution.Source: §4.4 · p. 35

5.1 RL-LSTM; 5.2 RL-DNC; 5.3 RL-MEM

• RL-LSTM retains MERLIN’s encoder and policy LSTM, RL-DNC substitutes a Differentiable Neural Computer, and RL-MEM adds direct memory writing without retroactive updates.RL-LSTM has separate policy and value heads; RL-DNC is used only for the short-episode memory game because its memory scaling has quadratic computational complexity. RL-MEM reads memory as in MERLIN’s policy.Source: §5.1 · p. 35, §5.2 · p. 35, §5.3 · pp. 35–36

6 Implementation and Optimisation; 7 Tasks; 7.8 Task Specific Parameters

• Training uses 192 asynchronous workers, truncated backpropagation, bootstrapped returns, and separate ADAM optimisers for the predictor and policy.Each model was trained five times with different random seeds, and learning curves show standard errors across runs. Task-specific settings include two-layer 256-unit LSTMs for DM Lab and one-layer 50-unit LSTMs for Memory; MERLIN uses three predictor read keys and one policy read key, while RL-MEM uses three.Source: §6 · p. 36, §7 · p. 36, §7.8 · p. 40

7.1 Memory Game

• The memory game requires matching eight Omniglot card pairs on a 4×4 grid within 24 moves despite changing card views and blank observations at cleared locations.Cards receive random affine transformations each time they are viewed; clearing the board also yields a bonus point at each timestep.Source: §7.1 · p. 36

7.2 Navigation Tasks

• Navigation tasks test goal-finding with distal and proximal cues, proximal cues alone, larger arenas, and randomly changing doors that require replanning.Four tasks are assigned across workers; goal-finding uses 3–5 rooms in an 11-unit-squared arena, while the large-environment variant uses a 15-unit-squared grid with up to seven rooms.Source: §7.2 · p. 38

7.3 Arbitrary Visuomotor Mapping; 7.4 Rapid Reward Valuation

• Arbitrary visuomotor mapping tests remembering stimulus-to-gaze directions as the query set grows, while rapid reward valuation tests learning object values and reusing them after objects move.In the mapping task, a new image joins the pool after k consecutive correct responses. In valuation, eight objects with fixed +2 or −1 values reappear in permuted locations after positive objects are collected.Source: §7.3 · p. 38, §7.4 · p. 39

7.5 Episodic Water Mazes; 7.6 Executing Transient Instructions

• The episodic water-maze task requires remembering three platform locations distinguished by wall colour, while transient-instruction trials require selecting an object matching a briefly presented object-colour and room-colour instruction.The instruction is shown for the first five time steps, and the episode ends when an object is collected; correct and incorrect selections yield +10 and −1, respectively. The task adapts an instruction-following task by Hermann et al. (49) and models work-dog instruction-following (48).Source: §7.5 · p. 39, §7.6 · p. 39

7.7 Latent Learning in a T-Maze

• The latent-learning T-maze separates exploration from exploitation with an extended delay, then tests whether the agent uses earlier experience to choose between two objects.After exploring objects in phase 1 and collecting apples during a corridor phase, the agent returns to the maze and chooses between objects; the displayed object earns +10 and the other −3.Source: §7.7 · pp. 39–40

8 Lesion Comparison Experiments

• Lesion comparisons removed external memory, observation reconstruction objectives, the return decoder, or retroactive memory updates.The no-memory variant retained separate recurrent components; the return-decoder-only variant retained future-return prediction while removing observation decoders and reconstruction costs.Source: §8 · p. 41

9.1 Time-to-Goal Analysis; 9.2 Goal Location Decoding

• Time-to-goal was compared across MERLIN, RL-LSTM, and RL-MEM, and goal-location decoders were evaluated using recurrent-memory or image-encoder inputs.Time-to-goal used a common set of 200 test mazes and included episodes where all three models reached the goal over 10 times. Decoding analyses considered allocentric and egocentric goal locations.Source: §9.1 · p. 41, §9.2 · p. 42

9.3 Value Function Event-Triggered Analysis; 9.4 Variance Explained Analysis

• Value-function error was aligned to goal or object acquisition events, and variance explained in z was decoded from estimated value functions across six return-loss weights.The event-triggered analysis averaged absolute value-function error over qualifying acquisitions; the variance analysis used 10-fold cross-validation.Source: §9.3 · p. 42, §9.4 · pp. 42–43

9.5 Return Prediction Saliency

• Return-prediction saliency was visualized from pixel gradients, and a comparison model initially prevented return-prediction error from modifying representations beyond z.The gradient maps used the L2 norm across colour channels and Gaussian smoothing; the comparison model’s backpropagation restriction was lifted for subsequent analysis.Source: §9.5 · p. 43

9.6 Read Head Analysis in Navigation

• One read head focused near the agent, while two read from positions nearer the goal; one read from just closer to the goal than the agent.The analysis covered 200 test mazes and included runs in which the agent took at least 50 steps to reach the goal.Source: §9.6 · p. 43

9.7 Rapid Reward Valuation Analyses

• 93% versus 63%: classifiers using retroactive memories predicted upcoming object reward valence more accurately than classifiers using the state variable.The state-variable classifier’s accuracy exceeded 50% because performant agents consumed good objects more often; retroactive memory included a discounted sum of future state variables.Source: §9.7 · p. 44

9.8 Episodic Water Mazes; 9.9 Executing Transient Instructions; 9.10 Latent Learning in a T-Maze

• The task analyses also tracked first versus later platform visits, decoded transient instructions over time, and tested whether memory reads encoded a T-maze goal location.The transient-instruction decoders used MBP state and reads, or each independently; the T-maze classifiers predicted left- versus right-arm goals from MBP memory reads.Source: §9.8 · p. 44, §9.9 · p. 44, §9.10 · p. 44

9.11 Spatial Receptive Fields of Memory Access

• Memory rows accessed by the third MBP read key tended to be most active near the spatial location where each row was initially written.The analysis used rows written after the first goal acquisition and showed the 10% most active grid locations for the most active rows.Source: §9.11 · p. 45

10 Human Testing

• Two experienced game testers practiced each navigation task for 15 minutes, then recorded 15 consecutive 90-second runs; the task creator performed six 90-second blocks for each visuomotor image set.Navigation scores are the testers’ average; the visuomotor sets were naturalistic images and a synthetic transfer set.Source: §10 · p. 45

Extended Figures

• Logistic regression decoded egocentric goal position (d, θ) more accurately from MERLIN’s MBP representations than from RL-model representations or any convnet encoder layer.Source: Extended Figures · p. 46
• MERLIN performed at or above human level after joint training across goal-finding tasks.The extended comparison also reports that the UNREAL agent performed worse than MERLIN and the comparison models.Source: Extended Figures · pp. 46, 54
• At low return-prediction cost coefficients, MBP priors predicted visuals reasonably and matched sharper posteriors; at higher coefficients, prior predictions deteriorated.The agent’s performance was not highest for the model producing the best visual samples.Source: Extended Figures · p. 48
• Across tested tasks, lesions disrupted MERLIN’s asymptotic performance or learning speed, with memory and joint return-observation decoding especially important.Removing retroactive memory updating was least harmful, modestly lowering final performance while primarily increasing learning time.Source: Extended Figures · p. 50
• In Episodic Water Mazes, MERLIN’s time to find the platform decreased dramatically with successful visits in each maze colour.In one environment, most MBP memory rows were read near where they were first written.Source: Extended Figures · pp. 51, 57
• On the Goal Finding Task, changing gradient paths did not obviously affect performance, and combining the policy with MBP inputs functioned similarly to the default architecture.Including or excluding gradients between memory-based prediction and policy made very little difference to performance.Source: Extended Figures · pp. 52–53
• Longer backpropagation windows improved RL-LSTM and RL-MEM on task-battery and navigation tasks, but they remained below MERLIN.On Episodic Water Mazes, increasing the window did not help; policy-gradient RL generally failed to store information when decision delays exceeded τ.Source: Extended Figures · pp. 55–56