Original: R OBUST E XPLORATION VIA C LUSTERING - BASED O NLINE D ENSITY E STIMATION
Revised: R OBUST E XPLORATION VIA C LUSTERING - BASED O NLINE D ENSITY E STIMATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Intrinsic motivation is a critical ingredient in reinforcement learning to enable progress when rewards are sparse.
Revised: Intrinsic motivation is a critical ingredient in reinforcement learning to enable progress when rewards are sparse.

Original: However, many existing approaches that measure the novelty of observations are brittle, or rely on restrictive assumptions about the environment which limit generality.
Revised: However, many existing approaches that measure the novelty of observations are brittle, or rely on restrictive assumptions about the environment which limit generality.

Original: Wepropose to decompose the exploration problem into two orthogonal sub-problems: (i) finding the right representation (metric) for exploration (ii) estimating densities in this representation space. To address (ii), we introduce Robust Exploration viaClustering-based Online Density
Revised: We introduce Robust Exploration via

Original:  Estimation ( RECODE ), a non-parametric method that estimates visitation counts for clusters of states that are similar according to the metric induced by any arbitrary representation learning technique.
Revised: Clustering-based Online Density Estimation ( RECODE ), a non-parametric method that estimates visitation counts for clusters of states that are similar according to the metric induced by a specified representation learning technique.

Original: We adapt classical clustering algorithms to  design a new type of memory that allows RECODE to keep track of the history of interactions over thousands of episodes, thus effectively tracking global visitation counts.
Revised: We adapt classical clustering algorithms to the online setting to design a new type of memory that allows RECODE to efficiently track global visitation counts over thousands of episodes.

Original: This is in contrast to existing non-parametric approaches, that can only store the recent history, typically the current episode.

Original: The generality of RECODE allows us to easily address (i) by leveraging both off-the-shelf and novel representation learning techniques.
Revised: RECODE can easily leverage both off-the-shelf and novel representation learning techniques.

Original: In particular, we introduce a novel generalization of the action-prediction representation that leverages  multi-step predictions and that we find to be better suited to a suite of challenging 3D-exploration tasks in DM-HARD-8 .
Revised: We introduce a novel generalization of the actionprediction representation that leverages transformers for multi-step predictions, which we demonstrate to be more performant on a suite of challenging 3Dexploration tasks in DM-HARD-8 .

Original: We show experimentally that our approach can work with a variety of RL agents, and obtain state-of-the-art performance on Atari and DM-HARD-8 .
Revised: We show experimentally that our approach can work with a variety of RL agents, obtaining state-of-the-art performance on Atari and DM-HARD-8 , and being the first agent to reach the end-screeen in Pitfall!

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Exploration mechanisms are a key component of reinforcement learning (RL, Sutton & Barto, 2018) agents, especially in sparse-reward tasks where long sequences of actions need to be executed before collecting a reward.
Revised: Exploration mechanisms are a key component of reinforcement learning (RL, Sutton & Barto, 2018) agents, especially in sparse-reward tasks where long sequences of actions need to be executed before collecting a reward.

Original: The exploration problem has been studied theoretically (Kearns & Singh, 2002; Azar et al., 2017; Brafman & Tennenholtz, 2003;
Revised: The exploration problem has been studied theoretically (Kearns & Singh, 2002; Azar et al., 2017; Brafman & Tennenholtz, 2003;

Original: Auer et al., 2002; Agrawal & Goyal, 2012;
Revised: Auer et al., 2002; Agrawal & Goyal, 2012;

Original: Audibert et al., 2010; Jin et al., 2020) in the context of bandits (Lattimore & Szepesv ´ ari, 2020) and
Revised: Audibert et al., 2010; Jin et al., 2020) in the context of bandits (Lattimore & Szepesv´ari, 2020) and

Original: Markov Decision Processes (MDP, Puterman, 1990; Jaksch et al., 2010).
Revised: Markov Decision Processes (MDP, Puterman, 1990; Jaksch et al., 2010).

Original: Among those theoretical works, one simple and theoretically-sound approach to perform exploration efficiently in MDPs is to use a decreasing function of the visitation counts as an exploration bonus (Strehl & Littman, 2008;
Revised: Among those theoretical works, one simple and theoretically-sound approach to perform exploration efficiently in MDPs is to use a decreasing function of the visitation counts as an exploration bonus (Strehl & Littman, 2008;

Original: Azar et al., 2017).
Revised: Azar et al., 2017).

Original: However, this approach is intractable with large or continuous state spaces, as generalization between states becomes essential.
Revised: However, this approach is intractable with large or continuous state spaces, as generalization between states becomes essential.

Original: Several experimental works have tried to come up with ways to estimate visitation counts/densities in complex environments where counting is not trivial.
Revised: Several experimental works have tried to come up with ways to estimate visitation counts/densities in complex environments where counting is not trivial.

Original: Two partially successful approaches have emerged to empirically estimate visitation counts/densities in deep RL: (i) the parametric approach that uses neural networks and (ii) the nonparametric approach that uses a slot-based memory to store representations of visited states, where the representation learning method serves to induce a more meaningful metric 1 between states.
Revised: Two partially successful approaches have emerged to empirically estimate visitation counts/densities in deep RL: (i) the parametric approach that uses neural networks and (ii) the nonparametric approach that uses a slot-based memory to store representations of visited states, where the representation learning method serves to induce a more meaningful metric 1 between states.

Original: Parametric methods either explicitly estimate the visitation counts using density models (Bellemare et al., 2016; Ostrovski et al., 2017) or implicitly estimate the counts using e.g., Random Network
Revised: Parametric methods either explicitly estimate the visitation counts using density models (Bellemare et al., 2016; Ostrovski et al., 2017) or implicitly estimate the counts using e.g., Random Network

Original: Distillation ( RND , Burda et al., 2019; Badia et al., 2020b).
Revised: Distillation ( RND , Burda et al., 2019; Badia et al., 2020b).

Original: Non-parametric methods rely on a memory to store encountered state representations (Badia et al., 2020b) and representation learning to construct a metric that differentiates states meaningfully (Pathak et al., 2017).
Revised: Non-parametric methods rely on a memory to store encountered state representations (Badia et al., 2020b) and representation learning to construct a metric that differentiates states meaningfully (Pathak et al., 2017).

Original: Parametric methods do not store individual states explicitly and as such their capacity is not directly bound by memory constraints; but they are less well suited to rapid adaptation on short timescales (e.g., within a single episode).
Revised: Parametric methods do not store individual states explicitly and as such their capacity is not directly bound by memory constraints; but they are less well suited to rapid adaptation on short timescales (e.g., within a single episode).

Original: To bring the best of both worlds into a single method, Never Give Up ( NGU , Badia et al., 2020b) combines a short-term novelty signal based on an episodic memory and a long-term novelty via RND , into a single intrinsic reward.
Revised: To obtain the best of both worlds, Never Give Up ( NGU , Badia et al., 2020b) combines a short-term novelty signal based on an episodic memory and a long-term novelty via RND , into a single intrinsic reward.

Original: However, this approach also naturally inherits the disadvantages of RND ; in particular, susceptibility to uncontrollable or noisy features (see Section 5).
Revised: However, this approach also naturally inherits the disadvantages of RND ; in particular, susceptibility to uncontrollable or noisy features (see Section 5), and being difficult to tune.

Original: More details on related works are provided in App.
Revised: More details on related works are provided in App.

Original: B.
Revised: C.

Original: In this paper, we propose to decompose the exploration problem into two orthogonal sub-problems.
Revised: In this paper, we propose to decompose the exploration problem into two orthogonal sub-problems.

Original: First, (i) Representation Learning which is the task of learning an embedding function on observations or trajectories that encodes a meaningful notion of similarity.
Revised: First, (i) Representation Learning which is the task of learning an embedding function on observations or trajectories that encodes a meaningful notion of similarity.

Original: Second, (ii) Density Estimation which is the task of estimating smoothed visitation counts to derive a novelty-based exploration bonus.
Revised: Second, (ii) Density Estimation which is the task of estimating smoothed visitation counts to derive a novelty-based exploration bonus.

Original: The first contribution of this paper is a general solution to (ii).
Revised: We first present a general solution to (ii) which is computationally efficient and scalable to complex environments.

