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 .ae 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 aate-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 RI. In particular, Deep Recurrent Q- Network (DRQN; [17]) replaces the feed-forward value network with LSTM counterparts, aiming to solve Partially-Obses Observabble 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: 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 (-1. rr) to C. Sample (sy, av) from B and optimize Cir wrt and (Eq. 1)
if = then
18:
for each 7, C do
19:
Compute V (7t) =E 2 Write M <- write
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.NH 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-parmettric 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-the-fly 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', rr) 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 differentt 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
Single
-160
-160
-180
-200
-000
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: (60::::666.400 1606,04460, 2016.
[5] Maxime Chevalier- com/ com/maximecto wawccb/ggggg. 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. /4443, 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, 1633116, 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": "bxxU6Msxw6", "reviewer": "Reviewer_m2e7", "summary": "This paper proposes to leverage episodic memory for improving the performance of control tasks. The proposed method uses estimated trajectory values for efficient policy learning. On top of the episodic memory control, they built a hybrid control system to combine the current learning algorithms. The experimental results are better than a few baselines in traditional control tasks and some of the atari benchmarks. ", "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": "The originality of this paper is good. This paper clearly clarifies the difference between their method with previous methods in this domain. \nThe paper is well written. \nThe significance of this paper is reasonable. Previous methods, even those value-based RL algorithms, might ignore the episodic nature of most of the benchmarks. In this paper, they utilize this fact and improve upon previous methods with an augmented trajectory memory.\nWhile the experiments still need a bit of improvement, I would say the idea itself is valuable and inspiring for the field. So I would recommend for acceptance if a bit more evidence are provided to show the effectiveness of the proposed method.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The experimental section missed some strong model-based baselines  --- dreamer v1/v2,  SIMPLE and strong model-free baselines such as rainbow.  While this method might be developed from DQN and it already shows improvement upon it, it is still important to test the same idea on sota or more advanced methods since many ideas only work on some specific backbone algorithms. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "Y8gUq0SyXqg", "reviewer": "Reviewer_bLjg", "summary": "This paper presents a novel method for episodic control (EC) in reinforcment learning, combining several improvements aimed at targeting weakensses of current EC methods. These improvements are a trajectory-based (rather than state-based) representation for storage in memory, a model-based 1-step planning component for decicion making (including learning a reward model), and a dynamic combination of the non-parametric (memory-based) and parametric value estimate based on the current trajectory representation. Combined, these methods result in MBEC++. This method is compared against previous EC methods from the literature. First they evaluate methods in a variety of toy environments with specific properties to demonstrate that MBEC handles noise in the reward and transition function better that previous work and that the planning component is useful. They then test on the ALE, training for 10M steps as is standard practice in EC works, and show better results in terms of both median and mean performance, outperforming previous EC methods that were trained for 40M steps. Several ablations show the importance of the different components of MBEC++.", "questions": "", "limitations": "", "rating": 6, "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": "## Overview\n\nOverall, I believe the paper is a novel combination of new techniques for improving episodic control work. The problems with existing methods are presented and tackled by the new techniques, motivating them well. The experiments empirically demonstrate the superior performance of MBEC, and the ablations (including those in the appendix) provide insight into the importance of the different components. The main weaknesses of the paper are the writing clarity (there are numerous small grammatical errors, and some unclear writing and structure); the multitude of hyperparameters for their method, with relatively little intuition given about what good settings these might be when tackling new problems; and the lack of a focused related work section on previous EC methods.\n\n## Strengths\n\nOriginality: I believe the paper presents an original method, combining original architectural components in a way that enables a novel planning algorithm. Particularly the formulation of the trajectory representation model to allow for 1-step lookahead planning for value estimation is novel and interesting. While dynamic modulation between parametric and non-parametric components has been done before (for example see [Adaptive Semiparametric Language Models](https://arxiv.org/abs/2102.02557)), it hasn't been done in the context of EC methods.\n\nQuality: The method presented seems to be technically sound and is described in sufficient detail in the appendix and the main paper to be reproduced. The performance claims are empirically validated. \n\nClarity: The method presented is complicated, which makes it difficult to present a simple explanation. I believe the explanation of the method presented in the paper is good but not great. Further details in the appendix clarified any confusions or uncertainties I had. There are multiple small grammatical errors which I haven't listed here, but which could probably be fixed by passing the paper through a rigorous online grammar checker (e.g. Google Docs, Grammarly).\n\nSignificance: I think the demonstration of the combination of model-based planning with model-free value estimation in an episodic control setting is useful and important, and it seems likely that future work in episodic control will use ideas from this paper or build upon it's method, possibly to integrate further model-based components or methods. I think episodic control in model-based RL is underexplored, and this paper is a first step in that direction.\n\n## Weaknesses\n\nAs described, the main weaknesses I see with the paper are (1) writing clarity, (2) many hyperparameters, (3) lack of EC related work section.\n\nI've made several suggestions for improvements below regarding (1). I also believe many grammatical errors could be fixed by passing the text through a grammar-checking software.\n\n(2) is a more difficult weakness to address. I believe the paper is still sufficiently useful despite the large amount of additional hyperparameters, but the paper would be improved if there was more intuition and justification for how to choose these hyperparameters in different settings based on characteristics of the environment. For example, if I was to try this method in a novel environment, what heuristics could I use to pick good initial values for these hyperparameters? Which hyperparameters do you expect to be agnostic to the environment, and which should be changed per environment? Perhaps linking certain hyperparameters together (e.g. given a chosen K or L, what would be reasonable choices for N). These insights could probably be drawn from the existing experiments and ablations in the paper and appendix, or new experiments could be run to support any claims made.\n\nFor (3), I think an additional subsection in the Background section focusing on previous EC works would better situate the work in the existing literature. I think less time could be spend discussing DQN methods, given it's of relative little relevance to the rest of the paper, compared with previous EC methods. This section could also stress the weaknesses described in the paragraph starting on L28.\n\n### Smaller Suggestions for improvement\n\n  - I think it would be beneficial to clarify your usage of \"episodic memory\" as not a memory that's within an episode, but rather memory that's maintained across a lifetime. There's confusion in the literature on this terminology, as there is prior work that uses \"episodic memory\" to refer to within-episode memory (for example [Generalization of Reinforcement Learners with Working and Episodic Memory](https://arxiv.org/abs/1910.13406) and [Rapid Task-Solving in Novel Environments](https://arxiv.org/abs/2006.03662) both use Episodic Memory to refer to a within-episode memory component), and also the papers cited in your work which use \"episodic memory\" or \"episodic control\" to refer to an across-lifetime memory. Perhaps in one of the first few paragraphs this could be clarified with a short sentence, so future readers are clear from the outset what you mean by \"episodic memory\".\n  - L28-34: It would motivate your argument more if you could cite or reference experimental evidence of these weaknesses in a more explicit way. For (a) I believe your own experiments support this claim, and possibly the same is true for (b) and (c). If that's the case I think it's worth rephrasing this paragraph to say that you demonstrate these are problems with current EC methods, which your method then addresses. If evidence is in prior literature then explaining how that is demonstrated, rather than just citing papers where the problems have been demonstrated, would be beneficial.\n  - L130: It would clarify the algorithm explanation to describe what estimations of the true trajectories values are used here. (It is explained later on as well, but mentioning you used Monte Carlo estimations of the value here would improve the ease of understanding your method).\n  - L139: I think it would be clearer to denote the estimated value by $\\hat V(\\vec\\tau)$ (i.e. V rather than R). $R$ is often used as a reward function, and $R(\\vec\\tau)$ I'd intuitively read as the reward from the trajectory $\\vec\\tau$, whereas in the paper it's the remaining reward after completing trajectory $\\vec\\tau$\n  - L148: The sentence starting \"Following [22]...\" is slightly confusing. I think it would be better as: \"Following [22], we choose to write the trajectory representation of every L'th time-step (rather than every time-step) to save compuation while still maintaining good memorizaton.\" Otherwise it implies that you write to the memory every L time-steps, but you could still be writing every trajectory representation, just saving them up for L time-steps (as is done for example in Neural Episodic Control, where every time-step is written, but only at the end of the episode)\n  - L151: The name \"update\" for the **memory update** function isn't the most informative, as it does a memory write just the same as the \"write\" function. I'd recommend changing it to \"bootstrap_write\", as the value you're writing is calculated with the 1-step bootstrap value estimate (in equation (3)).\n  - L153: \"As such, at an arbitrary timestep...\". At this point I think it would be worth clarifying how do you choose the time-step? I believe the answer is you randomly decide to update with probability $p_u = 0.1 $, from the appendix section B.1 (L513)\n  - Eq. 6 and 3 are basically identical - is there a way you could unify them to make it clear you're doing the same operation in both cases? I think this would clarify the algorithm further\n  - L193: Without an explanation of the reward function for this environment (which is in the appendix but not the main body of the paper), understanding the optimal behaviour isn't possible. For example it's unclear why staying in the maze is beneficial (see L214) unless you know that touchingthe walls leads to negative reward. In the interest of making the paper self-contained, I think a brief description of the standard reward function for the maze environments in the main body of the paper would be beneficial.\n  - L196: It's needlessly confusing to define the variable $n_e$ if the only value it takes in your paper is 3. I think you should just say you work in 3x3 mazes.\n  - L275: It would further justify your method if you highlighted results on atari games that (a) were stochastic (and hence MFEC/NEC doesn't cope well but MBEC++ does), (b) where you believe the state-action-value storage is slowing does learning, (c) where the dynamic combination of parametric and non-parametric value estimates is useful and (d) where you think the model-based planning component is providing improvements.\n  - L284: I believe you meant \"state-action-based value estimations fall short...\"?\n  - L291, 293, Table 1: Ephemeral Value Adjustment is EVA, not EEVA.\n  - Algorithm 1\n      - L12: this and L13 should be indented to match the scope of the **for** loop.\n      - L14: this line could be removed and lines 15-16 unindented for clarity (as it only happens at the end of each episode).\n  - Figure 2\n      - (b, upper): The logarithmic scales makes it unclear the actual improvement between MBEC and DQN here, which seems minimal visually but from the text of the paper and the graph below is substantial. Perhaps a different scale should be used.\n      - (c): Why is the x-axis only a fragment of the full training run? This should be mentioned in the caption (to avoid confusion).\n\n### Clarifications\n\n  - L137: Instead of alternating between these reading at random, did you try mixing them? If not, do you have justification for why randomly alternating between them is the best choice?\n  - L209: Does falling into the trap terminate the episode?\n  - L216: If you're using a much larger memory for MFEC than you do for MBEC, then it's possible this late-stage degradation in performance is only due to this hyperparameter change, rather than your method. Also, do you have evidence that this is the cause of the worse performance for MFEC? For example, the average reward plateus at 45k steps for MFEC in Fig 2 (b, upper) - is this when the memory fills up?\n  - On Atari results, you say you train for 10M frames. Do you mean 10M timesteps (which is 40M in-game frames, due to frame stacking?) If so, I think you should use \"timesteps\" rather than \"frames\"\n\n## Conclusion\n\nOverall, I believe as it stands the paper is marginally below acceptance (5). Making improvements on clarity or related work would move my score to marginally above acceptance (6). Further making improvements on intution and heuristics behind choosing hyperparameters, to compensate for the large amount of additional hyperparameters, would move my score to accept (7).\n\n### Edit\n\nOverall, I appreciate the authors efforts to improve the clarity of the paper, and working on the assumption that many of the comments or confusions in mine and other reviews will be addressed, I will raise my score to a 6. I think the current heurstics for picking hyperparameters aren't sufficiently justified to raise my score to a 7.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Societal impact is addressed (if only briefly). My concerns on the limitations of the paper are mainly centred on the large number of hyperparameters, as discussed in the main review above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "O0o7OKRPKdw", "reviewer": "Reviewer_Crem", "summary": "This paper proposed a model-based approach to leverage episodic memory in reinforcement learning problems.\nIt aims to tackle three fundamental issues of memory-based methods, i.e., the near-deterministic assumption, sample inefficiency of memory construction, and the inflexible connection between parametric and non-parametric modules.\nThe experiments show that the proposed method outperforms baselines over a variety of environments.", "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": "**Originality**: To my knowledge, the proposed method is novel and the motivation is clear and convincing. I agree this paper aims to tackle some fundamental problems in EC-based approaches. The proposed method has good contributions for this track of research.\n\n**Quality and Clarity**: The structure of this paper is well-organized. The proposed method is well-motivated and self-contained. The description of implementation details is acceptable.\n\n**Significance**: This paper aims to address some critical issues in episodic control algorithms. The design of the proposed approach is reasonable. The experiment results are strong and significant. My concerns are discussed in the next section.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "**Main Concerns**:\n\n(1) Given the trajectory representation is changing during learning, how to ensure the key stored in the memory is consistent with the current trajectory encoding? What will happen if we find incorrect neighboring records?\n\n(2) In some Atari games, the trajectory length might be very long, e.g., more than $10^4$ steps in a single trajectory. In this case, does the claim *\"from that memory, we must be able to recall any past event\"* still hold? How about the relation between model accuracy and learning performance in Atari games?\n\n(3) If the above two problems matter, I am curious about to what extent the error will lead the memory to find incorrect neighbors.\n\n**Minor Questions**: (I did not penalize my evaluation according to these points)\n\n(1) Maybe there are some typos in Algorithm 1. I guess that line 12-16 should be included in the *for*-loop at line 4. In the current version, these two *if*-blocks are executed after the *for*-loop.\n\n(2) A concurrent work [1] aims to address the second issue considered by this paper in continuous control settings. It is not required to compare performance with very recent papers, but a short discussion would be helpful.\n\nReference:\n\n[1] Hu et al., Generalizable Episodic Memory for Deep Reinforcement Learning, ICML 2021.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "EY37XIRjguC", "reviewer": "Reviewer_R2hp", "summary": "This paper aims to combine episodic memory (where we store previous RL traces to act in the environment) with model-free/model-based RL (where we learn a global policy or value function to act in the environment). Their main contributions are that they 1) learn encodings for the trajectories, and 2) use the trajectories within a model-based planning procedure, and 3) adaptively integrate episodic memory with model-free RL based on a learned, tunable parameter. The paper has extensive experiments in small tasks, Atari 2600 and a first-person view POMDP, showing better performance than previous episodic control methods. ", "questions": "", "limitations": "", "rating": 5, "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": "Update after rebuttal:\nI thank the authors for their rebuttal, but my main problem, about the clarity of the paper, unfortunately remains. I will therefore stick to my vote. \n\nStrong:\n* Interesting and promising direction to combine episodic control with common model-free/model-based RL techniques.\n* Extensive experiments. I also like the small-scale experiments with interpretability. Authors also include several ablation experiments. \n* I like the dynamic integration of model-free and episodic memory, i.e., learning the trade-off in a separate neural network. It indeed makes sense to adapt the contribution of episodic memory over training, since it can likely be high initially, but be decreased near convergence of the solution approximation.\n\nWeak/possible improvements: \n* Clarity: I think your writing could be much clearer. I will give some example from the introduction:\n    *You repeatedly mention ‘habitual’, but it took me some time to understand that you mean ‘model-free’ here, or ‘fast prediction of policy/value with function approximation. \n    * L30: “storing state-action value which demands experiencing all actions” → This is only true for tabular representations right, not if you profit from generalization?\n    * L36 “the model” → What model? The method you introduce? I also do not really see in you explanation how this solves the stochasticity issue, you do not get back to it all? \n    * In L39: what is the difference between ‘the world’ and ‘the trajectories’. I can sample trajectories from a world model right? Explain such things.\n    * L40-47: Again, I have a lot of trouble understand why this solves problem b, and what you are exactly proposing.\n\n\n* Related work: Your related work is good when it comes to memory for partial observability. But, as you indicate as well, episodic memory is different. RNN’s for partial observability try to propagate/memorize relevant state representations, while episodic memory tries to also mimic previous actions, i.e., it is not just a state estimation technique. You mention this in line 99, but you do not explain the difference, nor give some related work in the other direction (although it is present in the introduction).\n* Methods (Sec 3): At some points I have trouble following what happens in Sec 3. For example, section 3.2 introduces memory read, write and update procedures, but it is not clear too me when and why they are necessary. Overall, I often miss motivation why what you propose is necessary. I like Figure 1, but I think it needs more motivation, and then in your text you could get back to this figure and explain where each step happens, why it is relevant, and what considerations were important for your solution.\n* Noisy transition (L232 and onwards): I am not convinced of your choice of a noisy transition, where you with random probability do not execute the action. This means that you still stay on the exact same trace, and you can simply try the action again at the next step. Thereby, it is arguably the minimal and simplest form of stochasticity. These experiments would be much stronger with a truly stochastic domain, not a deterministic one you slightly modified.\n* Sec 4.4.: When you discuss results in your main paper, the results should also be visible there. Otherwise, state something like: “see appendix .. for additional experiments on a partially observable first-person view environment”. I think these experiments are really interesting and valuable, but then you should reserve the required space for them in the main paper. \n* Discussion: I miss a discussion. There is a one sentence remark about future work in the Conclusion, but no discussion of limitations of this work at all.  \n\nConclusion: I like the novel idea of this paper, which I think is relevant to the community. The experiments are also good and extensive. However, on the downside I do not think the paper is very clear (see the comments above). I especially sometimes miss a clear story line, why certain choices are made, how they solve existing problems, etc. This especially applies to the introduction and methodology section. I do think the paper could be improved a lot in this direction, and I think the paper deserves a discussion. I will therefore vote borderline here, although I do think the idea has much potential. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I could not find a section on this. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: I recommend accepting this paper.

