Momentum Boosted Episodic Memory for Improving Learning in Long-Tailed RL Environments
Authors: Dolton Milagres Fernandes
Pramod Kaushik
Harsh Shukla
Bapi Raju Surampudi
Venue: Submitted to ICLR 2023 (2023)
Keywords: long tail distribution
reinforcement learning
representation learning
contrastive learning
complementary learning system
hippocampus

Abstract: Conventional Reinforcement Learning (RL) algorithms assume the distribution of the data to be uniform or mostly uniform. However, this is not the case with most real-world applications like autonomous driving or in nature, where animals roam. Some objects are encountered frequently, and most of the remaining experiences occur rarely; the resulting distribution is called \emph{Zipfian}. Taking inspiration from the theory of \emph{complementary learning systems}, an architecture for learning from Zipfian distributions is proposed where long tail states are discovered in an unsupervised manner and states along with their recurrent activation are kept longer in episodic memory. The recurrent activations are then reinstated from episodic memory using a similarity search, giving weighted importance. The proposed architecture yields improved performance in a Zipfian task over conventional architectures. Our method outperforms IMPALA by a significant margin of 20.3\% when maps/objects occur with a uniform distribution and by 50.2\% on the rarest 20\% of the distribution.

Under review as a conference paper at ICLR 2023
MOMENTUM BOOSTED EPISODIC MEMORY FOR IM- PROVING LEARNING IN LONG-TAILED RL ENVIRON- MENTS
Anonymous authors Paper under double-blind review
ABSTRACT
Conventional Reinforcement Learning (RL) algorithms assume the distribution of the data to be uniform or mostly uniform. However, this is not the case with most real-world applications like autonomous driving in nature, where animals roam. Some objects are encountered frequently, and most of the remaining experiences occur rarely; the resulting distribution is called Zipfian. Taking inspiration from the theory complementary learning systems, an architecture for learning from Zipfian distributions is proposed where long tail states are discovered in unsu- pervised manner and states along with their recurrent activation are kept longer episodic memory. The recurrent activations are then reinstated from episodic memory using a similarity search, giving weighted importance. The proposed architecture yields improved performance in Zipfian task over conventional ar- chitectures. Our method outperforms IMPALA by a significant margin of 20.3% when maps/objects occur with a uniform distribution and by 50.2% on the rarest 20% of the distribution.
INTRODUCTION
Humans and animals roam around in environments that are unstructured in nature. However, existing algorithms in reinforcement learning are built around the assumption that environments are mostly uniform. Most of the time, small subset of experiences frequently recur while many important experiences occur only rarely (Zipi 20133) Smith et al. 2018). For example, imagine a deer trying survive in an environment with predators. It is drinking water from a source and it escapes from potential predator, the deer cannot afford to learn slowly from multiple such experiences to learn to avoid dangerous places corresponding to water sources. It needs learn that experience quickly and generalize across similar instances. Similarly, in autonomous driving, experiences are not uniform, and usually, the rare instances where there is an accident or unusual experiences are more critical in real-world settings. This is the fundamental premise learning systems (McClelland et al. 1995 Kumaran et al. (0166) is proposed. In this framework, an intelligent agent needs have fast learning system and slow learning system operating together to restructure the statistics of the environment for better survival internally and not be naive in expecting uniform environments. In the brain, this is hypothesized through the interplay between the hippocampus, a fast learning system, and the neocortex which is a slow learning system, and together they manage generalize and retain experiences crucial the goals of the organism (Kumaran et al., 20166). It is also true that to achieve their respective learning outcomes, both the systems need each other (Botvinick al., 2019). The hippocampus able to achieve fast learning through its reliance on the slow learning system of the cortex where high dimensional data coming from the sensory systems are converted to low dimensional representations which can be operated on by the hippocampus. Such top-down modulation from the cortex influences the processing in the hippocampus (Kumaran & Maguire 2007). Similarly, the slow structured learning of the cortex happens through interleaved learning by replaying experiences stored in the hippocampus (eill et al. 2010).
which the theory of complementary
Here, particularly look at this interplay between a fast learning and slow learning system and apply this to solve the long-tailed phenomena. The reinforcement learning algorithm (Sutton & Barto 2018h Espeholt al. 2018) uses the episodic buffer to generalize across experiences, and

---

Under review as a conference paper at ICLR 2023
familiarity memory prioritizes long-tail data from the outputs generated by the RL algorithm. This prioritization of samples happens through a contrastive learning-relate momentum loss which en- ables the unsupervised discovery of long-tailed data from the stream of experiences (Zhou et al. 2022). These prioritized samples are kept longer in memory and then, the hidden activations corre- sponding these prioritized samples reinstated in the recurrent layers of the RI. network.
Our Main Contributions are:
Proposing first solution for the problem of navigating to objects occurring with long tail distribution using deep reinforcement learning.
Application contrastive momentum loss for unsupervised discovery of long tail states in the context of reinforcement learning.
Novel method prioritize long-tail states in the buffer then reinstating hidden activations in recurrent layers.
BACKGROUND
2.1 MARKOV DECISION PROCESSES
Let's assume environment & which provides the agent with observation St, the agent selects an action A,, and then the environment responds by providing the agent with the next state S The interactions between the agent and environment
formalized by MDPs which reinforcement defined by the tuple A, R, T.7 > where, S represents the set of states, A the of actions, R : S x A R denotes the reward function. T : S x A + Dist(S) represents the transition function mapping state-action pairs distribution over next states (S) and y [0, 1] the discount factor.
learning tasks that satisfy Markov property (Sutton & Barto 2018a).
2.2 IMPALA
Previous works the actor-critic framework have been based on single learner multiple actor architectures where the communications are based on gradients with respect to policy parameters (Mnih et al., 2016b). IMPALA is distributed off-policy actor-critic framework (Espeholt et al., 2018); here actors communicate a sequence of trajectories to the learners which give it very high throughput. We consider IMPALA base architecture and build top of it.
2.3 MEMORY SYSTEMS RL
Memory systems in humans allow them to retrieve the relevant set of experiences for decision- making in case of unseen circumstances. neuroscience, some of the types of memories studied are Working Memory and Episodic Memory. Working memory is short-term temporary storage while episodic memory is a non-paramen non-parametric or semi-parametric long-term storage memory. Deep Rein- forcement Learning agents with episodic memory, in particular a combination of non-parametric and parametric networks have shown improved sample efficiency and are suitable for decision socision-making in rare events. [Blundell al. (2016) used a non- non-parametrie model to keep the best Q-values in tabular memory. Pritzel et al. (2017) in Neural Episodic Control proposed to keep the representations and Q-values in semi-tabular form. [ansen al.] (2018) took up trajectory-centric approach to model such systems. Our work looks up to Fortunato et al. (2019)'s 's
state-centric formulation of Episodic Memory (MEM) which implements working memory using a latent recurrent neural network and an episodic memory.
2.4
SELF SUPERVISED LEARNING
Self Supervised Learning has been widely used in Reinforcement Learning to learn better represen-
tations for learning actions (Anand et al. 2019 Laskin al. 2020). Here we look at the Contrastive Learning framework (Chopra et al. 2005; Chen al. 2020 He et al. 2019) which uses similarity constraints learn representations. The representation in pixel space is learned by bringing differ- ent views of the same images together and vice-versa. Usually, Self-supervised long-tailed learning methods have been developed mainly from loss perspective or model perspective. Focal loss
2

---

Under review as a conference paper at ICLR 2023
in hard-example mining (Lin et al. 2017 Liu al. 2021 belongs to the former perspective of loss re-weighting. These methods are however limited by long tail sample discovery because they require labels to work. SSL from model perspective requires highly specific model design (Tian et 2021; Tiang et al. 202T). However, [Zhou et al. (2022) takes up a different approach to long tail learning which is based on the data perspective. /Zhan et al. (202T); Arpit al. (2017); [Feldman (2019) has studied the memorization effect in deep networks, and it has been observed that easy patterns are usually memorized before the hard patterns, and similar behavior can be observed in long-tail learning. Zhou et al. (2022) extends this by distinguishing between the long tail and head samples with respect to memorization an unsupervised setting. They propose Boosted Contrastive Learning method which works with momentum loss that enables the unsupervised discovery of long-tail data.
TASK
Our task is inspired by Zipf"s Gridworld task in Chan et al. (0022). The task involves navigating to a target object in a partially observable gridworld environment. There are in total of 20 maps, and 20 objects in each of the maps are placed at random positions. The positions of these objects in these maps do not change during trials. The trial starts with the agent placed in the center. Given target object, the agent has to find the object in the partially observable gridworld environment by taking at most 100 steps. If the object selected at the end is the correct object, the trial ends, and the agent gets a positive reward. the object selected incorrect or the number of steps exceeds the limit, the agent gets no reward.
Due to the unavailability of code for the environment during the time of our experimentation, we designed and implemented similar Zipfian Gridworld environment with slight differences from the original environment proposed in [haa et al. (2022) but maintaining the core principles. Instead of embedding the target object in the top left corner of the view of the agent, we embed it in the top center with a black background. Also, we have just four actions (up, down, left, right) which can be used to navigate the agent is able to reach any object present on the map using the set of four moves. An example of the agent's partial view be seen in Figure [b. There maps, we have 10 objects. We chose the number 10 primarily because of the vast computational resources required to process 20 grid maps with 20 objects each.
any cell in the grid. It is ensured that all the objects in map distinct and
total of 10 maps, and in each of these
1/k"
zp(k, n, e) = ΣT (1/ir)
The probability of occurrence of the maps is governed by Zipf's power law (Equation []. where the number of maps/objects, i is the map/object index (1 <= k <= n) and is the Zipf's exponent. The same skew can be seen for target object selection in each map as shown in Figure D To solve the task, the agent not only needs to explore the environment intelligently but also needs to memorize the path if the agent solves the trial correctly.
4 ARCHITECTURE
Given an image observation (im), IMPALA's (Espeholt et al. 2018) feature extractor gives pixel input embedding (p) which is further passed an LSTM network with the hidden state (h) to get the new hidden state, policy, and value. As done in (Fortunato et al., 2019), we have MEM module (orange-red buffer in Figure 33 that is used to find the memory (m), which then additionally fed as input to the LSTM network along with the pixel input embedding p. We only keep rare states in the MEM by introducing an additional 'familiarity' buffer (light cyan buffer in Figure (3) that uses boosted contrastive learning to prioritize and filter rare states for storing in MEM.
4.1 STATE FAMILIARITY USING BOOSTED CONTRASTIVE LEARNING
A 'familiarity' buffer is a circular buffer that contains states over which we can get a prediction on whether state is rare or how rare the states are. To achieve this, we take inspiration from Zhou al.
(2022), where they improve performance on a long- tailed self-supervised learning task by proposing momentum loss that can predict which samples among the dataset long-tail samples.
3

---

Under review as a conference paper at ICLR 2023
Map 1
Map
Map 10
p: 64.53%
p:
0.64%
Agent's View
Target object to find
Agent in initial position
Door to enter next room
Object
(a)
Probability
(b)
Figure 1: Zipf's Gridworld Task: Contains 10 maps, each with 10 objects placed at random loca- tions. The location of these objects does not change during trials. The agent (white square) starts at the center in each trial and has to navigate towards the target object shown in the center at the top. The agent's view shown the bottom image in (a). The target object to find, the agent's initial position, door, and object can be seen in the annotated image. The value 'p' below each map shows the probability of occurrence of the map in a trial, highlighting the skew the distribution. A similar skew occurs for the distribution of objects in these maps. We can see this in the below part (b), which shows the distribution objects for the first map (most common). Adapted from Chan
et al. (2022)
(a) Input
(b) Gaussian Noise
(c) Random Cutout
(d) Final Augmented Image
Figure 2: Image augmentations for contrastive learning: (a) Shows downsampled input image for trial. (b) Input image after adding gaussian noise to it. (c) Input image after applying random cutout augmentation. The black rectangle near the agent's position is the area cutout. (d) Final augmented image after adding gaussian noise and random cutout.

---

Under review as a conference paper at ICLR 2023
h
hp TL, V
MEM
h-11
LSTM
kg
Pi
h2
[P mj]
KNN fetch
and
Pt
FF
CNN
im;
kg
Pi
Transfer "4
rarest states
im,
Input
Familiarity Buffer
The figure shows momentum-boosted episodic memory archi- tecture pipeline. The IMPALA backbone consists of a CNN feature extractor followed by a Feed Forward layer that gives the embedding. This embedding is concatenated with the one hot action encoding, reward & memory to get pixel embedding p; and then given to the LSTM network for fur- ther processing with working memory. The LSTM network additionally takes the past hidden state hi-1 as input. During training, the input image, pixel embedding, LSTM hidden states, and keys are stored the familiarity buffer. The momentum loss tracked on this buffer during contrastive learning is then used to prioritize long-tail states. The MEM is then periodically updated with top i states from the familiarity buffer. The memory (m²) is computed from the MEM using a weighted sum (+) after fetching using KNN similarity search on the keys present in the MEM (Equation
Figure 3: Model Architecture:
Our 'familiarity' buffer (light cyan buffer in Figure contains the input image (im). key (k), pixel input embedding (p) and LSTM hidden state (h). We will define these terms in the following section. In each learning step of IMPALA (Espeholt et al. 2018), the batch of trajectories is sent to this buffer. The entire trajectory isn't sent to the buffer, but subset (states hop of hp from the start of the trajectory) is sent. For a trajectory Tr (81, $2, sk) of size k consisting of states st. the subset of trajectory is defined
Trgubset = (51,, $(1+h)) (1+1
(2)
The same feature extractor of IMPALA is trained using an additional auxiliary contrastive loss (Equation (6) on states present in the familiarity buffer. Once the circular buffer is full of states, the contrastive learning process starts to help the learning process. For each image sample im; in the buffer, we find its pixel input embedding p; after IMPALA's feature extractor. The same feature extractor is used on the augmented image to get another embedding p.ng The image im; is aug- mented using two augmentations, namely gaussian noise (Boyat & Joshi 2015) and random cutout (DeVries & Taylor 2017). We use these two augmentations because they are simple and don't af- fect the task. For example, using something like random conv (Laskin et al. 2020) will change the color of all objects in the observation, and this might result in the agent confusing to be image (observation) from some other trial. For adding gaussian noise to the image, we first generate an image of the same dimension as that of the original image that is filled with random numbers from
5

---

Under review as a conference paper at ICLR 2023
normal distribution with mean 0 and variance 1. This image is then added the original image after amplifying by a factor of sigma. For random cutout, we take a random location in the image and cut a rectangular area of random size. By cutting, we mean replacing the pixels in that rectangular region with black pixels.
During the training process, consider the embeddings p; and p; positive pairs for contrastive learning. The NT-Xem loss (Sohn (216)), Equation 6 is used to calculate the loss per sample. For each sample i in the circular buffer, we track its momentum loss following [Zhou et al.) (2022). The momentum loss assists in knowing which samples in the circular buffer are long-tail samples. For sample if for consecutive epochs the contrastive losses are {ET,, C2...... G. then the moving average momentum loss is defined as follows:
Em m T.F:
Em. = 3lm (1 B)
where / is hyperparameter that controls the degree smoothed by the historical losses. The final normalized momentum used to find the familiarity states is defined as
l Emp
M max{[em]] Em, 1bj=I.....
+1
(4)
m where Em, is the average momentum loss of the dataset at the 4th training step of the algorithm and N is the number of samples. The higher the momentum value M2,-- the higher is the rareness of the sample in the circular buffer. The model is trained end end by optimizing both IMPALA's loss and the auxiliary contrastive loss. Let the loss given by IMPALA be Limpala and that given by the contrastive learning branch be Lcontrastives then we define the final loss to be the one shown in Equation 5 below.
L= Cimpals + 7 Coontrastive
(5)
where y is hyperparameter. The contrastive loss is given by:
P..P.
exp
Ceontrastive
log Σ Exeex' exp
(6)
where is the number of samples, X' represents X U {p.00), (p.. au9) is the positive sample sample pair, X- is the negative sample set of p and is the temperature.
4.2 COMBINING FAMILIARITY WITH EPISODIC MEMORY
The base architecture we use is IMPALA with MEM proposed in Fortunato et al. (2019). Note that we do not use the one-step contrastive predictive coding (CPC) module they introduced. In the paper, an Episodic Memory (MEM) is introduced on top of the IMPALA architecture, which is circular buffer that stores the pixel embedding (p), LSTM hidden state (h), and the key (k), which is calculated using
(7) where W and b are Icarnable parameters and [p, h] denotes the concatenation of p and h along the dimension axis. When IMPALA+MEM comes across situations with a similar context, it uses a method to learn how save summaries of previous experiences and extract crucial data. The neu- ral network (controller), which generates the model predictions, also receives extra inputs from the readings from memory. Learning long-term dependencies, which can be challenging when depend- ing solely on backpropagation in recurrent architectures, is made simpler by successfully enhancing the controller's working memory capacity with experiences from various time scales received from the MEM. In the original IMPALA+MEM algorithm, the pixel embeddings p.. LSTM hidden states h4, and keys k; were computed after the feature extraction layer of IMPALA and then added to the MEM buffer. But, we instead add this to the familiarity buffer, which computes the familiarity of states based on the momentum loss and then periodically sends it to the MEM buffer to facilitate learning during rare trials. The agent is able to perform well on frequent states even if the MEM module is not used and only the simple IMPALA algorithm is used. Hence, we only maintain rela- tively rare states in the MEM buffer help learn on rare states that actually require the help of an external episodic memory module. The MEM gets ta most rare states from the familiarity buffer after every ly training epochs of contrastive learning, where lg and I are hyperparameters.
[p. h] +

---

Under review as a conference paper at ICLR 2023
The overall architecture can be seen in Figure 3 For stimulus pa and previous hidden state h the agent chooses the most pertinent events to give as input m, to the LSTM network using a type of dot-product attention (Bahdanau al. 2015) over its MEM. Using the key kt, formed by Pt and h4-1 using Equation 7 a K Nearest Neighbour (KNN) search done over the keys in MEM to find the most relevant K keys. The hidden states for these K relevant items combined using the below-wett weighted sum (Equation g to get additional input my to be given to the LSTM network.
Ei weh;
(8)
> Wi
1/
W Pa.
h4 b13 +
where is a small constant and
| represents the squared L2 norm of æ.
Algorithm Pseudocode for our algorithm
1: Inputs:
fm, mem
// Familiarity memory and MEM
t
// Transfer frequency
lg
// Number of rare instances transfer
T
Number of IMPALA training epochs
hp
// Within trajectory hop
|| Contrastive loss weight
2: Initialize:
fm + 0 // Clear FM buffer mem.bu, fer
Clear MEM buffer
3: 1= do
4:
trajectories ppe_iimala.batchit
5: im, k. p. h, impala +- tmmall..trainttrjjettttrr hp, mem)
6:
fm.add (im, k p. h)
// subset added according Equation
7:
contrastive loss + fm.contrastive train() if 1%1, 0 then
// loss calc. as shown Equation
8:
9:
momenluum_vlaue + tennee/lwa..ueeeeeeemmmdddddddde
// Equation
rare_experiences fm..et.rore.ktte
momentum.baluess
11:
mem.add(rare experiences)
l'nall loss impala loss 7 contraslive loss final loss.backwad backkrak)
// Eqqation
13:
EXPERIMENTS AND RESULTS
We compare the results of our method with mainly four different types of architecture. The first
architecture that we compare is IMPALA (Espeholt et al.) 2018). which is off-policy actor-critic
framework and has shown substantial improvements over baselines like Clemente et al. (011))
Mnih al. (2016a). The second one is the IMPALA+MEM introduced in Fortunato et al. 2019
The third experiment that we do is IMPALA with visual reconstruction. Chan et al. (2022) have experimented with visual reconstruction (Hill et al. 202T), and it has been seen that it helps in improving performance on rare trials by small margin (see the second row of Table D In this, we add extra task for visual reconstruction on top of IMPALA with CNN-based autoencoder. The fourth one is where we included contrastive learning to learn good embeddings. It is to be noted that this approach (see the fourth row of Table does not find the rare states the familiarity buffer, it samples k states uniformly randomly from the familiarity buffer. In contrast in the proposed ap- proach, we pass the rare i states to MEM from the familiarity buffer (see the last row of results in [ From these experiments, we can see clearly that contrastive learning (feature representa- tion learning) alone cannot give good performance and we also need add a familiarity buffer to prioritize rare states and pass them on to MEM. The training curves (Training/Zipfian Accuracy) and different ablations can be seen in Section A.2 We can see that the Zipfian mean episode return for method increases faster than all the other methods in the initial phase of training and also converges later to have the highest accuracy (*=9..9%. The hyperparameters used for our model listed in Section A.3 The results Deepmind Zipf's Gridworld are shown in Section A.1

---

Under review as a conference paper at ICLR 2023
(A) IMPALA
(b) IMPALA-MEM
103
30
64
34
-03
-46
Object Ininstiiinng order neenon
Ours
02
[c) Visual Reconstruction
[d] Confrastive Learning
Figure 4: Performance plots: (a) Performance of IMPALA agent each map and object. The y-axis denotes the map axis, and the x-axis denotes the object axis. Value at (i. j) shows the per- formance (0-1 scale) of the agent the trial where the object with ID is chosen at the map with ID i. An increase in i and means an increase in the rareness the map and object respectively according the Zipf's distribution (Equation T (b) Performance of IMPALA with MEM added. We can see that are some medium-rare trials in which the agent has learned to navigate and learn the task. (c) IMPALA with Visual Reconstruction using CNN-based autoencoder. of IMPALA+MEN with contrastive learning. (e) Performance of our agent consisting of familiarity buffer that highlights long tail samples for MEM using modified boosted contrastive learning.
(d) Performance
Table I: Evaluation Performance: We compare four different methods with algorithm namely IMPALA, IMPALA+ Visual Reconstruction using a simple CNN-based autoencoder, IM- PALA+MEM, and IMPALA MEM with only contrastive learning. We report median results across three runs (# absolute median deviation across runs) with distinct random seeds for models trained Learning+Rare State Prioritization
for 4 x 10" steps. Our method IMMALA+MEENCCntrasttt using Familiarity Buffer) beats the remaining methods by large margin on all three evaluation met- rics.
Accuracy (%)
Rare
Uniform
Method
Zipfian
(Rarest 20% objects
(All maps and objects)
on rarest 20% maps)
IMPALA
95.0 + 0.03
64.1 + 0.02
0 + 0.00
IMPALA + Vis recon
95.9 + 0.02
68.0 + 0.01
0 + 0.00
IMPALA + MEM
97.2 + 0.03
72.2 + 0.04
25.1 + 0.01
IMPALA + MEM Contr. Learning
98.1 + 0.02
73.1 + 0.01
25.3 - 0.02
Ours
99.9 + 0.01
84.4 + 0.01
50.2 - 0.02
Having higher training accuracy is not what we are looking for; instead, we want to have good accuracy when tested uniformly or rare instances. Figures fa-e) show the performance of the five architectures all the maps and objects of our environment. For each (map, object) combination, we plot the average performance across trials. Figure 4 shows the performance the IMPALA agent. We can that the agent is not able to learn the extreme rare trials and also some medium rare trials are not learned by the agent. Figures 4 (b) (d) show the performance IMPALA+MEM with and without contrastive learning respectively. In the case of 1MPALA+MEN with contrastive
8

---

Under review as a conference paper at ICLR 2023
learning, the familiarity buffer is sampled uniformly randomly to fetch states for the MEM. We that the performance is almost similar, but still, the medium-rare trials are not learned. Figure 4 shows the performance of IMPALA with added visual reconstruction using CNN based autoencoder. This performs slightly better than the baseline (IMPALA) but fails to match the performance of other agents. Lastly, in Figure 4e we can see the performance of our agent with the familiarity module. The medium-rare trials are being learned, and also some very rare trials are been done successfully by our agent. Table gives more insights into our agent's performance. Our agent is able to beat other compared agents by large margin on all of the three evaluation metrics (Training/Zipfian accuracy, Uniform accuracy, and Rare accuracy).
DISCUSSION
This paper deals with the problem of long-tailed distribution in reinforcement learning. Inspired by the theory of complementary systems, which states that an intelligent agent requires fast and slow learning system acting complementary to each other. Here, the momentum loss of contrastive learning provides a mechanism to detect long tail samples in unsupervised manner. These samples are then prioritized to be stored in a separate buffer that stores hidden activation associated with such states. When rare sample comes, similarity search is done find relevant keys, and the corresponding hidden activations are merged to be reinstated in the recurrent layers.
This architecture relates to how the hippocampus which is a fast learning system acts tandem with the slow learning cortex of the organism to store relevant experiences and replay them to overcome the statistics of the environment the organism subjected to (O (O'Neill al.) 2010). The episodic memory relies on the network to discover long tail data from the incoming data stream. The network relies on the episodic memory for identifying the relevant memory of the long tail data in order to reinstate it in the recurrent layers of the working memory system to execute the episodic sequence. Similarly, the brain could reinstate episodic sequences from the hippocampus to the working mem- ory when animals execute task.
Furthermore, dopamine neurotransmitter has been found to detect novel states and relay them to the hippocampus (Duszkiewiez et al.) (211)). The firing of dopamine neurons is also related to curiosity and learning progress which is analogous to momentum loss here (Ten et al. 2021 Gruber & |Ranganath] 2019). The neural network prioritizes its representations and persistence of memory based on how informative the state is for executing the task. This architecture can also potentially execute episodic generalization proposed by Fortunato al. (2019) but adapted for long-tailed data. The resulting embedding could be different, but the underlying dynamics of executing the task could be the same. Future work could look how we could extend this more realistic 3D environments and include systematic generalization to compose temporally separated rare states to solve tasks (Fodor & Pylyshyn 1988)
CONCLUSION
This paper attempts to overcome the problem of long-tailed data for reinforcement learning which conventional architectures do not address well owing to their underlying assumptions. An unsuper- vised long-tail discovery method using self-supervised momentum loss is used to prioritize long-tail data. Using this prioritization, an episodic storing of hidden activations is done to be later rein- stated in the recurrent layers so that rare trajectories are executed. Both of these proposed features are crucial in enabling the network to perform better than conventional architectures on a long- tail dataset. We hope that this work will encourage the development of new RL methods in such data distributions and finally enable the development of agents capable of learning from a lifetime of non-uniform experience.
8 ACKNOWLEDGEMENT
To be added later.
9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and R. Devon Hjelm. Unsupervised state representation learning in atari. CoRR, abs/1906.08226, 2019. URL
http://arxiv. wttp://r..a.b..//ooe http://aaxxx..orrr/////15666.. 08226
Devansh Arpit, Stanislaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxin- der S. Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, and Simon Lacoste-Julien. A closer look at memorization in deep networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 233-242. PMLR, 06-11 Aug 2017. URI. https:
www.inng..................me.m.. ww.. www.s.pess//eee wwt.ppesss//o/rppt11111. html
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Con- ference Track Proceedings, 2015. http:://acxx..or////aaa///a
0473
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control, 2016. URL https: 04460
Matthew Botvinick, Sam Ritter, Jane X Wang, Zeb Kurth-Nelson, Charles Blundell, and Demis Hassabis. Reinforcement learning, fast and slow. Trends in cognitive sciences, 23(5):408-422 2019.
Ajay Kumar Boyat and Brijendra Kumar Joshi. A review paper: noise models in digital image processing. arXiv preprint arXiv: arXiv: 1505.03489, 2015.
Stephanie C. Chan, Andrew Kyle Lampinen, Pierre H. Richemond, and Felix Hill. Zip- fian environments for reinforcement learning. 1st Conference on Lifelong Learning Agents, abs/2203.08225 08222, 2022.
Ting Chen, Simon Korbllihh, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. CoRR, abs/2002.05709, 2020. URL https:
//arxiv.org///// /a.ob/s/222..00 05709
S. Chopra, R. Hadsell, and Y. LeCun. Learning similarity metric discriminatively, with application to face verification. In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'0S) '05), volume 1, pp. 539-5399-6 vol. 2005. doi: 10.1109/CVPR.2005.200
Alfredo V Clemente, Humberto N Castejón, and Arjun Chandra. Efficient parallel methods for deep reinforcement learning. arXiv preprint arXiv:1705.04880 2017.
Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXivv1708.04533 2017.
Adrian J Duszkiewicz, Colin G McNamara, Tomonori Takeuchi, and Lisa Genzel. Novelty and dopaminergic modulation of memory persistence: a tale of two systems. Trends in neurosciences, 42(2):022-11,, 2019.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan. Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In International conference machine learning, pp. 1407-1416. PMLR, 2018.
Vitaly Feldman. Does learning require memorization? A short tale about a long tail. CoRR, abs/1906.05220 05271, 2019, UtL www.vorg/aall. halgg//ab//i 1906. 0tte//9900000777111
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: critical analy- sis. Cognition, 28(1-2):3-71, 1988.
10

---

Under review as a conference paper at ICLR 2023
Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomênech Badia, Gavin Buttimore, Charles Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforeement learners with working and episodic memory. Advances in neural information processing systems, 32, 2019.
Matthias J Gruber and Charan Ranganath. How curiosity enhances cenpecaampusdeppeeeee mem- ory: the prediction, appraisal, curiosity, and exploration (pace) framework. Trendy in cognitive sciences, 23(12):1014-1023 2019.
Steven Hansen, Pablo Sprechmann, Alexander Pritzel, André Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past, 2018. URL https://
wwx.l.org//bb//lBa00 08163
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. CoRR, abs/1911.05722, 2019.
URL http:
05722
Felix Hill, Olivier Tieleman, Tamara von Glehn, Nathaniel Wong, Hamza Merzic, and Stephen Clark. Grounded language learning fast and slow. 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
https://spenreviee net. forum www.wddwwwwwwyyee Sts
Ziyu Jiang, Tianlong Chen, Bobak J Mortazavi, and Zhangyang Wang. Self-damaging contrastive learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Con- ference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4927 4939. PMLR, 18 24 Jul 2021. URL htttps://procedd ngs. mlr ptess/v139/ jiang2ia. html
Dharshan Kumaran and Eleanor A Maguire. Which computational mechanisms operate in the hip- pocampus during novelty detection? Hippocampus, 17(9):735-743 2007.
Dharshan Kumaran, Demis Hassabis, and James McClellndd. What learning systems do intelligent agents need? complementary learning systems theory updated. Trends in cognitive sciences, 20 (7):512-534, 2016
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. CoRR, 12.22004.19990 2020. URL https:
Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection, In 2017 IEEE International Conference on Computer Vision (ICCV), pp. 2999-3007,
2017. doi: 10. 109/ICCV2001.324
Hong Liu, Jeff Z. HaoChen, Adrien Gaidon, and Tengyu Ma. Self-supervised learning is more robust to dataset imbalance, 2021. URL htRhstts::/// http:////aaaiii..orrbb//////11/1 05025
James L McClelland, Bruce L McNaughton, and Randall C C''R C'RReily. Why there are complementary learning systems in the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937. PMLR, 2016a.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mchdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. CoRR, (obb/1602.0178.. 2016b. URL
/wtp:./612...111 01783
Joseph O'Neill, Barty (heydell-Boveeri David Dupret, and Jozsef Csicsvari. Play again: reacti- vation of waking experience and memory. Trends in neurosciences, 33(5):220-229 2010.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. CoRR, 2017. http:htt://arxi... abs 703. 01 988

---

Under review as a conference paper at ICLR 2023
Linda B Smith, Swapnaa Jayaraman, Elizabeth Clerkin, and Chen Yu. The developing infant creates a curriculum for statistical learning. Trends cognitive sciences, 22(4):325-333 2018
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. Advances in neural information processing systems, 29, 2016.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. A Bradford
Book, Cambridge, MA, USA, 2018a. ISBN 0262039249.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018b.
Alexandr Ten, Pramod Kaushik, Pierre Pierre-Yves Yves Oudeyer, and Jacqueline Gottlieb. Humans monitor learning progress in curiosity-driven exploration. Nature communications, 12(1):1-10, 2021.
Yonglong Tian, Olivier J. Hénaff, and Aäron den Oord. Divide and contrast: Self-supervised learning from uncurated data. CoRR, (oss22055000054 2021. URL https://arxiv.orr
abs/ 2oss2205... 08054
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM, 64(3):107-115, feb 2021. ISSN 0001-0782. doi: 10.1145/3446776.
org 11t.//...111/4344771111
Zhihan Zhou, Jiangchao Yao, Yan-Feng Wang, Bo Han, and Ya Zhang. Contrastive learning with boosted memorization. In International Conference on Machine Learning, pp. 27367-27377. PMLR, 2022.
George Kingsley Zipf. The psycho-biology of language: An introduction to dynamic philology. Routledge, 2013.
A APPENDIX
A.1 RESULTS ON DEEPMIND'S GRIDWORLD
We evaluate various baselines on Deepmind Zip's Gridworld Environment (10 maps and 10 objects) and report results in Table
Table 2: Evaluation Performance: Deepmind Zipf's Gridworld Environment.
Accuracy (%)
Rare
Uniform
Method
Zipfian
(Rarest 20% objects on rarest 20% maps)
(All maps and objects)
IMPALA
88.3 + 0.02
41.1 + 0.01
= 0.00
IMPALA + Vis recon
90.2 + 0.02
45.9 + 0.01
0 + 0.00
IMPALA MEM
92.9 + 0.03
51.2 0.02
25.0 + 0.02
IMPALA MEM Contr. Learning
94.8 + 0.02
52.3 0.01
25.1 0.02
Ours
98.5 + 0.02
66.3 + 0.01
25.2 0.01
A.2 SUPPLEMENTARY ANALYSES
The training curve for the main experiments can be seen in Figure
The effects of changing different hyperparameters of our architecture are explained here.
We look at the 'K' value of K. Nearest Neighbour, Trajectory Hop (hp). Rare State Transfer amount (tx), and Rare State Transfer Frequency (ty) for our ablation study. The effect on the training of these hyperparameters can be seen in Figure 6
12

---

Under review as a conference paper at ICLR 2023
Training Curve (Zipfian Accuracy)
1.0
Mean Episode Return
0.8
0.6
0.4
Ours
Contrastive Learning
IMPALA+MEM
0.2
IMPALA
IMPALA+ Vis recon
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
Steps
1e7
Figure 5: Performance plots: Training curves for different experiments.
Table 3: Effect of KNN 'K' value,
Accuracy (%)
Rare
Uniform
K Value
Zipfian
(Rarest 20% objects on rarest 20% maps)
(All maps and objects)
99.4
79.0
98.5
67.6
16
99.9
84.4
49.9
32
97.3
74.9
24.9
Table 4: Effect of trajectory hop (hp).
Accuracy (%)
Rare
Uniform
(hp)
Zipfian
(Rarest 20% objects
(All maps and objects)
rarest 20% maps)
99.1
78.4
24.8
99.3
79.1
24.9
99.9
84.4
49.9
99.1
77.7
25.2
13

---

Under review as a conference paper at ICLR 2023
Training Curve (Zipfian Accuracy)
Training Curve (Zipfian Accuracy)
2.0
1.0
Episode Return
0.0
Mean Episode Return
0.8
0.4
2pm2
Apm8
k-16
0.0
1o-32
Jpm32
0.0
6.5
1.0
15
2.0
2.5
20
3.5
4.0
0.0
0.5
2.0
1.5
2.0
2.5
3.0
25
4.0
Steps
Strps
le7
(a) Effect KNN 'K' value
(b) Effect trajectory hop (hp)
Training Curve (Zipfian Accuracy)
Training Curve (Zipfian Accuracy)
2.0
1.0
Retur
Mean Episode Retum
0.01
0.6
Red
0.4
At-37
1/-4
512
32
0.0
2.0
Steps
Steps
(c) Effect of transfer amount (tc)
(d) Effect of transfer frequency (fr)
Figure Training Performance: Effect of changing different hyperparameters on training.
Table 5: Effect transfer amount (4x).
Accuracy (%)
Rare
Uniform
(lx)
Zipfian
(Rarest 20% objects on rarest 20% maps)
(All maps and objects)
99.2
81.1
24.9
16
99.9
84.4
49.8
32
99.2
76.0
0
64
98.9
77.9
0
128
98.7
69.8
0
Table 6: Effecc of transfer frequency (ty).
Accuracy (%)
Rare
Uniform
(ty)
Zipfian
(Rarest 20% objects rarest 20% maps)
(All maps and objects)
99.1
79.9
25.0
99.9
84.4
49.8
99.3
82.0
49.9
99.1
76.9
74.8
14

---

Under review as a conference paper at ICLR 2023
A.3
EXPERIMENT HYPERPARAMETERS
Zipf's Gridworld
Image Width
84
Image Height
84
Action Repeats
1
Unroll Length
32
Discount (7)
0.99
Baseline loss scaling
0.5
Entropy cost
0.01
Action Repeats
1
Optimizer
RMSProp
Learning rate
3e 4
Number of training steps Maximum steps in a trial
4e7
100
Additional Parameters
Zipf's Exponent (e)
Number of Actors
50
Trajectory Hop (hp)
Average Momentum Beta (8)
0.97
Loss Gamma (7)
0.5
MEM Buffer capacity
1024
Familiarity Memory Buffer capacity
1024
Rare State Transfer Amount (tx)
512
Rare State Transfer Frequency (ty)
8
KNN (K)
16
Epsilon (c)
1e - 3
Sigma (sigma)
0.05
15


Reviews:
[{"review_id": "XBuOvQ21ugI", "reviewer": "Reviewer_AqsP", "summary": "", "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": "", "summary_of_the_paper": "This paper proposes a method for tackling RL with long-tailed distributions. It uses a self-supervised contrastive method to estimate state familiarity, and then adds the rare states to an episodic memory module based on prior work. The paper considers one task domain, in which it demonstrates compelling improvements from the proposed approach.", "summary_of_the_review": "\nI think this paper offers an interesting idea and some compelling preliminary results. I think it would be a great workshop paper now, but I feel it would need experiments on one more domain and some more ablations to accept as a main conference paper.\n", "strength_and_weaknesses": "Strengths:\nThe idea is interesting and compelling.\nThe results on the environment considered are impressive.\n\nWeaknesses:\nUsing only a single environment makes the results somewhat limited. It would be nice to see experiments on another domain, to verify whether the results generalize. \nIt would be very useful to consider more ablations to the procedure. \nIs the momentum in the familiarity necessary, or would using a simple average, or even just a single value of the loss give similar performance? While Zhou et al. presumably explored this in the image setting, it’s not clear that the results would be the same in RL.\nWhy these particular data augmentations, as opposed to others e.g. the others used in methods like SimCLR [Chen et al., 2020]?\n\nComments / notes:\nThe idea of identifying state familiarity through loss on an unsupervised method on those states is related to ideas that have been used to measure state novelty in exploration, such as in RND [Burda et al., 2018]. It might be worth making this link in the paper, because it suggests the possibility of more fertile overlap between these seemingly-disparate research streams in the future.\n“Due to the unavailability of the code for the environment …” — The code for Zipf’s Gridworld task has been released (at least now) at https://github.com/deepmind/zipfian_environments/tree/main/gridworld — in fact, I’m using it in some present work. Might be worth updating this comment. The other environments could potentially be used to address the weakness noted above. \n\nReferences\n------ \n\nBurda, Y., Edwards, H., Storkey, A., & Klimov, O. (2018, September). Exploration by random network distillation. In International Conference on Learning Representations.\n\nTing Chen, Simon Kornblith, Mohammad Norouzi, Geoffrey Hinton Proceedings of the 37th International Conference on Machine Learning, PMLR 119:1597-1607, 2020.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "X-hMZ_AX2R", "reviewer": "Reviewer_uaJ9", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "The paper proposes a mechanism to keep long-tail events in episodic memory to improve RL performance in Zipfian-like environments. Through contrastive learning, observations' rarity score is estimated and stored in a familiarity buffer. After certain training epochs, only the top rarest events from the buffer transit to the episodic memory, providing long-term context to the policy network following IMPALA architecture. The proposed method is tested in Zipf's Gridworld task and outperforms several variants of the IMPALA baseline, especially in handling long-tail maps. ", "summary_of_the_review": "Although the motivation of the paper is good, its method is incremental and the result is insufficient. The writing needs major improvement. ", "strength_and_weaknesses": "### Strength: \nEncouraging rare events in episodic memory is a reasonable idea. It fits well with the problem the paper aims to solve: Zipfian-like environments. Using contrastive learning to discover rare samples is well-motivated by long-tail learning literature. The experimental results show apparent performance gain in the 20% rarity case. Some baselines are wisely chosen to verify the contribution of the proposed idea.  \n\n### Weakness:\nThe method is relatively incremental. It merely combines components such as IMPALA, episodic memory and contrastive learning. The new module is the familiarity buffer, whose design is not surprising given the paper's main idea (storing the rarest events). This idea is not new either and can date back to this paper [1]. \n\nAnother problem is the writing and presentation. The introduction provides little information on the operation of the proposed method, which creates a gap in understanding the claimed contributions. The notations are confusing and sometimes overlapping. The algorithm is hard to read; one may not implement this algorithm to reproduce the method. \n\nLast but not least, the experiment is weak. Good results in a toy task and for only one setting (10 maps+10 objects) are not enough to validate the method. There are many additional hyperparameters. Since the performance changes as they vary (Table 2-5), it may require a lot of tuning to make this method work. \n\n### Detailed comments and questions:\n\n- The second paragraph in the introduction seems to be about other papers (Sutton & Barto, 2018b; Espeholt et al., 2018). Please rewrite it to clarify your contribution. \n- Symbols like $tk$ can be written as $t_k$ for readability. T is used in both Eq. 2 and 3; are they different? What is N in Eq.4? What is v in Eq. 7?\n- Sec. 4.2: IMPALA+MEM part should be described in detail in the background.\n- The literature review is missing. Your references (e.g. Pritzel et al. (2017); Blundell et al. (2016)) are not up-to-date. It is better to have a separate section covering recent advances in episodic reinforcement learning. \n- Algorithm: Please link the functions to the section you describe them. \n- Fig. 1:  There should be k_t in the KNN part\n- Fig. 4: There are some unexplainable cases. d) Contrastive learning can solve the rarest one (9-9) while your method fails. All methods fail in (5-1), which is far from the hardest. Why?\n- Table 1: Why is your method better in the Uniform case? Is there any advantage? \n- Hyperparameter: Does changing the memory/familiarity buffer size matter?\n- To strengthen the experiment, please consider the following:\n\n   * Test in other environments that also have rare events or sparse rewards. Many Atari and Minigrid environments exhibit these properties. Although they may not be precisely Zipfian, the need for long-tail discovery remains.  \n   * Test in different configurations of your toy environment by varying the environment's hyperparameters\n   * Compare your contrastive learning with simple long-tail discovery methods, such as using VAE for novelty/outliner detection. \n   * Consider some other episodic control methods as the baselines.\n\n[1] Frank, Jordan & Mannor, Shie & Precup, Doina. (2008). Reinforcement learning in the presence of rare events.  ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "DQOS6_t1JJ", "reviewer": "Reviewer_cQdQ", "summary": "", "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": "", "summary_of_the_paper": "This works presents a memory based method to deal with learning in Zipfian distributions, i.e., when some of the training data is rarely visited. To do so, authors start from an existing IMPALA + an Episodic Memory (MEM) method, and introduce a \"familiarity buffer\" that acts as a long-term memory and selects the top k rarest events to send them to the MEM, that acts like a shorter-spam or working memory.\n\nExperiments shows that adding this familiarity buffer does indeed boost performance specially with rare events", "summary_of_the_review": "This work tackles a complex problem and a very interesting research line. The method seems promising but the works is still incomplete with several points to be tackled for the community understand the solution presented in this work.", "strength_and_weaknesses": "Authors make a great work introducing the topic of research motivations and goals. Empirical results evidence that the presented method boosts the ability of agents to track long-tailed events.\n\nI believe there are two major weaknesses in this work. \n\nFirst, is the \"method\" section (Sec. 4), which is sometimes disjointed and could also improve with giving some intuition at some points. It also presents plenty of hyperparameters for the method, without detailing much about how sensible the proposed solution is to the hyperparameters. I'll detail further in the box below\n\nSecond, the experiment section is too narrow. In the sense that:\n* It only contrast their extension of IMPALA+MEM against IMPALA and IMPALA+MEM, it would have been a great addition to compare with a completely different kind of memory-based solution like HCAM [1] to contrast. It would also greatly impact the significance of this work if authors demonstrated that the familiarity buffer can be integrated with other methods than MEM\n* It would really benefit from ablation experiments. This would let the community know how influential is each element of your proposed solution, also it would clarify how sensible it is to the different hyperparameters.\n\n[1] Lampinen, Andrew, et al. \"Towards mental time travel: a hierarchical memory for reinforcement learning agents.\" Advances in Neural Information Processing Systems 34 (2021): 28182-28195.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0JxvarcG45", "reviewer": "Reviewer_h8Xx", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "This paper proposes a method for training RL agents in settings where the environment distribution is long-tailed, i.e. an agent must solve N tasks, but some of them occurs much more frequently than others. Examples of such a setting include the Zipfian RL environments introduced by [1]. This paper proposes a method that uses contrastive learning to get a familiarity score based on the contrastive loss (smoothed using momentum). This familiarity score is then combined with an episodic memory to provide additional learning for trajectories that belong to rare tasks. \n\nThe method is evaluated on a single environment, which is a modified version of the Zipfian gridworld introduced in [1]. It shows improvement over standard IMPALA, as well as different ablated versions of the proposed method, when the distribution of tasks is uniform rather than long-tailed, and the improvement is especially pronounced for rare tasks. \n\n\n\n\n\n\n\n\n[1] Chan et al., \"Zipfian environments for Reinforcement Learning\". ", "summary_of_the_review": "Overall, I think this paper proposes a promising idea, but in its current form it is not ready for publication at ICLR. I would highly encourage the authors to extend the empirical evaluation as described in my comments. If the improvements reported here hold up to a stronger empirical evaluation, I think this would make a nice paper at a top conference. ", "strength_and_weaknesses": "Strengths:\n- The setting is interesting and the general idea of considering heavily skewed task distributions is well-motivated\n- The method does seem to show improvement, although it's only evaluated on a single task\n\nWeaknesses:\n- The main weakness of this paper is the experimental evaluation. First, only a single environment is used - this is insufficient to showcase the generality of the method. It could be that this particular approach is overfit to this task. I would suggest the authors test their methods on several more environments. \n- For some reason, the authors do not use the original version of the Zipfian Gridworld task from [1] and make their own modified version of it (which is also easier, 10 objects/tasks instead of 20). They claim in their paper that this is because the code of [1] is not available, but it is: https://github.com/deepmind/zipfian_environments. It has been available long before the ICLR submission deadline. Since the environments from [1] are open source, I think they should be used instead to facilitate consistency across publications and to ensure comparability of results. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject