State Entropy Maximization with Random Encoders for Efﬁcient Exploration
Abstract
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.
Source code and videos are available at https://sites.google.com/view/re3-rl .
Introduction
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 [4, 32, 45] and prediction errors [5, 17, 33, 41].
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.
[24] and Hazan et al.
[15] 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.
[30] estimate state entropy by measuring distances between states and their k -nearest neighbors.
To extend this approach to highdimensional environments, recent works [2, 26, 46] 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 [46], inverse dynamics prediction [2], and contrastive learning [26].
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. 𝑘 = 𝑘 = 𝑘 = Source TrueState
Random Encoder 𝑘 = 𝑘 = 𝑘 = Source TrueState Random Encoder
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.
[48] and Caron et al.
[6] 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 [47] and MiniGrid [7] 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 [5, 33].
• RE3 is compute-efﬁcient because 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.
Related Work
Exploration in reinforcement learning.
Exploration algorithms encourage the RL agent to visit a wide range of states by injecting noise to the action space [25] or parameter space [10, 34], maximizing the entropy of the action space [13, 52], and setting diverse goals that guide exploration [8, 9, 31, 35].
Another line of exploration algorithms introduce intrinsic rewards proportional to prediction errors [5, 17, 33], and count-based state novelty [4, 32, 45].
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. [Figure] If a reward signal is given by the environment Maximize
Distance
State entropy maximization.
Most closely related to our work are methods that maximize the entropy of state distributions.
Hazan et al.
[15], Lee et al.
[24] proposed to maximize state entropy estimated by approximating the state density distribution.
Instead of approximating complex distributions
, Mutti et al.
[30] proposed to maximize a k -nearest neighbor state entropy estimate from on-policy transitions.
Recent works extend this method to environments with high-dimensional observations.
Tao et al.
[46] employ model-based RL techniques to build a representation space for the state entropy estimate that measures similarity in dynamics, and Badia et al.
[2] proposed to measure similarity in the representation space learned by inverse dynamics prediction.
The work closest to ours is Liu & Abbeel [26], which uses off-policy RL algorithms to maximize the k -nearest neighbor state entropy estimate in contrastive representation space [40] 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 [12] termed the Gamba perceptron by Minsky & Papert [27].
Moreover, nice properties of random projections are commonly exploited for low-rank approximation [36, 49].
These ideas have since been extended to deep convolutional networks, where random weights are surprisingly effective at image generation and restoration [48], image classiﬁcation and detection [6], and fast architecture search [37].
In natural language processing,
Wieting & Kiela [50] demonstrated that learned sentence embeddings show marginal performance gain over random embeddings.
In the context of RL, Gaier & Ha [11] showed that competitive performance can be achieved by architecture search over random weights without updating weights, and Lee et al.
[23] 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.
Method 3.1 Preliminaries
We formulate a control task with high-dimensional observations as a partially observable Markov decision process (POMDP) [20, 43], 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 [28], we reformulate the POMDP as an MDP [43] 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.
Algorithm
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 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 [3].
However, since it is difﬁcult to estimate p with high-dimensional data, particle-based k -nearest neighbors ( k -NN) entropy estimator [39] 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.2 Random Encoders for Efﬁcient Exploration
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 [26] 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 } .
Here, our intuition is that measuring the 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 high-dimensional 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 [30], or samples from a minibatch [26].
Our scheme enables stable, precise entropy estimation in a compute-efﬁcient manner.
[Figure]
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 pretraining , 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 ﬁne-tuning.
Formally, we introduce a policy π φ , parameterized by φ , that maximizes the expected return
[Figure]
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 RE3 with on-policy RL in Algorithm 2.
Experiments
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 4)?
• How does RE3 compare to state entropy maximization schemes that involve representation learning (see Figure 5) and other exploration schemes that introduce additional models for exploration (see Figure 6)?
• Can RE3 further improve the sample-efﬁciency of off-policy RL algorithms by unsupervised pre-training (see Figure 7 and Figure 8)?
• How compute-efﬁcient is RE3 (see Figure 9)?
• Can RE3 also improve the sample-efﬁciency of on-policy RL in discrete control tasks (see Figure 10)?
4.1 DeepMind Control Suite Experiments
Setup.
To evaluate the sample-efﬁciency of our method, we compare to Dreamer [14], a state-of-theart model-based RL method for visual control; and two state-of-the-art model-free RL methods, RAD [22] and DrQ [21].
For comparison with other exploration methods, we consider RND [5] and ICM [33] 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 A. Source code is available at https://sites.google.com/ view/re3-rl . [Figure] [Figure]
Comparative evaluation.
Figure 4 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 algorithms.
We provide additional experimental results on more tasks in Figure 12.
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 5.
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 [16] encoder pre-trained on ImageNet dataset (RAD + SE w/ ImageNet) and an ATC [42] 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 A. 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
[Figure] [Figure]
t
[Figure]
n t r op y E s t
i m a t e [Figure]
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 6, RAD + RE3 consistently exhibits superior sample efﬁciency in most 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]
We compare the diversity of visited states resulting from pre-training for 500K steps with APT [26] and RE3 (ours).
We provide corresponding videos in our website.
[Figure]
Estimating state entropy with a random encoder is signiﬁcantly more computeefﬁcient than with a encoder trained with representation learning. 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 computeefﬁ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 D. Figure 9 shows the FLOPS used by each agent to achieve its ﬁnal performance in Hopper Hop.
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.
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 8.
One can see that the pre-trained policy using RE3 exhibits more diverse behaviors compared to random exploration or APT [26], 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.e., proprioceptive) state space in the Hopper environment in Figure 7a.
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 7b and 7c.
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 A 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 pretraining a policy only for 500K environment steps, while previous work [26] reported results by training for 5M environment steps.
[Figure]
4.2 MiniGrid Experiments
[Figure]
Setup.
We evaluate our method on MiniGrid [7], a gridworld environment with a selection of sparse reward tasks.
We consider the following setups where the agent obtains a reward only by reaching the green goal square: Empty, a large room with the goal in the furthest corner; DoorKey, where the agent must collect a key and unlock a door before entering the room containing the goal.
The tasks are shown in Figure 11.
For evaluation, we consider two exploration methods, RND and ICM. For our method and other exploration methods, we use Advantage Actor-Critic (A2C)
[29] as the underlying RL algorithm.
the agent has access to a compact 7 × 7 × 3 embedding of the 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 C for more details).
Comparison with other exploration methods.
Figure 10 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 10 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 13 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.
Discussion
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 compute-efﬁ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 pre-trained ﬁ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 RE would facilitate future research by providing a simple-to-implement, stable, and compute-efﬁcient module that can be easily combined with other techniques.
Acknowledgements
This research is supported in part by ONR PECASE N000141612723, Tencent, and Berkeley Deep Drive. We would like to thank Seunghyun Lee, Colin Li, Sangwoo Mo, Aravind Srinivas, and Mandi
Zhao for providing helpful feedbacks and suggestions. We would also like to thank Hao Liu for providing helpful implementation details, and Adam Stooke for providing pre-trained parameters of
ATC. Appendix A Details on DeepMind Control Suite Experiments A.1 Environments
We evaluate the performance of RE3 on various tasks from DeepMind Control Suite [47]. 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.
[38] 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.
A.2 Implementation Details for Model-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 3 ,
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 [26], 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 [1].
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 [40].
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 fully-connected 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 [16] 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 [42] 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.
[42] 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 [5] and ICM [33]:
• RND.
For RND, we introduce a random encoder f θ whose architecture is same as in Yarats et al.
[51], 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.
[51], and introduce a inverse dynamics predictor h ψ with 2 fully-connected layers withunits.
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.
[33], we minimize 0 .
2 · L forward + 0 . {+
 · L+} inv .
We perform hyperparameter search over the weight β ∈ { 0 . 05 , 0 .
and report the best result on each environment.
[Table]
A.3 Implementation Details for Model-based RL.
For Dreamer [Table] we use the publicly available released implementation repository ( https://github. com/danijar/dreamer ) from the authors 4 .
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 [26] 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.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 [1].
Finally, tanh nonlinearity is added to the 50-dimensional output of the fully-connected layer.
[Table] [Figure]
B Additional Experimental Results on DeepMind Control Suite
We provide additional experimental results on various tasks from DeepMind Control Suite [47].
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.
[Figure]
C Details on MiniGrid Experiments
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]
D Calculation of Floating Point Operations
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.
[18] and Jeong & Shin [19] 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/
[Table]
E Comparison to State Entropy Maximization with Contrastive Encoder for MiniGrid Pre-training
We show a comparison to state entropy maximization with contrastive learning (i.e., A2C + SE w/
Contrastive + PT) in Figure 13.
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 [26], but found that using intrinsic reward during ﬁne-tuning results in very similar performance. As discussed in Section 4.
[Figure] 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 [Figure] 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]
F RE3 with On-policy Reinforcement Learning
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