Optimizing Agent Behavior over Long Time Scales by Transporting Value

The Reconstructive Memory Agent

• RMA alone learned correct P1-to-P3 color matching; unlike LSTM and LSTM+Mem, it remained effective across 0–60-second distractors and with CIFAR images.Its reconstruction objective helps build and stabilize perceptual features for later memory retrieval.Source: The Reconstructive Memory Agent · pp. 7–8

Temporal Value Transport

• TVT alone solved Active Visual Match with large P2 rewards, learning to seek the P1 cue while standard RMA behaved randomly.It learned faster than RMA exposed to no distractor reward.Source: Figure 2 · p. 9, Temporal Value Transport · p. 12
• TVT links a later memory retrieval to the time the memory was encoded, transporting value back as a fictitious reward weighted by memory attention.Splices require read strength above threshold; parallel reads can trigger them, and value is not sent to events within one half-life of retrieval.Source: Temporal Value Transport · p. 11
• In Key-to-Door, TVT learned key pickup while comparison agents failed reliably; its bootstrapped-return variance was nearly two orders of magnitude lower than undiscounted return in P1.Memory retrieval at the door triggered value transport to earlier behavior that led to key pickup, although P3 policy actions did not require recalling P1.Source: Temporal Value Transport · p. 12, Figure 4 · p. 13

Discussion

• TVT was generally on par with RMA on standard RL tasks, slightly worse on one, and also solved a task where actions avert more-negative distal outcomes.The authors describe TVT as heuristic and distinguish it from Sparse Attentive Backtracking [27] and RUDDER [28]: it uses reconstructive content-addressed memory and modifies rewards for the full policy gradient.Source: Temporal Value Transport · p. 14, Discussion · p. 17

Temporal Value Transport

• TVT reliably solved KtDtM and alone solved Latent Information Acquisition, learning exploratory coverage that touched all three objects while RMA touched about one.In KtDtM, transported value supported key pickup in P1 and the subgoal of seeing the colored square in P3.Source: Figure 5 · p. 15, Temporal Value Transport · pp. 16–17, Figure 6 · p. 16

Discussion

• Human cognition is often conjectured to be more model-based than current reinforcement-learning agents, including TVT, and to provide consciously available causal explanations [31, 32].The authors say long-term credit assignment may recruit planning, abstract, symbolic, and logical reasoning, commitment to indefinite goals, and language.Source: Discussion · p. 18
• People can conceptualize and plan for horizons without an upper limit, while societal attitudes and norms around savings and investment vary [33].The authors caution that some human abilities may require explanation at a different level of inquiry.Source: Discussion · p. 18

• The supplement contents list places “Psychlab Arbitrary Visuomotor Mapping” at section 7.11.3, page 45.Source: Transporting Value · p. 24

1 Agent Model

• The agent combines an observation encoder, memory-augmented recurrent network, policy, and decoder for observation and value prediction.Source: §1 · p. 26

1.1 Encoder

• The encoder independently processes the current image, previous action, and previous reward, then concatenates their outputs into a vector.Source: §1.1 · p. 26

1.1.1 Image Encoder

• The image encoder maps 64 × 64 × 3 inputs through six ResNet blocks to 8 × 8 × 64 features, then projects to 500 dimensions.Source: §1.1.1 · p. 26

1.1.2 Action Encoder

• Previous actions use a one-hot identity encoding, and previous rewards use an identity encoder.Source: §1.1.2 · p. 27, §1.1.3 · p. 27

1.2 Decoder

• The decoder has separate subnetworks for image, action, and reward reconstruction, plus a value-function decoder.Source: §1.2 · p. 27

1.2.1 Image Decoder

• The image decoder reverses the encoder architecture, using transposed convolutions and output channels that parameterize the reconstruction likelihood.Source: §1.2.1 · p. 27

1.2.2 Action and Reward Decoders

• Linear layers map the state variable to a scalar reward and to the action-space cardinality.Source: §1.2.2 · p. 27

1.2.3 Value Function Predictor

• The value predictor is an MLP over the state and action logits, with gradients stopped at the logits so value learning does not modify the policy.It uses one 200-unit tanh hidden layer and a one-dimensional output.Source: §1.2.3 · pp. 27–28

