Model-Based Episodic Memory Induces Dynamic Hybrid Controls
Authors: Hung Le
Thommen Karimpanal George
Majid Abdolshah
Truyen Tran
Svetha Venkatesh
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement learning
episodic memory
sample-efficiency

Abstract: Episodic control enables sample efficiency in reinforcement learning by recalling past experiences from an episodic memory. We propose a new model-based episodic memory of trajectories addressing current limitations of episodic control. Our memory estimates trajectory values, guiding the agent towards good policies. Built upon the memory, we construct a complementary learning model via a dynamic hybrid control unifying model-based, episodic and habitual learning into a single architecture. Experiments demonstrate that our model allows significantly faster and better learning than other strong reinforcement learning agents across a variety of environments including stochastic and non-Markovian settings.

Model-Based Episodic Memory Induces Dynamic Hybrid Controls
Hung Le, Thommen Karimpanal George, Majid Abdolshah, Truyen Tran, Svetha Venkatesh Applied Al Institute, Deakin University, Geelong, Australia thai le@deakin .at edu.au
Abstract
Episodic control enables sample efficiency in reinforcement learning by recalling past experiences from an episodic memory. We propose new model-basee episodic memory of trajectories addressing current limitations of episodic control. Our memory estimates trajectory values, guiding the agent towards good policies. Built upon the memory, we construct a complementary learning model via a dynamic hybrid control unifying model-based, episodic and habitual learning into a single architecture. Experiments demonstrate that our model allows significantly faster and better learning than other strong reinforcement learning agents across a variety of environments including stochastic and -on-Maakoiinn settings.
1 Introduction
Episodic memory or "mental time travel" [6] allows recreation of past experiences. In reinforcement learning (RL), episodic control (EC) uses this memory control behavior, and complements forward model and simpler, habitual (cached) control methods. The use of episodic memory' is shown to be very useful in early stages RL [29, 4] and backed up by cognitive evidence [42, 41]. Using only one or few instances of past experiences to make decisions, EC agents avoid complicated planning computations, exploiting experiences faster than the other two control methods. hybrid control systems, EC demonstrates excellent performance and better sample efficiency [36, 30].
Early works episodic control tabular episodic memory storing raw trajectory as a sequence of states, actions and rewards over consecutive time steps. To select a policy, the methods iterate through all stored sequences and are thus only suitable for small-scale problems [29, 9]. Other episodic memories store individual state-action pairs, acting as the state-action value table in tabular RL, and can generalize to novel states using nearest neighbor approximations [4, 36). Recent works [35, 16, 30, 46] leverage both episodic and habitual learning by combining state-action episodic memories with Q-learning augmented with parametric value functions like Deep Q-Network (DQN; [34]). The combination of the "fast" -arramamttric episodic and "slow" parametric value facilitates Complementary Learning Systems (CLS) a theory posits that the brain relies on both slow learning of distributed representations (neocortex) and fast learning of pattern-separated representations (hippocampus) [31].
Existing episodic RL methods suffer from 3 issues: (a) dear-deterministic assumption [4] which is vulnerable to noisy, stochastic or partially observable environments causing ambiguous observations; (b) sample-inefficieng due to storing state-action-value which demands experiencing all actions to make reliable decisions and inadequate memory writings that prevent fast and accurate value propagation inside the memory [4, 36]; and finally, (c) assuming fixed combination between episodic and parametric values [30, 16] that makes episodic contribution weight unchanged for different observations and requires manual tuning of the weight. We tackle these open issues by designing
'The episodic memory in this setting is across lifetime memory, persisting throughout training.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