Original: We introduce Robust Exploration via Clustering-based Online Density Estimation ( RECODE ), a non-parametric method that estimates visitation counts for clusters of states that are similar according to the metric induced by any arbitrary representation.
Revised: We introduce Robust Exploration via Clustering-based Online Density Estimation ( RECODE ), a non-parametric method that estimates visitation counts for clusters of states that are similar according to a metric induced by any arbitrary representation.

Original: We leverage classical clustering algorithms to design a new type of memory that allows RECODE to keep track of histories of interactions spanning thousands of episodes.
Revised: We adapt classical clustering algorithms to an online setting, resulting in a new type of memory that allows RECODE to keep track of histories of interactions spanning thousands of episodes.

Original: This is in contrast to existing non-parametric exploration methods, which store only the recent history and in practice usually only account for the current episode.
Revised: This is in contrast to existing non-parametric exploration methods, which store only the recent history and in practice usually only account for the current episode.

Original: The resulting exploration bonus is principled, simple, and matches  state-of-the-art exploration results on Atari.
Revised: The resulting exploration bonus is principled, simple, and matches or exceeds state-of-the-art exploration results on Atari; being the first agent to reach the end-screen in Pitfall! .

Original: In the presence of noise, we show that it strictly improves over state-of-the-art exploration bonuses such asRND or NGU  .
Revised: In the presence of noise, we show that it strictly improves over state-of-the-art exploration bonuses such as NGU or RND .

Original: The generality of RECODE also allows us to easily leverage both off-the-shelf and novel representation learning techniques, leading to our second contribution.
Revised: The generality of RECODE also allows us to easily leverage both off-the-shelf and novel representation learning techniques, which leads in to our second contribution.

Original: Specifically, we generalize the action-prediction representation learning method (Pathak et al., 2017), used in several state-of-the-art exploration agents,  to a multi-step prediction setting that is better-suited to 3D and partially-observable domains.
Revised: Specifically, we generalize the action-prediction representations (Pathak et al., 2017), used in several state-of-the-art exploration agents, by applying transformers to masked trajectories of state and action embeddings for multi-step action prediction. We refer to this method as CASM for Coupled

Original:  In particular, we show that it can yield significant performance gains insome hard 3D-exploration tasks included in the DM-HARD-8 suite, and achieves a new state of the art in the single-task setting.
Revised: Action-State Masking . In conjunction with RECODE , CASM can yield significant performance gains in hard 3D-exploration tasks included in the DM-HARD-8 suite; achieving a new state of the art in the single-task setting.

Original:  B ACKGROUND AND NOTATION
Revised:  B ACKGROUND AND NOTATION

Original: In this section, we provide the necessary background and notation to understand our method (see
Revised: In this section, we provide the necessary background and notation to understand our method (see

Original: Sec. 3).
Revised: Sec. 3).

Original: First, we present a general setting of interaction between an agent and its environment.
Revised: First, we present a general setting of interaction between an agent and its environment.

Original: Second, we define the terms embeddings, atoms and memory.
Revised: Second, we define the terms embeddings, atoms and memory.

Original: Third, we present our notation for visitation counts.
Revised: Third, we present our notation for visitation counts.

Original: Finally, we show how we derive intrinsic rewards from visitations counts.
Revised: Finally, we show how we derive intrinsic rewards from visitations counts.

Original: Interaction Process between an Agent and its Environment.
Revised: Interaction Process between an Agent and its Environment.

Original: We consider a discrete-time interaction process (McCallum, 1995; Hutter, 2004; Hutter et al., 2009; Daswani et al., 2013) between an agent and its environment where, at each time step t ∈ N , the agent receives an observation o t ∈ O and generates an action a t ∈ A .
Revised: We consider a discrete-time interaction process (McCallum, 1995; Hutter, 2004; Hutter et al., 2009; Daswani et al., 2013) between an agent and its environment where, at each time step t ∈ N , the agent receives an observation o t ∈ O and generates an action a t ∈ A .

Original: We consider an environment with stochastic dynamics p : H ×
Revised: We consider an environment with stochastic dynamics p : H ×

Original: A → ∆ O 2 that maps a history of past observations-actions and a current action to a probability distribution over future observations.
Revised: A → ∆ O 2 that maps a history of past observations-actions and a current action to a probability distribution over future observations.

Original: More precisely, the space of past observations-actions is H =
Revised: More precisely, the space of past observations-actions is H =

Original: (cid:83) t ∈
Revised: (cid:83) t ∈

Original: N H t where H 

Revised: N H t where H

Original: = O and ∀ t ∈ N ∗ , H t +1 =
Revised: = O and ∀ t ∈ N ∗ , H t +1 =

Original: H t × A × O .
Revised: H t × A × O .

Original: We consider policies π : H → ∆ A that maps a history of past observations-actions to a probability distribution over actions.
Revised: We consider policies π : H → ∆ A that maps a history of past observations-actions to a probability distribution over actions.

Original: Finally, an extrinsic reward function r e : H ×
Revised: Finally, an extrinsic reward function r e : H ×

Original: A → R maps a history to a scalar feedback.
Revised: A → R maps a history to a scalar feedback.

Original: Embeddings, Atoms and Memory.
Revised: Embeddings, Atoms and Memory.

Original: An embedder is a parameterized function f θ : H → E where E is an embedding space.
Revised: An embedder is a parameterized function f θ : H → E where E is an embedding space.

Original: Typically, the embedding space is the vector space R N where
Revised: Typically, the embedding space is the vector space R N where

Original: N ∈ N ∗ is the embedding size.
Revised: N ∈ N ∗ is the embedding size.

Original: Therefore, for a given time step t ∈ N , an embedder is a function f θ that associates to any history h t ∈ H t a vector e t = f θ ( h t ) called an embedding.
Revised: Therefore, for a given time step t ∈ N , an embedder is a function f θ that associates to any history h t ∈ H t a vector e t = f θ ( h t ) called an embedding.

Original: There are several ways to train an embedder f θ such as using an auto-encoding loss of the observation o t (Burda et al., 2018), using an inverse dynamics loss (Pathak et al., 2017) or using a multi-step predictionerror loss at the latent level (Guo et al., 2020; 2022).
Revised: There are several ways to train an embedder f θ such as using an auto-encoding loss of the observation o t (Burda et al., 2018a), using an inverse dynamics loss (Pathak et al., 2017) or using a multi-step prediction-error loss at the latent level (Guo et al., 2020; 2022).

Original: Those techniques are referred as representation learning methods.
Revised: Those techniques are referred as representation learning methods.

Original: An atom f ∈ E is a vector in the embedding space that is contained in a memory
Revised: An atom f ∈ E is a vector in the embedding space that is contained in a memory

Original: M = { f i ∈ E} | M | i =1 which is a finite slot-based container, where | M | ∈ N ∗ is the memory size.
Revised: M = { f i ∈ E} | M | i =1 which is a finite slot-based container, where | M | ∈ N ∗ is the memory size.

Original: The memory M is updated at each time step t by a non-parametric function of the memory M and the embedding e t .
Revised: The memory M is updated at each time step t by a non-parametric function of the memory M and the embedding e t .

Original: In the simplest case, the memory is filled in a first-in first-out (FIFO) manner along the interactions (Badia et al., 2020b;a) and atoms are simply the embeddings themselves.
Revised: In the simplest case, the memory is filled in a first-in first-out (FIFO) manner along the interactions (Badia et al., 2020b;a) and atoms are simply the embeddings themselves.

Original: However, more complex mechanisms than FIFO can be considered to fill/update a memory.
Revised: However, more complex mechanisms than FIFO can be considered to fill/update a memory.

Original: For instance given a memory M and an embedding e t , the embedding e t can be inserted in the memory if and only if it is  different from the other atoms in the memory and the memory is not at capacity.
Revised: For instance given a memory M and an embedding e t , the embedding e t can be inserted in the memory if and only if it is sufficiently different from the other atoms in the memory and the memory is not at capacity.

Original: The update rule that defines the atoms in the memory is a key component of our method.
Revised: The update rule that defines the atoms in the memory is a key component of our method.

Original: Visitation Counts.
Revised: Visitation Counts.

