PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalization of Reinforcement Learners with Working and Episodic Memory

Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charlie Deck, Joel Z Leibo, Charles Blundell

arXiv:1910.13406v2cs.LGcs.AIstat.ML

TL;DR

The paper addresses limited memory-specific evidence on how reinforcement-learning agents generalize to holdout data. It proposes memory-focused tasks and a combined-memory agent, which outperforms baselines across the suite, though challenging tasks remain difficult.

  • Problem

    Reinforcement-learning evaluation has commonly relied on training data, while existing generalization evaluations have not specifically tested memory systems.

    Source: §1 · p. 2

  • Method

    The paper constructs tasks with memory-relevant holdout splits and develops an agent combining working-memory control, external episodic memory, and auxiliary unsupervised learning.

    Source: §1 · p. 2, §3 · p. 6

  • Results

    The Memory Recall Agent outperforms baselines across the suite; episodic memory and learned representations are crucial in some cases and sometimes work synergistically.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The task suite and benchmark agent provide a testbed for further work on memory-enabled agents, while the trickiest tasks leave room for improvement.

    Source: §5 · p. 8

    Because experiments were computationally demanding, hyperparameter tuning used only small variations for each task.

    Source: §4 · p. 6

Abstract

from arXiv · show

Memory is an important aspect of intelligence and plays a role in many deep reinforcement learning models. However, little progress has been made in understanding when specific memory systems help more than others and how well they generalize. The field also has yet to see a prevalent consistent and rigorous approach for evaluating agent performance on holdout data. In this paper, we aim to develop a comprehensive methodology to test different kinds of memory in an agent and assess how well the agent can apply what it learns in training to a holdout set that differs from the training set along dimensions that we suggest are relevant for evaluating memory-specific generalization. To that end, we first construct a diverse set of memory tasks that allow us to evaluate test-time generalization across multiple dimensions. Second, we develop and perform multiple ablations on an agent architecture that combines multiple memory systems, observe its baseline models, and investigate its performance against the task suite.

1 Introduction

The paper addresses gaps in evaluating memory-specific generalization in reinforcement learning and proposes a task suite and agent architecture for studying it. Its Memory Recall Agent outperforms baselines across the suite, while component effects vary and challenging holdout tasks remain difficult.

1 Introduction

  • Working memory is a limited-capacity temporary store, whereas episodic memory stores experience over longer timescales and can compile event sequences into episodes.The paper situates these as two commonly studied memory types.Source: §1 · p. 1
  • Prior reinforcement-learning evaluation commonly used training data alone, and existing generalization evaluations were not specific to memory.The paper identifies this as a gap in testing agent generalization.Source: §1 · p. 2
  • The proposed approach tests working and episodic memory with environments whose holdout sets differ from training along memory-relevant dimensions.The dimensions include task scale and the precise objects used in the environment.Source: §1 · p. 2
  • The Memory Recall Agent combines working-memory control, external episodic memory, and long-term representational credit assignment, and outperforms baselines across the task suite.The architecture uses an auxiliary unsupervised loss and backpropagation through time that can jump across several time steps.Source: §1 · p. 2
  • Episodic memory and good representations are crucial to performance in some cases, and their combination can be synergistic.The paper also reports that none of its models fully generalize beyond training on the more challenging tasks, with greater degradation under extrapolation.Source: §1 · p. 2

2 Task suite overview

The task suite comprises 13 memory tasks with train-test splits designed to assess generalization across multiple dimensions. It combines psychology-inspired tasks with 3D environments that test scene comparison, goal navigation, and transitive inference.

2 Task suite overview

  • The suite defines 13 tasks with train-test splits that test memory and generalization across multiple dimensions.It adapts tasks from PsychLab and DMLab and adds Unity 3D tasks.Source: §2 · p. 2

2.1 PsychLab

  • PsychLab tasks present one or more consecutive images in trials and test object-response mapping, image recognition, change detection, or digit-location recall.Episodes contain multiple trials.Source: §2.1 · pp. 2–3

