E PISODIC R EINFORCEMENT L EARNING WITH A SSOCIATIVE M EMORY
Anonymous authors
Paper under double-blind review
A BSTRACT
Sample efﬁciency has been one of the major challenges for deep reinforcement learning.
Non-parametric episodic control has been proposed to speed up parametric reinforcement learning by rapidly latching on previously successful policies.
However, previous work on episodic reinforcement learning neglects the relationship between states and only stored the experiences as unrelated items.
To improve sample efﬁciency of reinforcement learning, we propose a novel framework, called Episodic Reinforcement Learning with Associative Memory (ERLAM), which associates related experience trajectories to enable reasoning effective strategies.
We build a graph on top of states in memory based on state transitions and develop an efﬁcient reverse-trajectory propagation strategy to allow rapid value propagation through the graph.
We use the non-parametric associative memory as early guidance for a parametric reinforcement learning model.
Results on Atari games show that our framework has signiﬁcantly higher sample efﬁciency and outperforms state-of-the-art episodic reinforcement learning models.
I NTRODUCTION
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 method, which usually depends on a parametrized value function.
The parametric approaches are quite sample inefﬁcient and requires several orders of magnitude more training samples than a human.
This is because gradient-based updates are incremental and slow, and has global impacts on parameters, which leads to catastrophic inference issue.
Recently, episodic reinforcement learning has attracted much attention for improving sample efﬁciency 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 polices when encountering with states similar to past experiences.
However, most of the existing breakthroughs have been focusing on episodic memory and leave the association of memory largely unstudied.
Previous work usually use 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 found in hippocampus plays an important role in human activities, which associates past experiences by remembering relationship between them.
Inspired by this, we propose a novel associative memory based reinforcement learning framework to improve the sample-efﬁciency of reinforcement learning, called
Episodic Reinforcement Learning with Associative Memory (ERLAM), which associates related experience trajectories 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 efﬁcient reverse-trajectory propagation strategy to allow the values of new experiences to rapidly propagate to all memory items through the graph.
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.
[Figure]
To illustrate the superiority of the associative memory in reinforcement learning, consider a robot exploring in a maze to seek out the apple (at place G), as shown in Figure 1.
It collects two trajectory experiences starting from place A and B, respectively.
All the states of trajectory A (the top blue dotted line) receive no reward because the agent terminates at a state with a non-zero reward (at place C), while in trajectory B (the bottom blue dotted line) the ﬁnal non-zero reward of catching an apple (at place G) back-propagates through the whole path.
Episodic memory keeps a high value at the intersection of two trajectories (the door) when taking actions toward lower-right corner while recording zero values at the other states in trajectory A.
If an episodic memory based robot starts from place A, it will wander around A because there are no positive values indicating the way to goal.
Thus based on the episodic memory, the robot may eventually take a policy like the green line after multiple attempts.
However, if the robot adopts associative memory, the high value in the door collected from trajectory B will be further propagated to the start point A and thus the robot can correctly take the red-line policy.
To some extent, our associative memory is equivalent to automatic augmentation of counterfactual combinatorial trajectories in memory.
Thus, our framework signiﬁcantly improves the sampleefﬁciency of reinforcement learning.
Comparisons with state-of-the-art episodic reinforcement learning methods show that ERLAM is substantially more sample efﬁcient 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 researches on associative memory based reinforcement learning.
B ACKGROUND
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 deﬁned by a tuple ( S , A , P, R , γ ) , where S is a ﬁnite set of states, A is a ﬁnite set of actions available to the agent, P : S ×
A × S → R deﬁnes 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 s t ∈ S , select an action a t ∈ A according to its policy π : S → A , and receives a scalar reward r t .
In the setting of ﬁnite horizon, the accumulated discounted return is calculated as,
R t = (cid:80) Tk =0 γ k r t + k where T is the episode length and goal of the agent is to maximize the expected return for each state s t .
The state-action value function Q π ( s, a ) =
E [ R t | s t =
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 s t at time step t .
The parameters of the value network θ are optimized by minimizing the L 2 difference between the networks output
Q θ ( s, a ) and the Q-learning target y t =
r t + γ max a Q ˆ θ ( s t +1 , a t ) , 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 ( s t , a t , r t , s t +1 ) tuple from a replay buffer for training.
DQN as a typical parametric reinforcement learning method suffers from sample inefﬁciency 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 the 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 polices when encountering with similar states, instead of waiting for many steps of optimisation.
R ELATED W ORK
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 conduct parametric control.
Following this line, a large number of extensions have been proposed to improve the learning efﬁciency 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 to 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 improvement, many work have been proposed to accelerate reward propagation and back up mechanism.
Optimality Tightening method(He et al., 2016) combines the strength of DQN with a constrained optimization approach to rapidly propagate closeby rewards.
Q ∗ ( λ ) (Harutyunyan et al., 2016) and Retrace( λ ) (Munos et al., 2016) incorporate onpolicy 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.
Different from these work, 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-nearestneighbours lookup in the memory table.
Recently, a number of 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 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, 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 objective function and signiﬁcantly boosts up the performance of DQN.
Different from these prior work which adopt either tabular memory or semi-tabular memory, our work builds a graph on memory items based on their relationship to from a associative memory.
Graph Based Methods in Deep Reinforcement Learning Recently, several work 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 run graph search to ﬁnd the sequence of waypoints, leading to many easier sub-tasks and thus improve learning efﬁciency.
Huang et al.
(2019) abstracts state space as a small-scale map which allows it to run high-level planning using pairwise shortest path algorithm.
Different from these prior work which use graph for planning, our method reorganize 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.
E PISODIC R EINFORCEMENT L EARNING WITH A SSOCIATIVE M EMORY 4. A SSOCIATING EPISODIC MEMORY AS A GRAPH
Similar with previous episodic reinforcement learning, we adopt an episodic memory to maintain the historically highest values Q EC ( φ ( s ) , a ) of each state-action pair, where φ is an embedding function and can be implemented as a random projection or variational auto-encoders (VAE) (Kingma &
Welling, 2013).
When receiving a new state, the agent will look up in the memory and update the values of states according to the following equation,
[Equation]
However, episodic memory stores states as unrelated items and does not make use of relationship between these items.
To fully exploit information in episodic memory, we further build a directed graph G on top of items in the episodic memory to form an associative memory, as shown in Figure 2.
In this graph, each node corresponds to a memory item that records the embedded vector of a state φ ( s ) and we leverage transitions of states to bridge the nodes.
The graph are deﬁned as,
G
[Equation]
E =
[Equation]
s →
[Equation]
( s, a, s
[Equation]
) is stored in memory } .
[Equation]
Given a sampled trajectory, we temporarily add each state to the graph.
We add directed edges from the given state to every other previously memorized state that is successor of it under a certain action.
Our associative memory reorganizes the episodic memory and connects these fragmented states that previously yielded high returns by a graph.
We rewrite these stored values Q EC ( φ ( s ) , a ) as
Q G ( φ ( s ) , a ) in our graph augmented episodic memory.
In addition, we adopt a strategy of discarding the least recently used items when the memory is full.
Experiences
[Figure]
4. P ROPAGATING VALUES THROUGH ASSOCIATIVE MEMORY
Typical deep RL algorithms sample experience tuples uniformly from the replay buffer to update value function.
However, the way of sampling tuples neglects the trajectory nature of an agent’s experience (i.e., one tuple occurs after another and so information of the next state should be quickly propagated into the current state).
EVA (Hansen et al., 2018) encourages faster value propagation by introducing trajectory-centric planning (TCP) algorithm.
Nonetheless, EVA only propagates value through the current episode, which we refer to as intra-episode propagation.
Our insight here is that a same state might appear in different trajectories and such join points can help connect different various trajectories.
Therefore, we explicitly build the graph between states from different trajectories in memory and thus allows inter-episode value propagation.
Since the graph over states is complicated (e.g., not a tree structure), value propagation over such graph is always slow.
To accelerate the propagating process, we propagate values using the sequential property.
The pseudo code of value propagation is shown in Algorithm 1.
Our general idea is to update the values of the graph in the reverse order of each trajectory.
Speciﬁcally, when adding to a new state to the memory, we record the sequential step ID t of the state at current trajectory.
For memory associating, we ﬁrst sort the elements in memory by their sequential step ids in descending order, and propagate the value from states with large sequential step ID to small one for several iterations until Q G values converge.
At each update, we get all successor state-action pairs ( s (cid:48) , a (cid:48) ) of the current one ( s, a ) and current reward r
according to the graph G and apply max operation on successor action a (cid:48) to propagate the values to current state-action pair.
Formally, our graph augmented memory is updated as follow:
[Equation]
Since most of states at the beginning are similar across different episodes, our reverse order updating strategy can efﬁciently propagate all the values of the graph.
In addition, we prove that our graphbased value propagation algorithm can converge to unique optimal value, as shown in Appendix
A.
In the previous episodic reinforcement learning with no graph built, only the values of exactly the same or similar states can be updated.
This is because in the typical update rule of episodic memory as shown in Eq.
1, the relationship between states has been neglected.
That is, episodic memory does not leverage the information of edges E in our graph G .
Consequently, stored values in episodic memory often violate bellman equation.
On the contrary, our associative memory allows efﬁcient value propagation through the edges of the graph to compute the more accurate values for each state.
G ← Sort nodes in graph G by sequential step ID t in descending order repeatfor
|G| do
(cid:48) using graph G .
4. L EARNING WITH ASSOCIATIVE MEMORY
Building associative memory can be viewed as a way of augmenting counterfactual experiences.
As shown in Figure 2, the same states might appear in N > 1 trajectories.
Vanilla episodic memory maps such states to highest values among N trajectories, while our associative memory regards such states as join points to connect different trajectories, leading to totally N 2 trajectories.
This is equivalent to sample more combinatorial trajectories from environments, and thus can signiﬁcantly improve sample efﬁciency of RL algorithms.
Our associative memory can be applied to both learning and control phase.
In this paper, we use our associative memory as a guidance for the learning of Q function.
Speciﬁcally, we use the associative memory as regularization term of objective function to supervise the learning of Q-network.
The
Q-network is learned by minimized the following objective function:
[Equation]
where λ is the weight of the regularization term, θ represents parameters of parametric Q-network.
Similar with DQN (Mnih et al., 2015), we also adopt a target network parameterized by ˆ θ to stabilize the learning process.
Through the combination of parametric and non-parametric term, we can efﬁciently guide the learning of a conventional Q-network by the fast-adjusted high values in associative memory so that the agent can rapidly latch on strategies that previously yield high returns instead of waiting for many steps of slow gradient update.
The pseudo code of our method is shown in Algorithm 2.
for Episode number e = 1 . . .
E do for t = 1 . . .
T e do
Receive initial observation s t from environment with state embedding h t = φ ( s t ) a t ← (cid:15) -greedy
policy baesd on Q θ ( s t , a )
Take action a t , receive reward r t and next state s t +
Append ( s t , a t , r t , s t +1 ) to D if t mod update freq =
=
then
Sample training experiences ( s, a, r, t ) from D
Retrieve Q G ( φ ( s ) , a ) from associative memory
Update parameter θ using Eq.end if end for for t =
T e . . .
1 do
R t ← r t + γR t
+1 , if t < T e ; R t ← r t , if t = T e
Append ( h t , a t , r t , t, R t ) to G if ( h t , a t ) / ∈ G
Update Q G using Eq.1 if ( h t , a t ) ∈ G end for if e mod K == 0 then
Run Algorithm 1 to update Q G end if end for
4. C ONNECTION TO GRAPH - BASED DEEP REINFORCEMENT LEARNING
When the general RL setting used in our approach degenerates to a setting of navigation-like task that is usually adopted by goal-conditional RL (Kaelbling, 1993; Schaul et al., 2015a), the update target of associative memory in Eq.
3, y = r
+ γ max a (cid:48)
Q G ( φ ( s (cid:48) ) , a (cid:48) ) can be rewritten as, y = (cid:26) r , if s (cid:48) is a terminal state , γ max a (cid:48) Q G ( φ ( s (cid:48) ) , a (cid:48) ) , otherwise .
[Equation]
Optimizing with target in Eq.
5 is equivalent to ﬁnding the shortest path in the graph of all states.
In this case, algorithm 1 is analogous to Bellman-Ford algorithm (Bellman, 1958), which is proved that the value can converge in limited iterations.
In the context of goal-conditional RL, some graphbased methods (Huang et al., 2019; Eysenbach et al., 2019) also calculated shortest path.
They focus on a graph of waypoints learned by goal-conditional RL instead of memorized states that previously yield high returns.
In addition, they use a parametric approach for value approximation, while we develop a non-parametric approach to improve sample efﬁciency of a parametric RL agent.
E XPERIMENTS 5. E XPERIMENT SETTING
To evaluate the sample efﬁciency of our algorithm, we conduct experiments on the benchmark suite of Atari games from the Arcade Learning Environment (ALE;Bellemare et al.
(2013)), which offer various scenes to test RL algorithms over different settings.
We follow the same setting for network architecture and all hyper-parameters as DQN (Mnih et al., 2015).
The raw images are resized to an 84 × 84 grayscale image s t and 4 consecutive frames are stacked into one state.
The Q value network alternates convolutions and ReLUs followed by a 512-unit fully connected layer and an output layer whose size is equal to the number of actions in each game.
Denote Conv( W , F , S ) as the convolutional layer with the number of ﬁlters W , kernel size F and stride S .
The 3 convolutional layers can be indicated as Conv(32,8,4), Conv(64,4,2), and Conv(64,3,1).
We used the RMSProp algorithm (Tieleman & Hinton, 2012) with learning rate α = 0 .
00025 for gradient descent training.
[Figure]
The discount factor γ is set to 0.99 for all games.
We use annealing (cid:15) -greedy policies from 1.0 to 0.1 in training stage while ﬁxing (cid:15)
= 0 .
05 during evaluation.
For associative memory, we also use the same settings for all game.
We set the value of λ as 0.
and associate frequency K as 50.
The memory size is set as 1 million.
We use random projection technique and project the states into vectors with dimension of d = 4 .
For efﬁcient table lookup, we build a kd-tree for these low-dimension vectors.
5. R ESULTS ON A TARI GAMES
We largely follow the training and evaluation protocol as (Mnih et al., 2015).
We train our agents for 10 epochs, each containing 1 million frames, thus 10 million frames in total.
For each game, we evaluate our agent at the end of every epoch for 0.5 million million frames, with each episode up to 18000 frames and start the game with up to 30 no-op actions to provide random starting positions for the agent.
In our experiments, we compare ERLAM with episodic reinforcement learning baselines, MFEC (Blundell et al., 2016), NEC (Pritzel et al., 2017), EMDQN (Lin et al., 2018), EVA (Hansen et al., 2018), as well as an ablation (i.e., DQN with no associative memory).
MFEC directly use the non-parametric episodic memory for action selection, while NEC, EMDQN and EVA combine nonparametric episodic memory and a parametric Q-network.
Different from previous work, ERLAM adopts associative memory to guide the learning of a Q-network.
We tested ERLAM on 25 popular and challenging Atari games.
To evaluate our approach, we follow
Wang et al.
and measure improvement in percentage in score over the better of human and
DQN agent scores for both ERLAM and EMDQN:
[Equation]
To test the sample efﬁciency of our method, we limit our training data in 10 million frames and compare with state-of-the-art results on episodic RL (i.e., EMDQN (Lin et al., 2018)) , which are trained with 40 million frames and reported in their original paper.
The results are shown in Figure 3.
We found that even though our agent uses 4 time less training samples than EMDQN, ERLAM still outperforms EMDQN on 17 games.
It demonstrates that associative memory can efﬁciently guide the learning of a parametric RL agent and our framework of combining associative memory with parametric RL can achieve signiﬁcantly better sample efﬁciency than existing RL algorithms.
We also compare the overall performance (mean and median) of ERLAM with other methods in
Table 1.
We found that ERLAM signiﬁcantly outperforms all baselines over most games.
For the
[Table]
games where ERLAM does not perform very well we summarize the reasons as follows.
First,
ERLAM is good at improving the sample-efﬁciency in near-deterministic environments but may suffer from overestimation in highly stochastic environments, such as Tutankham.
Second, since representations learning is not the focus of this paper, we simply use the naive random projection as the state representations in memory.
Random projection is only used for dimension reduction and does not contain useful high-level feature or knowledge (e g., objects and relations).
Thus in some games with rare revisited states, there are not enough joint nodes in our graph and our algorithm does not perform well, such as FishingDerby and Jamesbond.
To gain better understanding of our superior performance we further plot learning curves (Figure 4) on four games, which include three general good cases (Atlantis, BattleZone, StarGunner) and a bad case (BankHeist) to demonstrate when associative memory works extremely well and when it is not particularly effective.
In addition, we plot the average values of states in memory (Figure 4) for better revealing the performance difference on game scores.
Across most games, ERLAM is signiﬁcantly faster at learning than EMDQN and DQN, but ERLAM only has a slightly better performance than EMDQN on BankHeist.
The reasons lie in two folds.
Firstly, there are more crossed experiences on Atlantis, BattleZone, StarGunner than BankHeist, and thus on the ﬁrst three games the values computed by associative memory are signiﬁcantly larger than those in episodic memory.
Secondly, we observe that the background objects in BankHeist have abnormally changeable appearance and complex behaviors, which are intractable for memory based methods (e.g., MFEC,
NEC, EMDQN and ERLAM) especially with a simple random projection embedding function for state feature abstraction (we also discussion this in Conclusion Section).
It also accounts for the reason why ERLAM and EMDQN have similar performance with DQN.
We also add experiments to verify our superior performance beneﬁts from associative memory rather than representations (e.g., random projection).
As shown in Appendix Figure 5, DQN with only random projections as inputs has a much worse performance than ERLAM and the vanilla DQN, which suggests that it is associative memory that matters.
C ONCLUSION
In this paper, we propose a biologically inspired sample efﬁcient reinforcement learning framework, called Episodic Reinforcement Learning with Associative Memory (ERLAM).
Our method explicitly organizes memorized states as a graph.
We develop an efﬁcient reverse-trajectory propagation strategy to allow the values of new experiences rapidly propagate to all memory items through the graph.
Experiments in Atari games demonstrate that our proposed framework can signiﬁcantly improve sample efﬁciency of current reinforcement learning algorithms.
In the future, there are some interesting research directions that can be pursued within our proposed framework.
Firstly, in this paper, following the work of Blundell et al.
(2016) and Lin et al.
(2018), our state embedding function φ is implemented as random projection.
It is possible to incorporate advanced representation learning approaches that can capture useful features into our framework to support more efﬁcient memory retrieval and further boost up the performance.
Secondly, existing episodic reinforcement learning algorithms mainly focus on value-based methods.
It will be an interesting future work to extend the episodic memory to policy-based methods.
Thirdly, we instantiate our associative memory in the learning phase in this paper.
However, associative memory can also be used in explicit episodic control to further enhance exploitation.
Fourthly, as a episodic
RL approach, ERLAM is also good at improving the sample-efﬁciency in only near-deterministic environments.
For a completely stochastic environment, our model can be extended by storing distribution of Q values (Bellemare et al., 2017; Dabney et al., 2018) instead of the maximum Q value in the associative memory.
A T HEORETICAL C ONVERGENCE
Proof.
Note that our graph-based value propagation is similar to the proof of value iteration (Bellman, 1966; Bertsekas et al., 1995; Sutton & Barto, 2011).
For any estimate of our graph-based action-value function ˆ Q G , we deﬁne the Bellman backup operator B : R |S| → R |S|
[Equation]
where P G ( s (cid:48) | s, a ) deﬁnes transition probability given graph G .
For any action-value function estimates ˆ
[Equation]
So the contraction property of Bellman operator holds:
[Equation]
For the ﬁxed point Q ∗G
[Equation]
we have:
[Equation]
Therefore, we prove that our graph-based value propagation algorithm can converge to unique optimal value
[Equation]
B R AW S CORES ON A TARI G AMES
[Table]
2693.96 StarGunner 604.
9432.97 Tutankham 148.
146.94 YarsRevenge 7614.