Density-Based Bonuses on Learned Representations for Reward-Free Exploration in Deep Reinforcement Learning
Abstract
In this paper, we study the problem of representation learning and exploration in reinforcement learning.
We propose a framework to compute exploration bonuses based on density estimation, that can be used with any representation learning method, and that allows the agent to explore without extrinsic rewards.
In the special case of tabular Markov decision processes (MDPs), this approach mimics the behavior of theoretically sound algorithms.
In continuous and partially observable MDPs, the same approach can be applied by learning a latent representation, on which a probability density is estimated.
1. Introduction
Exploration is a major challenge in reinforcement learning (RL), where an agent must continually gather information about its environment in order to learn how to act optimally .
In an online setting and when rewards are available, the performance of an RL algorithm is usually deﬁned in terms of its regret, which is the difference between the rewards gathered by an optimal policy and the rewards gathered by the algorithm.
In a reward-free setting, the performance can be evaluated by the number of trajectories (i.e., the sample complexity) required for the algorithm to build a dataset that can be used by a planning algorithm to compute a nearoptimal policy for any possible reward function (Jin et al., 2020a).
In ﬁnite MDPs, the regret, the sample complexity, and the computational complexity of near-optimal algorithms 1 scale with the number of states and actions (Jaksch et al., 2010;
DeepMind D Domingues [Table]
Hence, beyond small ﬁnite environments, we are faced with the challenge of designing provably efﬁcient algorithms that are, at the same time computationally efﬁcient and able to handle arbitrarily large state spaces.
Under structural assumptions on the MDP that allow generalization, it is possible to derive algorithms whose performance bounds depend on the dimension of the state-action space with respect to some representation, instead of its cardinality .
The representation is often in the form of a metric, kernel, or linear function approximation (Ortner & Ryabko, 2012; Chowdhury & Gopalan, 2019; Jin et al., 2020b; Song & Sun, 2019;
Sinclair et al., 2019; Yang & Wang, 2020;
Yang et al., 2020; Domingues et al., 2020; Chowdhury &
Oliveira, 2020; Wang et al., 2020).
Although such algorithms present theoretical guarantees, they may suffer from two main drawbacks: a high computational complexity, preventing them from being applicable to large-scale MDPs; and the requirement that a good representation should be provided, e.g. in the form of a feature map or a kernel function.
Deep reinforcement learning algorithms have shown success in tackling large-scale reinforcement learning problems, and allow combining representation learning methods and
RL in an end-to-end fashion.
However, deep RL methods may have an excessively large sample complexity in hard exploration problems, and we have a limited theoretical understanding of such methods, especially in terms of how to perform exploration optimally.
In this paper, we focus on the problem of reward-free exploration in deep RL, and study an exploration method that draws inspiration from theoretically grounded algorithms for ﬁnite or low-dimensional
MDPs.
Theoretically-inspired exploration bonuses In ﬁnite
MDPs where a reward function is available, algorithms relying on exploration bonuses scaling with (cid:112) 1 /n ( s, a ) , where n ( s, a ) is the number of visits to a state-action pair ( s, a ) , have been shown to achieve near-optimal performance (Strehl & Littman, 2008; Azar et al., 2017).
Such bonuses depending on n ( s, a ) have inspired many algorithms that improve exploration in deep RL.
For instance, Bellemare et al.
(2016) propose a method to compute pseudo-counts approximating n ( s, a ) using density estimation on images, and Tang et al.
(2017) use locality-sensitive hashing to map continuous states to discrete embeddings, where explicit counts are computed.
A common property among all these approaches is that the more a state-action pair ( s, a ) is visited, the smaller the bonus at ( s, a ) .
This property is also satisﬁed for algorithms such as Random Network Distillation (RND) (Burda et al., 2019) and Never Give Up (NGU) (Badia et al., 2020), that manage to tackle many hard exploration problems.
Reward-free exploration In the absence of rewards from the environment, a major question is to decide what the agent should maximize .
Hazan et al.
and (Guo et al., 2021) propose algorithms that search a policy maximizing the entropy of its induced state-space distribution.
One disadvantage of this approach is that the learned policy may avoid regions of the state-space that do not increase the entropy, but may be important for downstream tasks.
Another approach is to evaluate the agent by its ability to explore and build a dataset of transitions , which can then be used to build a model and compute a near-optimal policy for any possible reward function (Jin et al., 2020a).
Kaufmann et al.
(2021) and M ´ enard et al.
(2020) show that under this criterion, near-optimal algorithms can be obtained using bonuses of the form (cid:112)
1 /n ( s, a ) or 1 /n ( s, a ) .
This motivates the use of such exploration bonuses both in the presence and in the absence of extrinsic rewards.
A key difference between the dataset approach and entropy maximization, is that the former will recompute a policy in every episode, aiming to visit different regions of the state space, whereas the latter aims to ﬁnd a single maximum-entropy policy that may fail to cover some regions of the state-space.
Representation Learning When dealing with environments with inﬁnitely many states and partial observability, learning a good state representation is crucial.
In general, this representation must satisfy at least two properties: (i) it should allow generalization across similar states; and (ii) it should be a sufﬁcient statistic for the past transitions and discard irrelevant information.
However, whether a representation is good or not will depend on the purpose for which we design the learning agent.
For instance, we might look for representations that ignore uncontrollable features of the environment, that are useful to build world models, or that allow generalization across tasks (Pathak et al., 2017;
Azar et al., 2019; Guo et al., 2020).
Given the vast amount of representation learning methods, we believe that it is im- portant to look for a uniﬁed exploration strategy that allows the agent to integrate different methods to learn representations and that do not rely on extrinsic rewards from the environment, which is the problem that we investigate in this paper.
Our contribution We study a framework in which exploration is encouraged by intrinsic rewards (or exploration bonuses) computed as a function of a density estimation on top of learned representations, which are inspired by the 1 /n ( s, a ) bonuses used in reward-free ﬁnite MDPs.
We argue that the advantages of this approach are:
(i)
It mimics the behavior of near-optimal algorithms in ﬁnite or low-dimensional MDPs
and it generalizes naturally to large scale problems.
(ii) It allows us to integrate and combine different representation learning methods, while keeping the same strategy for computing exploration bonuses.
(iii)
It enables exploration in the absence of extrinsic rewards.
2. Method
We consider an agent interacting with an environment in episodes of length H , such that, at every time t , the agent chooses an action a t ∈ A and it receives an observation o t +1 ∈ O .
Let h t
=
( o i , a i − 1 ) i ≤ t ∈ H be the history of observations and actions up to time t , which are stored in a replay buffer.
We propose the following method for reward-free exploration with learned representations:
• Given a history h t , we use a representation function f to compute an embedding x t
= f ( h t ) , where f is optimized according to any representation learning algorithm trained with data from the replay buffer; • We compute p ( x t , a t ) representing the probability of observing the embedding x t and the action a t in the replay buffer, where p is learned using any density estimation method and we deﬁne the intrinsic reward at ( h t , a t ) as r ( x t , a t ) =
1 /p ( x t , a t ) ; • A reinforcement learning agent is trained to maximize the sum of intrinsic rewards.
Below, we argue that in the special case of ﬁnite environments, this algorithm becomes similar to the RF-Express algorithm of M´enard et al.
(2020) that has sample complexity guarantees for reward-free exploration.
Then, we propose a simple method to estimate the inverse density 1 /p ( x t , a t ) that is applicable to general environments and evaluate it experimentally combining different representation learning methods.
Fully observable ﬁnite MDPs In this case, the observation o t completely describes the state of the environment, and we can take x t = f ( h t ) =
o t ∈ O .
Then, for any ( x, a ) =
( o, a ) ∈
O ×
A we can compute the probability density after T transitions as:
[Equation]
where n T ( o, a ) = (cid:80)
Ti =1 I { ( o i , a i ) =
( o, a ) } is the number of times that the observation-action pair ( o, a ) appears in the history of T transitions.
Consequently, the intrinsic reward at ( x, a ) =
( o, a ) becomes r ( x, a ) =
1 /p ( x, a ) =
T/n T ( o, a ) .
An optimal policy π T is invariant to the scaling of the rewards: we can multiply the intrinsic reward by any constant and π T will remain optimal.
This implies that π T also maximizes the sum of rewards r ( x, a ) /T , which is equivalent to using 1 /n T ( o, a ) as intrinsic reward.
Hence, an algorithm using intrinsic rewards deﬁned as 1 /p ( x, a ) is expected to have a behavior similar to that of RF-Express (M ´ enard et al., 2020) that uses bonuses proportional to 1 /n T ( o, a ) and has a near-optimal sample complexity.
Notice that, however, the scaling 1 /T is timedependent, which is not an issue for RF-Express that recomputes a policy in every episode, but in deep RL we need to be careful when optimizing non-stationary rewards.
Kernel density estimation (KDE) When a kernel function k is available, such that k
(( x, a ) , ( x (cid:48) , a (cid:48) )) represents the similarity between the state-action pairs ( x, a ) and ( x (cid:48) [Equation]
Domingues et al.
(2020) show that exploration bonuses of the form
[Equation]
allows us to obtain a sample efﬁcient exploration algorithm for continuous MDPs, where β > 0 is a regularization factor.
We can deﬁne a regularized kernel density estimator as
[Equation]
and notice that the bonus (1) is proportional to (cid:112) 1 /p ( x, a ) .
By analogy with ﬁnite MDPs, we propose the use of 1 /p ( x, a ) bonuses also for continuous MDPs.
Fitted KDE using learned representations As a simple method for density estimation, we propose the use of a learned version of (2).
We deﬁne a kernel function based on the learned representation f as follows:
[Equation]
For simplicity, we ignore the actions a and a (cid:48) in the deﬁnition of the kernel 3 , and we consider β = 0 .
Now, due to computational issues, we need to avoid the sum over all T past transitions in (2).
To do so, let J be a random variable uniformly distributed on { 1 , . .
, T } , and notice that, for [Equation] 0 , we have
[Equation]
by Jensen’s inequality.
Consequently, we can learn a function g ( x, a ) to estimate an upper bound on 1 /p ( x, a ) 4 by minimizing the following loss: [Equation] 
[Equation]
 where B is a batch of histories and actions ( h j , a j ) sampled uniformly from the replay buffer and x j = f ( h j ) .
Similarly, to compute the expectation over ( x, a ) , we take samples ( h, a ) uniformly from the replay, and set x = f ( h ) .
Finally, we propose to use the following exploration bonus at any state-action pair ( x, a ) :
[Equation]
where the term added to g ( x, a ) increases the bonus to compensate the approximation error that is made when learning the kernel density estimator (2), and is computed by sampling ( x J , a J ) uniformly from the replay buffer.
Although we propose the bonuses based on ﬁtted KDE as a ﬁrst method to test our approach, other density estimation methods can be used, such as normalizing ﬂows (Rezende & Mohamed, 2015).
Also, notice the bonuses computed ﬁtted KDE have a similar idea as the ones used by Never Give Up (NGU) (Badia et al., 2020), but NGU uses only data from a single episode for its kernel-based bonuses, whereas ﬁtted KDE uses all data from the past stored in the replay buffer.
Another similarity is that
NGU includes a multiplicative term in the bonuses using prediction errors from RND, whereas we use an additive term representing the error in predicting the KDE.
[Figure]
r
[Figure]
f v i s i t e d s t a t
[Figure]
r
[Figure]
f v i s i t e [Figure] t a t [Figure]
3. Experiments
To evaluate our framework, we consider four environments: (i) a fully observable GridWorld with 9 rooms, where the state and the observations are one-hot encodings of the position of the agent; (ii) a partially observable version
GridWorld, where a velocity is included in the state, but the observation only shows the position; (iii) Acrobot 5 ; and (iv) a partially observable version of Acrobot, where only the positions, and not the velocities, are observed.
For representation learning, we combined three methods: (i) action prediction (Pathak et al., 2017; Badia et al., 2020); (ii) adjacency regularization (Guo et al., 2021); and (iii) observation prediction, where the embedding of the past f ( h t ) is used to predict the next observation o t +1 .
The latter was only used in the partially observable enviroments.
For density estimation, we used the Fitted KDE method presented in
Section 2, and Real NVP (Dinh et al., 2017).
Finally, as reinforcement learning algorithm, we used Soft-DQN (Haarnoja et al., 2018; Vieillard et al., 2020).
Additionally, as a baseline, we consider RND (Burda et al., 2019) that computes intrinsic rewards based on the error incurred when predicting the output of a random neural network.
We study the case where the predictions are made using the observations o t only, and also the case where the predictions are made on top of the embeddings x t = f ( h t ) .
Notice that the latter case is a way to generalize RND to environments where a representation needs to be learned.
We evaluate the algorithms with respect to how quickly they are able to discover new states in the environment.
6 Figureshows the number of visited states versus the number of training steps with our framework, compared to the RND baselines.
We observe that Fitted KDE allows us to discover at least as many states as the baselines in the GridWorld, and more states are discovered in the Acrobot environments.
4. Conclusion
We propose a uniﬁed framework to compute exploration bonuses in reward-free exploration for RL using density estimation.
This framework is inspired by provably efﬁcient algorithms in tabular MDPs and we generalize it to arbitrary environments by allowing the algorithm to integrate any set of representation learning and density estimation methods.
We illustrate its empirical effectiveness in continuous and partially observable problems, by proposing a simple method to compute density based bonuses, Fitted
KDE, that can be related to the intrinsic rewards used by state-of-the-art methods, such as RND (Burda et al., 2019) and NGU (Badia et al., 2020).
Interesting directions for future work include testing our approach in more complex environments, such as hard exploration games in Atari, and investigate the impact of different representation methods.
The authors thank Bilal Piot, Alaa Saade and Daniel Guo for their very useful feedback and support during this work.