2.2 3D tasks

  • The 3D tasks test spotting differences between rooms, navigating to a goal, and choosing the higher-valued object in a pair.Spot the Difference includes variants in object count and motion pattern; Goal Navigation re-randomizes the goal location each episode.Source: Figure 1 · p. 3, §2.2 · p. 3, Figure 2 · p. 3

Transitive Inference:

  • Transitive Inference tests whether the agent can learn an overall ordering from ordered pairs of adjacent objects.Source: Transitive Inference · p. 3

2.3 Scale and Stimulus Split

  • Holdout levels differ from training along scale and stimulus dimensions, with scale tests using interpolation between training sizes and extrapolation beyond the larger size.The stimulus split is intended to guard against overfitting to particular visual inputs.Source: §2.3 · pp. 3–4

3 The Memory Recall Agent

The Memory Recall Agent combines working and episodic memory with auxiliary representation learning and a training regime that lets gradients jump across time. Its components pair a recurrent controller with retrieval from stored experiences.

3 The Memory Recall Agent

  • MRA combines a pixel-input network, working memory, slot-based episodic memory, a contrastive auxiliary loss, and jumpy backpropagation on top of IMPALA.The pixel network uses convolutional and residual layers.Source: §3 · p. 4
  • The LSTM receives pixel embeddings and retrieved memory, then produces policy and value outputs from its hidden state.That hidden state also contributes to episodic-memory entries.Source: §3 · p. 4
  • The episodic memory stores past pixel and LSTM activations in slots, retrieves relevant experiences, and supplies them to the controller to support long-term dependencies.Its fixed-size buffer overwrites the oldest entries when full and is cleared between episodes.Source: §3 · pp. 4–5
  • A query derived from the pixel embedding and previous LSTM state selects nearest memory keys; inverse-distance-weighted values are concatenated with the pixel embedding for the LSTM.Each stored key is computed from its pixel embedding and LSTM state.Source: §3 · p. 5, Figure 4 · p. 5
  • Jumpy backpropagation stops gradients at memory writes, while allowing key-projection parameters to learn and skipping intermediate recurrent steps in gradient updates.The method stores pixel embeddings, hidden states, and keys so it can recompute keys for retrieved memories.Source: §3 · p. 6
  • MRA adds a contrastive predictive coding loss to the reinforcement-learning losses, using the working-memory state to predict future pixel encodings.The number of CPC steps is typically 10 or 50 in the experiments.Source: §3 · p. 6
  • The reconstruction baseline minimizes L2 distance between predicted and true rewards and pixel inputs using the working-memory state.The paper describes further baseline details in Appendix C.4.Source: §3 · p. 6

4 Experiments

Across the task suite, LSTM + MEM + CPC achieved the best overall scores on training and holdout, though different memory systems favored different tasks. The results also show that generalization gaps and memory-component benefits varied across tasks and holdout levels.

4 Experiments

  • The experiments ran 10 MRA ablations across training and two holdout levels, varying working memory, episodic memory, auxiliary loss, and jumpy backpropagation.Because the experiments were computationally demanding, hyperparameter variations for each task were small.Source: §4 · p. 6

4.1 Full comparison

  • The heatmap compares ablations by task across Train, Holdout-Interpolate, and Holdout-Extrapolate, with tasks sorted by normalized score.Source: Figure 5 · p. 7