Original: Let M = { f l ∈ E} | M | l =1 be a slot memory containing atoms. The vanilla visitation count, N δ ( M, e ) , for a given embedding e ∈ E with respect to the memory M is:
Revised: The exact visitation count, N δ ( M, e ) , for a given embedding e ∈ E with respect to the memory M can be written as:

Original: [Equation]
Revised: [Equation]

Original: δ ( f l , e ) , where δ ( f, e ) : ( e, f ) ∈ E 2 = (cid:26) 1 , if e = f 0 , otherwise , ( [Equation] When the embedding space is complex and too large, the vanilla visitation countbecomes uninformative because each embedding is potentially different from all the atoms.
Revised: However, when the state-space is very large or continuous, the exact visitation count is often uninformative since the same embedding may rarely be encountered twice.

Original: To overcome this problem,  soft-visitation counts are computed.
Revised: To overcome this problem, we can instead compute soft-visitation counts N K ( M, e ) :

Original: The soft-visitation count, N K ( M, e ) , for a given state e with respect to the memory M is N K ( M, e )

Original: = (cid:80) | M | l =1 K ( f l , e ) , where K ∈ R E 2{+
+} + is a definite positive kernel.
Revised: = (cid:80) | M | l =1 K ( f l , e ) , where K ∈ R E {+
+} + is a definite positive kernel.

Original: Different choices of kernel can be made such as a Gaussian kernel
Revised: Different choices of kernel can be made such as a Gaussian kernel

Original: K ( f, e ) = exp( −∥ e − f ∥ 22 ) or an inverse kernel K ( f, e ) = 1 1+ ∥
Revised: K ( f, e ) = exp( −∥ e − f ∥ 2 2 ) or an inverse kernel K ( f, e ) = 1 1+ ∥

Original: e − f ∥ 22 where ∥ .
Revised: e − f ∥ 2 2 where ∥ .

Original: ∥ 2 is the Euclidean distance in the embedding space E .
Revised: ∥ 2 is the Euclidean distance in the embedding space E .

Original: Finally, we can compute a weighted soft visitation count N K ( M, e, { w l } | M | l =1 ) =
Revised: Finally, we can compute a weighted soft visitation count N K ( M, e, { w l } | M | l =1 ) =

Original: (cid:80)
Revised: (cid:80)

Original: | M | l =1 w l K ( f l , e ) , where the weights w l ∈ R + are positive real numbers.
Revised: | M | l =1 w l K ( f l , e ) , where the weights w l ∈ R + are positive real numbers denoting the count at each atom.

Original: Note that the vanilla visitation count can be recovered from the more general weighted soft visitation count as a special case by setting the weights w l = 1 and K = δ .
Revised: Note that the exact visitation count can be recovered from the more general weighted soft visitation count as a special case by setting the weights w l = 1 and K = δ .

Original: Intrinsic Rewards from Visitation Counts.
Revised: Intrinsic Rewards from Visitation Counts.

Original: It ispretty straightforward to define an intrinsic reward from visitation counts.
Revised: It is straightforward to define an intrinsic reward from visitation counts.

Original: Indeed, as the goal of an exploratory agent is to go to less-visited states,then an intrinsic reward can be any decreasing function of the visitation counts.
Revised: Indeed, as the goal of an exploratory agent is to go to less-visited states, an intrinsic reward can be any decreasing function of the visitation counts.

Original: In the literature, the inverse of the square root of the visitation counts is known to be theoretically sound (Azar et al., 2017).
Revised: In the literature, the inverse of the square root of the visitation counts is known to be theoretically sound (Azar et al., 2017).

Original: Therefore, for a given time t , the intrinsic reward associated to the transition ( o t , a t , o t +1 ) , is defined as:
Revised: Therefore, for a given time t , the intrinsic reward associated to the transition ( o t , a t , o t +1 ) , can be defined as:

Original: [Equation]
Revised: [Equation]

Original: where a small constant c ∈ R + increases stability, e t +1 = f θ ( h t +1 ) and h t +1 = { o 0 , a 0 , · · · , o t +1 } .
Revised: where a small constant c ∈ R + is added to avoid numerical instability.

Original:  RECODE
Revised:  RECODE

Original: In this section, we introduce our Robust Exploration via Clustering-based Online Density Estimation ( RECODE ) approach that computes intrinsic rewards for exploration.
Revised: In this section, we introduce our Robust Exploration via Clustering-based Online Density Estimation ( RECODE ) approach that computes intrinsic rewards for exploration.

Original: At a high level, RECODE stores a fixed number of weighted atoms (typically 5 · 10 4 or 2 · 10 5 depending on the domain) that are interpreted in the following as cluster-centers, along with their counts.
Revised: At a high level, RECODE stores a fixed number of weighted atoms (typically 5 · 10 4 or 2 · 10 5 depending on the domain) that are interpreted in the following as cluster-centers, along with their counts.

Original: The update rule of our memory M for each new embedding e ∈ E and the computation of the intrinsic reward are detailed in Algorithm.
Revised: The update rule of our memory M for each new embedding e ∈ E and the computation of the intrinsic reward are detailed in Algorithm.

Original: To each atom/cluster-center f l ∈ E , we associate a count c l ∈ N (all initialized to 0 ) that are updated after each new embedding e is observed. 
Revised: To each atom/cluster-center f l ∈ E , we associate a count c l ∈ N (initialized to 0 ) that is updated after each new embedding e is observed. When a new embedding is inserted we probabilistically choose either to add it as a new cluster (if it is far away from existing clusters relative to an adaptive threshold) or increment the count of its nearest neighbor and adjust the cluster center toward the new embedding.

Original: The update rule has a close connection to the DP-means algorithm of Kulis & Jordan (2011), with two key differences:
Revised: The update rule has a close connection to the DP-means algorithm of Kulis & Jordan (2011), with two key differences:

Original: • the counts of the cluster-centers are discounted at each step, allowing our approach to deal with the non-stationarity of the data due both to changes in the policy and the embedding function, effectively reducing the weight of stale cluster-centers in the memory, • when creating a new cluster-center, we remove an underpopulated one, so as to keep the size of the memory constant.
Revised: • the counts of the cluster-centers are discounted at each step, allowing our approach to deal with the non-stationarity of the data due both to changes in the policy and the embedding function, effectively reducing the weight of stale cluster-centers in the memory, • when creating a new cluster-center, we remove an underpopulated one, so as to keep the size of the memory constant.

Original:  A theoretical analysis is sketched in Appendix C. To help build some intuition about the quality of our density estimation, we illustrate on Figure 1 the result of Algorithm 1 on a toy example witha non-stationary embedding distribution. We find in particular that tuning the discount allows to
Revised: We compare the qualitative behavior of different removal strategies in Fig. A theoretical analysis is sketched in Appendix D. To help build some intuition about the quality of our density estimation, we illustrate in Figure 1 the result of Algorithm 1 on a toy example with

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original:  smoothly interpolate between short-term and long-term memory, a property that will prove crucial to achieve the strong experimental results of Section 5.
Revised: left a non-stationary embedding distribution. We find in particular that tuning the discount allows to smoothly interpolate between short-term and long-term memory, a property that will prove crucial to achieve the strong experimental results of Section 5.

Original: Note that contrary to the usual episodic memory used in Badia et al.
Revised: Note that contrary to the usual episodic memory used in Badia et al.

Original: (2020b;a) , the memory is never reset, andit is shared between all actors when using a distributed RL agent.
Revised: (2020b;a), the memory is never reset, and is shared between all actors in our distributed RL agent.

Original: We use the following kernel function defined for ( e, f ) ∈ E 2 by:
Revised: We use the following kernel function defined for ( e, f ) ∈ E 2 by:

Original: [Equation]
Revised: [Equation]

Original: where ϵ ∈ R + , d 2 m is an estimate of the squared distance between an embedding and its nearest neighbors in the memory (see Algorithm. 1)) and 1 { . } is the indicator function.
Revised: where ϵ ∈ R + , d 2 m is an estimate of the squared distance between an embedding and its nearest neighbors in the memory (see Algorithm. 1)) and 1 { . } is the indicator function.

Original: Compared to Badia et al.
Revised: Compared to Badia et al.

Original: (2020b;a), our smoothed visitation-counts computes the normalized distances from a given embedding e to all its neighbors within a d 2 m -ball, instead of only the k -nearest neighbors.
Revised: (2020b;a), our smoothed visitation-counts computes the normalized distances from a given embedding e to all its neighbors within a d 2 m -ball, instead of only the k -nearest neighbors.