1.3 Memory-Augmented RNN

• The memory-augmented RNN uses a deep LSTM and slot-based external memory, with memory readouts carried as recurrent state and memory initialized blank each episode.Source: §1.3 · p. 28

1.3.1 State Variable Generation

• State generation combines the encoded current observation with the previous LSTM state and memory readout through an MLP.Source: §1.3 · p. 28, §1.3.1 · p. 28

1.3.2 Deep LSTMs

• The controller is a two-layer deep LSTM whose layer outputs are concatenated to form its output state.Source: §1.3.2 · pp. 28–29

1.3.3 LSTM Update

• At each step, the LSTM receives the state variable concatenated with the previous memory readouts and produces its output.Source: §1.3.3 · p. 29

1.3.4 External Memory Reading

• Memory reading uses content-based cosine similarity to select top-K memory rows, computes weighting vectors, and forms readouts; unwritten zero rows may be selected.Source: §1.3.4 · pp. 29–30

1.3.5 External Memory Writing

• Memory writing appends the current state variable to the time-indexed row after the read operation.Source: §1.3.5 · p. 30

1.4 Policy

• A 200-unit tanh MLP maps the state, LSTM output, and memory readouts to action logits, then samples an action from a multinomial softmax.Source: §1.4 · p. 30

2 Loss Functions

• RMA combines a policy-gradient loss with reconstruction objectives for decoding observations and a specific loss that regularizes memory use for TVT.The reconstruction losses decode observations, while the additional memory regularizer is detailed in Section 5.3.Source: §2 · p. 30

2.1 Reconstruction Loss

• The reconstruction objective is a negative conditional log-likelihood of observations and returns, with separate action, reward, value, and image losses.It uses softmax cross-entropy for actions, mean-squared error for rewards and values, and Bernoulli cross-entropy for image pixels; image loss is divided by the number of pixel-channels.Source: §2.1 · p. 31
• For standard RL control tasks, truncated return windows bootstrap from a predicted value when the window ends before the episode; a full-episode window uses τwindow = T.For other tasks, the target return is the discounted sum through the end of the episode.Source: §2.1 · p. 31

2.2 Policy Gradient

• The policy advantage uses generalized advantage estimation with discount γ and bootstrapping parameter λ, alongside an entropy-increasing loss for the action distribution.The TVT-specific regularization term is described separately in Section 5.3.Source: §2.2 · p. 32, §2.3 · p. 32

3 Comparison Models

• The two comparison models are the LSTM+Mem Agent and the LSTM Agent.The LSTM+Mem model retains memory, whereas the LSTM Agent additionally removes the external memory system.Source: §3 · p. 32, §3.1 · p. 32, §3.2 · p. 33

3.1 LSTM+Mem Agent

• LSTM+Mem resembles RMA but has no reconstruction decoders or losses, and predicts value with a one-hidden-layer MLP of 200 units.Its value estimate is conditioned on the state variable and stopped-gradient policy log probabilities.Source: §3.1 · p. 32

3.2 LSTM Agent

• The LSTM Agent has no external memory system and uses a deep two-layer LSTM with policy and value heads, each a one-hidden-layer MLP of 200 units.It shares RMA’s encoder networks and is described as essentially the same design as A3C [38].Source: §3.2 · p. 33

Implementation and optimisation

• Training used 384 parallel worker threads, each computing gradients after a truncation window; most main-paper experiments used a window spanning the full episode.Standard RL control experiments were excluded from the full-episode-window setting.Source: §4 · p. 33
• Workers sent gradients to an asynchronous parameter server, and optimization used ADAM with β1 = 0.9 and β2 = 0.999.The parameter server applied an optimization step as each gradient arrived.Source: §4 · p. 33
• The experiments used Sonnet and its default network-parameter initialization.Worker pseudocode is presented in Algorithm 2.Source: §4 · p. 33

5 Temporal Value Transport

• TVT has two stages: it identifies significant memory reads as splice events, then transports value predictions from those reads back to the time points read from.The transported values modify rewards and thereby the reinforcement-learning updates.Source: §5 · p. 33

5.1 Splice Events

• Splice-event selection excludes reads too near the past, using a 1/(1 − γ) cutoff, and selects a maximum-strength point within identified windows.The passages describe setting the read strength to zero for points inside the near-past cutoff and choosing tmax by the maximum β(i).Source: §5.1 · p. 35

