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: The paper proposes a new method for organizing episodic memory in with deep Q-networks. It organizes the memory as a graph in which nodes are internal representations of observed states and edges link state transitions. Additionally, nodes from different episodes are merged into a single node if they represent the same state, allowing for inter-episode value propagation of stored rewards. The authors experimentally evaluate the method against reasonable baselines and show improved performance in the majority of tasks.