S TATE E NTROPY M AXIMIZATION WITH R ANDOM E NCODERS FOR E FFICIENT E XPLORATION
A BSTRACT
Recent exploration methods have proven to be a recipe for improving sampleefﬁciency in deep reinforcement learning (RL).
However, efﬁcient exploration in high-dimensional observation spaces still remains a challenge.
This paper presents
Random Encoders for Efﬁcient Exploration (RE3), an exploration method that utilizes state entropy as an intrinsic reward.
In order to estimate state entropy in environments with high-dimensional observations, we utilize a k -nearest neighbor entropy estimator in the low-dimensional representation space of a convolutional encoder.
In particular, we ﬁnd that the state entropy can be estimated in a stable and compute-efﬁcient manner by utilizing a randomly initialized encoder, which is ﬁxed throughout training.
Our experiments show that RE3 signiﬁcantly improves the sample-efﬁciency of both model-free and model-based RL methods on locomotion and navigation tasks from DeepMind Control Suite and MiniGrid benchmarks.
We also show that RE3 allows learning diverse behaviors without extrinsic rewards, effectively improving sample-efﬁciency in downstream tasks.
I NTRODUCTION
Exploration remains one of the main challenges of deep reinforcement learning (RL) in complex environments with high-dimensional observations.
Many prior approaches to incentivizing exploration introduce intrinsic rewards based on a measure of state novelty.
These include count-based visitation bonuses (Bellemare et al., 2016; Ostrovski et al., 2017; Tang et al., 2017) and prediction errors (Stadie et al., 2015;
Houthooft et al., 2016; Burda et al., 2019; Pathak et al., 2017).
By introducing such novelty-based intrinsic rewards, these approaches encourage agents to visit diverse states, but leave unanswered the fundamental question of how to quantify effective exploration in a principled way.
To address this limitation, Lee et al.
(2019) and Hazan et al.
(2019) proposed that exploration methods should encourage uniform (i.e., maximum entropy) coverage of the state space.
For practical state entropy estimation without learning density models, Mutti et al.
(2020) estimate state entropy by measuring distances between states and their k -nearest neighbors.
To extend this approach to high-dimensional environments, recent works (Badia et al., 2020; Tao et al., 2020; Liu & Abbeel, 2021) have proposed to utilize the k -nearest neighbor state entropy estimator in a low-dimensional latent representation space.
The latent representations are learned by auxiliary tasks such as dynamics learning (Tao et al., 2020), inverse dynamics prediction (Badia et al., 2020), and contrastive learning (Liu & Abbeel, 2021).
However, these methods still involve optimizing multiple objectives throughout RL training.
Given the added complexity (e.g., hyperparameter tuning), instability, and computational overhead of optimizing auxiliary losses, it is important to ask whether effective state entropy estimation is possible without introducing additional learning procedures.
[Figure]
In this paper, we present RE3: R andom E ncoders for E fﬁcient E xploration, a simple, computeefﬁcient method for exploration without introducing additional models or representation learning.
The key idea of RE3 is to utilize a k -nearest neighbor state entropy estimator in the representation space of a randomly initialized encoder, which is ﬁxed throughout training.
Our main hypothesis is that a randomly initialized encoder can provide a meaningful representation space for state entropy estimation by exploiting the strong prior of convolutional architectures.
Ulyanov et al.
and
Caron et al.
(2018) found that the structure alone of deep convolutional networks is a powerful inductive bias that allows relevant features to be extracted for tasks such as image generation and classiﬁcation.
In our case, we ﬁnd that the representation space of a randomly initialized encoder effectively captures information about similarity between states, as shown in Figure 1.
Based upon this observation, we propose to maximize a state entropy estimate in the ﬁxed representation space of a randomly initialized encoder.
We highlight the main contributions of this paper below:
• RE3 signiﬁcantly improves the sample-efﬁciency of both model-free and model-based RL methods on widely used DeepMind Control Suite (Tassa et al., 2020) and
MiniGrid (Chevalier-Boisvert et al., 2018) benchmarks.
• RE3 encourages exploration without introducing representation learning or additional models, outperforming state entropy maximization schemes that involve representation learning and exploration methods that introduce additional models for exploration (Burda et al., 2019; Pathak et al., 2017).
• RE3 is compute-efﬁcient as it does not require gradient computations and updates for additional representation learning, making it a scalable and practical approach to exploration.
• RE3 allows learning diverse behaviors in environments without extrinsic rewards; we further improve sample-efﬁciency in downstream tasks by ﬁne-tuning a policy pre-trained with the RE3 objective.
R ELATED W ORK
Exploration in reinforcement learning.
Exploration algorithms encourage the RL agent to visit a wide range of states by injecting noise to the action space (Lillicrap et al., 2015) or parameter space (Fortunato et al., 2018;
Plappert et al., 2018), maximizing the entropy of the action space (Ziebart, 2010; Haarnoja et al., 2018), and setting diverse goals that guide exploration (Florensa et al., 2018; Nair et al., 2018; Colas et al., 2019; Pong et al., 2020).
Another line of exploration algorithms introduce intrinsic rewards proportional to prediction errors (Houthooft et al., 2016; Pathak et al., 2017; Burda et al., 2019), and count-based state novelty (Bellemare et al., 2016;
Ostrovski et al., 2017; Tang et al., 2017).
Our approach differs in that we explicitly encourage the agent to uniformly visit all states by maximizing the entropy of the state distribution, instead of depending on metrics from additional models.
State entropy maximization.
Most closely related to our work are methods that maximize the entropy of state distributions.
Hazan et al.
and Lee et al.
(2019) proposed to maximize state entropy estimated by approximating the state density distribution.
Instead of approximating complex distributions
, Mutti et al.
(2020) proposed to maximize a k -nearest neighbor state entropy estimate
[Figure]
from on-policy transitions.
Recent works method to environments with high-dimensional observations.
Tao et al.
(2020) employ model-based RL techniques to build a representation space for the state entropy estimate that measures similarity in dynamics, and Badia et al.
(2020) proposed to measure similarity in the representation space learned by inverse dynamics prediction.
The work closest to ours is Liu & Abbeel (2021), which uses off-policy RL algorithms to maximize the knearest neighbor state entropy estimate in contrastive representation space (Srinivas et al., 2020) for unsupervised pre-training.
We instead explore the idea of utilizing a ﬁxed random encoder to obtain a stable entropy estimate without any representation learning.
Random encoders.
Random weights have been utilized in neural networks since their beginnings, most notably in a randomly initialized ﬁrst layer (Gamba et al., 1961) termed the Gamba perceptron by Minsky & Papert (1969).
Moreover, nice properties of random projections are commonly exploited for low-rank approximation (Vempala, 2005; Rahimi & Recht, 2007).
These ideas have since been extended to deep convolutional networks, where random weights are surprisingly effective at image generation and restoration (Ulyanov et al., 2018), image classiﬁcation and detection (Caron et al., 2018), and fast architecture search (Saxe et al., 2011).
In natural language processing,
Wieting & Kiela (2019) demonstrated that learned sentence embeddings show marginal performance gain over random embeddings.
In the context of RL, Gaier & Ha (2019) showed that competitive performance can be achieved by architecture search over random weights without updating weights, and Lee et al.
(2020) utilized randomized convolutional neural networks to improve the generalization of deep RL agents.
Building on these works, we show that random encoders can also be useful for efﬁcient exploration in environments with high-dimensional observations.
M ETHOD 3. P RELIMINARIES
We formulate a control task with high-dimensional observations as a partially observable Markov decision process (POMDP; Kaelbling et al.
1998; Sutton & Barto 2018), which is deﬁned as a tuple ( O , A , p, r e , γ ) .
Here, O is the high-dimensional observation space, A is the action space, p
( o (cid:48) | o ≤ t , a t ) is the transition dynamics, r e : O × A → R is the reward function that maps the current observation and action to a reward r e t =
r e ( o ≤ t , a t ) , and γ ∈ [0 , 1) is the discount factor.
By following common practice (Mnih et al., 2015), we reformulate the POMDP as an MDP (Sutton & Barto, 2018) by stacking consecutive observations into a state s t = { o t , o t − 1 , o t − 2 , ... } .
For simplicity of notation, we redeﬁne the reward function as r e t = r e
( s t , a t ) .
The goal of RL is to learn a policy π ( a t
| s t ) that maximizes the expected return deﬁned as the total accumulated reward.
k -nearest neighbor entropy estimator.
Let X be a random variable with a probability density function p whose support is a set X ⊂ R q .
Then its differential entropy is given as H ( X ) =
− E x ∼ p ( x ) [log p ( x )] .
When the distribution p is not available, this quantity can be estimated given
N i.i.d realizations of { x i } Ni =1 (Beirlant et al., 1997).
However, since it is difﬁcult to estimate p with high-dimensional data, particle-based k -nearest neighbors ( k -NN) entropy estimator (Singh et al., 2003) can be employed:
[Equation]
where x k -NN i is the k -NN of x i within a set { x i } Ni =1 , C k = log k − Ψ
( k ) a bias correction term, Ψ the digamma function, Γ the gamma function, q the dimension of x , (cid:98) π ≈ 3 . 14159 , and the transition from (1) to (2) always holds for q > 0 .
3. R ANDOM E NCODERS FOR E FFICIENT E XPLORATION
We present Random Encoders for Efﬁcient Exploration (RE3), which encourages exploration in high-dimensional observation spaces by maximizing state entropy.
The key idea of RE3 is k -nearest neighbor entropy estimation in the low-dimensional representation space of a randomly initialized encoder.
To this end, we propose to compute the distance between states in the representation space of a random encoder f θ whose parameters θ are randomly initialized and ﬁxed throughout training.
The main motivation arises from our observation that distances in the representation space of f θ are already useful for ﬁnding similar states without any representation learning (see Figure 1).
State entropy estimate as intrinsic reward.
To deﬁne the intrinsic reward proportional to state entropy estimate by utilizing (2), we follow the idea of Liu & Abbeel (2021) that treats each transition as a particle, hence our intrinsic reward is given as follows:
[Equation]
where y i =
f θ ( s i ) is a ﬁxed representation from a random encoder and
y k -NN i is the k -nearest neighbor of y i within a set of N representations { y 1 , y 2 , ..., y N } .
Our intuition is that measuring distance between states in the ﬁxed representation space produces a more stable intrinsic reward as the distance between a given pair of states does not change during training.
To compute distances in latent space in a compute-efﬁcient manner, we propose to additionally store low-dimensional representations y in the replay buffer B during environment interactions.
Therefore, we avoid processing highdimensional states through an encoder for obtaining representations at every RL update.
Moreover, we can feasibly compute the distance of y i to all entries y ∈ B , in contrast to existing approaches that utilize on-policy samples (Mutti et al., 2020), or samples from a minibatch (Liu & Abbeel, 2021).
Our scheme enables stable, precise entropy estimation in a compute-efﬁcient manner.
The RE3 objective.
We propose to utilize the intrinsic reward r i for (a) online RL , where the agent solves target tasks guided by extrinsic reward r e from environments, and (b) unsupervised pre-training , where the agent learns to explore the high-dimensional observation space of environments in the absense of extrinsic rewards, i.e., r e = 0 .
This exploratory policy from pre-training, in turn, can be used to improve the sample-efﬁciency of the agent in downstream tasks by ﬁnetuning.
Formally, we introduce a policy π φ , parameterized by φ , that maximizes the expected return
E π φ (cid:104)(cid:80) ∞
[Equation]
=
γ j r total j (cid:105) , where the total reward r total j is deﬁned as:
[Equation]
where β t ≥ 0 is a hyperparameter that determines the tradeoff between exploration and exploitation at training timestep t .
We use the exponential decay schedule for β t throughout training to encourage the agent to further focus on extrinsic reward from environments as training proceeds, i.e., β t = β 0 (1 − ρ ) t , where ρ is a decay rate.
While the proposed intrinsic reward would converge to 0 as more similar states are collected during training, we discover that decaying β t empirically stabilizes the performance.
We provide the full procedure for RE3 with off-policy RL in Algorithm 1 and RE with on-policy RL in Algorithm 2.
1: Initialize parameters of random encoder θ , policy φ , replay buffer B ← ∅ 2: for each timestep t do 3:
Collect a transition τ t = ( s t , a t , s t +1 , r e t ) using policy π φ // C OLLECT T RANSITIONS 4:
Get a ﬁxed representation y t
= f θ ( s t ) 5:
B ← B ∪ { ( τ t , y t ) } 6:
Sample random minibatch { ( τ j , y j ) } Bj =1 ∼ B // C OMPUTE INTRINSIC REWARD 7: for j = 1 to B do 8:
Compute the distance ||
y j −
y
|| 2 for all y ∈ B and ﬁnd the k -nearest neighbor y k -NN j 9:
Compute r i j ← log( || y j − y k -NN
j
||
+ 1) 10:
Update β t ← β 0 (1 − ρ ) t 11:
Let r total j ← r e j + β t · r i j 12: end for 13:
Update φ with transitions { ( s j , a j , s j +1 , r total j ) } Bj = //
U PDATE POLICY 14: end for
E XPERIMENTS
We designed experiments to answer the following questions:
• Can RE3 improve the sample-efﬁciency of both model-free and model-based RL algorithms (see Figure 3)?
• How does RE3 compare to state entropy maximization schemes that involve representation learning (see Figure 4) and other exploration schemes that introduce additional models for exploration (see Figure 5)?
• Can RE3 further improve the sample-efﬁciency of off-policy RL algorithms by unsupervised pre-training (see Figure 6 and Figure 7a)?
• How compute-efﬁcient is RE3 (see Figure 7b)?
• Can RE3 improve the sample-efﬁciency of on-policy RL in discrete control (see Figure 8)?
D EEP
M IND C ONTROL S UITE E XPERIMENTS
Setup.
To evaluate the sample-efﬁciency of our method, we compare to Dreamer (Hafner et al., 2020), a state-of-the-art model-based RL method for visual control; and two state-of-the-art modelfree RL methods, RAD (Laskin et al., 2020) and DrQ (Kostrikov et al., 2021).
For comparison with other exploration methods, we consider RND (Burda et al., 2019) and ICM (Pathak et al., 2017) that introduce additional models for exploration.
For RE3 and baseline exploration methods, we use RAD as the underlying model-free RL algorithm.
To further demonstrate the applicability of
RE3 to model-based RL algorithms, we also consider a combination of Dreamer and RE3.
For random encoders, we use convolutional neural networks with the same architecture as underlying
RL algorithms, but with randomly initialized parameters ﬁxed during the training.
As for the newly introduced hyperparameters, we use k = 3 , β 0 = 0 .
05 , and ρ ∈ { 0 . 0 , 0 .
We provide more details in Appendix B and source code in Appendix A.
Comparative evaluation.
Figure 3 shows that RE3 consistently improves the sample-efﬁciency of RAD on various tasks.
In particular, RAD + RE3 achieves average episode return of 601.6 on
Cheetah Run Sparse, where both model-free RL methods RAD and DrQ fail to solve the task.
We emphasize that state entropy maximization with RE3 achieves such sample-efﬁciency with minimal cost due to its simplicity and compute-efﬁciency.
We also observe that Dreamer + RE3 improves the sample-efﬁciency of Dreamer on most tasks, which demonstrates the applicability of RE3 to both model-free and model-based RL.
We provide experimental results on more tasks in Figure 9.
Effects of representation learning.
To better grasp how RE3 improves sample-efﬁciency, we compare to state entropy maximization schemes that involve representation learning in Figure 4.
Speciﬁcally, we consider a convolutional encoder trained by contrastive learning (RAD + SE w/
Contrastive) and inverse dynamics prediction (RAD + SE w/ Inverse dynamics).
We also consider two pre-trained encoders whose parameters are ﬁxed throughout training: a ResNet-50 (He et al.,
[Figure] [Figure]
2016) encoder pre-trained on ImageNet dataset (RAD + SE w/ ImageNet) and an ATC (Stooke et al., 2020) encoder pre-trained by contrastive learning on pre-training datasets that contain a wide range of states from the environments (RAD + SE w/ ATC), where more details are available at Appendix B.
We found that our method (RAD + SE w/ Random) exhibits better sample-efﬁciency than approaches that continually update representations throughout training (RAD + SE w/ Contrastive,
RAD + SE w/ Inverse dynamics).
This demonstrates that utilizing ﬁxed representations helps improve sample-efﬁciency by enabling stable state entropy estimation throughout training.
We also observe that our approach outperforms RAD + SE w/ ImageNet, implying that it is not necessarily beneﬁcial to employ a pre-trained encoder, and ﬁxed random encoders can be effective for state entropy estimation without having been trained on any data.
We remark that representations from the pre-trained ImageNet encoder could not be useful for our setup, due to the different visual characteristics of natural images in the ImageNet dataset and image observations in our experiments.
However, our approach also achieves comparable sample-efﬁciency to RAD + SE w/ ATC, which shows that random encoders can provide a stable state entropy estimation even competitive with encoders pre-trained on data from the environments.
Comparison with other exploration methods.
We also compare our state entropy maximization scheme to other exploration methods combined with RAD, i.e., RAD + RND and RAD + ICM, that learn additional models to obtain intrinsic rewards proportional to prediction errors.
As shown in
Figure 5, RAD + RE3 consistently exhibits superior sample efﬁciency in most tasks (see Figurefor results on more tasks).
While RND similarly employs a ﬁxed random network for the intrinsic reward, it also introduces an additional network which requires training and therefore suffers from instability.
This result demonstrates that RE3 can improve sample-efﬁciency without introducing additional models for exploration, by utilizing ﬁxed representations from a random encoder for stable state entropy estimation.
[Figure] [Figure]
[Figure]
Evaluation of unsupervised pre-training.
To evaluate the effectiveness of RE3 for learning diverse behaviors in the pre-training phase without extrinsic rewards, we ﬁrst visualize the behaviors of policies pre-trained for 500K environment steps in Figure 7a.
One can see that the pre-trained policy using RE3 exhibits more diverse behaviors compared to random exploration or APT (Liu & Abbeel, 2021), where a policy is pre-trained to maximize state entropy estimate in contrastive representation space.
To further evaluate the diversity of behaviors quantitatively, we also show the state entropy estimate in the ground-truth (i
[Figure]
e., proprioceptive) state space in the Hopper environment in Figure 6a.
We observe that RE3 exhibits a higher ground-truth state entropy estimate than state entropy maximization schemes that use contrastive learning and inverse dynamics prediction during pre-training.
This implies that RE3 can effectively maximize the ground-truth state entropy without being able to directly observe underlying ground-truth states.
Fine-tuning in downstream tasks.
We also remark that the diversity of behaviors leads to superior sample-efﬁciency when ﬁne-tuning a pre-trained policy in downstream tasks, as shown in Figure 6b and 6c.
Speciﬁcally, we ﬁne-tune a pre-trained policy in downstream tasks where extrinsic rewards are available, by initializing the parameters of policies with parameters of pre-trained policies (see
Appendix B for more details).
We found that ﬁne-tuning a policy pre-trained with RE3 (RAD + SE w/ Random + PT) further improves the sample-efﬁciency of RAD + SE w/ Random, and also outperforms other pre-training schemes.
We emphasize that RE3 allows learning such diverse behaviors by pre-training a policy only for 500K environment steps, while previous work (Liu &
Abbeel, 2021) reported results by training for 5M environment steps.
[Figure]
[Figure]
[Figure]
Compute-efﬁciency.
We show that RE3 is a practical and scalable approach for exploration in
RL due to its compute-efﬁciency.
In particular, RE3 is compute-efﬁcient in that (a) there are no gradient updates through the random encoder, and (b) there are no unnecessary forward passes for obtaining representations at every update step since we store low-dimensional latent representations in the replay buffer.
To evaluate compute-efﬁciency, we show the ﬂoating point operations (FLOPs) consumed by RAD, RAD + SE w/ Random (ours), RAD + SE w/ Contrastive, and RAD + SE w/
Inverse dynamics.
We account only for forward and backward passes through neural network layers.
We explain our full procedure for counting FLOPs in Appendix E. Figure 7b shows the FLOPS used by each agent to achieve its ﬁnal performance in Hopper Hop (see Figure 4 for corresponding learning curves).
Estimating state entropy with a random encoder is signiﬁcantly more computeefﬁcient than with a encoder trained with representation learning.
One important detail here is that
RAD + SE w/ Random has comparable compute-efﬁciency to RAD.
Therefore, we improve the sample-efﬁciency of RAD without sacriﬁcing compute-efﬁciency.
4. M INI G RID E XPERIMENTS
Setup.
We evaluate our method on MiniGrid (Chevalier-Boisvert et al., 2018), a gridworld environment with a selection of sparse reward tasks.
The tasks we consider are shown in Figure 11.
For evaluation, we use Advantage Actor-Critic (A2C; Mnih et al. 2016)
as the underlying RL algorithm, and consider two exploration methods, RND and ICM.
The agent has access to a compact 7 × 7 ×embedding of the 7 × 7 grid directly in front of it, making the environment partially-observable.
To combine RE3 with A2C, an on-policy RL method, we maintain a replay buffer of 10K samples solely for computing the RE3 intrinsic reward, and compute k -NN distances between the on-policy batch and the entire replay buffer.
For RND and ICM, the intrinsic reward is computed using the on-policy batch.
For RE3, ICM, and RND, we perform hyperparameter search over the intrinsic reward weight and report the best result (see Appendix D for more details).
Comparison with other exploration methods.
Figure 8 shows that RE3 is more effective for improving the sample-efﬁciency of A2C in most tasks, compared to other exploration methods,
RND and ICM, that learn additional models.
In particular, A2C + RE3 achieves average episode return of 0.49 at 2.4M environment steps in DoorKey-8x8; in comparison, A2C + ICM achieves a return of 0.20 and A2C + RND and A2C both fail to achieve non-trivial returns.
These results demonstrate that state entropy maximization with RE3 can also improve the sample-efﬁciency of on-policy RL algorithms by introducing only a small-size replay buffer.
Fine-tuning in downstream tasks.
To evaluate the effectiveness of RE3 for unsupervised pretraining in MiniGrid tasks, we ﬁrst pre-train a policy in a large vacant room (Empty-16 × 16) to maximize RE3 intrinsic rewards for 100K environments steps.
Then, we ﬁne-tune the pre-trained policy in downstream tasks by initializing a policy with pre-trained parameters and subsequently training with A2C + RE3.
Figure 8 shows that A2C + RE3 + PT signiﬁcantly improves the sampleefﬁciency of A2C, which demonstrates that the ability to explore novel states in a large empty room helps improve sample-efﬁciency in DoorKey tasks, which involve the added complexity of additional components, e.g., walls, doors, and locks.
We show a comparison to state entropy maximization with contrastive learning in Figure 12 and observe that it does not work well, as contrastive learning depends on data augmentation speciﬁc to images (e.g., random shift, random crop, and color jitter), which are not compatible with the compact embeddings used as inputs for MiniGrid.
We remark that RE3 eliminates the need for carefully chosen data augmentations by employing a random encoder.
D ISCUSSION
In this paper, we present RE3, a simple exploration method compatible with both model-free and model-based RL algorithms.
RE3 maximizes a k -nearest neighbor state entropy estimate in the ﬁxed representation space of a randomly initialized encoder.
Our experimental results demonstrate that RE3 can encourage exploration in widely-used benchmarks, as it enables stable and computeefﬁcient state entropy estimation.
Here, we emphasize that our goal is not to claim that representation learning or additional models are not required for exploration, but to show that ﬁxed random encoders can be useful for efﬁcient exploration.
For more visually complex domains, utilizing pretrained ﬁxed representations for stable state entropy estimation could be more useful, but we leave it to future work to explore this direction further because this would require having access to environments and a wide distribution of states for pre-training, which is itself a non-trivial problem.
Another interesting direction would be to investigate the effect of network architectures for state entropy estimation, or to utilize state entropy for explicitly guiding the action of a policy to visit diverse states.
We believe RE3 would facilitate future research by providing a simple-to-implement, stable and compute-efﬁcient module that can be easily combined with other techniques
Appendix A S OURCE C ODE AND V IDEOS
We provide source code that reproduces our experimental results and videos in Figure 7a in https: //anonymous.4open.science/r/3f955657-271b-4e5d-9a66-f77ecc3708de/ .
B D ETAILS ON D EEP M IND C ONTROL S UITE E XPERIMENTS B. E NVIRONMENTS
We evaluate the performance of RE3 on various tasks from DeepMind Control Suite (Tassa et al., 2020).
For Hopper Hop, Quadruped Run, Cartpole Swingup Sparse, Pendulum Swingup, we use the publicly available environments without any modiﬁcation.
For environments which are not from the publicly available released implementation repository ( https://github.com/deepmind/ dm_control ), we designed the tasks following Seyde et al.
(2021) as below:
• Walker/Cheetah Run Sparse : The goal of Walker/Cheetah Run Sparse task is same as in Walker/Cheetah Run, moving forward as fast as possible, but reward is given sparsely until it reaches a certain threshold: r = r original · 1 r original > 0 .
25 , where r original is the reward in Walker/Cheetah Run from DeepMind Control Suite.
B. I MPLEMENTATION D ETAILS FOR M ODEL - FREE RL
For all experimental results in this work, we report the results obtained by running experiments using the publicly available released implementations from the authors 2 ,
i.e., RAD ( https://github. com/MishaLaskin/rad ) and DrQ ( https://github.com/denisyarats/drq ).
We use random crop augmentation for RAD and random shift augmentation for DrQ.
We provide a full list of hyperparameters in Table 1.
Implementation details for RE3.
We highlight key implementation details for RE3:
• Intrinsic reward.
We use r i ( s i ) :
= ||
y i − y k -NN
i || 2 for the intrinsic reward.
We got rid of log from intrinsic reward in (3) for simplicity in DeepMind Control Suite experiments, but results using log are also similar.
To make the scale of intrinsic reward r i consistent across tasks, following Liu & Abbeel (2021), we normalize the intrinsic reward by dividing it by a running estimate of the standard deviation.
As for newly introduced hyperparameters, we use k = 3 and β 0 = 0 .
05 , and perform hyperparameter search over ρ ∈ { 0 . 00001 , 0 . 000025 } .
• Architecture.
For all model-free RL methods, we use the same encoder architecture as in Yarats et al.
Speciﬁcally, this encoder consists of 4 convolutional layers followed by ReLU activations.
We employ kernels of size 3 × 3 with 32 channels for all layers, and 1 stride except of the ﬁrst layer which has stride 2.
The output of convolutional layers is fed into a single fully-connected layer normalized by LayerNorm (Ba et al., 2016).
Finally, tanh nonlinearity is added to the 50-dimensional output of the fully-connected layer.
• Unsupervised pre-training.
For unsupervised pre-training, we ﬁrst train a policy to maximize intrinsic rewards in (3) without extrinsic rewards for 500K environment steps.
For ﬁne-tuning a policy in downstream tasks, we initialize the parameters using the pre-trained parameters and then learn a policy to maximize RE3 objective in (4) for 500K environment steps.
To stabilize the initial ﬁne-tuning phase by making the scale of intrinsic reward consistent across pre-training and ﬁne-tuning, we load the running estimate of the standard deviation from pre-training phase.
Implementation details for representation learning baselines .
We highlight key implementation details for state entropy maximization schemes that involve representation learning, i e., contrastive learning and inverse dynamics prediction and that employs a pre-trained ImageNet Encoder:
• Contrastive learning.
For state entropy maximization with contrastive learning, we introduce a separate convolutional encoder with randomly initialized parameters learned to minimize contrastive loss (Srinivas et al., 2020).
Note that we use the separate encoder to investigate the independent effect of representation learning.
Then we compute intrinsic reward r i using representations obtained by processing observations through this separate encoder.
• Inverse dynamics prediction.
For state entropy maximization with inverse dynamics prediction, we introduce a separate convolutional encoder with randomly initialized parameters learned to predict actions from two consecutive observations.
Speciﬁcally, we introduce 2 fullyconnected layers with ReLU activations to predict actions on top of encoder representations.
We remark that this separate encoder is separate from RL.
Then we compute intrinsic reward r i using representations obtained by processing observations through this separate encoder.
• Pre-trained ImageNet encoder.
For state entropy maximization with pre-trained ImageNet encoder, we utilize a ResNet-50 (He et al., 2016) encoder from publicly available torchvision models ( https://pytorch.org/vision/0.8/models.html ).
To compute intrinsic reward r i , we utilize representations obtained by processing observations through this pre-trained encoder.
• Pre-trained ATC encoder.
For state entropy maximization with pre-trained ATC (Stooke et al., 2020) encoder, we utilize pre-trained encoders from the authors.
Speciﬁcally, these encoders are pre-trained by contrastive learning on pre-training datasets that contain samples encountered while training a RAD agent on DeepMind Control Suite environments (see Stooke et al.
(2020) for more details).
We use pre-trained parameters from Walker Run, Hopper Stand, and Cheetah Run for RAD + SE w/ ATC on Walker Run Sparse, Hopper Hop, and Cheetah Run Sparse, respectively.
To compute intrinsic reward r i , we utilize representations obtained by processing observations through this pre-trained encoder.
Implementation details for exploration baselines.
We highlight key implementation details for exploration baselines, i.e., RND (Burda et al., 2019) and ICM (Pathak et al., 2017):
• RND.
For RND, we introduce a random encoder f θ whose architecture is same as in Yarats et al.
(2019), and introduce a predictor network g φ consisting of a convolutional encoder with the same architecture and 2-layer fully connected network with 1024 units each.
Then, parameters φ of the predictor network are trained to predict representations from a random encoder given the same observations,
i
e., minimize (cid:15) =
||
f θ ( s i ) − g φ ( s i ) ||
We use prediction error (cid:15) as an intrinsic reward and learn a policy that maximizes r total =
r e + β · r
i .
We perform hyperparameter search over the weight β ∈ { 0 . 05 , 0 .
and report the best result on each environment.
• ICM.
For ICM, we introduce a convolutional encoder g φ whose architecture is same as in Yarats et al.
(2019), and introduce a inverse dynamics predictor h ψ with 2 fully-connected layers with 1024 units.
These networks are learned to predict actions between two consecutive observations, i.e., minimize L inv =
|| a t − h ψ ( g φ ( s t ) , g φ ( s t +1 ))
||
We also introduce a forward dynamics predictor f Ψ that is learned to predict the representation of next time step,
i
e., minimize L forward =
12 || g φ ( s t +1 )
− f Ψ ( g φ ( s t ) , a t )
|| 22 .
Then, we use prediction error as an intrinsic reward and learn a policy that maximizes r total =
r e + β · r i .
For joint training of the forward and inverse dynamics, following Pathak et al.
(2017), we minimize 0 .
·L forward +0 . 8 ·L inv .
We perform hyperparameter search over the weight β ∈ { 0 . 05 , 0 .
and report the best result on each environment
B. I MPLEMENTATION D ETAILS FOR M ODEL - BASED RL.
For Dreamer, we use the publicly available released implementation repository ( https:// github.com/danijar/dreamer ) from the authors 3 .
We highlight key implementation details for the combination of Dreamer with RE3:
• Intrinsic reward.
We use r i ( s i ) :
= ||
y i − y k -NN
i || 2 for the intrinsic reward.
We got rid of log from intrinsic reward in (3) for simplicity in DeepMind Control Suite experiments, but
results using log are also similar.
To make the scale of intrinsic reward r i consistent across tasks, following Liu & Abbeel (2021), we normalize the intrinsic reward by dividing it by a running estimate of the standard deviation.
Since Dreamer utilizes trajectory segments for training batch, we use large value of k = 50 to avoid ﬁnd k -NN only within a trajectory segment.
We also use β 0 = 0 . 1 , and ρ = 0 . 0 ,
i
[Table]
e., no decay schedule.
We also remark that we ﬁnd k -NN only within a minibatch instead of the entire buffer as in model-free RL, since the large batch size of Dreamer is enough for stable entropy estimation.
• Architecture.
We use the same convolutional architecture as in Dreamer.
Speciﬁcally, this encoder consists of 4 convolutional layers followed by ReLU activations.
We employ kernels of size 4 × 4 with { 32 , 64 , 128 , 256 } channels, and 2 stride for all layers.
To obtain low-dimensional representations, we additionally introduce a fully-connected layer normalized by LayerNorm (Ba et al., 2016).
Finally, tanh nonlinearity is added to the 50-dimensional output of the fully-connected layer.
[Table]
[Table]
C A DDITIONAL E XPERIMENTAL R ESULTS ON D EEP M IND C ONTROL S UITE
Comparative evaluation.
We provide additional experimental results on various tasks from DeepMind Control Suite in Figure 9.
We observe that RE3 improves sample-efﬁciency in several tasks, e
[Figure]
g., Reacher Hard and Hopper Stand, while not degrading the performance in dense-reward tasks like Cartpole Balance.
This demonstrates the simple and wide applicability of RE3 to various tasks.
Comparison with other exploration methods.
Figure 10 shows that RAD + RE3 consistently exhibits superior sample efﬁciency to other exploration methods (i
[Figure]
e., RAD + RND, RAD + ICM) on additional tasks from DeepMind Control Suite
[Figure] [Figure]
[Figure]
D D ETAILS ON M INI G RID E XPERIMENTS
For our A2C implementation in MiniGrid, we use the publicly available released implementation repository ( https://github.com/lcswillems/rl-starter-files ) and use their default hyperparameters.
We provide a full list of hyperparameters that are introduced by our method or emphasized for clarity in Table 3.
We highlight some key implementation details:
• We use r
i
( s i ) :
= log( || y
i −
y k -NN
i ||
+ 1) for the intrinsic reward.
The additional 1 is for numerical stability.
• We use the average distance between y i and its k nearest neighbors (i.e., y 2 -NN
i , · · · , y k -NN i ) for the intrinsic reward, instead of the single k nearest neighbor.
This provides a less noisy state entropy estimate and empirically improves performance in MiniGrid environments.
• We perform hyperparameter search over β ∈ { 0 .
05 , 0 . 1 } for RE3, ICM, and RND and report the best result.
• We do not change the network architecture from the above publicly available implementation.
We use the same encoder architecture as the RL encoder for state entropy maximization.
This encoder architecture consists of 3 convolutional layers with kernel 2, stride 1, and padding 0, each followed by a ReLU layer.
The convolutional layers have 16, 32, and 64 ﬁlters respectively.
The ﬁrst ReLU
is followed by a two-dimensional max pooling layer with kernel 2.
The actor and critic MLPs both contain two fully-connected layers with hidden dimension 64, with a tanh operation in between.
[Table]
E C ALCULATION OF F LOATING P OINT O PERATIONS
We explain our FLOP counting procedure for comparing the compute-efﬁciency of RAD, RAD +
SE w/ Random (ours), RAD + SE w/ Contrastive, and RAD + SE w/ Inverse dynamics in Figure 7.
We consider each backward pass to require twice as many FLOPs as a forward pass, as done in https://openai.com/blog/ai-and-compute/ .
Each weight requires one multiply-add operation in the forward pass.
In the backward pass, it requires two multiply-add operations: at layer i , the gradient of the loss with respect to the weight at layer i and with respect to the output of layer ( i − 1 ) need to be computed.
The latter computation is necessary for subsequent gradient calculations for weights at layer ( i − 1 ).
We use functions from Huang et al.
and Jeong & Shin (2019) to obtain the number of operations per forward pass for all layers in the encoder (denoted E ) and number of operations per forward pass for all MLP layers (denoted M ).
We assume that (1) the number of updates per iteration is 1, (2) the architecture of the encoder used for state entropy estimation is the same as the RL encoder used in RAD, and (3) the FLOPs required for computations (e.g., ﬁnding the k -NNs in representation space) that are not forward and backward passes through neural network layers is negligible.
We denote the number of forward passes per training update F , the number of backward passes per training update B , and the batch size b (in our experiments b = 512 ).
Then, the number of FLOPs per iteration of RAD is:
[Equation]
where the last term is for the single forward pass required to compute the policy action.
Speciﬁcally, RAD + SE w/ Random only requires E extra FLOPs per iteration to store the ﬁxed representation from the random encoder in the replay buffer for future k -NN calculations.
In comparison, RAD + SE w/ Contrastive requires 4 bE extra FLOPs per iteration, and RAD + SE w/
Inverse dynamics requires more than 6 bE extra FLOPs.
F C OMPARISON TO S TATE E NTROPY M AXIMIZATION WITH C ONTRASTIVE E NCODER FOR M INI G RID P RE - TRAINING
We show a comparison to state entropy maximization with contrastive learning (i.e., A2C + SE w/
Contrastive + PT) in Figure 12.
In the results shown, we do not use state entropy intrinsic reward during the ﬁne-tuning phase for A2C + SE w/ Contrastive + PT, following the setup of Liu & Abbeel (2021), but found that using intrinsic reward during ﬁne-tuning results in very similar performance.
As discussed in Section 4.2, we observe that the contrastive encoder does not work well for state entropy estimation, as contrastive learning depends on data augmentation speciﬁc to images (e.g., random shift, random crop, and color jitter), which are not compatible with the compact embeddings used as inputs for MiniGrid.
We re-mark that RE3 eliminates the need for carefully chosen data augmentations by employing a random encoder.
[Figure]
G RE3 WITH O N - POLICY R EINFORCEMENT L EARNING
We provide the full procedure for RE3 with on-policy RL in Algorithm 2.
1: Initialize parameters of random encoder θ , policy φ 2: Initialize replay buffer B ← ∅ , step counter t ←3: repeat 4: t start ← t // C OLLECT TRANSITIONS 5: repeat 6:
Collect a transition τ t = ( s t , a t , s t +1 , r e t ) using policy π φ 7:
Get a ﬁxed representation y t
= f θ ( s t ) 8:
B ← B ∪ { ( τ t , y t ) } 9: t ← t +10: until terminal s t or t − t start =
t max 11: for j = t − 1 to t start do // C OMPUTE INTRINSIC REWARD 12:
Compute the distance ||
y j −
y
|| 2 for all y ∈ B and ﬁnd the k -nearest neighbor y k -NN j 13:
Compute r i j ← log( || y j − y k -NN
j
||
+ 1) 14:
Let r total j ← r e j + β · r i j 15: end for 16:
Update φ with transitions { ( s j , a j , s j +1 , r total j ) } t − 1 j = t start 17: until t > t max // U PDATE POLICY