Both the pre-discussion scores (5, 5, 6, 6) and post-discussion scores (5, 6, 6, 6) are borderline.

Reviewers agreed that there was a clear contribution, but the primary shared concerns were (1) writing clarity and (2) hyperparameter tuning:
* Reviewer R2hp gave the lowest score of 5, and wrote in the discussion: "Although the rebutall does point to some locations in the paper where unclarity might be resolved, I still think this was and is a major factor in my overall rating. The authors used much terminology which felt unnecessary, and I had trouble sticking to the main storyline." But the same reviewer also observed, "I like the novel idea of this paper, which I think is relevant to the community. [...] However, on the downside I do not think the paper is very clear. [...] I will therefore vote borderline here, although I do think the idea has much potential."
* Reviewer bLjg raised their score from 5 to 6, assuming the clarity revisions discussed will be implemented ("on the assumption that many of the comments or confusions in mine and other reviews will be addressed"), and underscored that hyperparameter tuning prevented a higher score ("the current heurstics for picking hyperparameters aren't sufficiently justified to raise my score to a 7").
* Reviewer m2e7 was concerned about strong baselines beyond just EC methods, and not satisfied with the author rebuttal, but these concerns were not shared by the other reviewers, and even with these concerns the reviewer leaned towards acceptance ("the idea itself is valuable and inspiring for the field", "this paper has some genuine contributions").
* Reviewer Crem did not participate in the discussion but the main concerns raised were in my estimation addressed in the rebuttal.

To summarize, while reviewers agreed on shortcomings in the paper which should be addressed (especially clarity), there is consensus on the paper having interesting contributions and an aggregate lean towards acceptance. For that reason I recommend acceptance, and also encourage the authors to implement the improvements discussed.
Decision: Accept (Poster)