5.2 Reward Modification

• For each selected tmax, rewards are modified at time points more than 1/(1 − γ) steps earlier.TVT sends back the value prediction at tmax+1, the first one incorporating information from the read; multiple read processes make independent additive reward changes.Source: §5.2 · p. 35

5.3 Reading Regularization

• A small regularization cost is imposed whenever a memory-read strength exceeds a threshold, to prevent extraneous TVT triggers.This cost is added to the other loss terms.Source: §5.3 · p. 35

6.1 Undiscounted Case

• The analysis defines signal as the squared norm of the expected P1 policy change induced by the policy gradient.The returns are baseline-subtracted for the calculation.Source: §6.1 · p. 36
• Noise is defined as the trace of the policy-gradient variance, with the P2 distractor contributing variance under the stated independence approximation.The analysis assumes P1 and P2 are approximately independent and neglects their weak dependence through the agent’s internal state and parameters.Source: §6.1 · p. 37
• The signal calculation treats the task as having no P2 distractor phase, while the noise calculation includes P2 reward variance.The approximation neglects changes to the system’s memory state caused by P2 experience.Source: §6.1 · pp. 37–38
• Low signal-to-noise ratio can negatively affect stochastic-gradient convergence, and averaging episodes only postpones the difficulty when P2 reward variance is sufficiently high.The analysis notes that some distractor reward variance can match or overwhelm the variance reduction from averaging.Source: §6.1 · p. 38

7.4 Task Phases

• The standardized P2 distractor is a 30-second apple-collection phase on an 11 × 11 map, with independently spawned apples worth 5 reward each.The duration is varied in some experiments; the map uses the Tetris texture set unless noted otherwise.Source: §7.4 · p. 40

7.6.2 Active Visual Match

• Active Visual Match changes P1 to a larger maze with a randomized target location, requiring active search for the color later matched in P3.The task uses the same basic matching structure as Passive Visual Match.Source: §7.6.2 · p. 41

7.6.3 Key-to-Door

• Key-to-Door tests whether the agent remembers acquiring a key in P1 to open a door and reach a rewarded goal in P3.The key is placed randomly and gives no immediate reward; reaching the goal after opening the door yields 10 points.Source: §7.6.3 · pp. 41–42

7.6.4 Key-to-Door-to-Match

• Key-to-Door-to-Match combines key acquisition, two distractor phases, a later color cue, and a final matching choice across five phases.Its two distractor phases last 15 seconds each, and the intervening color-cue phase has no reward.Source: §7.6.4 · p. 42

7.6.5 Two Negative Keys

• Two Negative Keys tests choosing between remembered key identities, with different negative outcomes for failing or choosing the blue rather than red key.The P3 rewards are -20 for failing to collect the goal, -10 after the blue key, and -1 after the red key.Source: §7.6.5 · pp. 42–43

7.6.6 Latent Information Acquisition

• Latent Information Acquisition tests learning which objects are good or bad without P1 reward, then using that information when objects are repositioned in P3.Good objects yield 20 in P3 and bad objects yield -10; P1 cues last 0.25 seconds.Source: §7.6.6 · p. 43

7.7–7.12 Task Variants and Controls

• Distractor variants manipulate apple rewards or counts, while control tasks test memory access in goal-finding, naturalistic mushroom collection, and arbitrary image-to-direction mapping.Fixed-count apples preserve the standard expected reward under an optimal collection policy but remove reward variance; variable apple rewards preserve the optimal expected return while increasing variance with reward value. Experiments use DM Lab, 15-frame-per-second agent observations, and a common six-action set except for the four-action visuomotor task.Source: §7.9 · p. 44, §7.10 · p. 44, §7.11.1 · p. 45

8.2 Active Visual Match

• In Active Visual Match, RMA learning was already significantly disrupted when each apple gave reward 1, so higher-variance examples were not reported.The analysis compares zero apple reward with reward 1.Source: §8.2 · p. 47

8.3 Key-to-Door

