Optimizing Agent Behavior over Long Time Scales by Transporting Value

TL;DR

The paper tackles the challenge of long-horizon credit assignment in reinforcement learning by introducing Temporal Value Transport (TVT), a memory-attentive framework that learns to encode, store, recall, and revalue past events to credit actions across long delays. The authors implement a Reconstructive Memory Agent (RMA) with a differentiable external memory and train it alongside TVT, LSTM+Mem, and LSTM baselines to test memory-based credit assignment in a suite of long-delay tasks. TVT demonstrates faster learning and substantially reduced gradient variance, solving information- and causation-based tasks that impair traditional discounting approaches. The work provides a mechanistic link between episodic memory processes and reward-based learning, with potential implications for neuroscience-inspired AI and theories of intertemporal choice in economics.

Abstract

Humans spend a remarkable fraction of waking life engaged in acts of "mental time travel". We dwell on our actions in the past and experience satisfaction or regret. More than merely autobiographical storytelling, we use these event recollections to change how we will act in similar scenarios in the future. This process endows us with a computationally important ability to link actions and consequences across long spans of time, which figures prominently in addressing the problem of long-term temporal credit assignment; in artificial intelligence (AI) this is the question of how to evaluate the utility of the actions within a long-duration behavioral sequence leading to success or failure in a task. Existing approaches to shorter-term credit assignment in AI cannot solve tasks with long delays between actions and consequences. Here, we introduce a new paradigm for reinforcement learning where agents use recall of specific memories to credit actions from the past, allowing them to solve problems that are intractable for existing algorithms. This paradigm broadens the scope of problems that can be investigated in AI and offers a mechanistic account of behaviors that may inspire computational models in neuroscience, psychology, and behavioral economics.

Figures (19)

• Figure 1: Task Setting and Reconstructive Memory Agent.a. The three phase task structure. In phase 1 (P1), there is no reward, but the agent must seek information or trigger an event. In phase 2 (P2), the agent performs a distractor task that delivers reward. In phase 3 (P3), the agent can acquire a distal reward, depending on its behavior in P1. At each time step, the RL agent takes in observations $o_t$ and produces actions $a_t$, and passes memory state to the next time step. b. The Passive Visual Match task: the agent passively observes a colored square on the wall in P1 (gray here), consumes apples in P2, and must select from a lineup the previously observed square from P1. The agent and colored square are indicated by the yellow and red arrow, respectively. c. The Reconstructive Memory Agent (RMA) takes in observations, $o_t$, encodes them, $e_t$, compresses them into a state variable $z_t$, and decodes from $z_t$ the observations and value prediction $\hat{V}_t$. The state variable is also passed to an RNN controller $h_t$ that can retrieve (or read) memories $m_t$ from the external memory $M_t$ using content-based addressing with search keys $k_t$. $z_t$ is inserted into the external memory at the next time step, and the policy $\pi_t$ stochastically produces an action $a_t$ as a function of $(z_t, m_t, h_t)$ (only $z_t$ shown). d. The RMA solves the Passive Visual Match, achieving better performance than a comparable agent without the reconstruction objective (and decoders), LSTM+Mem, and better than an agent without an external memory, LSTM. Here and henceforth, all learning curves show standard errors about the mean, computed over 5 independent runs. e. The RMA uses its attentional read weight on time step 526 in P3 to retrieve the memories stored on the first few time steps in the episode in P1, when it was facing the colored square, to select the corresponding square and acquire the distal reward, worth 10 points.
• Figure 1: Passive Image Match Learning.Left. Full episode score. Right. P2 score. ($\gamma=0.96$ for all models.)
• Figure 2: Temporal Value Transport and Type 1 Information Acquisition Tasks.a. First-person (upper row) and top-down view (lower row) in Active Visual Match task while the agent is engaged in the task. In contrast to Passive Visual Match, the agent must explore to find the colored square, randomly located in a two-room environment. The agent and colored square are indicated by the yellow and red arrow, respectively. b. Without rewards in P2, RMA models with large discount factors (near 1) were able to solve the task; the RMA with $\gamma=0.998$ exhibited retarded but definite learning with modest P2 reward (1 point per apple). c. Cartoon of the Temporal Value Transport mechanism: the distractor interval is spliced out, and the value prediction $\hat{V}_{t_3}$ from a time point $t_3$ in P3 is directly added to the reward at time $t_1$ in P1. d. The TVT agent alone was able to solve Active Visual Match with large rewards during the P2 distractor, and faster than agents exposed to no distractor reward. The RMA with discount factor $\gamma=0.96$ was able to solve a greater than chance fraction because it could randomly encounter the colored square in P1 and retrieve its memory in P3.
• Figure 2: Passive Image Match with Varying Delay Period. All models learned to retrieve the P3 reward with no P2 delay, but performance is hampered for longer delays for models with no reconstructive loss.
• Figure 3: Analysis of Agent in Active Visual Match.a. In P1, TVT trained on Active Visual Match, actively sought out and oriented to the colored squared. RMA meandered randomly. b. Its attentional read weights focused maximally on the memories from time points when it was facing the colored square. c. With statistics gathered over 20 episodes, TVT's average value function prediction in P1 (blue) was larger than the actual discounted reward trace (green) -- due to the transported reward. Difference shown in gray. The RMA value function in contrast matched the discounted return very closely. d. The P3 rewards for TVT rose during learning (upper panel) after the maximum read strength per episode first crossed threshold on average (lower panel, red line).