Original: This change prevents some undesirable properties of the k -NN approach where inserting a cluster with unit-count close to the embedding e might reduce its pseudo-count instead of increasing it if the k -th nearest neighbor of e had a large count.
Revised: This change avoids an undesirable property of the k -NN approach where inserting a cluster with unit- count close to the embedding e might reduce its pseudo-count instead of increasing it if the k -th nearest neighbor of e had a large count.

Original: We now detail how RECODE can be integrated in a typical distributed RL agent (Espeholt et al., 2018; Kapturowski et al., 2018) that comprises several processes that run in parallel and interact with each other.
Revised: We now detail how RECODE can be integrated in a typical distributed RL agent (Espeholt et al., 2018; Kapturowski et al., 2018) that comprises several processes that run in parallel and interact with each other.

Original: Classically, a Learner performs gradient steps to train a policy π θ and an embedding (representation) function f θ , forwarding the parameters θ to an Inference Worker.
Revised: Classically, a Learner performs gradient steps to train a policy π θ and an embedding (representation) function f θ , forwarding the parameters θ to an Inference Worker.

Original: A collection of independent Actors query the inference worker for actions that they execute in the environment and send the resulting transitions to the Learner, optionally through a (prioritized) Replay (Mnih et al., 2015; Schaul et al., 2015).
Revised: A collection of independent Actors query the inference worker for actions that they execute in the environment and send the resulting transitions to the Learner, optionally through a (prioritized) Replay (Mnih et al., 2015; Schaul et al., 2015).

Original: When using RECODE , the Actors additionally communicate with a shared Memory implementing Algorithm 1: at each step t , they query from the Inference Server an embedding f θ ( h t ) of their history and send it to the shared Memory which returns an intrinsic reward r t that is then added to the extrinsic reward to train the policy in the Learner process.
Revised: When using RECODE , the Actors additionally communicate with a shared Memory implementing Algorithm 1: at each step t , they query from the Inference Server an embedding f θ ( h t ) of their history and send it to the shared Memory which returns an intrinsic reward r t that is then added to the extrinsic reward to train the policy in the Learner process.

Original: In practice, we normalize the intrinsic reward by a running estimate of its standard-deviation as in
Revised: In practice, we normalize the intrinsic reward by a running estimate of its standard-deviation as in

Original: Burda et al.
Revised: Burda et al.

Original: A diagram giving an overview of a distributed agent using RECODE is given in
Revised: A diagram giving an overview of a distributed agent using RECODE is given in

Original: Fig.
Revised: Fig.

Original:  R EPRESENTATION L EARNING M ETHODS
Revised:  R EPRESENTATION L EARNING M ETHODS

Original: The Learner process uses the data sent by the actors to train the policy π θ via an RL algorithm and to train the embedder f θ via a representation learning loss.
Revised: The Learner process uses the data sent by the actors to train the policy π θ via an RL algorithm and to train the embedder f θ via a representation learning loss.

Original: As noted in Section 2, the choice of the embedding function f θ : H → E induces a metric in the embedding space E allowing to compare histories.
Revised: As noted in Section 2, the choice of the embedding function f θ : H → E induces a metric in the embedding space E allowing to compare histories.

Original: Many different representation learning techniques have been studied in the context of exploration in RL (Burda et al., 2018;
Revised: Many different representation learning techniques have been studied in the context of exploration in RL (Burda et al., 2018a;

Original: Guo et al., 2020; 2022; 2021; Erraqabi et al., 2021).
Revised: Guo et al., 2020; 2022; 2021; Erraqabi et al., 2021).

Original: In the following, we focus on action prediction embeddings, introducing first the standard 1 -step prediction formulation (Pathak et al., 2017; Badia et al., 2020b;a).
Revised: In the following, we focus on action prediction embeddings, introducing first the standard 1 -step prediction formulation (Pathak et al., 2017; Badia et al., 2020b;a).

Original: Our embedding function f θ is parameterized as a feed-forward neural network taking o t , the observation at time t , as inputs.
Revised: Our embedding function f θ is parameterized as a feed-forward neural network taking o t , the observation at time t , as inputs.

Original: We further define a classifier g ϕ that, given the embeddings of two consecutive observations f θ ( o t ) , f θ ( o t +1 ) , outputs an estimate p θ,ϕ ( a t | o t , o t +1 )
Revised: We further define a classifier g ϕ that, given the embeddings of two consecutive observations f θ ( o t ) , f θ ( o t +1 ) , outputs an estimate p θ,ϕ ( a t | o t , o t +1 )

Original: =
Revised: =

Original: g ϕ
Revised: g ϕ

Original: ( f θ ( o t ) , f θ ( o t +1 )) of the probability of taking an action given two consecutive observations ( o t , o t +1 ) .
Revised: ( f θ ( o t ) , f θ ( o t +1 )) of the probability of taking an action given two consecutive observations ( o t , o t +1 ) .

Original: Both f θ and g ϕ are then jointly trained by minimizing an expectation of the following loss:
Revised: Both f θ and g ϕ are then jointly trained by minimizing an expectation of the following loss:

Original: [Equation]
Revised: [Equation]

Original: where L ( θ, ϕ )( a t ) is the negative log likelihood and a t is the true action taken between o t and o t +1 .
Revised: where L ( θ, ϕ )( a t ) is the negative log likelihood and a t is the true action taken between o t and o t +1 .

Original: These embeddings have been shown to be helpful in environment with many uncontrollable features in the observation (Badia et al., 2020b), such as the game of Pitfall!
Revised: These embeddings have been shown to be helpful in environments with many uncontrollable features in the observation (Badia et al., 2020b), such as the game of Pitfall!

Original: in Atari, where they might result in spurious sources of novelty even when the agent is standing still.
Revised: in Atari, where they might result in spurious sources of novelty even when the agent is standing still.

Original: We note however that RECODE can be used with an arbitrary embedding function, e.g. one tailored for the domain of interest.
Revised: We note however that RECODE can be used with an arbitrary embedding function, e.g. one tailored for the domain of interest.

Original: One downside of the standard, 1 -step action-prediction method is that the simplicity of the prediction task may only require highly localized and low-level features to be learned for its solution, which may not be informative of more geometrical or topological notions of environment structure, that partially-observable or 3D-exploration tasks might require.
Revised: One downside of the standard, 1 -step action-prediction method is that the simplicity of the prediction task may only require highly localized and low-level features to be learned for its solution, which may not be informative of more geometrical or topological notions of environment structure, that partially-observable or 3D-exploration tasks might require.