4.2 Results

  • LSTM + MEM + CPC achieved the best overall normalized scores across training and holdout, although different memory systems worked best for different tasks.Figure 6 averages normalized scores across tasks.Source: §4.2 · p. 7, Figure 6 · p. 7
  • The generalization gap was minimal for AVM and Continuous Recognition but larger for What Then Where and Spot the Difference: Multi-Object, tending to widen with task difficulty.In PsychLab, tasks scaled by trial count seemed easier than those scaled by delay duration; the hypothesized gap was observed on tasks much harder for agents than humans.Source: §4.2 · p. 7, §4 · p. 7, Figure 7 · p. 8
  • MEM improved scores on 9 training tasks, 6 holdout-interpolate tasks, and 6 holdout-extrapolate tasks, including recognition, navigation, and inference tasks.Gains covered all Spot the Difference tasks except Multi-Object and all Goal Navigation tasks except Visible Goal Procedural Maze.Source: §4.2 · pp. 7–8
  • MEM alone noticeably boosted holdout performance on Transitive Inference, AVM, and Continuous Recognition, but had little or no effect alone on Change Detection and Multi-Object.On Change Detection and Multi-Object, combining MEM with CPC or REC provided a noticeable boost.Source: §4.2 · p. 8
  • Adding MEM + CPC or MEM + REC to working memory appeared to improve holdout generalization on average; MEM + CPC synergy with LSTM occurred on 7, 6, and 6 tasks across the three levels.Synergy across training and both holdout levels occurred on Visible Goal Procedural Maze, Visible Goal with Buildings, Spot the Difference: Basic, and Spot the Difference: Multi-Object.Source: §4.2 · p. 8
  • CPC outperformed REC on Spot the Difference and the two harder PsychLab tasks, while REC outperformed CPC on all Goal Navigation tasks except Invisible Goal Empty Arena.Averaged across tasks, REC was more useful with FF working memory, whereas CPC was more useful with LSTM.Source: §4.2 · p. 8

5 Discussion & Future Work

The paper introduces memory-focused environments and an agent to study generalization on holdout sets. It reports that auxiliary unsupervised losses can improve holdout generalization with external episodic memory, while the hardest tasks remain open challenges.

5 Discussion & Future Work

  • The testbed combines diverse memory environments based on tasks designed to identify working and episodic memory in humans with an agent demonstrating many of these abilities.The authors propose it as a testbed and benchmark for further work on agents with memory.Source: §5 · p. 8
  • The agent fared relatively poorly on the suite’s trickiest tasks, leaving room for improvement.The supplied discussion begins its specific example with Spot the Difference: Motion.Source: §5 · p. 8
  • Auxiliary CPC or reconstruction losses improve holdout generalization with external episodic memory, sometimes synergistically.The paper suggests this combination could also benefit agents such as DNC and NEC.Source: §5 · p. 9
  • Spot the Difference: Motion might require a generative model for forward planning that imagines how future motion unfolds.The authors cite Racanière et al. (2017) as an example.Source: §5 · p. 9

A Level descriptions and further experimental findings

The task suite separates training from holdout evaluation across scale and stimulus dimensions. Training samples small and large levels, while holdouts test interpolation and extrapolation; PsychLab tasks require retaining associations or information across trials.

A Level descriptions and further experimental findings

  • Each task has small and large training levels plus holdout-interpolation and holdout-extrapolation levels; training uniformly samples the two training levels.Interpolation uses an intermediate scale, while extrapolation uses the extralarge scale; dashed settings without a clear generalization interpretation are omitted.Source: A Level descriptions and further… · p. 11

Scale

  • The level structure varies scale and stimulus, with scale defined by corridor delay or room object count in Spot the Difference tasks.Stimuli can vary by color set or image set in PsychLab, and by color set or motion pattern in Spot the Difference.Source: Scale · p. 11, Table 1 · p. 11, Table 2 · p. 11
  • PsychLab episodes use repeated single-image trials: agents learn image-action associations, retain patterns across delays, recognize previously seen images, or recall a digit’s location.In What Then Where, the agent must recall the digit shown in the ‘what’ phase, locate it in the ‘where’ phase, and respond with that position.Source: Figure 8 · p. 12

A.1.1 PsychLab: main experimental findings

PsychLab results vary by task and controller: episodic memory is especially useful for some association and recognition tasks, whereas CPC leads on change detection and the LSTM baseline leads on What Then Where.

