Episodic Reinforcement Learning with Associative Memory

1 INTRODUCTION . Deep reinforcement learning ( RL ) has achieved remarkable performance on extensive complex domains ( Mnih et al. , 2015 ; Lillicrap et al. , 2016 ; Silver et al. , 2016 ; Schulman et al. , 2017 ) . Deep RL research largely focuses on parametric methods , which usually depend on a parametrized value function . The model-free approaches are quite sample inefficient and require several orders of magnitude more training samples than a human . This is because gradient-based updates are incremental and slow and have global impacts on parameters , leading to catastrophic inference issues . Recently , episodic reinforcement learning has attracted much attention for improving sample efficiency of deep reinforcement learning , such as model-free episodic control ( MFEC ) ( Blundell et al. , 2016 ) , neural episodic control ( NEC ) ( Pritzel et al. , 2017 ) , ephemeral value adjustments ( EVA ) ( Hansen et al. , 2018 ) , and episodic memory deep q-networks ( EMDQN ) ( Lin et al. , 2018 ) . Episodic control is inspired by the psychobiological and cognitive studies of human memory ( Sutherland & Rudy , 1989 ; Marr et al. , 1991 ; Lengyel & Dayan , 2008 ; Botvinick et al. , 2019 ) and follows the idea of instance-based decision theory ( Gilboa & Schmeidler , 1995 ) . It builds a non-parametric episodic memory to store past good experiences and thus can rapidly latch onto successful policies when encountering states similar to past experiences . However , most of the current breakthroughs have focused on episodic memory and leave the association of memory largely unstudied . Previous work usually uses a tabular-like memory , and experiences are stored as unrelated items . Studies in psychology and cognitive neuroscience ( Kohonen , 2012 ; Anderson & Bower , 2014 ) discover that associative memory in the hippocampus plays a vital role in human activities , which associates past experiences by remembering the relationships between them . Inspired by this , we propose a novel associative memory based reinforcement learning framework to improve the sample-efficiency of reinforcement learning , called Episodic Reinforcement Learning with Associative Memory ( ERLAM ) , which associates related experience trajectories ∗Equal Contribution to enable reasoning effective strategies . We store the best historical values for memorized states like episodic memory , and maintain a graph on top of these states based on state transitions at the same time . Then we develop an efficient reverse-trajectory propagation strategy to allow the values of new experiences to propagate to all memory items through the graph rapidly . Finally , we use the fast-adjusted non-parametric high values in associative memory as early guidance for a parametric RL agent so that it can rapidly latch on states that previously yield high returns instead of waiting for many slow gradient updates . To illustrate the superiority of the associative memory in reinforcement learning , consider a robot exploring in a maze to seek out the apple ( place G ) , as shown in Figure 1 . It collects two trajectory experiences starting from place A and B ( i.e. , blue dash line A-D-C and B-D-G ) , respectively . All the states of trajectory A-D-C receive no reward because the agent terminates at a non-reward state ( place C ) . While in trajectory B-D-G , the final non-zero reward of catching an apple ( place G ) back-propagates to all the states of this trajectory . Episodic memory keeps a higher value of two trajectories at the intersection ( place D ) when taking actions toward the lower-right corner , but the other states in trajectory A-D are still 0 . If an episodic memory based robot starts from place A again , it will wander around A because there are no positive values indicating the way to the goal . Thus based on the episodic memory , the robot may eventually take a policy like the green line ( A-B-D-G ) after multiple attempts . However , if the robot adopts associative memory , the high value in the place D collected from trajectory B-D-G will be further propagated to the start point A , and thus the robot can correctly take the red-line policy ( A-D-G ) . To some extent , our associative memory is equivalent to automatic augmentation of counterfactual combinatorial trajectories in memory . Thus , our framework significantly improves the sampleefficiency of reinforcement learning . Comparisons with state-of-the-art episodic reinforcement learning methods show that ERLAM is substantially more sample efficient for general settings of reinforcement learning . In addition , our associative memory can be used as a plug-and-play module and is complementary to other reinforcement learning models , which opens the avenue for further research on associative memory based reinforcement learning . 2 BACKGROUND . In the framework of reinforcement learning ( Sutton & Barto , 1998 ) , an agent learns a policy to maximize its cumulative rewards by exploring in a Markov Decision Processes ( MDP ) environment . An MDP is defined by a tuple ( S , A , P , R , γ ) , where S is a finite set of states , A is a finite set of actions available to the agent , P : S × A × S → R defines the transition probability distribution , R is the reward function , and γ ∈ ( 0 , 1 ] is the discount factor . At each time step t , the agent observes state st ∈ S , selects an action at ∈ A according to its policy π : S → A , and receives a scalar reward rt . In the setting of finite horizon , the accumulated discounted return is calculated as , Rt = ∑T k=0 γ krt+k where T is the episode length and goal of the agent is to maximize the expected return for each state st . The state-action value function Qπ ( s , a ) = E [ Rt|st = s , a ] is the expected return for executing action a on state s and following policy π afterwards . DQN ( Mnih et al. , 2015 ) parameterizes this action-value function by deep neural networks Qθ ( s , a ) and use Q-learning ( Watkins & Dayan , 1992 ) to learn it to rank which action at is best to take in each state st at time step t. The parameters of the value network θ are optimized by minimizing the L2 difference between the networks output Qθ ( s , a ) and the Q-learning target yt = rt + γmaxaQθ̂ ( st+1 , at ) , where θ̂ are parameters of a target network that is a older version of the value network and updated periodically . DQN uses an off-policy learning strategy , which samples ( st , at , rt , st+1 ) tuple from a replay buffer for training . DQN , as a typical parametric reinforcement learning method , suffers from sample inefficiency because of slow gradient-based updates . Thus episodic reinforcement learning is proposed to speed up the learning process by a non-parametric episodic memory . Episodic reinforcement learning enables fast learning by modeling hippocampal instance-based learning . The key idea is to store good past experiences in a tabular-based non-parametric memory and rapidly latch onto past successful policies when encountering similar states instead of waiting for many steps of optimization . 3 RELATED WORK . Deep Reinforcement Learning Our method is closely related to DQN ( Mnih et al. , 2015 ) . As the seminal work of deep reinforcement learning , DQN learns a deep neural network for state-action value function by gradient back-propagation and conducts parametric control . Following this line , a large number of extensions have been proposed to improve the learning efficiency of the parametric model . Double DQN ( Van Hasselt et al. , 2016 ) alleviates the over-estimation issue of Q-Network . Dueling network ( Wang et al. , 2015 ) separates Q-Network into two streams which predict state value and advantage value respectively and achieves better generalization across actions . Prioritized experience replay ( Schaul et al. , 2015b ) changes the sampling priority of each training sample according to its learning error . Apart from these prior improvements , many algorithms have been proposed to accelerate reward propagation and backup mechanism . Optimality Tightening method ( He et al. , 2016 ) combines the strength of DQN with a constrained optimization approach to rapidly propagate close-by rewards . Q∗ ( λ ) ( Harutyunyan et al. , 2016 ) and Retrace ( λ ) ( Munos et al. , 2016 ) incorporate on-policy samples into off-policy learning targets . Noisy Net ( Fortunato et al. , 2017 ) adds noise to the parametric model during learning to improve the exploration ability . Distributional RL ( Bellemare et al. , 2017 ) learns the value function as a full distribution instead of a expected value . Unlike these works , we focus on combining non-parametric memory and parametric model in this paper . Thus our method is complementary to these prior extensions and can be combined with them seamlessly . Episodic Reinforcement Learning Our work is also related to episodic reinforcement learning . Model-free episodic control ( Blundell et al. , 2016 ) uses a completely non-parametric model that keeps the best Q values of states in a tabular-based memory and replays the sequence of actions that so far yielded the highest return from a given start state . At the end of each episode , the Q values in memory are updated by the greater of the existing values and the accumulated discounted returns in the current episode . In the execution stage , the agent selects actions according to a k-nearestneighbors lookup in the memory table . Recently , several extensions have been proposed to integrate episodic control with parametric DQN . Neural episodic control ( Pritzel et al. , 2017 ) develops end-to-end episodic control by a differentiable neural dictionary to generate semi-tabular representation as slow-changing keys and then retrieves fast-updating values by context-based lookup for action selection . To better leverage the trajectory nature of experience , ephemeral value adjustments method ( Hansen et al. , 2018 ) proposes to further leverage trajectory information from replay buffer to propagate value through time and produce trajectory-centric value estimates . Our method differs from EVA in that we associate memory by a graph , and thus we can leverage not only intra-episode but also inter-episode information . Episodic memory deep q-networks ( Lin et al. , 2018 ) distills the information of episodic memory into a parametric model by adding a regularization term in the objective function and significantly boosts up the performance of DQN . Unlike these prior works , which adopt either tabular memory or semi-tabular memory , our work builds a graph on memory items based on their relationship to form an associative memory . Graph Based Methods in Deep Reinforcement Learning Recently , several works have also been proposed to use graph for planning in deep reinforcement learning . Eysenbach et al . ( 2019 ) builds a directed graph directly on top of states in replay buffer and runs graph search to find the sequence of waypoints , leading to many easier sub-tasks and thus improve learning efficiency . Huang et al . ( 2019 ) abstracts state space as a small-scale map which allows it to run high-level planning using a pairwise shortest path algorithm . Unlike these prior works that use graphs for planning , our method reorganizes episodic memory by a graph to allow faster reward propagation . In addition , these graph-based models rely on goal-conditioned RL ( Kaelbling , 1993 ; Schaul et al. , 2015a ) and only demonstrate their performance in navigation-like problems , while our approach is intended for general RL settings . Exploration Efficient exploration is a long-standing problem in reinforcement learning . Prior works have proposed guiding exploration based on criteria such as intrinsic motivation ( Stadie et al. , 2015 ) , state-visitation counts ( Tang et al. , 2017 ) , Thompson sampling and bootstrapped models ( Chapelle & Li , 2011 ; Osband et al. , 2016 ) , optimism in the face of uncertainty ( Kearns & Singh , 2002 ) , parameter-space exploration ( Plappert et al. , 2017 ; Fortunato et al. , 2017 ) . Recently , Oh et al . ( 2018 ) proposed self-imitation learning ( SIL ) and found that exploiting past good experiences can indirectly drive deep exploration . In their work , the agent imitates its own decisions in the past only when such decisions resulted in larger returns than expected . Like SIL , EMDQN ( Lin et al. , 2018 ) learns from episodic memory to replay past best decisions , therefore incentivizing exploration . In our method , we build associative memory through a graph , which enhances the exploitation of past good experiences and thus can indirectly encourage deeper exploration than EMDQN ( Lin et al. , 2018 ) .

Summary: This paper proposes Episode Reinforcement Learning with Associative Memory (ERLAM), which maintains a graph based on the state transitions (i.e. nodes correspond to states, and edges correspond to transitions) and propagates the values through the edges in the graph in the reverse order of each trajectory. The learned associative memory is then used for the regularization loss for training Q-network. Experimental results show that ERLAM significantly improves the sample efficiency in Atari benchmarks.