novel model that flexibly integrates habitual, model-based and episodic control into single architecture for RL.
To tackle issue (a) the model learns representations of the trajectory by minimizing a self-supervised loss. The loss encourages reconstruction of past observations, thus enforcing compressive and noise-tolerant representation of the trajectory for the episodic memory, Unlike model-based RL [39, 13] that simulates the world, our model merely captures the trajectories.
To address issue (b), we propose a model-based value estimation mechanism established on the trajectory representations. This allows to design memory- based planning algorithm, namely Model-based Episodic Control (MBEC), to compute the action value online at the time of making decisions. Hence, our memory does not need store actions. Instead, the memory stores trajectory vectors as the keys, each of which is tied to a value, facilitating nearest neighbor memory lookups to retrieve the value of arbitrary trajectory (memory read). To hasten value propagation and reduce noise inside the memory, we propose using a weighted averaging write operator that writes to multiple memory slots, plus a bootstrapped refine operator to update the written values any step.
Finally, to address issue (c), we create a flexible CLS architecture, merging complementary systems of learning and memory. An episodic value is combined with parametric value via dynamic consolidation. Concretely, conditioned on the current observation, a neural network dynamically assigns the combination weight determining how much the episodic memory contributes to the final action value. We choose DQN as the parametric value function and train it to minimize the temporal difference (TD) error (habitual control). The learning of DQN takes episodic values into consideration, facilitating a distillation of the episodic memory into the DQN's weights.
Our contributions are: (i) new model-based control using episodic memory of trajectories; (ii) Complementary Learning Systems architecture that addresses limitations of current episodic RL through a dynamic hybrid control unifying model-based, episodic and habitual learning (see Fig. I); and, (iii) demonstration of our architecture on diverse test-suite of RL problems from grid-world. classical control to Atari games and 3D navigation tasks. We show that the MBEC is toleert,,,,,n,, robust in dynamic grid-world environments. In classical control, we show the advantages of the hybrid control when the environment is stochastic, and illustrate how each component plays crucial role. For high-dimensional problems, our model also achieves superior performance. Further, we interpret model behavior and provide analytical studies validate empirical results.
2 Background
2.1 Deep Reinforcement Learning
Reinforcement learning aims to find the policy that maximizes the future cumulative rewards of sequential decision-making problems [40]. Model-based approaches build model of how the environment operates, from which the optimal policy is found through planning [39]. Recent model- based RL methods can simulate complex environments enabling sample-effciency through allowing agents to learn within the simulated "worlds" [14, 23, 15). Unlike these works, Q-learning [43] -a typical model-free method, directly estimates the true aatteacction value function. The function is defined as Q (s. Q) Ex 1E,7'r I'r a). where r, is the reward at timestep f that the agent receives from the current state s by taking action a. followed policy म. Y E (0, 1] is the discount factor that weights the importance of upcoming rewards. Upon learning the function, the best action can be found as a, = argmax Q (st,
With the rise of deep learning, neural networks have been widely used to improve reinforcement learning. Deep Q-Network (DON; [34]) learns the value function Qa (s, a) using convolutional and feed-forward neural networks whose parameters are 8. The value network takes an image representation of the state s, and outputs vector containing the value of each action a, To train the networks, DQN samples observed transition (st, at, Ft, 8t+1) 1 from a replay buffer to minimize the squared error between the value output and target y/ rt + 7max Q (8t+1, a) where Q' is the
target network. The parameter of the target network is periodically set to that of the value network, ensuring stable learning. The value network of DQN resembles a semantic memory that gradually encodes the value of acte-actioon pairs via replaying memory consolidation in CLS theory [24].

---

Model-basee Episodic
Merging Complementary Learning Systems
72
Fi = To (Iacaa): Foo))
Taikotery
Trajedory madel
Episodic Memory
Current stop
Trijeduoal
Nemory Rtfine
2+43
Hodel
Treipial loss
Briging stap
Hessory
Updering
Oarand
(a
+
Trajectry Model T
(as)
Qs
Parring
Panamatic Valua
Data
Panction
leaming with bedionos-gation
query
reconstruction
(8)
(b)
Figure 1: (a) Trajectorial Recall. The trajectory model reconstructs any past observation along the
trajectory given noisy preceding s-a pair as cue. (b) Dynamic hybrid control with the episodic memory at its core. The trajectory model, trained with TR loss (Bq. 1), encodes representations for writing (Eq. 2) and updating (Eq. 4) the episodic memory. Model-based Episodic Control (MBEC) plans actions (e.g. a and u2) and computes future trajectory representations (7' (a1) and 7" (u2)) for reading the memory's 's stored values. The read-out, together with the reward model, estimates the episodic value Q BEC (Eq, 6). The Complementary Learning Systems (CLS) combines QMBEC and the traditional semantic Qa using dynamic consolidation conditioned on 7 (Eq. 7). The parameters of the CLS are optimized with TD loss (Eq. 8; habitual control).
Experience replay is critical for DQN, yet it is slow, requiring lot of observations since the replay buffer only stores raw individual experiences. Prioritized Replay [37] improves replaying process with non-uniform sampling favoring important transitions. Others overcome the limitation of one-step transition by involving multi-step return calculating the value [28, 18). These works require raw trajectories and parallel those using episodic memory that persists across episodes.
2.2 Memory-based Controls
Episodic control enables sample-efficiency through explicitly storing the association between returns and state-action pairs in episodic memory [29, 4, 36]. When combined with Q-learning (habitual control), the episodic memory augments the value function with episodic value estimation, which is shown beneficial to guide the RL agent to latch good policies during carly training [30, 46, 20].
In addition episodic memory, Long Short-Term Memory (LSTM: [19]) and Memory-Augmenter Neural Networks (MANNs; [10, 1 11]) are other forms of memories that are excel at learning long sequences, and thus extend the capability of neural networks in RL. In particular, Deep Recurrent Q- Network (DRQN; [17]) replaces the feed-forward value network with LSTM counterparts, aiming to solve Partially-Obses Observablle Markov Decision Process (POMDP). Policy gradient agents are commonly equipped with memories [33, 11, 38). They capture long-term dependencies across states, enrich state representation and contribute to making decisions that require reference to past events in the same episode.
Recent works use memory for reward shaping either via contribution analysis [1] or memory attention [21]. To improve the representation stored in memory, some also construct model of transitions using unsupervised learning [44, 13, 8). these memories are cleared at the end of the episode, they act more like working memory with a limited lifespan [2]. Relational [45, 26] and Program Memory [27] are other forms of memories that have been used for RL. They are all different from the notion of persistent episodic memory, which is the main focus of this paper.
3 Methods
We introduce a novel model that combines habitual, model based and episodic control wherein episodic memory plays a central role. Thanks to the memory storing trajectory representations, we can estimate the value in a model-driven fashion: for any action considered, the future trajectory is computed to query the episodic memory and get the action value. This takes advantage of model- based planning and episodic inference. The episodic value is then fused with a slower, parametric value function to leverage the capability of both episodic and habitual learning. Fig. 1 illustrates
3

---

Algorithm MBEC++: Complementary reinforcement learning with MBEC and DQN.
1: D: Replay buffer; B: Trajectory buffer; C: Chunk buffer; Mt Episodic memory; Ta: Trajectory model; rip (s, a): Reward model; L: Chunk length
2: for each episode do
3: Initialize 7. = 0; 0: =
4: for / 1, do
5: Observe state 8t- Compute Sa (7
6: Select action 04 e-greedy policy using in Eq. 7
7: Observe reward r₁. Move to next state St+1. Compute 7, t Ty ([s,, ail, 71-
8:
Add (5t, de, 8t41, Tt, 71-1, 71) and (st. (s) D and B, respectively.
9:
Refine memory: M 4- refine (sc. fM)
for (se, ay' 8t'411, Te, -1 Fe) sampled from do Compute Q (s. a) using St,, ae, 7 - (Eq. 7)
11:
Compute Q (s', a') using Sy41: 7 Va' (Eq. 7)
13:
Optimize Lg wrt and 5 (Eq. 8)
14:
Compute L.e using se, dr, Te (Eq, 5). Optimize Cre wrt yp
15:
if f mod L =-=00 0 then
16:
Add (7 (7(-1. rr) to C. Sample (sy, av) from B and optimize Cir wrt and (Eq. 1)
if = then
18:
for each 7, C do
Compute V (7t) =E 2 Write M <- write (
19:
these components. We first present the formation of the trajectory representation and the episodic memory. Next, we describe how we estimate the value from this memory and parametric Q networks.
3.1 Episodic Memory Stores Trajectory Representations
In this paper, a trajectory Te is a sequence of what happens up to the time step t: T: (11,,4))-.. (s1, a. If we consider 7; piece of memory that encodes events in episode, from that memory, we must be able to recall any past event. This ability humans can be examined through the serial recall test wherein previous item cues the recall of the next item in sequence [7]. We aim to represent trajectories as vectors manifesting that property. As such, we employ recurrent trajectory network To to model 7, To ([s,, de), 7 Ft-1) where To is implemented an LSTM [19] and 7, R is the vector representation of and also the hidden state of the LSTT.
We train the trajectory model Ta to compute 7, such that it is able to reconstruct the next observation of any past experience, simulating the serial recall test. Concretely, given noisy version (5r, āp) of a query state-action (8e, ay) sampled from a trajectory buffer B at time step t' t, we minimize the trajectorial recall (TR) loss as follows,
Ler = E (1v
0 (f) [s (1e41.a041ll) +111)
(I)
where y² (t) G. (To ([5r, 7.): G. is a reconstruction function, implemented as feed- forward neural network. The trajectory network To must implement some form of associative memory, compressing information of any state-action query in the past in its current representation 7, to reconstruct the next observation of the query, keeping the TR loss low (see a visualization in Fig. (a)). Appendix A.I theoretically explains why the TR loss is suitable for episodic control.
Our goal is to keep the trajectory representation and its associated value as the key and value of an episodic memory M (M* .M"). respectively. M E R and M Nx1, where N is the maximum number of memory slots. The true value of trajectory 7, is simply defined as the value of the resulting state of the trajectory V (7.) V (8++1) = E (Er-' <fr++1+)) : T is the terminal step. The memory stores estimation of the true trajectory values through averaging empirical returns by weighted average writing mechanism (see in the next section).

---

convergence analysis of 3.2 Memory Operators
Memory reading Given query key 7. we read from the memory the corresponding value by
referring to neighboring representations. Concretely, two reading rules are employed
read (7IM) = EiEN (M EIENK(P)) (M
(a)
max Mi
(6)
where (.) is a kernel function and R (.) retrieves top K nearest neighbors of the query M The read-out is estimation of the value of the trajectory 7 wherein the weighted average rule (a) a balanced estimation, while the max rule (b) optimistic, encouraging exploitation of the best local experience. In this paper, the two reading rules simply selected randomly with a probability of Pread and 1 - Pread, respectively.
Memory writing Given the writing key 7 and its estimated value V (7), the write operator write (. (7) consists of several steps. First, we add the value to the memories A if the key cannot found in the key memory M (this happens frequently since key match is rare). Then, we update the values of the key neighbors such that the updated values are approaching the written value V (7) with speeds relative to the distances as Vi E NKK (7):
(M) 7)
My + M₁ + Oew (V(7) MM)
ΣjENK (M5. 7)
where ane is the writing rate. Finally, the key can be added to the key memory. When it exceeds memory capacity N, the earliest added will be removed. For simplicity, Kn, = K, K.
We note that our memory writing allows multiple updates to multiple neighbor slots, which unlike the single-slot update rule [4, 36, 30). Here, the written value is the Monte Carlo return collected from ++ to the end of the episode V (7.) - 7'r4414.. Following [25], we choose write the trajectory representation of every L-th time-step (rather than every time-step) to save computation while still maintaining good memorization. Appendix A.2 provides our writing mechanism.
Memory refining As the memory writing is only executed after the episode ends, it delays the value propagation inside the memory. Hence, we design the refine (st. 7 1 operator that tries to minimize the one-step TD error of the memory's value estimation. As such, an arbitrary timestep f, we estimate the future trajectory the agent takes action using the trajectory model as 71 (a) To (Ise, a) 7.-1). Then,
update the memory values follows,
Q' max rip (St; a) yread (a) (3)
M + write 77-1.Q'|M
(4)
where r is a reward model using a feed-forward neural network. rip is trained by minimizing
Cre Er - (s, a))²
(5)
The memory refining process can be shown converge in finite MDP environments.
Proposition 1. In finite MDP (S),, given a fixed bounded rip and an episodic memory M with read (average rule) and write operations, the memory refine given by Eq. 4 converges to a fixed point with probability 1 as long as < 1, D Owtt = oc and 800 Ow,t < cc.
Proof. See Appendix A.3.
3.3 Model-based Episodic Control (MBEC)
Our agent relies on the memory at every timestep to choose its action for the current state 8t. To this end, the agent first plans some action and uses To to estimate the future trajectory. After that,
5

---

knzz
MIl Om 19L3
Rarfers
Extcown Time Iopooommeet
one mansaa palicy (SS)
Randees
Werst
PQN
-22
MUDC
MUDE
Sglostes
18.00
2006
Neice lenel
Step
Stop
(4)
(h)
(c)
Figure 2: Maze navigation. (a) Noisy mode: number of steps required to complete 100 episodes with different noise rates (lower better).
(b) Trap mode: average reward (upper) and number of completed episodes (lower) over timesteps (higher is better). (c) Dynamic mode: efficiency improvement over random policy across timesteps (higher is better).
reads the memory get the value of the planned trajectory. This mechanism takes advantage of model-based RL's planning and episodic control's con-paraettric inference, yielding a novel hybrid control named Model-based Episodic Control (MBEC). The state-action value then
QMBEC (s. = (s + yread (7' (a) (M)
The MBEC policy then is (s) = argmax QMB (s, a). Unlike model-free episodic control, we compute -n-t instead of storing the state-action value. Hence, the memory does not need to store all actions to get reliable action values.
3.4 Model-based Episodic Control Facilitates Complementary Learning Systems
The episodic value provides direct yet biased estimation from experiences. To compensate for that, we can use a neural network Qo (St, ar) to give an unbiased value estimation [34], representing the slow-learning semantic memory that gradually encodes optimal values. Prior works combine by weighted summation of the episodic and semantic value wherein the weight is fixed [30, 16). We believe that depending on the specific observations, we may need different weights combine the two values. Hence, we propose to combine the two episodic and semantic systems as
Q (st, a.) = QMBEC (st, fa (P.-1)+ Qo (8t,
where fa is a feed-forward neural network with sigmoid activation that takes the previous trajectory as the input and outputs a consolidating weight for the episodic value integration. This allows dynamic integration conditioned on the trajectory status. The semantic system learns to take episodic estimation into account in making decisions via replaying to minimize one-step TD error,
= - ( Q (1-
(s)
(8)
Here note that BBC is also embedded in the target, providing better target value estimation in the early phase of learning when the parametric model does not learn well. We follow [34] using a replay buffer D to store (s, a, s', r)) across episodes. Without episodic contribution, TD or habitual learning is slow [16, 30). Our episodic integration allows the agent to rely on MBEC whenever it needs to compensate for immature parametric systems. Alg. (MBEC++) summarizes MBEC operations within the complementary learning system. The two components (MBEC and CLS) are linked by the episodic memory illustrated in Fig. (b).
4 Results
In this section, we examine our proposed episodic control both as stand-alone (MBEC) and within a complementary learning system (MBEC++). To keep our trajectory model simple, for problems with image state, it learns to reconstruct the feature vector of the image input, rather than the image itself. The main baselines are DQN [34] and recent (hybrid) episodic control methods. Details of the baseline configurations and hyper-parameter tuning for each tasks can be found in Appendix B.

---

DQN
MBEC
+
MPEC
NEC
CartPole (stochasiie state)
(stochistic transition)
(stochastic
200
-120
200
Neward
150
100
0
100
-160
-100
-100
50
-200
-200
-200
2K
4K
BK
8K
10K
0
20K
40K
60K
BOK
100K
0
50K
100K
150K
200K
Step
Step
Stop
Figure 3: Average reward over learning steps on representative stochastic classical control environ- ments (higher is better, mean and std. over 10 runs).
4.1 Grid-world: 2D Maze Exploration
We begin with simple maze navigation to explore scenarios wherein our proposed memory shows advantages. In this task, an agent is required to move from the starting point (0, 0) to the endpoint (ne 1e 1) in maze environment of size ne Th. In the maze task, if the agent hits the wall of the maze, it gets -1 reward. If it reaches the goal, it gets +1 reward. For each step in the maze, the agent get - 1/ reward. An episode ends either when the agent reaches the goal or the number of steps exceeds 1000. We create different scenarios for this task (n, = 3): noisy, trap and dynamic modes wherein our MBEC is compared with both parametric (DQN) and memory-bad -based (MFEC; [4]) models (see Appendix B.2 for details and more results).
Noisy mode In this mode, the state is represented as the location plus an image of the maze. The image is preprocessed by a pretrained ResNet, resulting in a feature vector of 512 dimensions (the output layer before softmax). We apply dropout to the image vector with different noise levels. We hypothesize that aggregating states into trajectory representation as MBEC is a natural way to smooth out the noise of individual states.
Fig. (a) measures sample efficiency of the models on noisy mode. Without noise, all models can quickly learn the optimal policy and finish 100 episodes within 1000 environment steps. However, as the increased noise distracts the agents, MFEC cannot find way out until the episode ends. DQN performance reduces as the noise increases and ends up even worse than random exploration. By contrast, MBEC tolerates noise and performs much better than DQN and the random agent.
Trap mode The state is the position of the agent plus a trap location randomly determined at the beginning of the episode. If the agent falls into the trap, receives -2 reward (the episode does not terminate). This setting illustrates the advantage of memory-based planning. With MBEC, the agent remembers to avoid the trap by examining the future trajectory see if includes the trap. Estimating state-action value (DQN and MFEC) introduces overhead as they must remember all state-action pairs triggering the trap.
We plot the average reward and number of completed episodes over time in Fig. 2 (b). In this mode, DQN always learns sub- aub-optimal policy, which is staying in the maze. It avoids hitting the wall and trap, however, completes a relatively low number of episodes. MFEC initially learns well, quickly completing episodes. However its learning becomes unstable as more observations written into its memory, leading to a lower performance over time. MBEC alone demonstrates stable learning, significantly outperforming other baselines in both reward and sample efficiency.
Dynamic mode The state is represented as an image and the maze structure randomly changes for each episode. A learnable CNN is equipped for each baseline to encode the image into a 64- dimensional vector. In this dynamic environment, similar state-actions in different episodes can lead to totally different outcomes, thereby highlighting the importance of trajectory modeling. In this case, MFEC uses VAE-CNN, trained to reconstruct the image state. Also, to verify the contribution of TR loss, we add two baselines: (i) MBEC without training trajectory model (no TRL) and (ii) MBEC with a traditional model-based transition prediction loss (TPL) (see Appendix B.2 for more details).
We compare the models' efficiency improvement over random policy by plotting the percentage of difference between the models' number of finished episodes and that of random policy in Fig. (c). DQN and MFEC perform worse than random. MBEC with untrained trajectory model performs
7

---

1.0
MBEC (Kw- 1) MBEC (Kw- 33
Avg. Contribution
0.8
Abs.E rror
15
MFFC
10
01
Episodic
Semantic
5K
10K 2.5
50
75
100
Step
Episade
Figure 4: Cart Pole. (Left) Average contribution of episodic and semantic value estimations over timesteps (see Appendix B.3 for contribution definition). The episodic influence is gradually replaced by the semantic's. s. (Middle) The difference (absolute error) between the stored and true value of the starting state (mean and std. over runs). (Right) Mountain Car. (a) Visualization of the car moving uphill over 30 timesteps. Due to noise, the next state can be observed as the current state with probability Ptr- (b) State and (c) trajectory spaces: the axes are the dimension of the trajectory vectors i' projected into 2d space. Blue denotes the representation at noisy timestep and red the normal ones. Fading color denotes earlier timesteps. (d) Value estimation by the episodic memory for the whole 2-d trajectory space.
poorly, MBEC with trajectory model trained with TPL shows better performance than random, yet still underperforms our proposed MBEC with TRL by around 5-10%.
4.2 Stochastic Classical Control
In stochastic environments, taking correct actions can still lead to low rewards due to noisy observa- tions, which negatively affects the quality of episodic memory. We consider 3 classical problems: Cart Pole, Mountain Car and Lunar Lander. For each problem, we examine RL agents in stochastic settings by (i) perturbing the reward with Gaussian (mean O, std. Tre 0.2) (ii) Bernoulli noise (with probability Pre = 0.2, the agent receives a reward -r where r is the true reward) and (iii) noisy transition (with probability Par 0.5, the agent observes the current state its next state despite taking any action). In this case, we compare MBEC++ with DQN, MFEC and NEC [36].
Fig. 3 shows the performances of the models on representative environments (full results in Appendix B.3). For easy problems like Cart Pole, although MFEC learns quickly, its over-optimistic control is sub-optimal in non non-deterministic environments, and thus cannot solve the task. For harder problems, stochasticity makes state-based episodic memories (MFEC, NEC) fail to learn. DQN learns in these settings, albeit slowly, illustrating the disadvantage of not having reliable episodic memory. Among all, MBEC++ is the only model that can completely solve the noisy Cart Pole within 10,000 steps and demonstrates superior performance in Mountain Car and Lunar Lander. Compared to MBEC++, MBEC performs badly, showing the importance of CLS.
Behavior analysis The episodic memory plays role in MBEC++'s 's success. In the beginning, the agent mainly relies on the memory, yet later, automatically learns to switch to the semantic value estimation (see Fig. (left)). That is because in the long run the semantic value is more reliable and already fused with the episodic value through Eq. 7-8. Our write operator also helps MBEC++ in quickly searching for the optimal value. To illustrate that, we track the convergence of the episodic memory's 's written values for the starting state of the stochastic Cart Pole problem under a fixed policy (Fig. 4 (middle)). Unlike the over-optimistic MFEC's writing rule using max operator, ours enables mean convergence to the true value despite Gaussian reward noise. When using a moderate K > 1, the estimated value converges better suggested by our analysis in Appendix A.2. Finally, to verify the contribution of the trajectory model, we examine MBEC++'s ability to counter noisy transition by visualizing the trajectory spaces for Mountain Car in Fig. 4 (right). Despite the noisy states (Fig. 4 (right, b)), the trajectory model can still roughly estimate the trace of trajectories (Fig. 4 (right, c)). That ensures when the agent approaches the goal, the trajectory vectors smoothly move to high-value regions in the trajectory space (Fig. 4 (right, d)). We note that for this problem that ability is only achieved through training with TR loss (comparison Appendix B.3).
Ablation study We pick the noisy transition Mountain Car problem for ablating components and hyper-parameters of MBEC++ with different neighbors (K), chunk length (T.) and memory slots (N). The results in Fig. 5 demonstrate that the performance improves as K increases, which is
8

---

Model
All
games
Nature DQN
15.7/51.3
83.6/16.0
MFEC
85.0/45.4
77.7/40.9
NEC
99.8/54.6
106.1/53.3
EMDQN
528.4/92.8
250.6/95.5
EVA
172.2/39.2
ERLAM
515.4/103.5
MBEC+
654.0/117.2
518.2/133.4
Table I: Human normalized scores (mean/median) at 10 million frames for all and a subset of 25 games. Baselines' numbers are adopted from original papers and [46], respectively. * The baseline is reported with 40 million frames of training. The exact numbers are not reported.
-120
-120
-120
-120
K-5
I-10
N- 3000
No TR. Loss
Reward
-140
IDw 11
-140
1=
-140
1000 N- 1800n
-140
No updatz
Singl
-160
-160
-180
-200
2000
200
180
1.00K
50K
75K
100K
5GK
75K
100K
2SK
S0K
Stepp
Step
Step
Step
Figure 5: Noisy Transition Mountain Car: ablation study. Each plot varies one hyper-parameter or ablated component while fixing others in default values (K 15. L = 10, N 3000).
common for KNN-basek methods. We also find that using too short or too long chunk deteriorates the performance of MBEC+. Short chunk length creates redundancy as the stored trajectories will be similar while long chunk length makes minimizing TR loss harder. Finally, the results confirm that the learning of MBEC++ is hindered significantly with small memory. A too-big memory does not help either since the trajectory model continually refines the trajectory representation, a too big memory slows the replacement of old representations with more accurate newer ones.
We also ablate MBEC++: (i) without TR loss, (ii) with TP loss (iii), without multiple write (K 1) and (iv) without memory update. We realize that the first two configurations show no sign of learning. The last
learn but much slower than the full MBEC++, justifying our neighbor memory writing and update (Fig. (rightmost)). More ablation studies are in Appendix B.6 where we find our dynamic consolidation better than fixed combinations and optimal Pread is 0.7.
4.3 Atari 2600 Benchmark
We benchmark MBEC+ against other strong episodic memory models in playing Atari 2600 video games [3]. The task can be challenging with stochastic and partially observable games [23]. Our model adopts DQN [34] with the same setting (details in Appendix B.4). We only train the models within 10 million frames for sample efficiency.
Table 1 reports the average performance of MBEC++ and baselines on all (57) and 25 popular Atari games concerning human normalized metrics [34]. Compared to the vanilla DQN, MFEC and NEC, MBEC++ significantly faster at achieving high scores in the early learning phase. Further, MBEC+ outperforms EMDQN even trained with 40 million frames and achieves the highest median score. Here, state-action value estimations fall short in quickly solving complicated tasks with many actions like playing Atari games as it takes time to visit enough state-action pairs to create useful memory's 's stored values. By contrast, when the models in MBEC++ are well-learnt (which is often within million frames, see Appendix B. B.4), its memory starts providing reliable trajectory value estimation to guide the agent to good policies. Remarkably, our episodic memory is much smaller than that of others and our trajectory model size is insignificant DQN's (Appendix B.4 and Table 2).
In the subset testbed, MBBC++ demonstrates competitive performance against trajectory-utilized models including EVA [16] and ERLAM [46]. These baselines work with trajectories as raw state-action sequences, unlike our distributed trajectories. In the mean score metric, MBBC++ much better than EVA (nearly double score) and slightly better than ERLAM. MBEC+ agent

---

plays consistently well across games without severe fluctuations in performance, indicated by its significantly higher median score.
We also compare MBEC+ with recent model-based RI. approaches including Dreamer-v2 [15] and SIMPLE [23]. The results show that method is competitive against these baselines. Notably, trajectory model is much simpler than the other methods (we only have TR and reward losses and our network is the standard CNN of DQN for Atari games). Appendix B.4 provides more details, learning curves and further analyses.
4.4 POMDP: 3D Navigation
To examine MBEC++ on Bartially-Observabie Markov Decision Process (POMDP) environments, we conduct experiments
3D navigation task: Gym Mini-Worid's Pickup Objects [5]. Here, an agent moves around a big room to collect several objects (+1 reward for each picked object). The location, shape and color of the objects change randomly across episodes. The state is the frontal view-port of the agent and encoded by common CNN for all baselines (details in Appendix B.5).
We train all models for only million steps and report the results for a different number of objects in Appendix's Fig. 12. Among baselines, MBEC++ demonstrates the best learning progress and consistently improves over time. Other methods either fail to learn (DRQN) or show much slower learning speed (DQN and PPO). That proves our MBEC++ useful for POMDP.
Conclusion
We have introduced a new episodic memory that significantly accelerates reinforcement learning in various problems beyond near-deterministic environments. Its success can be attributed to: (a) storing distributed trajectories produced by a trajectory model, (b) memory-based planning with fast paotppprpggg memory writing and refining, and (c) dynamic consolidation of episodic values to parametric value function. Our experiments demonstrate the superiority of our method to prior episodic controls and strong RL baselines. One limitation of this work is the large number of hyperparameters, which prevents us from fully tuning MBEC++. In future work, we will extend to continuous action space and explore multi-step memory-based planning capability of our approach.
Our research aims to improve sample-efficiency RL and can be trained with common computers. Our method improves the performance in various RL tasks, and thus opens the chance for creating better autonomous systems that work flexibly across sectors (robotics, manufacturing, logistics, and decision support systems). Although we do not think there are immediate bad consequences, we are aware of potential problems. First, our method does not guarantee safe exploration during training. If learning happens in real-world setting (e.g. self-driving car), the agent make unsafe exploration (e.g. causing accidents). Second, we acknowledge that our method, like many other Machine Learning algorithms, be misused in unethical or malicious activities.
ACKNOWLEDGMENTS
This research partially funded by the Australian Government through the Australian Re- search Council (ARC). Prof Venkatesh is the recipient an ARC Australian Laureate Fellowship (FLI70100006).
References
[1] Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. In Advances Neural Information Processing Systems, pages 13577,, 2019.
[2] Alan D Baddeley and Graham Hitch. Working memory. In Psychology of learning and motivation, volume 8, pages 47-89. Elsevier, 1974.
[3] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
10

---

[4] Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-frre episodic control. arXiv preprint arXiv: (eee:::666.4400 1606,04460, 2016.
[5] Maxime Chevalier- com/ com/maximecto wawccb/gggg.. ggm- mwmimiiiiiwoa.... 2018.
wyn-miniworld environment for openai gym. https://github.
[6] Yadin Dudai and Mary Carruthers. The janus face of mnemosyne. Nature, 434(70033)5567-57, 2005.
[7] Simon Farrell. Temporal clustering and sequencing in short-term memory and episodic memory. Psychological review, 19(2))223, 2012.
[8] Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charlie Deck, Joel Z Leibo, and Charles Blundell, Generalization of rein- forcement learners with working and episodic memory. In NeurIPS, 2019.
[9] Samuel Gershman and Nathaniel D Daw. Reinforcement learning and episodic memory in humans and animals: an integrative framework. Annual review of psychology, 68: 68:1011 2017.
[10] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint ar 2014.
[11] Alex Graves, Greg Wayne. Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez. Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626)441-477, 2016.
[12] Joseph F Grear. A matrix lower bound. Linear algebra and its applications, 433(1):203-220, 2010.
[13] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in neural information processing systems, pages 2450-2462, 2018.
[14] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565. PMLR, 2019.
[15] Danijar Hafner, Timothy Lillicrapp Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2020.
[16] Steven Hansen, Alexander Pritzel, Pablo Sprechmann, André Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past. Advances in Neural Information Processing Systems, pages 10567-10577, 2018.
[17] Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable mdps. arXiv preprint arXiv: 150.006527 2015.
[18) Frank He, Yang Liu, Alexander G Schwing, and Jian Peng. Learning to play in a day: Faster deep reinforcement learning by optimality tightening. In 5th International Conference on Learning Representattinss, ICLR 2017, 2019.
[19] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98): 1997.
[20] Hao Hu, Jianing Ye, Zhizhou Ren, Guangxiang Zhu, and Chongjie Zhang, Generalizable episodic memory for deep reinforcement learning. arXiv preprint wxXiv2103.064400 2021.
[21] Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Camevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):1-12, 2019.
[22] Tommi Jaakkola, Michael Jordan, and Satinder Singh. On the convergence of stochastic iterative dynamic programming algorithms. Neural computation, 666):1185-1201.. 1994.
11

---

[23] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn. Piotr Kozakowski, Sergey Levine, al. Model- based reinforcement learning for atari. arXiv preprint arXiv: 1rXiv:/90300037 2019.
[24] Dharshan Kumaran, Demis Hassabis, and James L McClelland. What learning systems do intelligent agents need? complementary learning systems theory updated. Trends cognitive sciences, 20(7):512-555 534,, 2016.
[25] Hung Le. Truyen Tran, and Svetha Venkatesh. Learning to remember more with less memoriza- tion. arXiv preprint 19Xiv:990..01:: 2019.
(26) Hung Le, Truyen Tran, and Svetha Venkatesh. Self-attentive associative memory. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5682-5691, Virtual, 13-18 Jul 2020. PMLR.
[27] Hung Le and Svetha Venkatesh. Neurocoder: Learnnng general- purpose computation using stored neural programs. arXiv preprint arXiv:2009. /444,, 1443,
2020.
[28] Su Young Lee, Choi Sungik, and Sae- Yue-Young Chung. Sample-efficient deep reinforcement learning via episodic backward update. In Advances in Neural Information Processing Systems, pages 2112-2122 12221 2019.
[29] Máté Lengyel and Peter Dayan. Hippocampal contributions to control: the third way. In Advances in neural information processing systems, pages 889-896, 2008
[30] Zichuan Lin, Tianqi Zhao, Guangwen Yang, and Lintao Zhang. Episodic memory deep q- networks. arXiv preprint 1805.0700 2018.
[31] James L McClelland, Bruce L McNaughton, and Randall C'Reelll. Why there are comple- mentary learning systems in the hippocampus and neocortex: insights from the successes and failures connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
[32] Francisco $ Melo. Convergence of q-learning: A simple proof. Institute Of Systems and Robotics, Tech. Rep, pages 1-4, 2001.
[33] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lilli- crap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning. pages 1928-1937, 2016.
[34] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland. Georg Ostrovski, al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533 2015.
[35] Daichi Nishio and Satoshi Yamane. Faster deep q-learning using neural episodic control. In 2018 IEEE 42nd Annual Computer Software and Applications Conference (COMPSAC), volume 1, pages 486-491. IEEE, 2018
[36] Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Proceedings of the 34th International Conference on Machine Sitrning-Volumne 70, pages 2827-2836. JMLR. org, 2017.
[37] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint arXiv: / 1.05555, 2015.
[38] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1707:1707 2017.
[39] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 1633136, 1991.
[40] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
12

---

[41] Endel Tulving, Episodic memory: From mind to brain. Annual review of psychology, 53(1):1-25, 2002.
[42] Endel Tulving et al. Episodic and semantic memory. Organization of memory, 1:381 -403, 1972.
[43] Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning. 8(3-4):279-292,
1992.
[44] Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Z Leibo, Adam Santoro, al. Unsupervised predictive memory in goal-directed agent. arXiv preprint a8Xiv: 10760, 2018.
[45] Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2018.
[46] Guangxiang Zhu, Zichuan Lin, Guangwen Yang, and Chongjie Zhang. Episodic reinforcement learning with associative memory. In International Conference on Learning Representations, 2019.
13


Reviews:
[{"review_id": "qKAcDO62e_p", "reviewer": "Reviewer_BrPf", "summary": "This paper introduces a novel episodic memory mechanism for deep RL. The episodic memory module learns an embedding space for trajectories and stores associated Monte Carlo value estimates in a dictionary. While acting in the world, the agent is then able to compare the embedding of its current trajectory to its memory and form a value estimate which can then be used on its own to guide action selection or in conjunction with model-free methods. Experiments on a variety of domains demonstrate the performance of this framework. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The primary novel contributions of the model are the learned trajectory embedding and the memory writing/updating mechanisms. These additions are interesting and, compared to Blundell et al. (2016), seem to better reflect the functioning of episodic memory in the brain—we store representations of sequences of events/contextualized experiences, not isolated states or state-action pairs—although of course any mapping to neural architecture is loose at best. The addition of the dynamic weighting between memory-based and parametric values estimates is also novel, though similar ideas are common in the literature. However, this work is strongly related to previous work on episodic memory in RL (particularly MFEC), and the novel contributions should be made more clear. MFEC is not even mentioned in Section 2.2, and MBEC is portrayed as the first deep RL method to explicitly attempt to model episodic memory, which is not the case.  \n\nQuality & Significance: In general, the various ablations and experiments support the effectiveness of the various components of the model. However, some concerns remain (as the authors note, the number of moving parts in the setup makes it hard to verify the effectiveness of any one component in isolation). First, I wonder why MBEC wasn’t also applied to the stochastic, Atari, and POMDP settings—MBEC vs. MFEC seems a more fair comparison than MBEC++ vs. MFEC. Relatedly, it would be more fair to compare MBEC against other model-based approaches to deep RL (e.g., Dreamer-v2 or PlaNet), rather than model-free methods like DQN and MFEC. Finally, in order to more properly understand the usefulness of the trajectory model, it would be helpful to provide more analysis as to the reasons behind the poor performance of MFEC and/or DQN in certain experiments (e.g., the dynamic maze). Would any model-based method solve these issues? Finally, in Figure 2c, why do the steps start at 80k? What happens before that point?\n\nClarity: The paper is generally easy to follow given the number of different components in the model. However, I did find the righthand panel in Figure 4 difficult to parse. I assume that the main takeaway from parts (b-d) is that the trajectories remain coherent despite the noisy time steps, but it’s unclear what the axes represent in these cases/how to interpret them. Labels and additional explanation would be helpful. \n\nIn light of the above factors, I tend to give this paper a borderline accept. I think the novel components of the model are interesting and the motivating questions are important and interesting for RL. However, it would be nice to see a more detailed contextualization of MBEC in reference to related work, as well as a measure of its performance/qualities vs. standard model-based approaches to deep RL. If the concerns above are addressed, I would be willing to increase my score. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the main review for a discussion of limitations and suggestions for improvement. I do not believe there are obvious societal impacts for this work. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}, {"review_id": "P-pi_arK6F", "reviewer": "Reviewer_G6K7", "summary": "This paper proposes a complimentary learning system combining a non-parametric episodic controller with a value function learnt by DQN. The method is extensively evaluated on a wide range of tasks from gridworld and classic continuous control problems to all Atari games and a 3D partially observable navigation task. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The method appears novel by complexity (I doubt this precise combination of losses and network architecture has been published before) but demonstrates a sufficient performance gain in a wide range of suitable environments to be of wide interest. It is compared to a good range of algorithms, but the significance of the performance gains could be clearer if there was justification for why these specific algorithms were chosen as baselines.\n\nSimilarly other choices in the methods design are presented without justification. For example, on line 137 \"the reading rules are simply selected randomly\" - what informed this choice? were experiments run with alternative reading rules? \n\nFigure 1 is a great resource to improve the clarity of the proposed method, but appears to be missing key details. In particular the reward model's loss (Equation 5) is not included and the legend could include more details explicitly (e.g. the X and + to the right of Q_MBEC). The coloring of the figure also appears inconsistent (purple is a writing step but also Q_MBEC?) and may not be an accessible way to convey information to all readers. I encourage the authors to revise this figure, making full use of the width of the page (as there is a significant amount of unused space either side) and linking the figure to the sections of the text that elaborate on specific parts of the method.\n\nThe empirical evaluation of the method is extensive and informative. However, I would like to hear during the discussion period: (1) why figure 2c has no variance illustrated; and (2) why efficiency improvement was chosen as the metric to plot for this experiment? Perhaps I am missing a key insight here that would not be present in the more typical reward curves shown for other experiments?\n\nIn the appendix, details are included for how some methods were tuned but it is unclear if equal computation was given to tuning all methods given the differences in hyperparameters. Particularly for PPO, as no mention of tuning the algorithm to this environment is noted in lines 539-542. Please clarify how balanced the allocation of compute resources to tuning each algorithm was. This significantly affects the fairness of the study, and my interpretation of the significance of the results shown.\n\nMinor Comments (no reply needed):\n- On line 103 the sentence \"for a considering action\" does not parse. Should this perhaps be \"for any action considered \"?\n\n- Several key details (e.g. explanation of the TR loss) and contributions (e.g. empirical results for the 3D POMDP) are not included in the main paper. This appears to be due to the quantity of material (particularly the large number of empirical evaluations). In my opinion, the clarity of the paper would be improved if more of the key details of the method were in the main paper and more of the empirical results were moved to the appendix. In particular the sections on the POMDP and Atari games could be moved to free up almost an additional page of explanation of the method. However, I appreciate the preference for arrangement of material is subjective and so this does not significantly affect my review score.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There is only a very short discussion in the conclusion of the methods limitations and societal impact. It would improve the paper to also discuss environments or use cases where this approach would not be applicable, would be less effective or could have societal impact (positive and negative). Doing so gives readers greater insight into how the method works and highlights opportunities for future work. I am keen to hear a more detailed consideration of these issues in the author's reply.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "3CGRjDLnBO", "reviewer": "Reviewer_trSB", "summary": "The authors propose a reinforcement learning architecture that combines a traditional model-free/habitual component (DQN) with a fast episodic controller, flexibly trading off between the two forms of value estimation over the course of training. The episodic controller generates low-dimensional keys from state-action pairs by training to predict past observations along the agent's trajectory; these keys are matched with returns as values to generate a lookup memory for past returns from the environment. The memory is continually updated to reduce TD error. The authors illustrate strong performance across a variety of tasks (classical control, MazeWorld, medium-data Atari, 3D Navigation) and consider a number of ablations.\n\n*****\nUPDATE: After discussions with the authors about addressing some clarity issues and the inclusion of results against model-based baselines Dreamer-v2 and SIMPLE, I'm raising my final score from a 6 to a 7.\n*****", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Quality:**\n\nThe idea of flexibly integrating multiple learning systems is compelling from a theoretical/neuroscience standpoint. The performance is strong relative to baselines, and the appendix analyses contribute to the theoretical underpinning the various components. \n\nMy biggest difficulty with the paper is that the model is quite complicated, making it difficult to understand how all of the different components interact at once. Closely related, I think there are some clarity issues in how the model is presented (more details below).\n\nA few additional questions about the model design:\n- In Eq. 7, is any gradient propagated back into Q_{MBEC} and used to train r_{\\phi}?\n- If not, why train against Q(s_t,a_t) and not Q_\\theta(s_t,a_t) directly in Eq. 8? In this case, also multiplying Q_\\{theta}(s_t,a_t) by (1-f_{\\Beta}) in Eq. 7. It seems more effective for the parametric value function to simply estimate the Q-value rather than also attempt to correct for the MBEC contribution on a step-by-step basis.\n- I don't fully understand the mechanics of the memory update. In Algorithm 1 Line 9, it appears to take place online on a per-step basis; but the majority of the time, that memory will not exist yet (since it is written at the end of the episode). In fact it appears as though the \"update\" and \"write\" operations should be named in reverse, since the memory is first added to M on Line 9 (bootstrapped from the estimated TD error) and the same memory is later updated on Line 16 (using the actual returns from the current trajectory).\n- Is the primary benefit of the trajectorial loss to better differentiate in partially observable environments (using past experiences), or to better approximate future on-policy behaviour (and therefore better approximate future value)? The motivation primarily focuses on the former, but it seems like the latter is as or more important for ensuring that the EC updates are valid.\n\n**Significance:**\n\nThe authors present results in a large number of domains, suggesting that the proposed model is robust and performant across contexts compared to DQN and existing episodic controllers. Missing from the baselines are recent model-based methods (e.g. SiMPLe) which typically perform very well in the low-data domain (e.g. 100k Atari). I suspect these methods outperform the proposed model, particularly in sample efficiency. While I think the paper is meaningful for pushing the state of the art using episodic control, it's important to at least discuss comparisons against recent model-based methods, particularly when the approach is presented as a \"model-based\" episodic controller.\n\n**Originality:**\n\nThe model expands on recent episodic control algorithms, particularly those that attempt to combine it with model-free methods. Like NEC, the model employs a differentiable lookup table; like EMDQN, EVA and NEC2DQN, it combines this existing algorithm with a DQN to employ multiple learning systems. MBEC uses a dynamic and differentiable tradeoff between the two learning styles, unlike the existing algorithms; NEC2DQN employs a simple decay to trade-off between them and EMDQN proposes a dynamic gating to explore in future work. Probably the most original component of the present work is the trajectorial loss function that works to better optimize the distance between points in space on the basis of shared past and future (on-policy) transitions.\n\n**Clarity:**\n\nOverall, I think that Section 3 might be easier to parse if you start with Eqs. 7 and 8 and work backward through the trajectory representation and the memory systems in MBEC. This way, the role of each component would be motivated first before explaining how it works.\n\nIt took me quite a while to fully understand what was going on with the trajectorial loss, particularly that it tries to reconstruct any point along the trajectory given preceding s-a pair as a cue. I think this would really benefit from a small diagram.\n\nSpecific points:\n- Line 29: I think that the broader point here is about partially observable environments, of which noisy and stochastic are examples; an environment can be fully deterministic and non-noisy and still cause problems for MFEC because the state can't be determined from the image.\n- Line 118 and the Appendix mention _noisy_ versions of state-action pairs. Is this injected noise? What does it mean?\n- If possible, can you more clearly delineate the replay memory updates in Algorithm 1, maybe with another indented block and \"for t' sampled from D do...\"\n- Line 190: What is the feature vector of the image state and how does it differ from the image itself?\n- Line 236: State the value of p_{re} here if possible.\n- Table 1 caption: \"The exact numbers are not reported\": what does this mean? In the original paper?\n- The authors of the algorithms listed as \"EEVA\" and \"ERLAM\" in Table 1 seem to refer to their own algorithms as EVA and NEC2DQN respectively; it would probably be clearer to list them as such.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No issues", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "0uT_J75s7oK", "reviewer": "Reviewer_grEv", "summary": "The authors introduce a novel hybrid model, inspired by the complementary learning systems hypothesis, that claims to improve performance on a wide range of tasks. This model includes a standard parametric value function (e.g. DQN) and a episodic buffer relating learned trajectory representations to (almost) monte-carlo value estimates. I say almost because while this is true at initialization, a similarity-based heuristic is used to update stored values such that similar vectors end up with similar values.\n\nThe trajectory representations are learned via an auto-regressive prediction process that (I believe) utilizes a pretrained network to avoid reconstruction of high-dimensional observations.\n\nThe two value estimates (parametric and non-parametric) are combined via a weighting function which attempts to minimize the temporal difference error of the aggregate value function.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work is full (perhaps too full) of unique modifications to the basic idea of combining parametric and non-parametric value estimates for improved data efficiency, and the results do seem quite impressive. But a proper evaluation of this algorithm is greatly hindered by persistent issues involving clarity and a lack of sufficient detail.\n\nThe most glaring example is the issue of using a pretrained network to yield targets for the TR/TP losses. What loss is this network trained on? What is its architecture? Is it trained in tandem with the other components, or is there a separate pre-training phase? I poured over both the paper and appendix and can't find any answers. And these answers should greatly impact the overall merit of this paper. If everything is trained in tandem, then the empirical results are very solid. But if a long-pretraining period and/or highly curated data are used, then all of the experimental results are comparing apples to oranges.\n\nThe TR/TP distinction is confusing. It's mentioned in the main text, but only explained in the appendix. The explanation is that TP uses time t to predict time t+1, whereas TR uses t-1 to predict time t (line 561). But for the sorts of environments under consideration, I can't see how this distinction matters. On line 118, it is mentioned that a \"noisy\" version of the state-action tuple is used in the TR loss, perhaps this isn't the case for the TP loss? Relatedly, how is this \"noisy\" version constructed?\n\nThe weighting function is another place where more explanation is needed. As mentioned in my summary, I believe its trained to minimize the TD error of the aggregate value function (Equation 8). But why is this a reasonable thing to do, considering your changing the mixture weight rather than the value estimates directly? For example, the parametric values will initially collapse to zero on many environments with sparse rewards. In which case, it'd be optimal to weigh this value highly since it minimizes TD error, but this feels like exactly the case where you'd want non-parametric estimates to be up-weighted.\n\nI like the explicit list of open problems this works resolves (line 28), but the proposed fix for the problem of needing deterministic environments doesn't make sense (line 36). MFEC requires deterministic environments because it behaves based on sampled experiences with the highest return. Thus, a single winning lottery ticket would forever force MFEC to buy more tickets, since it only considers the best return rather than the expected. MBEC's reconstruction loss doesn't impact this at all. Rather, MBEC is likely more robust to stochasticity due to it averaging over nearest neighbors rather that simply taking the max. Though it still does that sometimes (Line 133), and to the extent that it does, it also is biased towards deterministic environments.\n\nThis is a bit of a quibble, but the psychology angle hurts clarity and doesn't sufficiently connect to the literature enough to make up for it. Parametric vs non-parametric is a lot more explicit than episodic vs semantic. Semantic and habitual also seem to be used synonymously, which isn't the case in psychology nor machine learning. Learning a forward model makes sense on its own, and relating it to a serial recall test doesn't better motivate it (line 118). I'd suggest all text relates the method to complementary learning systems hypothesis be relegated to the discussion section.\n\nThe experimental methods for Atari are quite lacking. Rather than report any notion of spread, only the best seeds appear to be reported (line 626), and it is unclear how many seeds were run.\n\nThe ablations that are there are appreciated, but this paper adds many more distinction components that aren't ablated, e.g. is switching between two reading rules really necessary? I believe the ablation in the appendix (Figure 13) is showing the effect of dynamic weighting between the two value functions (again, this is unclear due to lack of supporting text). If so, there doesn't seem to be any support for it outside of a single Atari game. If an ablation shows a component isn't needed, eliminate it from the method.\n\nEDIT: several of my points were misunderstandings, and my others were largely addressed through the additional baselines and clarifications. I'm conditionally voting to accept under the assumption that the final draft is explicit about experimental seeds and other sources of confusion.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper improves on recent works that leverage both episodic and habitual learning by combining state-action episodic memories with parametric value functions like Deep Q-Network.

The strengths of the paper include:  The idea is compelling from a theoretical/neuroscience standpoint (i.e., complementary learning systems); The proposed trajectory embedding and memory mechanisms are novel; The proposed dynamic weighting between memory-based and parametric values estimates is also novel; The performance is strong.  On the other hand, the reviewers have concerns regarding the complexity of the proposed approach.

Decision: Accept (Poster)