A.1.1 PsychLab: main experimental findings

  • For AVM, MEM was the most useful component; auxiliary loss helped with an FF controller but not an LSTM, and CPC versus REC made little difference.The task requires remembering associations between images and movement patterns.Source: §A.1.1 · p. 12
  • For Continuous Recognition, MEM helped with an LSTM but not an FF controller alone; MEM plus CPC or REC substantially boosted both controllers.The task asks whether an image has appeared before in the episode.Source: §A.1.1 · p. 12
  • For Change Detection, CPC gave the largest benefit, while MEM slightly hurt the FF baseline and made no difference for LSTM.The agent judges whether two images separated by a delay differ.Source: §A.1.1 · p. 13
  • In What Then Where, the LSTM baseline beat all other models, while REC reduced performance below random.The task requires recalling which image appeared in the ‘what’ phase and where it appeared during the ‘where’ phase.Source: §A.1.1 · p. 13

A.2 Spot the Difference (SD)

  • Spot the Difference tasks vary corridor-delay duration or room object count, use either a color set or motion-pattern set, and last 120 seconds except Motion, which has a 240-second timeout.Source: §A.2 · p. 13

A.2.1 Spot the Diff: main experimental findings

Spot the Difference results depend on task conditions: MEM plus CPC is useful in several settings, but motion and multi-object variants remain difficult. The navigation environments vary arena size and goal spawn region.

A.2.1 Spot the Diff: main experimental findings

  • In Basic, MEM alone or with REC had little effect, while MEM plus CPC provided the biggest gain and was synergistic.CPC also helped when added to an LSTM.Source: §A.2.1 · p. 13
  • In Passive, MEM alone made the biggest positive difference; auxiliary-loss effects depended on controller and combination.CPC helped with FF plus MEM but hurt LSTM alone; REC helped less than CPC with FF plus MEM and hurt with LSTM plus MEM.Source: §A.2.1 · p. 13
  • Nothing performed well on Spot the Difference: Motion, which requires memorizing and comparing 3D motion patterns traced by multiple objects.The authors suggest hyperparameter tuning or further architectural improvements could potentially improve results.Source: §A.2.1 · p. 13
  • No models fared well on holdout-interpolation or holdout-extrapolation in Multi-object, the family’s hardest task.MEM alone had no effect or hurt performance; MEM plus CPC was most useful with an LSTM, though overall performance remained poor.Source: §A.2.1 · p. 14

A.3 Navigate to Goal

  • Navigate to Goal tasks vary the square arena’s in-game size and the goal spawn region.Visible Goal Procedural Maze is a modification of DMLab’s Explore Goal Locations task; its episodes last 120 seconds, compared with 200 seconds for the other tasks in this family.Source: §A.3 · p. 14

A.3.1 Navigate to Goal: main experimental findings

Across navigation tasks, memory and auxiliary losses helped in task-dependent ways, while reconstruction was especially useful for the goal family. In transitive inference, MEM with an auxiliary loss was crucial, with the best auxiliary loss depending on controller type.

A.3.1 Navigate to Goal: main experimental findings

  • An auxiliary unsupervised reconstruction loss was the most useful component for the goal task family.Source: §A.3.1 · p. 14
  • The successful LSTM + MEM + CPC agent learned the invisible goal’s location rather than merely memorizing a route, and its time to reach the goal shortened after respawns.Figure 9 shows example routes and time-to-goal across training, Holdout-Interpolate, and Holdout-Extrapolate.Source: §A.3.1 · p. 14, Figure 9 · p. 15
  • MEM + REC was the most successful model for Visible Goal Procedural Maze and Visible Goal With Buildings, while MEM was slightly more helpful than CPC with an LSTM in the latter.MEM + CPC also had a synergistic effect with an LSTM in Visible Goal With Buildings.Source: §A.3.1 · p. 14
  • For Invisible Goal With Buildings, MEM + REC was most useful with both FF and LSTM controllers; in the empty arena, MEM alone helped slightly and CPC edged REC with an LSTM.For an FF controller in the empty arena, REC helped more than CPC.Source: §A.3.1 · p. 14

A.4 Transitive Inference

  • Transitive inference used scrambled demonstrations of adjacent object pairs followed by a challenge pair, with chain length and color set varied across training and holdout.The challenge pair tested an unobserved relation between the second and second-to-last objects in the chain.Source: §A.4.1 · p. 16, §A.4 · p. 16