Original: Other popular forms of representation learning such as Contrastive Predictive Coding (CPC, Oord et al.
Revised: Other popular forms of representation learning such as Contrastive Predictive Coding (CPC, Oord et al.

Original: or Predictions of Bootstrapped Latents (PBL, Guo et al.
Revised: or Predictions of Bootstrapped Latents (PBL, Guo et al.

Original: (2020)) utilize temporally-extended prediction tasks but do not enforce any notion of controllability.
Revised: (2020)) utilize temporally-extended prediction tasks but do not enforce any notion of controllability.

Original: We now present a novel generalization of action-prediction embeddings that we show to yield strong results on DM-HARD-8 in section 5.
Revised: We now present a novel generalization of action-prediction embeddings that we show to yield strong results on DM-HARD-8 in section 5.

Original: Perhaps the most straightforward generalization of 1-step action prediction is to predict sequences of actions between observations o t and o t + k , but in general there may be many such sequences which are possible, besides the one which is obtained by the behavior policy.
Revised: A straightforward generalization of 1-step action prediction is to predict sequences of actions between observations o t and o t + k , but in general there may be many such sequences which are possible, besides the one which is obtained by the behavior policy.

Original: This introduces an additional policy-dependent non-stationarity in the prediction task which could potentially hinder learning efficiency and stability.
Revised: This introduces an additional policydependent non-stationarity in the prediction task which could potentially hinder learning efficiency and stability.

Original: To counteract this problem we couldinstead provide as context the sequence ( o t , a t , a t +1 . . .
Revised: To counteract this problem we could provide all but one of the intervening actions to our context ( o t , a t , a t +1 . . .

Original: a t + k − 2 , o t + k ) for the prediction of a t + k − 1 .
Revised: a t + k − 2 , o t + k ) , and predict only a t + k − 1 .

Original: However, in partially observed domains it is possible that these two observations alone are insufficient to accurately localize the agent’s state, and therefore it may be beneficial to provide additional context before o tor after o t + k .
Revised: However, in partially observed domains it is possible that these two observations alone are insufficient to accurately localize the agent’s state, and therefore it may be beneficial to provide additional context before o t .

Original: Concretely, we propose to apply a causally-masked transformer to sequences of observation and action embeddings, such that at each timestep t exactly one of o t and a t is provided.
Revised: Concretely, we propose to apply a causally-masked transformer to sequences of observation and action embeddings, such that at each timestep t exactly one of o t and a t is provided.

Original: The transformer output is then projected down to the size of the embedding  f θ ( o t ) , and the difference between the two is input into a final MLP classifier, g ϕ .
Revised: The transformer output is then projected down to the size of the embedding ( dim z t = dim f θ ( o t ) ), and the difference between the two is input into a final MLP classifier, g ϕ .

Original: During training, we randomly sample N = 4 masks per trajectory to help reduce gradient variance.
Revised: During training, we randomly sample N = 4 masks per trajectory to help reduce gradient variance.

Original: Note that we use e t = f θ ( o t ) , the transformer inputs , as the embeddings for RECODE in order to avoid leaking information about the agent’s trajectory.
Revised: Note that we use e t = f θ ( o t ) , the transformer inputs , as the embeddings for RECODE in order to avoid leaking information about the agent’s trajectory.

Original: As with 1 -step action prediction, we train the representation using maximum likelihood.
Revised: As with 1 -step action prediction, we train the representation using maximum likelihood.

Original: We refer to this approach as Coupled Action-State Masking ( CASM ) in the following.
Revised: We refer to this approach as Coupled Action-State Masking ( CASM ) in the following.

Original: Figure 2 shows a diagram of the architecture.
Revised: Figure 2 shows a diagram of the architecture.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we experimentally validate the generality of our approach by applying it to several domains with distinct properties.
Revised: In this section, we experimentally validate the generality of our approach by applying it to several domains with distinct properties.

Original: We first show that we can obtain state-of-the-art results on the hardest exploration games of the  Atari domain, and even improve upon existing approaches when noise is added to the environment.
Revised: We first show that we can match or improve state-of-the-art learning efficiency on the hardest exploration games contained in the well-established Atari benchmark.

Original: We then turn to DM-HARD-8 , a suite of partially-observable hard-exploration tasks, where we obtain results competitive with the recently proposed BYOLExplore (Guo et al., 2022) agent. 
Revised: We then turn to DM-HARD-8 , a suite of partially-observable 3D hard exploration tasks, where we obtain results that match or exceed the recently proposed BYOL-Explore (Guo et al., 2022) agent. All the candidates evaluated in the experiments in the main paper and the extended experiments in

Original: Unless otherwise noted, the agent we consider in this section is based on MEME (Kapturowski et al., 2022), a recent improvement over Agent57 (Badia et al., 2020a) that achieves much greater sample efficiency, and is the current state-of-the-art on Atari.
Revised: App. E are composed by three main modules: (1) a base agent, responsible for core RL aspects such as collecting observations and updating the policy, (2) an algorithm responsible to generate an exploration bonus, and (3) a representation learning algorithm responsible to generate latent representation of the observations that allow the exploration bonus to be meaningful. Each of the agent names in the experiments will reflect which specific algorithm is responsible for each of the three modules. For example, in our more detailed taxonomy the original MEME agent described in Kapturowski et al.

Original: Like its predecessors, MEME  uses the NGU intrinsic reward, combining RND and Episodic Memory. We modify MEME by simply replacing the NGU reward by the simpler RECODE reward, and keep everything else fixed.
Revised: (2022) is denoted as the MEME - NGU - AP baseline. Similarly, our MEME - RECODE - AP uses the same base agent ( MEME ) and representation learning algorithm ( AP ), with the only modification being the intrinsic reward mechanism being swapped from NGU to RECODE . For a full description of all the baselines we refer the reader to App.

Original:  We use a memory comprised of 5 · 10 4 atoms for our Atari experiments and 2 · 10 5 atoms for our DM-HARD-8 experiments.
Revised: A. We use a memory comprised of 5 · 10 4 atoms for our Atari experiments and 2 · 10 5 atoms for our DM-HARD-8 experiments.

Original: We find that the resulting agent runs at roughly the same speed as the original MEME agent.
Revised: We find that the resulting agent runs at roughly the same speed as the original MEME agent.

Original: More hyperparameters values can be found in App A.
Revised: More hyperparameters values can be found in App. A.

Original: 5. A TARI
Revised: 5. A TARI

Original: The Atari Learning Environment (ALE, Bellemare et al., 2013) is one of the most used RL benchmarks for deep RL. It comprises 57 Atari games which are  2 -D,fully-observable and (fairly) deterministic environments.
Revised: The Atari Learning Environment (ALE, Bellemare et al., 2013) is one of the most used RL benchmarks for deep RL, comprising 57 Atari games which are mostly 2 -D, and have a limited degree of partial observability which can often be rectified by stacking a small number of frames.

Original: In addition, they have avery long optimization horizon (episodes can last more than 10000 steps), complex observations (preprocessed greyscale images which are 210 ×byte arrays) but a fairly simple discrete action space composed of 18 actions.
Revised: On the other hand, many games have a long optimization horizon (with episodes lasting up to 27000 steps using the standard action-repeat of 4), and rewards vary considerably in both scale and density.

Original: Among the 57 Atari games, only  fewgames are considered hard-explorationgames (Bellemare et al., 2016) such as
Revised: Among the 57 Atari games, only a few are considered hard-exploration (Bellemare et al., 2016) such as

Original: Montezuma’s Revenge , Pitfall and Private Eye .
Revised: Montezuma’s Revenge , Pitfall and Private Eye .

Original: For evaluation, we follow the classicalrandom no-ops evaluation regime (Mnih et al., 2015; Van Hasselt et al., 2016), and average performance over 3 seeds.
Revised: For evaluation, we follow the classicalrandom no-ops evaluation regime (Mnih et al., 2015; Van Hasselt et al., 2016), and average performance over 3 seeds.

Original: This evaluation regime does not use sticky actions (Machado et al., 2018).
Revised: This evaluation regime does not use sticky actions (Machado et al., 2018).

Original: We compare our approach with MEME  and its ablations on 8 of the hardest exploration games on Fig.
Revised: We compare our approach with MEME - NGU - AP and its ablations on 8 of the hardest exploration games on Fig.

Original: We find that  RECODE  matches the performance of  MEME  with a single, simpler intrinsic reward, achieving super-human performance on all 8 hard exploration games.
Revised: We find that MEME - RECODE - AP matches the performance of the original MEME agent with a single, simpler intrinsic reward, achieving super-human performance on all 8 hard exploration games.

Original: Indeed, as shown in the ablations reported in Appendix D, MEME  requires  both RND and episodic memory to solve all 8 games: RND on its own cannot solve Pitfall!
Revised: Indeed, as shown in Fig. 3, MEME - NGU - AP requires the full NGU exploration bonus (i.e., both RND and EMM ) to solve all 8 games: RND on its own cannot solve Pitfall!

Original: because of the many uncontrollable features of its observations, while Episodic Memory on its own cannot solve Montezuma’s Revenge because it requires long-term memory.
Revised: because of the many uncontrollable features in its observations, while EMM on its own cannot solve Montezuma’s Revenge because it requires long-term memory.

Original: Because RECODE estimates the visitation counts over many episodes using Action Prediction embeddings that discard uncontrollable dynamics, it is able to solve both games with a single intrinsic reward.
Revised: Because RECODE estimates the visitation counts over many episodes using AP embeddings that discard uncontrollable dynamics, it is able to solve both games with a single intrinsic reward.

Original: We probe how far back the memory of RECODE goes in Montezuma’sRevenge in Figure 4 and find that the distribution of the age of the clusters learned by
Revised: We probe how far back the memory of RECODE goes in Montezuma’s

Original:  RECODE exhibits a mode around 2 · 10 6 actor steps, which corresponds to hundreds of episodes, with a significant number of clusters ten times older than that.
Revised: Revenge in Figure 4 and find that the distribution of the age of the clusters learned by RECODE exhibits a mode around 2 · 10 6 actor steps, which corresponds to hundreds of episodes, with a significant number of clusters ten times older than that.

Original: In Appendix D, we also compare our approach with  a modification of RND built on top of Action Prediction embeddings, in an attempt to fix the aforementioned undesirable properties of RND .
Revised: We also compare our approach with MEMERNDonAP (detailed in Appendix E), a modification of RND built on top of AP embeddings, in an attempt to fix the aforementioned undesirable properties of RND .

Original: We find that this approach does not allow to solve some of the hardest games such as Montezuma’s Revenge or Pitfall!.
Revised: We find that this approach does not allow to solve some of the hardest games such as Montezuma’s Revenge or Pitfall!.

Original: One possible explanation of this failure is the fact that a large RND error can be caused by either the observation of a new state, or a drift in the representation of an already observed one.
Revised: One possible explanation of this failure is the fact that a large RND error can be caused by either the observation of a new state, or a drift in the representation of an already observed one.

Original: The failure of RND to disentangle these two effects results in poor exploration.
Revised: The failure of RND to disentangle these two effects results in poor exploration.

Original: 5. N OISY M ONTEZUMA ’ S R EVENGE
Revised: 5. N OISY M ONTEZUMA ’ S R EVENGE

Original: Atari Games areknown to be fairly deterministic  which isnot a property commonly encountered in real-world environments which have noisy observations  due to imperfect sensors.
Revised: Atari Games are deterministic with respect to the RAM state, which is a property that most realworld environments do not share (e.g.: observations may be noisy due to imperfect sensors).

Original: As RND relies on predicting a random embedding of the pixel observation to determine whether a state is new or not, it cannot learn to discard noisy features, hindering its ability to detect meaningful novelty in the presence of noise.
Revised: As RND relies on predicting a random embedding of the raw observation to determine whether a state is new or not, it cannot learn to discard noisy features, hindering its ability to detect meaningful novelty in the presence of noise.

Original: In this section, we show that MEME  inherits the limitations of RND .
Revised: In this section, we show that MEME - NGU - AP inherits the limitations of RND .

Original: To this end, we consider the game of Montezuma’s Revenge, for which effective exploration requires long-term memory as shown by the ablations of Section 5.1.
Revised: To this end, we consider the game of Montezuma’s Revenge, for which effective exploration requires long-term memory as shown by the ablations of Section 5.1.

Original: We challenge RND by concatenating the original grey-scale, 210 × 160 pixels frame with a noisy frame of the same shape, where each pixel is sampled uniformly at random from the range [0 , 255] .
Revised: We challenge RND by concatenating the original grey-scale, 210 × 160 pixels frame with a noisy frame of the same shape, where each pixel is sampled uniformly at random from the range [0 , 255] .

Original: This type of noise is commonly referred to as noisy TV (Pathak et al., 2017).
Revised: This type of noise is commonly referred to as noisy TV (Pathak et al., 2017).

Original: The results of our experiments on this environment are presented on Figure 5.
Revised: The results of our experiments on this environment are presented on Figure 5.

Original: Perhaps as expected, we find that the performance of MEME  is strongly deteriorated, since it relies on RND to build its exploration bonus, and reduced to that of a pure RL baseline without exploration bonus (Kapturowski et al., 2018).
Revised: Perhaps as expected, we find that the performance of MEME - NGU - AP is strongly deteriorated, since it relies on RND for the long-term component of its exploration bonus, and performs no better than a pure RL baseline without exploration bonus (Kapturowski et al., 2018).

Original: RECODE , on the other hand, relies on action-prediction embeddings to estimate the visitation-counts.
Revised: RECODE , on the other hand, relies on action-prediction embeddings to estimate global visitation-counts.

Original: In this embedding space, states that only differ in the noisy, uncontrollable part of the observation tend to be aliased together, so that the effect of the noise on exploration vanishes.
Revised: In this embedding space, states that only differ in the noisy, uncontrollable part of the observation tend to be aliased together, so that the effect of the noise on exploration vanishes.

Original: Indeed, the performance of RECODE on Montezuma’s Revenge is unchanged when adding noise.
Revised: Indeed, the performance of MEMERECODE - AP on Montezuma’s Revenge is unchanged when adding noise.

Original: 5. DM-HARD- 
Revised: 5. DM-HARD-

Original: DM-HARD-8 (Gulcehre et al., 2019) is a benchmark comprised of 8 hard exploration tasks, originally built to emphasize the difficulties encountered by an RL agent when learning from sparse rewards in a procedurally-generated 3 -D world with partial observability, continuous control, and highly variable initial conditions.
Revised: DM-HARD-8 (Gulcehre et al., 2019) is a benchmark comprised of 8 hard exploration tasks, originally built to emphasize the difficulties encountered by an RL agent when learning from sparse rewards in a procedurally-generated 3 -D world with partial observability, continuous control, and highly variable initial conditions.

Original: Each task requires the agent to interact with specific objects in its environment in order to reach a large apple that provides reward (see Figure 11  for an example).
Revised: Each task requires the agent to interact with specific objects in its environment in order to reach a large apple that provides reward (see Figure 11 in the Appendix for an example).

Original: Being procedurally-generated, properties such as object shapes, colors, and positions are different every episode.
Revised: Being procedurally-generated, properties such as object shapes, colors, and positions are different every episode.

Original: A recently proposed approach called BYOL-Explore (Guo et al., 2022) was shown to be effective on this domain, while previous successes were only achieved through the use of human demonstrations (Gulcehre et al., 2019).
Revised: A recently proposed exploration bonus called BYOL-Explore (Guo et al., 2022) was shown to be effective on this domain, while previous successes were only achieved through the use of human demonstrations (Gulcehre et al., 2019).

Original: We first assess the effect of using the RECODE bonus overthe NGUone by performing a drop-in replacement of the intrinsic reward in  MEME  as in Section 5.
Revised: We first assess the effect of using the RECODE bonus over NGU by performing a drop-in replacement of the intrinsic reward in the MEME base agent as in Section 5.

Original: We consider the performance of the resulting agent in the single-task version of DM-HARD-8 since  MEME  is not designed to work in the multi-task setting out-of-the-box.
Revised: We consider the performance of the resulting agent in the single-task version of DM-HARD-8 since the MEME base agent is not designed to work in the multi-task setting out-of-the-box.

Original: The results, presented on Fig.
Revised: The results, presented on Fig.

Original: 6, show that our approach can solve 4 out of 8 games, largely improving over  MEME  or Episodic Memory in this direct comparison.
Revised: 6, show that MEMERECODE - AP approach can solve 4 out of 8 games, strictly improving over the original MEME agent using either NGU or EMM in this direct comparison.

Original: Second, we turn to using the CASM  approach introduced in Section 4 as our embedding.
Revised: Second, we turn to using the CASM representation learning approach introduced in Section 4 for our embedding.

Original: We find that the resulting  agent attains state-of-the-art results on the DM-HARD-8 benchmark, reliably solvingout of 8 games, and improving over the  performance of BYOL-explore.
Revised: We find that the resulting MEMERECODE - CASM agent attains state-of-the-art results on the DM-HARD-8 benchmark, reliably solving 6 out of 8 games. This matches and often improves over the state-of-the-art performance of Guo et al.

Original:  Finally, to demonstrate the generality of our approach, we implemented RECODE with standard action-prediction embeddings in a VMPO-based agent similar to that used in Guo et al. 
Revised: (2022), which combined the BYOL-Explore exploration bonus with a VMPO -based base agent. Finally, to demonstrate the generality of our approach, we also test the combination of a VMPO base agent similar to that used in Guo et al. (2022) with RECODE and AP embeddings.

Original: and evaluate in the multi-task setting. The resulting agent  reliably solves 4 games with zero additional tuning (App. F), matching the result obtained with Human Demonstrations in Gulcehre et al.
Revised: The resulting agent is evaluated in the multi-task setting, reliably solving 4 games with zero additional tuning (App. G), matching the result obtained with Human Demonstrations in Gulcehre et al.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: We have introduced Robust Exploration via Clustering-based Online Density Estimation ( RECODE ), a principled yet simple exploration bonus for deep Reinforcement Learning (RL) agents that allows to perform robust exploration by computing visitation counts from a slot-based memory.
Revised: We have introduced Robust Exploration via Clustering-based Online Density Estimation ( RECODE ), a principled yet simple exploration bonus for deep Reinforcement Learning (RL) agents that allows to perform robust exploration by computing visitation counts from a slot-based memory.

Original: Contrary to previous work, where the memory was short-term (i.e. only able to attend to the current episode due to memory limit), our memory is able to model a wide range of timescales determined by the choice of discount.
Revised: Contrary to previous work, where the memory was short-term (i.e. only able to attend to the current episode due to memory limit), our memory is able to model a wide range of timescales determined by the choice of discount.

Original: This is made possible by the use of an online clustering algorithm that is able to approximate the density of visited states and from which we derive a simple  intrinsic reward.
Revised: This is made possible by the use of an online clustering algorithm that is able to approximate the density of visited states and from which we derive a simple and easy-to-tune intrinsic reward.

Original: We evaluate our exploration bonus on top of the recent MEME agent, which is a value-based agent making use of the NGU intrinsic reward, that achieves state-of-the-art results on the Atari domain.
Revised: We evaluate our exploration bonus on top of the recent MEME agent, which is a value-based agent making use of the NGU intrinsic reward, that achieves state-of-the-art results on the Atari domain.

Original: We show that we can replace the complex NGU intrinsic reward that combines RND and episodic memory with the simpler RECODE reward without loss of performance on the hardest exploration levels in Atari.
Revised: We show that we can replace the complex NGU intrinsic reward that combines RND and episodic memory with the simpler RECODE reward without loss of performance on the hardest exploration levels in Atari.

Original: Furthermore, we highlightsome important failure modes of NGU in the presence of noise, and show that RECODE ’s performance is unaffected.
Revised: Furthermore, we highlight important failure modes of NGU in the presence of noise, and show that RECODE ’s performance is unaffected.

Original: Similarly, we highlight the limitations of NGU in procedurally-generated 3D environments such as DM-HARD-8 , and demonstrate the improvements brought by using RECODE as an intrinsic reward.
Revised: Similarly, we highlight the limitations of NGU in procedurally-generated 3D environments such as DM-HARD-8 , and demonstrate the improvements brought by using RECODE as an intrinsic reward.

Original: Next we introduce a novel representation learning method better suited to 3D and partially observable domains, which provides a significant boost in performance on this task suite, enabling to achieve a new state-of-the-art on DM-HARD-8 in the single-task setting.
Revised: Next we introduce a novel representation learning method better suited to 3D and partially observable domains, which provides a significant boost in performance on this task suite, enabling to achieve a new state-of-the-art on DM-HARD-8 in the single-task setting.

Original: Importantly we note that RECODE is agnostic to the choice of embeddings f θ ( h t ) , and while CASM presents one compelling option to extend 1-step action prediction, we hypothesize that much progress could be made in challenging exploration problems by developing more sophisticated representations which can be used in conjunction with RECODE . 
Revised: Importantly we note that RECODE is agnostic to the choice of embeddings f θ ( h t ) , and while CASM presents one compelling option to extend 1-step action prediction, we hypothesize that much progress could be made in challenging exploration problems by developing more sophisticated representations which can be used in conjunction with RECODE . A H YPER - PARAMETERS We omit hypers which do not differ from the base agent, MEME. [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: B A GENT T AXONOMY

Revised: All methods evaluated in the experiments in the main paper and the extended experiments in App.

Revised: E are composed by three main components.

Revised: • A base agent that oversees the overall RL learning process (e

Revised: g., executing actions and collection observations, computing adjusted returns, updating the policy, ...).

Revised: We focus on MEME (Kapturowski et al., 2022), a recent improvement over Agent57 (Badia et al., 2020a) that achieves much greater sample efficiency and is the current state-of-the-art on Atari, and a VMPO-based agent (Guo et al., 2022) that is the current state-of-the-art on DM-HARD-8 .

Revised: An algorithm to generate intrinsic rewards.

Revised: In addition to RECODE , we also consider the recent BYOL-Explore Guo et al.

Revised: (2022), NGU Badia et al.

Revised: (2020b)

Revised: and NGU ’s two building blocks, RND Burda et al.

Revised: (2018b)

Revised: and Episodic Memory ( EMM ) Pritzel et al.

Revised: • A representation learning mechanism to generate observation embeddings which are fed to the intrinsic reward generator.

Revised: We consider both Action Prediction ( AP ) and CASM embeddings.

Revised: Note that some intrinsic reward modules cannot make effective use of the representation learning module (e

Revised: g., RND ), while others merge both second and third module in a single approach (e.g., BYOL-Explore )

Revised: For example, in our more detailed taxonomy the original MEME agent described in Kapturowski et al.

Revised: (2022) is denoted as the MEME - NGU - AP baseline, and compared against our novel MEME - RECODE - AP agent where the only modifications is the changed exploration reward.

Revised: Table E reports more details on all combinations available present in our experiments

Revised: [Table]

Revised: [Table]

Revised: the

Revised: [Table]

Revised: RND takes

Revised: [Table]

Revised: input raw

Revised: [Table]

Revised: [Table]

Revised: test RND ’s ability to cope with non-stationary representations, we train an AP encoder concurrently with the policy and use it to create embeddings of the observations that are fed in RND (i

Revised: [Table]

Revised: e., running RND on top of AP ).

Revised: [Table]

Revised: c

Revised: [Table]

Revised: The BYOL-Explore mechanism internally trains a neural network to predict the dynamical evolution of the observations.

Revised: This provides the agent with both a reward/novelty signal (prediction error) as well as an embedded representation of the observations (that can be extracted from the last few layers of the network).

Revised: C R ELATED W ORKS

Revised: In this section, we give a brief and non-exhaustive overview of past works computing visitation counts or estimating densities in RL.

Revised: We classify them as either parametric or non-parametric.

Revised: Parametric methods.

Revised: Bellemare et al.

Revised: (2016) and Ostrovski et al.

Revised: (2017) propose to compute pseudo-visitation counts using density estimators on images such as Context Tree Switching (CTS,Bellemare et al., 2014) or PixelCNN (Van den Oord et al., 2016).

Revised: On the other hand, Tang et al.

Revised: (2017) use locality-sensitive hashing to map continuous states to discrete embeddings, where explicit visitation counts are computed.

Revised: Some methods such as RND (Burda et al., 2019) can be interpreted as estimating implicitly the density of observations by training a neural network to predict the output of a randomly initialized and untrained neural network which operates on the observations.

Revised: Hazan et al.

Revised: ; Pong et al.

Revised: ; Lee et al.

Revised: ; Guo et al.

Revised: (2021) propose algorithms that search a policy maximizing the entropy of its induced state-space distribution.

Revised: In particular, the loss optimized by Guo et al.

Revised: (2021) allows to compute a density estimate as well as maximizing the entropy.

Revised: Finally, Domingues et al.

Revised: (2021b) computes a density estimation on top of learned representations, which are inspired by bonuses used in reward-free finite MDPs.

Revised: Non-parametric methods.

Revised: Non-parametric density estimates that we build on date back to Rosenblatt (1956); Parzen (1962)

Revised: (Parzen–Rosenblatt window) and are widely used in machine learning as they place very mild assumptions on the data distribution.

Revised: Non-parametric, kernel-based approaches have been already used in RL and shown to be empirically successful on smaller scale environments by Kveton & Theocharous (2012) and Barreto et al.

Revised: and are theoretically analyzed by Ormoneit & Sen (2002);

Revised: Pazis & Parr (2013); Domingues et al.

Revised: (2021a).

Revised: In NGU (Badia et al., 2020b), Agent57 (Badia et al., 2020a) and MEME (Kapturowski et al., 2022), a non-parametric approach is used to compute a short term reward at the episodic level.

Revised: Liu & Abbeel (2021) propose an unsupervised pre-training method for reinforcement learning which explores the environment by maximizing a non-parametric entropy computed in an abstract representation space.

Revised: The authors show improved performance on transfer in Atari games and continous control tasks.

Revised: Seo et al.

Revised: (2021) use random embeddings and a non-parametric approach to estimate the state-visitation entropy, but do not generalize to concurrently learned embeddings.

Revised: Tao et al.

Revised: (2020) show that K-NN based exploration can improve exploration and data efficiency in model-based RL.

Revised: While non-parametric methods are good models for complex data, they come with the challenge of storing and computing densities on the entire data set.

Revised: We tackle this challenge in Sec. 3 by proposing a method that estimates visitation counts over a long history of states, allowing our approach to scale to much larger problems than those considered in previous works, and without placing assumptions on the representation, that can be trained concurrently with the exploration process and doesn’t need to be fixed a priori.

Revised: D A CLUSTERING VIEW OF RECODE

Revised: The update rules of RECODE for its memory structure in Algorithm 1 can be interpreted as an approximate inference scheme in a latent probabilistic clustering model.

Revised: We explore this connection here as means to better understand and justify the proposed algorithm as a density estimator.

Revised: Briefly, our scheme is related to DP-means (Kulis & Jordan, 2011), with adaptations made to accommodate the additional complexities of our setting, which follows a streaming protocol (i.e., data must be explicitly consumed or stored as it arrives and data that are not stored cannot be accessed again) and is non-stationary (i.e., data are not assumed to be identically distributed as time advances).

Revised: The clustering algorithm resulting from these adaptations is shown in Algorithm 2.

Revised: We first address the modifications introduced to deal with the memory limitations of the streaming setting:

Revised: each

Revised: datum

Revised: (embedding e t in our notation

Revised: is

Revised: incorporated into

Revised: a

Revised: cluster distribution approximation

Revised: , once, then discarded,

Revised: the total number of clusters is stochastically projected down onto an upper limit on the number of clusters (otherwise they would grow without bound

Revised: albeit

Revised: progressively

Revised: more

Revised: slowly)

Revised: Both modifications

Revised: allow

Revised: our method

Revised: to maintain

Revised: constant space

Revised: com plexity

Revised: in the face of an infinite

Revised: stream of data.

Revised: The step-wise justification of the Algorithm 2 is relative straightforward.

Revised: At step t , for embedding e t , we show that the following objective is minimised:

Revised: [Equation]

Revised: Working backwards: Updating the cluster center of the closest cluster reduces the objective directly and will not violate the constraint (unless it was already in violation; this excluded in the precondition of this branch).

Revised: This accounts for the “else” branch.

Revised: The “if” branch introduces a new cluster center precisely at e t , thus equation 5 is minimised completely: it is zero for this branch.

Revised: Finally, selecting the index of the nearest cluster center directly minimises placement of the branch according to equation 5, ignoring the constraint (which is latest ensured by the “if/else”).

Revised: Note that the hard constraint of equation 5 takes the place of the soft cluster penalty of DP-means (Kulis & Jordan, 2011).

Revised: The updates to the cluster centers, unlike k-means and DP-means, are done in an exponentiallyweighted moving average of the embeddings, rather than as global optimisation step utilising all of the data.

Revised: Consequently, and importantly, what happens to equation 5 evaluated for e s where s ̸ = t , is of significant interest, as objectives for k-means and DP-means account for all data, rather than a single datum.

Revised: We now turn to this question, which will be answered by examining the modifications we introduced to deal with the non-stationarity of our data stream: 1) the cluster count decays, 2) merging two clusters to accommodate a new one, 3) the exponentially weighted moving average update of cluster centers.

Revised: In k-means, all of the data are retained.

Revised: This makes k-means costly: at each step of fitting the entire data set is examined to update the cluster assignments and update the cluster means.

Revised: Instead, we take a distributional approximation to the data associated with each cluster, and when re-adjusting cluster assignments according to equation 5, we do so in terms of this distributional approximation.

Revised: In particular, each cluster is approximated by a Gaussian distribution with precision 1 and whose mean is unknown but with prior zero and precision 1.

Revised: Specifically:

Revised: [Equation]

Revised: where N ( µ, τ ) denotes a Gaussian (or normal) distribution with mean µ and precision τ (precision is the inverse variance).

Revised: Since the prior on m l is conjugate to the likelihood on e i , we know that the posterior on m l will have the form N ( µ l , c l ) .

Revised: Updating this posterior with a single embedding e i has the form:

Revised: [Equation]

Revised: This is precisely the update in Algorithm 2.

Revised: Note that in this model, the counts c l are also the precision parameters of the distribution representing the inverse spread (or the concentration) of each cluster.

Revised: At each step of Algorithm 2 these counts are decayed.

Revised: Effectively, this causes the variance of the distribution representing each cluster to spread out: thus at each time step, each cluster becomes less concentrated and more uncertain about which data points belong to it.

Revised: The hyperparameter γ captures the rate of diffusion of all clusters in this manner.

Revised: This uncertainty increase applied at each step acts as a “forgetting” mechanism that helps the algorithm to deal with a changing data distribution.

Revised: Cluster re-sampling, as already justified for e t above in terms of equation 5, ensures that the number of clusters is bounded by L .

Revised: There are two details to examine: what is merged, and how it is merged.

Revised: As c j (cid:55)→ 0 , the probability assigned by the Gaussian likelihood of cluster j to any new datum approaches zero also, thus the cluster with the lowest counts is likely to have the least impact on future density estimates (as it is most diffuse).

Revised: When c j ≫ 0 , however, it is not so clear which cluster should be removed, therefore, we stochastically select which cluster to remove inversely proportional to the square of the counts (using the square of the counts emphasizes small differences in counts more than 1 /c j ).

Revised: The cluster could potentially be removed completely, but we instead choose to re-assign its counts to the nearest cluster as we experimentally found this strategy to be less sensitive to the choice of hyperparameters.

Revised: E RND ON TOP OF A CTION P REDICTION E MBEDDINGS

Revised: We adapt RND to leverage trained action-prediction embeddings, which we refer to as RNDonAP .

Revised: To that effect, we use a randomly initialize Multi-Layer Perceptron (MLP) to perform a random projection of the embedding, and use a second, trained MLP, to reconstruct this random projection.

Revised: The reconstruction error provides an intrinsic reward for exploration, which we normalize by a running estimate if its standard deviation as in Burda et al.

Revised: We find that the resulting agent is unable to solve some of the hardest exploration games such as Montezuma’s Revenge or Pitfall .

Revised: The results of this ablation is shown in Fig.

Revised: Preliminary experiments with pre-trained embeddings do seem to indicate that RNDonAP can obtain stronger performance in this setting, but the inability to concurrently train the embeddings greatly limits the general applicability of the method.

Revised: F A RCHITECTURE OF AN AGENT USING RECODE

Revised: Figure 10 shows the typical architecture of a distribued RL agent using RECODE .

Revised: G M ULTITASK E XPERIMENTS

Revised: We also implemented RECODE in a VMPO-based agent similar to the one used with BYOLExplore (Guo et al., 2022), and compared our performance with BYOL-Explore in the multi-task

Revised:            t =

Revised: P( j ) 1/ c

Revised: j          

Revised: 

Revised: P( j ) 1/ c 2 j            argmin j ( c j )            t = 50 k

Revised:                         

Revised: [Figure]

Revised: setting.

Revised: This experiment serves two different purposes.

Revised: First, this demonstrates the generality of our exploration bonus, that is thus shown to be useful in widely different RL agents, be they value-based or policy-based.

Revised: Second, we can do a direct comparison with the state of the art BYOL-Explore agent in the multi-task settings.

Revised: We however note that the representation learning technique used in this experiment, 1-step Action Prediction, is based on a feed-forward embedding that discards past history, and may therefore not be the best fit for exploration in Partially Observable MDPs (POMDPs).

Revised: Still, we show in Fig.

Revised: 12 that the performance of RECODE is competitive with that of

Revised: BYOL-Explore, with only one level missing to match its performance.

Revised: Improving this performance using better-suited representations, such as CASM , is left for future work.

Revised: H A GGREGATED R ESULTS

Revised: Here we present our main results aggregated over all environments in each task suite.

Revised: To ensure that no single environment dominates due to larger reward scales we use the Human Normalized

Revised: Score (Mnih et al., 2015) in each environment, and then cap scores above 100% prior to averaging.

Revised: As can be observed in Fig.

Revised: 13 (Left), the uncapped score can swing significantly over time, which in this case is simply an artifact the high variance present in Q*bert .