• For Key-to-Door, apple rewards of 1, 3, 6, and 10 corresponded to total P2 reward variances of 25, 100, 196, and 361, respectively.Episode scores indicate that all apples were usually collected in P2 at policy convergence.Source: §8.3 · p. 47
• For Key-to-Door, increasing P2 reward variance disrupts learning more when apple counts remain variable, even after the agent has learned the apple-collection policy.The comparison includes zero reward, fixed apple counts, and variable apple counts at the same expected return for the latter two conditions.Source: §8.3 · p. 47

8.4 Return Prediction Saliency

• The value-prediction saliency visualization was generated from gradients of the agent’s predictions with respect to input pixels, then smoothed and normalized.The estimates were smoothed with a 2-pixel-wide Gaussian filter and normalized using the 97th percentile across time and pixels.Source: §8.4 · p. 48

8.5 Learning Rate Analysis for High Discount Factor

• The default learning rate of 5 × 10−6 was best among the tested rates for the high-discount RMA and LSTM variance tasks.The tested rates were 3.2×10−7, 8×10−7, 2×10−6, 5×10−6, and 1.25×10−5.Source: §8.5 · p. 48

8.6 Behavioral Analysis of Active Visual Match

• The Active Visual Match behavioral analysis compared TVT and RMA trajectories under three fixed target-square locations and visualized TVT memory retrieval weights.Each fixed-location condition used 10 episodes.Source: §8.6 · p. 49

8.7 Behavioral Analysis of Latent Information Acquisition

• The latent-information behavioral analysis compared TVT and RMA exploration using position occupancy and the number of objects acquired across 50 trials.The visualization included 3,750 P1 position samples.Source: §8.7 · p. 49

8.8 Return Variance Analysis

• The Key-to-Door return-variance analysis compared undiscounted returns with recursively bootstrapped returns computed using TVT with γ = λ = 0.96.The comparison used trajectories from the same TVT agent over 20 trials.Source: §8.8 · p. 49

Supplementary Figure 1

• Supplementary Figure 1 reports Passive Image Match full-episode scores and P2 scores for models with γ = 0.96.Source: Figure 1 · p. 50

Supplementary Figure 2

• With no P2 delay, all models learned to retrieve the P3 reward, while longer delays hampered models without reconstructive loss.The figure compares delay periods for Passive Image Match.Source: Figure 2 · p. 50

Supplementary Figure 3

• When plotted by episodes to control for delay-induced extra steps, RMA learned in roughly the same number of episodes across delays from 0 to 60 seconds.Source: Figure 3 · p. 51

Supplementary Figure 4

• RMA performed Passive Image Match when CIFAR-10 images replaced the colored-square P1 and P3 images.Source: Figure 4 · p. 52

Supplementary Figure 5

• At higher P2 reward variance in Active Image Match, RMA models failed to solve the task while TVT was largely unaffected.The variance was introduced by varying apple-reward probability and value.Source: Figure 5 · p. 53

P3 Reward

• TVT solved Active Image Match with a 60-second P2 delay period.Source: Figure 6 · p. 54
• With a black P3 door, TVT solved Key-to-Door equally easily when the P1 key was blue, indicating retrieval was not based on surface color similarity.Source: Figure 7 · p. 55
• Across control tasks, TVT matched RMA on Natlab Varying Map Randomized, modestly reduced final performance on Explore Goal Locations Small, and slowed learning on Psychlab Arbitrary Visuomotor Mapping.Final performance is shown for five training runs.Source: Figure 8 · p. 55, Figure 9 · p. 56
• In Two Negative Keys, TVT selected the red key and received an average P3 reward of −1, rather than −10 or −20.Source: Figure 10 · p. 57
• Variable P2 reward was maximally detrimental to the LSTM agent’s performance on Key-to-Door compared with zero or fixed apple reward.The comparison used γ = 0.998; the variable-reward case had the same expected return as the fixed-reward case.Source: Figure 11 · p. 58
• Across the learning-rate analyses, the default rate of 5 × 10−6 performed best.The tested rates ranged from 3.2 × 10−7 to 1.25 × 10−5.Source: Figure 12 · p. 59
• TVT remained stable when P2 distractor reward was introduced in Active Image Match, while RMA performance degraded; larger discounts also worsened RMA performance on Passive Image Match.The figure compares performance with and without P2 apple reward.Source: Figure 13 · p. 60