A.4.1 Transitive Inference: main experimental findings

  • In transitive inference, stacking MEM with an auxiliary loss was crucial; CPC helped more than REC for FF, while REC helped more than CPC for LSTM.LSTM + MEM + REC was more data-efficient and took off earlier than LSTM + MEM + CPC, though their normalized scores were not far apart.Source: §A.4.1 · p. 16

A.5 Jumpy Backpropagation (JB) ablation

The JB ablation found improved training and holdout performance on some tasks and no harm on the remaining tasks. The appendix also presents model curves, ablation rankings, score computation, and episode rewards.

A.5 Jumpy Backpropagation (JB) ablation

  • Jumpy Backpropagation (JB) improved training and holdout performance on some tasks and did not hurt performance on the remaining Memory Task Suite levels.The authors conclude that JB is an important component of MRA.Source: §A.5 · p. 16, Figure 10 · p. 17, Figure 11 · p. 18

A.6 Agent Performance Curves

  • Figure 14’s heatmap includes standard errors and orders tasks by normalized training score and models by mean training score.The highest-mean training model is at the top and the highest-mean training task is at the left.Source: Figure 14 · p. 23

B Human-Normalized Scores and Episode Rewards

A.6 Agent Performance Curves

B.1 Human-Normalized Score Computation

  • Human-normalized scores use each seed’s smoothed, rolling training-reward maximum to select a snapshot for comparison on both holdout levels; Figure 3 ranks the results.The rolling reward window is 10 episodes, and the ranking table is sorted by overall task-averaged human-normalized score.Source: §B.1 · p. 24, Table 3 · p. 24

B.2 Episode Rewards

  • Tables 4–16 report absolute episode rewards by task and level for trained agents, random agents, and humans, with standard error bars.Source: §B.2 · p. 24

C.1 Importance Weighted Actor-Learner Architecture

The methods describe an actor-learner architecture and visual encoder, then explain how the agent predicts future representations and reconstructs prior inputs and outputs. These auxiliary objectives are combined into a reconstruction loss.

C.1 Importance Weighted Actor-Learner Architecture

  • IMPALA uses decoupled actors that send batched experience trajectories to a learner, which computes gradients and returns parameter updates.V-trace importance weighting corrects for policy lag between actors and learner.Source: §C.1 · p. 25

C.2 Residual Network Architecture

  • A residual network with three convolutional blocks and a 256-unit MLP converts pixel input into latent representations for working memory and the query network.Source: §C.2 · p. 25

C.3 Contrastive Predictive Coding

  • Contrastive Predictive Coding uses working-memory states to predict future latent representations, scoring each memory-state/future-latent pair with a log-bilinear function.For each trajectory, it evaluates up to N predictive steps and contrasts a positive sample with negative samples using categorical cross-entropy.Source: §C.3 · pp. 25–26

C.4 Reconstruction

  • The reconstruction objective combines squared-error losses for prior-step rewards and actions with sigmoid cross-entropy for image reconstruction, scaling each by its cost hyper-parameter.The experiments generally set image, action, and reward costs to 1.0, with specified exceptions for three tasks.Source: §C.4 · p. 26

D Hyper-parameter Tuning

Hyper-parameter tuning was preliminary and constrained by the computational scope of the experiments. PsychLab received manual sweeps, while tuning was minimal on other tasks and REC was not tuned.

D Hyper-parameter Tuning

  • All experiments used three seeds, and hyper-parameter tuning was preliminary rather than exhaustive given the scope of the experiments.Initial settings were inherited from IMPALA or chosen as reasonable first guesses, with shared tuning aimed at adequate training performance across model variations.Source: D Hyper-parameter Tuning · p. 26
  • PsychLab received manual sweeps when training performance stalled or was noisy, but tuning on other tasks was minimal and mostly used first-guess settings.For Goal Navigation and Transitive Inference, the discount rate was standardized at 0.99.Source: D Hyper-parameter Tuning · p. 26
  • The reported result is that REC was not tuned at any point in the study; specifically, the authors state they performed no tuning for REC throughout.Source: D Hyper-parameter Tuning · p. 26

Related papers

Loading 1910.13406v2…