A Mixture of Surprises for Unsupervised Reinforcement Learning
Abstract
Unsupervised reinforcement learning aims at learning a generalist policy in a reward-free manner for fast adaptation to downstream tasks.
Most of the existing methods propose to provide an intrinsic reward based on surprise.
Maximizing or minimizing surprise drives the agent to either explore or gain control over its environment.
However, both strategies rely on a strong assumption: the entropy of the environment’s dynamics is either high or low.
This assumption may not always hold in real-world scenarios, where the entropy of the environment’s dynamics may be unknown.
Hence, choosing between the two objectives is a dilemma.
We propose a novel yet simple mixture of policies to address this concern, allowing us to optimize an objective that simultaneously maximizes and minimizes the surprise.
Concretely, we train one mixture component whose objective is to maximize the surprise and another whose objective is to minimize the surprise.
Hence, our method does not make assumptions about the entropy of the environment’s dynamics.
We call our method a M ixture O f S urprise S (MOSS) for unsupervised reinforcement learning.
Experimental results show that our simple method achieves state-of-the-art performance on the URLB benchmark, outperforming previous pure surprise maximization-based objectives.
Our code is available at: https: //github.com/LeapLabTHU/MOSS .
Introduction
Humans can learn meaningful behaviors without external supervision, i e. , in an unsupervised manner, and then adapt those behaviors to new tasks [3].
Inspired by this, unsupervised reinforcement learning decomposes the reinforcement learning (RL) problem into a pretraining phase and a finetune phase [32].
During a pretraining phase, an agent prepares all possible tasks that a user might select.
Afterward, the agent tries to figure out the selected task as quickly as possible during finetuning [24].
Doing so allows solving the RL problem in a meaningful order [3], e.g. , a cook first has to look at what is in the fridge before deciding what to cook.
Unsupervised representation learning has shown great success in computer vision [27] and natural language processing [10]; however, one challenge is that RL includes both behavior learning and representation learning [32, 53].
Current unsupervised RL methods provide an intrinsic reward to the agent during a pretraining phase to tackle the behavior learning problem [32].
Intuitively, this intrinsic reward should incentivize the agent to understand its environment [15].
Current methods formulate the intrinsic reward as either maximizing or minimizing surprise [45, 6, 36, 57, 42, 43, 11, 35, 31, 17, 47, 46, 13], where knowledge-based methods quantify surprise as the uncertainty of a prediction model [11, 42, 43] and data-based methods measure surprise as an information-theoretic quantity [45, 6, 57, 36, 5] 3 .
Surprise maximization methods [24, 31, 36] formulate the problem as an exploration problem.
Intuitively, to prepare all possible downstream tasks, an agent has to explore the state space and figure out what is possible in the environment.
For instance, one instantiation is to maximize the agent’s state entropy [36].
In contrast to surprise maximization, another line of work takes inspiration from the free-energy principle [22, 21, 20] and proposes to minimize the surprise [6, 45].
In particular, these works argue that external perturbations naturally provide the agent with surprising events.
Hence, an agent can learn meaningful behaviors by minimizing its surprise, and minimizing this quantity requires gaining control over these external perturbations [45, 6].
For example, SMiRL [6] proposed to minimize the agent’s state entropy.
A closer look at these two approaches indicates a strong assumption.
Substantial external perturbations already naturally provide an agent with a high state entropy [6].
Therefore, surprise-seeking approaches assume that the environment does not provide significant external perturbations.
On the other hand, in an environment without any external perturbations, the agent already achieves minimum entropy by not acting [45].
Therefore, surprise minimization approaches assume that the environment provides external perturbations for the agent to control.
However, in real-world scenarios, it is often difficult to quantify the entropy of the environment’s dynamics beforehand, or the agent might face both settings.
For example, a butler robot performs mundane chores daily in a household.
However, the robot might also encounter surprising events ( e.g. , putting out a fire).
Therefore a fixed assumption on the entropy of the environment dynamics is often not possible.
In other words, choosing between surprise maximization or minimization for pretraining is a dilemma.
Simultaneously optimizing these two opposite objectives does not make sense.
Instead, in this paper, we show that competence-based methods [32] offer a simple yet effective way to combine the benefits of these two objectives.
In addition to conditioning the policy on the state, competence-based methods condition the policy on a latent vector [31, 35, 26, 55, 23, 1, 18].
Conditioning the policy offers an appealing way to formulate the unsupervised RL problem.
During a pretraining phase, the agent tries to learn a set of possible behaviors in the environment and distills them into skills.
Then, during finetuning, the agent tries to figure out the selected task as quickly as possible, using the repertoire of skills gathered during pretraining.
For example, as illustrated in Fig. 1, we view skill distributions as a mixture of policies instead of a single policy.
In particular, we train one set of skills whose objective is to maximize the surprise and another set of skills whose objective is to minimize the surprise.
Surprisingly, this paper shows that this simple approach works well in practice, which simultaneously optimizes two contradicting objectives.
Our primary contribution, presented in Section 4, is a simple intrinsic reward called MOSS that does not make assumptions about the entropy of the environment’s dynamics.
In Section 5, our experimental results on URLB [32] and
ViZDoom [30] show that, surprisingly, our MOSS methods achieve state-of-the-art results.
We organize the paper as follows.
Section 3 briefly analyzes previous unsupervised RL algorithms under the surprise framework.
Then in Section 4, we introduce our MOSS method.
Next, experimental results in Section 5 shows that on URLB [32] and ViZDoom [30], our MOSS method improves upon previous pure maximization and minimization methods.
Finally, we provide discussions and limitations in Section 6.
Preliminaries
Markov Decision Process.
Unsupervised RL methods studied in this paper operate under a Markov
Decision Process (MDP) [51].
In particular, we specify an MDP as a tuple M =
( S , A , T, r ext , ρ, γ ) where S is the state space and A is the action space of the environment.
T ( S ′ | S , A ) represents the state-transition dynamics, ρ : S → [0 , 1] is the initial state distribution, and γ ∈ [0 , 1) is the discount factor.
At each discrete time step t ∈
Z ∗ , the agent receives a state and performs an action, which we denote as S t ∈ S and A t ∈ A , respectively.
During pretraining, unsupervised RL algorithms compute an intrinsic reward r int ; during the finetune phase, the agent receives the extrinsic reward r ext given by the environment at each interaction.
Skill.
Intuitively, a skill is an abstraction of a specific behavior ( e.g. , walking), and in practice, a skill is a latent conditioned policy [17].
Given a latent vector z , we denote a skill as π θ ( a t | s t , z ) , where π θ is the policy parameterized by θ .
For instance, during pretraining, the latent vectors are sampled every n steps such that the latent vector z is associated with the behavior executed during the associated n steps.
Mutual Information.
Knowledge-based, data-based, and competence-based methods have different measures of surprise.
The study in this paper falls into the category of competence-based methods.
In particular, data-based and competence-based methods rely on an information-theoretic definition of surprise, i.e. , entropy.
Previous competence-based methods acquire skills by maximizing mutual information [16] between T and skills Z
[Equation]
where T can be the states S , the joint distribution of state-transitions ( S ′ , S ) , or the state-transitions ( S ′ | S ) .
In particular, these methods differ in how they decompose the mutual information.
Theoretically, these different decompositions are equivalent, i.e. , they all maximize the mutual information between states and skills.
However, the particular choice greatly influences the performance in practice as optimizing this objective relies on approximations.
To motivate the potential of competence-based methods over data-based or knowledge-based methods, we provide an intuitive understanding of Eq.
On the one hand, the entropy term says that we want skills in aggregate that explore the state space; we use it as a proxy to learn skills that cover the set of possible behaviors.
On the other hand, It is not enough to learn skills that randomly go to different places.
We want to reuse those skills as accurately as possible, meaning we need to be able to discriminate or predict the agent’s state transitions from skills.
To do so, we minimize conditional entropy.
In other words, appropriate skills should cover the set of possible behaviors and should be easily distinguishable.
Information-Theoretic Skill Discovery
Competence-based methods employ different intrinsic rewards to maximize mutual information: (1) discriminability-based and (2) exploratory-based intrinsic rewards.
For example, the former rewards the agent for discriminable skills.
In contrast, the latter rewards the agent for skills that effectively cover the state space using a KNN density estimator [48, 36] to approximate the entropy term.
Below we analyze both approaches.
3.1 Discriminability-based Intrinsic Reward
Previous work such as DIAYN [17] uses a discriminability-based intrinsic reward.
They use the decomposition in Eq.
(2) and given a variational distribution q ϕ , they optimize the following variational lower bound
[Equation]
where Z ∼ p
( Z ) is a discrete random variable.
In particular, the discriminator rewards the agent if it can guess the skill from the state, i.e. , r int ( s ) ≜ log q ϕ ( z | s ) − log p ( z ) .
These methods may easily run into a chicken and egg problem, where skills learn to be diverse using the discriminator’s output.
However, the discriminator cannot learn to discriminate skills if the skills are not diverse.
Hence, it discourages the agent from exploring.
Previous work [49] has tried to solve this problem by decoupling the aleatoric uncertainty from the epistemic uncertainty .
Furthermore, solving the chicken and egg problem is not enough since a state must map to a single skill in Eq.
Accordingly, the methods that use the decomposition in Eq.
(2) require the skill space to be smaller than the state space so that the skills are distinguishable.
Since previous work showed that a high-dimensional continuous skill space empirically performs better, the intrinsic reward should not rely on the decomposition in Eq.
Intuitively, a continuous skill space allows (1) to interpolate between skills and (2) to represent a more significant set of skills in a more compact representation.
Instead, in Eq.
(1), in contrast to skills that are predictable by states H
[ Z | T ] , we require that states are predictable by skills H [ T | Z ] .
Since any state should be predictable from a given skill, the skill space must be larger than the state space ( i.e. , we do not want a skill mapping to multiple states).
Therefore, another work [47] relies on a variational bound on Eq.
[Equation]
which uses a continuous skill space; however, it does not scale to high dimensions as the intrinsic reward r int ( s ′ , a , s ) ≜ log q ϕ ( s ′ | s , z ) p ( s ′ | s ) relies on an estimation of p ( s ′ | s ) .
In particular, they assume that p ( z | s ) =
p ( z ) .
Intuitively, given s , if we assume each element in the latent vector ( z ) i to be independent of each other, the error of this assumption will be scaled by the dimension of z .
3.2 Exploratory-based Intrinsic Reward
As aforementioned, discriminability-based intrinsic reward may easily run into a chicken and egg problem.
Hence, other work [31, 35] uses an exploratory-based intrinsic reward to address the chicken and egg problem.
In other words, they use the decomposition in Eq.
(1), where the intrinsic reward explicitly rewards the agent for exploring through H [ T ] .
Previous work [36] maximizes the state entropy, i.e. , T = S .
However, the authors in [4] argue that it often results in the discriminator simply memorizing the last state of each skill.
Instead, CIC [31] proposes to maximize the entropy of the joint distribution of state-transitions (which from now on we refer as joint entropy), i.e. , H [ S ′ , S ] .
Therefore, CIC [31] proposes to estimate the conditional entropy H [ S ′ , S | Z ] using noise contrastive estimation [31, 25] and the joint entropy H [ S ′ , S ] using a k -nearest neighbor estimator [48] to handle high dimensional skill space.
KNN-density estimation.
Previous works [36, 31] approximate the joint entropy H
[ S ′ , S ] using a k -nearest neighbor estimator [48].
Given a random sample of size N , { ( S ( i ) ′ , S ( i ) ) } Ni =1 ∼ P ( S ′ , S ) , the estimator computes a Monte Carlo estimate [7] of the entropy defined as
[Equation]
where ˆ p k ( s ( i ) ′ , s ( i ) )
=
k NV ( i ) k is an estimation of p ( s ( i ) ′ , s ( i ) ) , V ( i ) k is the volume of the hypersphere of radius R ( i ) k , and b ( k ) is a constant such that the estimator is unbiased.
In particular, R k is the distance from s to its k -nearest neighbor s ′ k where we assume a uniform distribution of points within the hypersphere (local uniformity assumption [37]).
In practice, to make the distance R k meaningful, we do not operate in the state space.
For instance,
CIC [31] trains an MLP f ψ : R |S| → R d that projects the state into a latent representation such that
R k = ∥
f ψ ( s ) − f ψ ( s ′ k ) ∥ where d is the skill dimension.
Moreover, in APT [36], it was found that averaging over all k -nearest neighbors leads to better results.
Hence, CIC [31] optimizes a quantity proportional to Eq.
[Equation]
where c = 1 is constant for numerical stability.
We view the joint entropy as an expected reward, and for a transition ( s ( i ) , s ( i ) ′ ) , the reward function is
[Equation]
Noise Contrastive Estimation.
To train parameters ψ , CIC [31] minimizes the conditional entropy H [ S ′ ; S | Z ] using a noise contrastive estimation (NCE) [25, 40].
Denote h = f ψ ( s ) and h ′ =
f ψ ( s ′ )
the projected representations of s and s ′ , respectively.
Furthermore, let g ϕ z : R d
→
R
d be an MLP that projects the skill into a latent representation and g ϕ s : R d × R d → R d be an MLP that projects h and h ′ into a latent representation.
We define the NCE loss
[Equation]
as an approximation to H [ S ′ ; S | Z ] , where we define cos ( · , · ) as the cosine similarity operator, ω as a temperature term, and c ( i ) j = cos ( g ϕ z ( z ( i ) ) , g ϕ s ( h ( j ) , h ( j ) ′ )) /ω .
As aforementioned, this term serves as representation learning for parameters ψ .
Intuitively, the NCE loss [25] pushes p ( s ′ | s , z ) to be a delta-like density function.
However, since we define behaviors with n steps, we want a wider distribution, so in practice, ω is set to 0 .
5 to smooth the distribution.
Mixture Of Suprises (MOSS)
Our goal is to provide an algorithm that does not build on any assumption about the entropy of the environment’s dynamics.
A straightforward way to achieve this is to maximize and minimize surprise simultaneously.
A benefit of such an objective is that it promotes exploration to find a more effective policy for minimizing the surprise [6].
In particular, in [6], they maximize surprise according to all prior experiences while minimizing surprise according to the current episode.
However, they find that it only helps in some cases.
The challenge is to find an effective way to optimize these two opposite objectives since if we add the two objectives, it would simply yield a linear interpolation of the two objectives, which still corresponds to either maximizing or minimizing.
A previous work [19], dubbed
Adversarial Surprise (AS), resolves this challenge using an adversarial game where two policies compete against each other: one maximizes surprise while the other minimizes surprise.
AS [19] requires training two policies that do not share parameters or data.
Instead, we propose a novel yet simple mixture of policies to maximize and minimize surprise simultaneously.
Our mixture of policies does not rely on an adversarial game [19] in which training is known to be challenging [38] and uses a single network for both objectives resulting in higher sample efficiency.
Mainly, our method only
Input: Environment, number of pretraining steps N T .
Initialize DDPG, and ϕ s , ϕ z , ψ for t = 1 to N T do if ( t % update skill every ) =
=
then
Sample m t ∼ p ( m )
▷ setting to m = 0 is equivalent to CIC[31]
Sample z t ∼ p ( z | m t ) end if
Take some action a t and observe s t +
Store ( s t , s t +1 , m t , z t ) into buffer B
if t ≥ 4000 then
Sample a batch { ( a ( i ) , s ( i ) , s ( i ) ′ , m ( i ) , z ( i ) ) }
Ni =1 and compute intrinsic reward r int using equation 9.
Update DDPG using intrinsic reward.
Update ϕ s , ϕ z , ψ using noise contrastive loss in equation 8.
end if end for requires a minor change to CIC [31] that does not involve any additional computational cost during training.
In particular, our mixture of policies builds on a competence-based approach [31], meaning we use two disjoint skill sets.
We separate the two opposite objectives using a different skill set for each objective: one skill set for surprise maximization and another for surprise minimization.
Specifically, we define the mixture of policies using a binary random variable M , where we associate each mixture component with a continuous uniform distribution U d ( a, b ) on an interval [ a, b ] .
Each uniform distribution constitutes a skill distribution.
Maintaining two different skill distributions allows us to define a different objective for each distribution.
In particular, we define two distributions
Z max ∼ P ( Z | M = 0) and Z min ∼ P ( Z | M = 1) .
Since we do not assume any priors on the entropy of the environment’s dynamics, we deterministically set M = 0 for the first half of the steps and M = 1 for the other half of the steps in the episode.
Compared to more sophisticated methods ( e.g. , appendix B) for setting M ( i.e. , switching the objective), such a heuristic adds no computational cost during training, requires no hyper-parameter tuning, and performs well.
Moreover, motivated by the discussions in Section 3.2, we build our method on top of the CIC 4 [31] framework.
This framework is appealing as it bypasses the chicken and egg problem and supports high-dimensional skills.
In addition, it significantly outperforms previous competence-based methods on the URLB benchmark [32].
Hence, following CIC [31], we use the decomposition in Eq.
I ( S ′ , S ; Z ) =
H [ S ′ , S ] − H
[ S ′ , S | Z ] .
In particular, we define surprise as the joint entropy, which corresponds to maximizing and minimizing H [ S ′ , S ] , and as in CIC [31] we use Eq.
(6) to approximate the joint entropy.
We approximate the conditional entropy H [ S ′ , S | Z ] using a noise contrastive estimation following Eq.
By minimizing the conditional entropy, we distill behaviors into skills.
Doing so also serves as representation learning for computing the joint entropy [31].
Implementation.
To ensure fairness, all hyper-parameters are kept as in CIC [31] and we use the same RL algorithm ( i.e. , DDPG [34] as implemented in DrQ-V2 [56]).
In practice, Z max ∼ P ( Z | M = 0) ≜
U d (0 , 1) and Z min ∼ P ( Z | M = 1) ≜
U d ( − 1 , 0) (we refer readers to appendix for other variants that we tried), where we set M = 0 for the first half of steps and M = 1 for the remaining half of steps in the episode.
Hence, the RL agent maximizes
F ( θ ) ≜ − E z max ∼U d (0 , 1) E p z max ( s ′ , s ) [log p z max ( s ′ , s )]
+ E z min ∼U d ( − 1 , 0)
E p z min ( s ′ , s ) [log p z min ( s ′ , s )] .
And the intrinsic reward is a function of the state-transition and the binary random variable M
[Equation]
As in CIC [31], we approximate the joint entropy using Eq.
MOSS’s implementation is straightforward; i.e. , at each time step, instead of sampling Z ∼ P ( Z ) , we sample Z ∼ P ( Z | M ) and replace the intrinsic reward with Eq.
Pseudocode for MOSS is provided in Alg.
In fact, fixing M = 0 yields CIC [31].
Experiments
Environments.
We present the main results of method by evaluating on the Unsupervised Reinforcement Learning Benchmark (URLB) [32], which is a standard unsupervised RL benchmark for continuous control.
URLB [32] has three different domains listed in order of increasing difficulty:
• Walker is a biped robot domain constrained to a 2D vertical plane.
The challenge for this domain is to learn how to balance and maneuver simultaneously.
• Quadruped is a four-legged robot in 3D space.
This domain has a higher dimensional state and action space than the Walker domain.
• Jaco is a 6-DOF robotic arm domain with a three-finger gripper in 3D space that has to learn action constraints and manipulation.
Furthermore, each of the three domains has four different downstream tasks with varying difficulties.
We refer readers to Tab. 1 for the complete list of 12 tasks.
Finally, in OpenAI gym environments [9], since an episode terminates as soon as a robot falls, the agent will naturally avoid falling to the ground as it will get a 0 reward.
Hence, during pretraining, the OpenAI gym environment [9] may leak some task information ( e.g. , standing) to the unsupervised learning agent [31].
Therefore, URLB [32] builds on top of the Deepmind Control Suite [52] to avoid leaking some task information.
The domains in URLB [32] use low-dimensional state information, deterministic transitions, and non-evolving environments.
In contrast, ViZDoom [30] uses raw pixel observations, stochastic transition, and evolving environments.
Furthermore, since the environment randomly spawns enemies that shoot fireballs, the agent faces external perturbations that surprise the agent.
We evaluate our method on the DefendTheLine and TakeCover maps.
Please refer to Figure 2 for environment renders.
[Figure]
Evaluation.
We follow the benchmark’s standard training procedure by pretraining the agent formillion steps in each of the three domains and then finetuning the pre-trained agent for 100 k steps with downstream task rewards.
Since AS trains two network-independent policies, we ensure both policies each perform 2 million environment steps during pretraining.
To ensure fairness, we pretrain and finetune each method with the same 12 seeds (0-11) using the code provided in URLB 5 [32].
We present the results in Fig.
3 for a total of 1584 ( = 11 algorithms × 12 tasks × 12 seeds) runs.
Our main evaluation metrics are interquartile mean (IQM) and Optimality Gap (OG) of normalized scores, where we used the score of a DDPG [56] agent trained from scratch for 2 million steps as the expert scores [32].
On the one hand, IQM is more robust to outliers than the mean, and IQM is significantly less biased than the median [2].
On the other hand, the Optimality Gap measures how far scores are from the expert scores [2].
Baselines.
Baselines include methods presented in the URLB benchmark [32] as well as CIC [31].
The baseline methods fall into knowledge-based methods, i.e. , ICM [42], Disagreement [43], RND [11], data-based methods, i.e. , APT [36], Proto-RL [57], AS [19], and competence-based methods, i.e. , SMM [33], DIAYN [17], APS [35], CIC [31].
We refer readers to the Appendix of URLB [32] for details on the different baselines.
Lastly, since the two policies (Alice and Bob) in AS are parameter-independent, we pretrain for 4 million steps, each policy trained using 2 million steps.
We evaluate our agent in ViZDoom [30] against CIC [31], and a modified version of CIC [31], dubbed NegativeCIC, where the intrinsic reward is the negative joint entropy.
Since the ViZDoom environment has high natural environment dynamics entropy, NegativeCIC should act as a strong baseline.
We adopt the same pretraining and finetuning procedure as in URBL for ViZDoom.
MOSS.
We keep all hyperparameters as those in CIC [31].
Specifically, unless specified otherwise, we set M = 0 for the first half of steps and M = 1 for the other half of steps in the episode.
We refer readers to the Appendix for more details.
(a) Aggregate Statistics .
[Table]
[Figure]
[Table]
5.1 Results
[Table]
To ensure the reliability of our main results, we report results as described in [2] to account for the uncertainty.
We present results in Fig.
3, where, surprisingly, our method obtains state-of-the-art performance on the URLB benchmark [32] on our main evaluation metrics [2].
In particular, the performance profile in Fig.
3(b) reveals that MOSS stochastically dominates all other methods.
We also present numerical results in Tab. 1 for a detailed breakdown over individual tasks.
For the IQM metric, we report an increase of 6% compared to CIC, which is the second leading method.
Furthermore, as for the Optimality Gap metric, we decreased 20% compared to CIC.
Furthermore, in Fig.
3(b) previous methods intersect at multiple points (making it hard to conclude whether a method is better than the other), while MOSS is always above all previous methods.
In addition, in Fig. 3(b), MOSS maintains all runs greater than a normalized score as high as 0 .
[Table]
We also evaluated our method on ViZDoom [30].
Since the TakeCover and DefendTheLine maps in ViZDoom provides natural perturbations, we use NegativeCIC as the baseline.
We present the numerical results in Tab. 2.
Since the environments in ViZDoom are highly stochastic, we remove the top and bottom 10% of scores and report the mean and standard error to account for any outliers.
Like in the main results presented earlier, MOSS consistently outperforms the baseline method under this completely new setting.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
5.2 Ablation
Since MOSS does not make assumptions about the entropy of the environment’s dynamics, we are interested in knowing whether adding stochasticity affects performance.
For this stochastic setup, during each interaction with the environment, we sample Y ∼ Bernoulli(0 . 3) and if Y = 1 , we add a Gaussian Noise N (0 , 0 . 2) to the agent’s intended action.
In Fig. 4(a), we observe that CIC’s
[31] performance drops while MOSS’ performance is almost unaffected.
This finding empirically supports previous works which claim that solely maximizing surprise for unsupervised RL agents is sometimes not enough, especially in stochastic environments [6, 45].
In MOSS, we maintain two different skill distributions, i.e. , Z max ∼ U d (0 , 1) and Z min ∼ U d ( − 1 , 0) .
In particular, setting Z min to the same distribution as Z max ∼ U d (0 , 1) is equivalent to directly optimizing CIC [31] with the two objectives simultaneously.
We dub MOSS with the same skill distribution, MOSS_SAME.
In Fig. 4(b), we observe that naively doing so can result in poor performance, showing the effectiveness of our method in mitigating the simultaneous optimization of two contradicting objectives.
Moreover, as shown in Fig. 4(b), NegativeCIC performs poorly on the URLB benchmark [32] since most tasks required structured exploration in a low entropy environment.
Discussions & Limitations
We presented MOSS, an unsupervised RL method that does not make assumptions about the environment dynamics’ entropy by simultaneously maximizing and minimizing surprises.
Our method falls into competence-based methods.
Leveraging a mixture of policies for more than surprise maximization and minimization is an interesting future research.
Limitations.
One limitation is that our method has a simple and effective deterministic rule for switching between the two objectives independent of state or history.
We did explore a more sophisticated method in Appendix B that had a better performance but required tuning and extra computation.
Therefore, a promising direction is to explore methods that take advantage of state information ( e.g. , using a meta-controller [29]) to output a more optimal switching mechanism that incentivizes deep structured exploration [44].
Moreover, we only considered competence-based methods.
Extending our method to other unsupervised RL methods would be interesting to investigate in future work.
Additionally, when finetuning, we transferred network weights which could lead to unlearning of already learned behaviors [12].
It would be meaningful to combine with works that circumvent this issue like [12] which utilized frozen pretrained policies to improve downstream finetuning further.
Lastly, we only considered two objectives, and including more meaningful objectives like empowerment or anxiety [39, 50] during pretraining could be an exciting direction.
Potential Negative Impact.
Since we operate under the unsupervised reinforcement learning domain, the agent receives rewards that incentives exploration.
Therefore, agents under this domain can be more unpredictable than agents that receive task rewards for intended behaviors.
Furthermore, there is no current safeguard for the agent to not act maliciously in its environment; the agent may potentially cause harm to properties situated in the environment, including the agent itself.
Therefore, unsupervised learning with safety constraints would be an exciting area of future research to mitigate this potential negative impact.
Acknowledgement.
This work is supported in part by the National Science and Technology Major
Project of the Ministry of Science and Technology of China under Grants 2018AAA0101604, the
National Natural Science Foundation of China under Grants 62022048 and the State Key Lab of
Autonomous Intelligent Unmanned Systems.
We also appreciate the generous donation of computing resources by High-Flyer AI
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
Please do not modify the questions and only use the provided macros for your answers.
Note that the Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See supplementary material and Section 5.
The code is in the supplementary material.
(b) Did you specify all the training details ( e g. , data splits, hyperparameters, how they were chosen)?
[Yes] See supplementary material and Section(c)
Did you report error bars ( e g. , with respect to the random seed after running experiments multiple times)?
[Yes] See Section
(d) Did you include the total amount of compute and the type of resources used ( e.g. , type of GPUs, internal cluster, or cloud provider)?
[Yes] See supplementary material.
If you are using existing assets ( e.g. , code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Yes code in the supplementary material (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Additional Implementation Details
[Table]  A.1 MOSS Implementation
We implement MOSS using JAX [8], Haiku [28].
We chose to build on top of JAX [8] as we observed a 2x speedup compared to PyTorch [41].
Moreover, we use PyTorch [41] and Reverb [Table] Episode partition  [14] to implement the replay buffer.
Tab.
3 details the hyper-parameters used in MOSS which are taken directly from CIC [31].
Since MOSS builds on CIC [31], we empirically verified that our implementation on top of JAX [8] matches the performance on top of PyTorch [41].
Training follows the URLB benchmark [32], where an agent is pretrained for 2 million steps and then finetuned on a downstream task for 100 k steps.
A.2 Baseline Implementation
For the baselines that were presented in URLB, we obtained the results by running the code from the
URLB GitHub repo: https://github.com/rll-research/url_benchmark .
All hyperparameters were kept the same as the original implementation.
A.3 Environment
For continuous control domains, we used the custom Deep Mind Control Suite [52] environments from the URL Benchmark GitHub repo: https://github.com/rll-research/url_benchmark .
For the VizDoom domain, we used code from their official GitHub repo: https://github.com/ mwydmuch/ViZDoom .
We include the environment renders in Figure 2.
A.4 Double DQN
We made modifications to MOSS to evaluate in discrete action settings.
Specifically, we adopted
Double-DQN [54] as the backbone reinforcement learning algorithm, which has the Q-value target as
[Equation]
where θ t , θ − t are online network parameters and the target network parameters at time t , respectively.
Furthermore, since ViZDoom [30] has a smaller action space, we provided the agent with a discrete number of skill embeddings, similar to DIAYN [17].
Moreover, since performing CPC loss [31] in high dimensional pixel space is not ideal, and to save the number of parameters, we use the CNN backbone of the agent’s DQN network to project observations into state vectors, then use an MLP identical to the continuous action setting to calculate the CPC loss [31] using the discrete set of skill embeddings.
Tab.
4 details the hyper-parameters used for Double DQN and MOSS in the ViZDoom environment.
A.5 Compute Resources
All experiments were run on an internal cluster with 8 NVIDIA A100 GPU and AMD EPYC64-Core Processor.
Pretraining MOSS takes roughly 5 hours.
While finetuning takes roughlymins.
A.6 Thing that did not work: Skill Distribution
In MOSS, we maintain two different skill distributions, i.e. , Z max ∼ P ( Z | M = 0) and Z min ∼
P ( Z | M = 1) .
In particular, we defined them as Z max ∼ U d (0 , 1) and Z min ∼ U d ( − 1 , 0) .
Given a skill of dimension d , another way to define P ( Z | M = 0) and P ( Z | M = 1) is to allocate half of a skill vector for m =
and the other half for m = 1 .
In other words, Z max and Z min are both drawn from U d (0 , 1) , however, when m = 0 , ( Z max ) : d/ 2 ∼ P ( Z
| M = 0) while ( Z max ) d/ 2: = 0 .
Conversely, when m = 1 , ( Z min ) d/ 2: ∼ P ( Z | M = 0) while ( Z max ) : d/ 2 = 0 .
B Objective Switching
[Table]
Since our framework has two objectives, the reinforcement learning agent requires collecting experience to train its conditional policy under both proposed objectives.
Moreover, the temporal structure of staying in different modes of behavior in humans and animals is not monolithic; therefore, designing when to switch modes for the agent might be an interesting area to consider.
For example, previous works investigated a similar setting and used a threshold based on Q-values [44].
However, in unsupervised reinforcement learning, without any sense of the possible downstream task distribution, we wish to design a mode switching mechanism that is both adaptive and scale-invariant across environments.
Literature in active learning proposes a strategy to query high uncertainty samples.
We adopt this framework to MOSS.
Intuitively, we wish to encourage the agent to collect more trajectories at places with higher epistemic uncertainty for the agent to learn more in a more unfamiliar situation.
Since we do not wish to train additional networks, we use the online critic network to output a sample variance to act as a proxy for uncertainty.
Concretely, for a given mode m , we sample a batch N of skill-vectors z ( i ) m ∼ p ( z | m ) .
We then calculate the sample variance as
[Equation]
where ¯ Q πm is the sample mean Q-value for the batch.
However, since variance is a measurement with units, and we do not wish to introduce additional environment-specific hyperparameters, we use the history of variance and the current Q-value variance as function inputs to select modes, bypassing the unit sensitivity problem across environments.
Concretely, we keep a record of the sample variance of both modes at each time step t , denoted as V t,M =0 for maximization skills and V t,M =1 for minimization skills.
At each skill switching interval, the agent selects the mode m t for the current time step by,
[Equation]
otherwise , where β is a coefficient greater than 1 that controls the threshold of switching modes, and k is the interval we use to switch skills and modes.
Intuitively, we wish to switch to the opposite mode if the opposite mode’s uncertainty at time t is greater than a coefficient times the last recorded uncertainty.
We show some preliminary results for the quadruped domain in Tab. 5.
The effectiveness of adaptive methods is demonstrated in the results, where MOSS adaptive was able to beat out MOSS.
However, this method requires memory and computational resources and we found that the hyperparameter β is somewhat sensitive across domains and requires tuning.
Therefore, we leave additional investigations for mode switching to future work because presenting an efficient and hyperparameter-insensitive method of maximization and minimization of surprise was the main scope of this paper.
C Additional Results C.1 Prior Distribution P ( M )
In MOSS [Table] we deterministically set M = 0 for the first half of the steps and M = 1 for the other half of the steps in the episode.
This corresponds to having 50% of maximization data and 50% of minimization data.
We report additional results on Walker for different ratios of maximization and minimization in Tab. [Table]
[Table]
We observe that a prior towards entropy maximization results in a better performance than a prior towards entropy minimization on Walker.
[Table]
C.2 Skill Prior
We present additional results on the dimensionality of the skill vector along with the event space (discrete vs continuous) in Tab. 7.
We find that a continuous skill performs better than a discrete skill but not by a large margin.
However, it seems that in general, a higher dimensional skill performs better, especially in Quadruped which has a higher dimensional state and action space than the Walker domain.
Our results suggest that the optimal skill vector may be task-dependent. [Table] [Table]
C.3 Zero-shot Results
We show zero-shot results of different URL methods in Tab. 8.
In addition, Tab. 9 provides zero-shot results for both Z max and Z min .
C.4 Finetune Learning Curves
We include the finetune learning curves in Figure 5 for MOSS and two methods that are most related to MOSS, namely, CIC and Adversarial Surprise.
We can see that these methods generally improved with finetuning with task reward.
However, we observe that the Quadruped domain had a noticeably flatter learning curve, meaning that methods benefited less from training than in other domains.