Contrastive Intrinsic Control for Unsupervised Reinforcement Learning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We introduce Contrastive Intrinsic Control (CIC), an unsupervised reinforcementlearning (RL) algorithm that maximizes the mutual information between statetransitions and latent skill vectors.
CIC utilizes contrastive learning betweenstate-transitions and skills vectors to learn behaviour embeddings and maximizesthe entropy of these embeddings as an intrinsic reward to encourage behaviouraldiversity.
We evaluate our algorithm on the Unsupervised RL Benchmark (URLB)in the asymptotic state-based setting, which consists of a long reward-free pretraining phase followed by a short adaptation phase to downstream tasks withextrinsic rewards.
We ﬁnd that CIC improves over prior exploration algorithms interms of adaptation efﬁciency to downstream tasks on state-based URLB.
Deep RL is a powerful approach toward solving complex control tasks in the presence of extrinsicrewards.
Successful applications include playing video games from pixels [1], mastering the game of
Go [2, 3], robotic locomotion [4, 5, 6] and dexterous manipulation [7, 8, 9] policies.
While effective,the above advances produced agents that are unable to generalize to new downstream tasks beyond theone they were trained to solve.
Humans and animals on the other hand are able to acquire skills withminimal supervision and apply them to solve a variety of downstream tasks.
In this work, we seekto train agents that acquire skills without supervision with generalization capabilities by efﬁcientlyadapting these skills to downstream tasks.
Over the last few years unsupervised RL has emerged as a promising framework for developing RLagents that can generalize to new tasks.
In the unsupervised RL setting, agents are ﬁrst pre-trained with self-supervised intrinsic rewards and then ﬁnetuned to downstream tasks with extrinsic rewards.
Unsupervised RL algorithms broadly fall into three categories - knowledge-based, data-based, andcompetence-based methods 2 .
Knowledge-based methods maximize the error or uncertainty of apredictive model [12, 13, 14].
Data-based methods maximize the entropy of the agent’s visitation [15,16].
Competence-based methods learn skills that generate diverse behaviors [17, 18].
This work fallsinto the latter category of competence-based exploration methods.
Unlike knowledge-based and data-based algorithms, competence-based algorithms simultaneouslyaddress both the exploration challenge as well as distilling the generated experience in the formof reusable skills.
This makes them particularly appealing, since the resulting skill-based policies(or skills themselves) can be ﬁnetuned to efﬁciently solve downstream tasks.
While there are manyself-supervised objectives that can be utilized, our work falls into a family of methods that learns skillsby maximizing the mutual information between visited states and latent skill vectors.
Many earlierworks have investigated optimizing such objectives [17, 18, 19, 20].
However, competence-basedmethods have been empirically challenging to train and have under-performed when compared toknowledge and data-based methods [21].
[Figure]
In this work, we take a closer look at the challenges of pre-training agents with competence-basedalgorithms.
We introduce Contrastive Intrinsic Control (CIC) – an exploration algorithm that usesa new estimator for the mutual information objective.
CIC combines particle estimation for stateentropy [22, 15] and noise contrastive estimation [23] for the conditional entropy which enables itto both generate diverse behaviors (explore) and discriminate high-dimensional continuous skills(exploit) .
To the best of our knowledge, CIC is the ﬁrst exploration algorithm to utilize noisecontrastive estimation to discriminate between state transitions and latent skill vectors.
Empirically,we show that CIC adapts to downstream tasks more efﬁciently than prior exploration approacheson the state-based Unsupervised Reinforcement Learning Benchmark (URLB).
CIC achieves 79%higher returns on downstream tasks than prior competence-based algorithms and 18% higher returnsthan the next-best exploration algorithm overall.
Background and Notation
Markov Decision Process: We operate under the assumption that our system is described by a
Markov Decision Process (MDP) [24].
An MDP consiss of the tuple ( S , A , P , r,  ) which has states 49 s 2 S , actions a 2 A , transition dynamics p
( s 0 | s, a ) ⇠ P , a reward function r , and a discount 50 factor  .
In an MDP, at each timestep t , an agent observes the current state s , selects an action froma policy a ⇠ ⇡ ( ·| s ) , and then observes the reward and next state once it acts in the {+environment: 52 r, s 
 ⇠ env.step ( a ) .+}
Note that usually r refers to an extrinsic reward.
However, in this work we willﬁrst be pre-training an agent with intrinsic rewards r int and ﬁnetuning on extrinsic rewards r ext .
For convenience we also introduce the variable ⌧ ( s ) which refers to any function of the states s . Forinstance ⌧ can be a single state, a pair of states, or a sequence depending on the algorithm. Ourmethod uses ⌧ =
( s, s 0 ) to encourage diverse state transitions while other methods have differentspeciﬁcations for ⌧ .
Importantly, ⌧ does not denote a state-action trajectory, but is rather shorthandfor any function of the states encountered by the agent.
In addition to the standard MDP notation, wewill also be learning skills z 2 Z and our policy will be skill-conditioned a ⇠ ⇡ ( ·| s, z ) .
Unsupervised Skill Discovery through Mutual Information Maximization: Most competencebased approaches to exploration maximize the mutual information between states and skills.
Our workand a large body of prior research [17, 20, 18, 25, 26, 27] aims to maximize a mutual informationobjective with the following general form:
[Equation]
Competence-based algorithms use different choices for ⌧ and can condition on additional informationsuch as actions or starting states.
For a full summary of competence-based algorithms and theirobjectives see Table 2 in Appendix D.
Lower Bound Estimates of Mutual Information: The mutual information I ( s ; z ) is intractable tocompute directly.
Since we wish to maximize I ( s ; z ) , we can approximate this objective by instead
[Figure]
maximizing a lower bound estimate.
Most known mutual information maximization algorithms usethe variational lower bound introduced in [28]:
[Equation]
The variational lower bound can be applied to both decompositions of the mutual information.
Thedesign decisions of a competence-based algorithm therefore come down to (i) which decompositionof I ( ⌧ ; z ) to use, (ii) whether to use discrete or continuous skills, (iii) how to estimate H ( z ) or H ( ⌧ ) ,and ﬁnally (iv) how to estimate H ( z | ⌧ ) or H ( ⌧ | z ) .
Motivation
Results from the recent Unsupervised Reinforcement Learning Benchmark (URLB) [21] show thatcompetence-based approaches underperform relative to knowledge-based and data-based baselineson DeepMind Control (DMC).
We argue that the underlying issue with current competence-basedalgorithms when deployed on harder exploration environments like DMC has to do with the currentlyused estimators for I ( ⌧ ; z ) rather than the objective itself.
To produce structured skills that lead todiverse behaviors, I ( ⌧ ; z ) estimators must (i) explicitly encourage diverse behaviors and (ii) havethe capacity to discriminate between high-dimensional continuous skills.
Current approaches do notsatisfy both criteria.
Competence-base algorithms do not ensure diverse behaviors: Most of the best known competencebased approaches [17, 18, 25, 26], optimize the ﬁrst decomposition of the mutual information
H ( z )  H ( z | ⌧ ) .
The issue with this decomposition is that while it ensures diversity of skill vectors 88 it does not ensure diverse behavior from the policy, meaning max
H ( z ) does not imply max H ( ⌧ ) .
Of course, if H ( z )  H ( z | ⌧ ) is maximized and the skill dimension is sufﬁciently large, then H ( ⌧ )will also be maximized implicitly.
Yet in practice, to learn an accurate discriminator q ( z | ⌧ ) , theabove methods assume skill spaces that are much smaller than the state space (see Table 2), and thusbehavioral diversity may not be guaranteed.
In contrast, the decomposition I ( ⌧ ; z )
= H ( ⌧ )  H ( ⌧ | z )ensures diverse behaviors through the entropy term H ( ⌧ ) .
Methods that utilize this decompositioninclude [27, 20].
Why it is important to utilize high-dimensional skills: Once a policy is capable of generating diversebehaviors, it is important that the discriminator can distill these behaviors into distinct skills.
If the setof behaviors outnumbers the set of skills, this will result in degenerate skills – when one skill mapsto multiple different behaviors.
It is therefore important that the discriminator can accommodatecontinuous skills of sufﬁciently high dimension.
Empirically, the discriminators used in prior workutilize only low-dimensional continuous skill vectors.
DIAYN [17] utilized 16 dimensional skills,
DADS [20] utilizes continuous skills of dimension 2  5 , while APS [27], an algorithm that utilizessuccessor features [30, 31] for the discriminator, is only capable of learning continuous skills withdimension 10 .
We show how small skill spaces can lead to ineffective exploration in a simplegridworld setting in Appendix H and evidence that skill dimension affects performance in Fig.
[Figure]
On the importance of benchmarks for evaluation: While prior competence-based approaches suchas DIAYN [17] were evaluated on OpenAI Gym [32], Gym environment episodes terminate whenthe agent loses balance thereby leaking some aspects of extrinsic signal to the exploration agent.
On the other hand, DMC episodes have ﬁxed length.
We show in Fig 3 that this small difference inenvironments results in large performance differences.
Speciﬁcally, we ﬁnd that DIAYN is able tolearn diverse skills in Gym but not in DMC, which is consistent with both observations from DIAYNand URLB papers.
Due to ﬁxed episode lengths, DMC tasks are harder for reward-free explorationsince agents must learn to balance without supervision.
Method
3.1 Contrastive Intrinsic Control
From Section 2 we are motivated to ﬁnd a lower bound for I ( ⌧ ; z ) with a discriminator that is capableof supporting high-dimensional continuous skills 3 .
Additionally, we wish to increase the diversity ofbehaviors so that the discriminator can continue learning new skills throughout training.
We choosethe forward decomposition of MI I ( ⌧ ; z ) =
H ( ⌧ )  H ( ⌧ | z ) similar to [27] and estimate the lowerbound with Eq. 2.
The entropy is estimated H ( ⌧ ) with a particle-based estimator similar to [15].
Assuch, the primary technical contribution of this work is a novel estimator for the discriminator.
To improve the discriminator, we propose to utilize noise contrastive estimation (NCE) [23] betweenstate-transitions and latent skills as a lower bound for I ( ⌧ ; z ) .
4 It has been shown previously thatsuch estimators provide a valid lower bound for mutual information [33].
However, to the best of ourknowledge, this is the ﬁrst work to investigate contrastive representation learning for intrinsic control.
Representation Learning:
Speciﬁcally, we propose to learn embeddings by parameterizing thediscriminator with a contrastive density estimator.
This is a novel choice that differs from prior workswhich utilize a classiﬁer [17] or non-contrastive density estimation [20].
[Equation]
where f ( ⌧ , z ) is any real valued function.
For our practical algorithm, we parameterize this function as f ( ⌧ , z ) =
g 1 ( ⌧ ) > g 2 ( z ) / k g 1 ( ⌧ ) kk g 2 ( z ) k T where ⌧ =
( s, s 0 ) is a transition tuple, g k are neural encoders, and T is a temperature parameter.
This inner product is similar to the one used in
SimCLR [34].
The representation learning loss backpropagates gradients from the NCE loss which maximizessimilarity between state-transitions and corresponding skills.
[Equation]
We provide pseudocode for the CIC representation learning loss:
[Figure]
[Equation]
[Figure]
PyTorch -like pseudocode for the CIC loss
[Equation]
[Figure]
{+5 def cic_loss(s, s_next , z, temp): 
+}
[Equation]
142states: s, s_next (B, D), skills: z (B, D) {+
+}
[Equation]
144 9 tau = concat(s, s_next , dim=1) 146 11 query = query_net(z) 148key = key_net(tau) 149 14 query =
normalize(query , dim=1)
151key = normalize(key , dim=1) 152 17 logits =
matmul(query , key.
T) / temp 154labels = arange(logits.shape [0]) # positives are on diagonal
[Figure]
loss = cross_entropy(logits , labels)
# negatives are off diagonal 157 22 return loss Listing 1: Pseudocode for the CIC loss
Intrinsic reward: Although we have a representation learning objective, we still need to specify theintrinsic reward for the algorithm for which there can be multiple choices.
Prior works considerspecifying an intrinsic reward that is proportional to state-transition entropy [15], the discriminator [17], a similarity score between states and skills [36], and the uncertainty of the discriminator [37].
We investigate each of these choices in Section 6 and ﬁnd that an intrinsic reward that maximizesstate-transition entropy coupled with representation learning via the CPC loss deﬁned in Sec.
3.1 isthe simplest variant that also performs well (see Table 1). [Figure]
For the intrinsic reward, we use a particle estimate [22, 38] as in [15] of the state-transition entropy.
Similar to [15, 16] we estimate the entropy up to a proportionality constant, because we want theagent to maximize entropy rather than estimate its exact value.
The APT particle entropy estimate is proportional to the distance between the current visited statetransition and previously seen neighboring points.
[Equation]
[Equation]
where h i is an embedding of ⌧ i shown in Fig.
h
⇤ i is a kNN embedding, N k is the number of kNNs,and N  1 is the number of negatives.
The total number of elements in the summation is N becauseit includes one positive.
Explore and Exploit: With these design choices the two components of the CIC algorithm canbe interpreted as exploration with intrinsic rewards and exploitation using representation learningto distill behaviors into skills.
The marginal entropy maximizes the diversity of state-transitionembeddings while the contrastive discriminator log q ( ⌧ | z ) encourages exploitation by ensuring thatskills z lead to predictable states ⌧ .
Together the two terms incentivize the discovery of diverseyet predictable behaviors from the RL agent.
While CIC shares a similar intrinsic reward structureto APT [15], we show that the new representation learning loss from the CIC estimator results insubstantial performance gains in Sec 6.
Practical Implementation
Our practical implementation consists of two main components: the RL optimization algorithm andthe CIC architecture.
For fairness and clarity of comparison, we use the same RL optimizationalgorithm for our method and all baselines in this work.
Since the baselines implemented in
URLB [21] use a DDPG 5 [29] as their backbone, we opt for the same DDPG architecture to optimizeour method as well (see Appendix B).
CIC Architecture: We use a particle estimator as in [15] to estimate H ( ⌧ ) .
To compute the variationaldensity q ( ⌧ | z ) , we ﬁrst sample skills from uniform noise z ⇠ p ( z ) where p ( z ) is the uniformdistribution over the [0 , 1] interval.
We then use two MLP encoders to embed g 1 ( ⌧ ) and g 2 ( z ) ,and optimize the parameters 1 , 2 with the CPC loss similar to SimCLR [34] since f ( ⌧ , z ) =
g 1 ( ⌧ ) T g 2 ( z ) .
We ﬁx the hyperparameters across all domains and downstream tasks.
We refer thereader to the Appendices E and F for the full algorithm and a full list of hyperparameters.
Adapting to downstream tasks: To adapt to downstream tasks we follow the same procedure forcompetence-based method adaptation as in URLB [21].
During the ﬁrst 4k environment interactionswe populate the DDPG replay buffer with samples and use the extrinsic rewards collected duringthis period to ﬁnetune the skill vector z .
While it’s common to ﬁnetune skills with Cross Entropy
Adaptation (CMA), given our limited budget of 4k samples (only 4 episodes) we ﬁnd that a simplegrid sweep of skills over the interval [0 , 1] produces the best results (see Fig. 5).
After this, we ﬁx theskill z and ﬁnetune the DDPG actor-critic parameters against the extrinsic reward for the remaining96k steps.
Note that competence-based methods in URLB also ﬁnetune their skills during the ﬁrst 4kﬁnetuning steps ensuring a fair comparison between the methods.
The full adaptation procedure isdetailed in Appendix E.
Experimental Setup
Environments We evaluate our approach on tasks from URLB, which consists of twelve downstreamtasks across three challenging continuous control domains for exploration algorithms – walker,quadruped, and Jaco arm.
Walker requires a biped constrained to a 2D vertical plane to performlocomotion tasks while balancing.
Quadruped is more challenging due to a higher-dimensionalstate-action space and requires a quadruped to in a 3D environment to learn locomotion skills.
Jacoarm is a 6-DOF robotic arm with a three-ﬁnger gripper to move and manipulate objects withoutlocking.
All three environments are challenging in the absence of an extrinsic reward.
Baselines: We implemented CIC using the URLB [21] codebase 6 and compare CIC to baselinesincluded in URLB across all three exploration categories.
Knowledge-based basedlines include
ICM [12], Disagreement [13], and RND [14].
Data-based baselines incude APT [15] and ProtoRL [16].
Competence-based baselines include DIAYN [17], SMM [26], and APS [27].
The closestbaselines to CIC are APT, which is similar to CIC but without state-skill CPC representation learning(no discriminator), and APS which uses the same decomposition of mutual information as CIC andalso uses a particle entropy estimate for H ( ⌧ ) .
The main difference between APS and CIC is that
APS uses successor features while CIC uses a contrastive estimator for the discriminator.
For furtherdetails regarding baselines we refer the reader to Appendix C.
Evaluation: We follow an identical evaluation to the 2M pre-training setup in URLB.
First, wepre-train each RL agent with the intrinsic rewards for 2M steps.
Then, we ﬁnetune each agent tothe downstream task with extrinsic rewards for 100k steps.
All baselines were run for 10 seeds perdownstream task for each algorithm using the code and hyperparameters provided by URLB [21].
Built on top of URLB, CIC is also run for 10 seeds per task.
A total of 1080 = 9 algorithms ⇥12 tasks ⇥ 10 seeds experiments were run for the main results.
Importantly, all baselines and CICuse a DDPG agent as their backbone.
To ensure that our evaluation statistics are unbiased we use stratiﬁed bootstrap conﬁdence intervals toreport aggregate statistics across M runs with N seeds as described in Rliable [35] to report statisticsfor our main results in Fig.
Our primary success metric is the interquartile mean (IQM) and the
Optimality Gap (OG).
IQM discards the top and bottom 25 % of runs and then computes the mean.
It is less susceptible to outliers than the mean and was shown to be the most reliable statistic forreporting results for RL experiments in [35].
OG measures how far a policy is from optimal (expert)performance.
To deﬁne expert performance we use the convention in URLB, which is the scoreachieved by a randomly initialized DDPG after 2M steps of ﬁnetuning (20x more steps than ourﬁnetuning budget).
Results
We investigate empirical answers to the following research questions: (Q1)
How does CIC adaptationefﬁciency compare to prior competence-based algorithms and exploration algorithms more broadly?(Q2)
Which intrinsic reward instantiation of CIC performs best?
(Q3)
How do the two terms in the
CIC objective affect algorithm performance?
(Q4)
How does skill selection affect the quality of thepre-trained policy?
(Q5)
Which architecture details matter most?
Adaptation efﬁciency of CIC and exploration baslines: Expert normalized scores of CIC andexploration algorithms from URLB are shown in Fig.
We ﬁnd that CIC substantially outperformsprior competence-based algorithms (DIAYN, SMM, APS) achieving a 79% higher IQM than the nextbest competence-based method (APS) and, more broadly, achieving a 18% higher IQM than the nextbest overall baseline (ProtoRL).
In further ablations, we ﬁnd that the contributing factors to CIC’sperformance are its ability to accommodate substantially larger continuous skill spaces than priorcompetence-based methods.
Intrinsic reward speciﬁcation: The intrinsic reward for competence-based algorithms can beinstantiated in many different ways.
Here, we analyze intrinsic reward for CIC with the formr int = H ( ⌧ ) + D ( ⌧ , z ) , where D is some function of ( ⌧ , z ) .
Prior works, select D to be (i)the discriminator [27], (ii) a cosine similarity between embeddings [36], (iii) uncertainty of thediscriminator [37], and (iv) just the entropy D ( ⌧ , z ) =
We run CIC with each of these variantson the walker and quadruped tasks and measure the ﬁnal mean performance across the downstreamtasks (see Tab. 1).
The results show that the entropy-only intrinsic reward performs best followed byan uncertainty-based intrinsic reward.
We hypothesize that the reason why a simple entropy-onlyintrinsic reward works well is that state-skill CPC representation learning clusters similar behaviorstogether.
Since similar behaviors are clustered, maximizing the entropy of state-transition embeddingsproduces increasingly diverse behaviors.
[Table]
The importance of representation learning: To what extent does representation learning with thestate-skill CIC loss affect the agent’s exploration capability?
To answer this question we train the CICagent with the entropy intrinsic reward with and without the representation learning auxiliary loss for2M steps.
The zero-shot reward plotted in Fig.
6 indicates that without representation learning thepolicy collapses.
With representation learning, the agent is able to discover diverse skills evidencedby the non-zero reward.
This result suggests that state-skill CPC representation learning is a criticalpart of CIC.
Qualitative analysis of CIC behaviors: Qualitatively, we ﬁnd that CIC is able to learn locomotionbehaviors in DMC without extrinsic information such as early termination as in OpenAI Gym.
Whilemost skills are higher entropy and thus more chaotic, we show in Fig 1 that structured behaviors canbe isolated by ﬁxing a particular skill vector.
For example, in the walker and quadruped domainsbalancing, walking, and ﬂipping skills can be isolated.
For more qualitative investigations we referthe reader to Appendix I.
Skill architecture and adaptation ablations:
We ﬁnd that projecting the skill to a latent spacebefore inputting it as the key for the contrastive loss is an important design decision (see Fig.
5a),most likely because this reduces the diversity of the skill vector making the discriminator task simpler.
We also ﬁnd empirically that the skill dimension is an important hyperparameter and that larger skillsresults in better zero-shot performance (see Fig.
5b), which empirically supports the hypothesis posedin Section 2 and Appendix H that larger skill spaces are important for internalizing diverse behaviors.
Interestingly, CIC zero-shot performance is poor in lower skill dimensions (e.g. dim ( z ) < 10 ),suggesting that when dim ( z ) is small CIC performs no better than prior competence-based methodssuch as DIAYN, and that scaling to larger skills enables CIC to pre-train effectively.
To measure the effect of skill ﬁnetuning described in Section 4, we sweep mean skill values along theinterval of the uniform prior [0 , 1] with a budget of 4k total environment interactions and read out the performance on the downstream task.
By sweeping, we mean simply iterating over the interval[0 , 1] with ﬁxed step size (e.g. v = 0 , 0 . 1 , . . .
and setting z i = v for all i .
This is not anoptimal skill sampling strategy but works well due to the extremely limited number of samples forskill selection.
We evaluate this ablation on the Quadruped Stand and Run downstream tasks.
Theresults shown in Fig.
5 indicate that skill selection can substantially affect zero-shot downstream taskperformance.
[-DADS maximizes I ( s 0 | s, z ) and estimates entropy H ( s 0 | s )by marginalizing over z such that H ( s 
-]
Limitations and Impact
While CIC achieves leading results on state-based URLB, we would also like to address its limitations.
First, in this paper we only consider MDPs (and not partially observed MDPs) where the full state isobservable.
We focus on MDPs because generating diverse behaviors in environments with largestate spaces has been the primary bottleneck for competence-based exploration.
Combining CICwith visual representation learning to scale this method to pixel-based inputs is a promising futuredirection for research not considered in this work.
One issue with unsupervised RL algorithms (and hence CIC) in terms of potentially negative societalimpact is that self-supervised exploration can be dangerous.
Since self-supervised agents maximizeintrinsic rewards, this can lead to destructive behavior.
For example, when deploying CIC on a Walkeror Quadruped robot it learns chaotic exploration behaviors 7 that would most likely break the robotin real-world settings.
Alignment of exploration agents to prevent them from learning dangerouspolicies is a promising direction for future work.
Conclusion
We have introduced a new competence-based algorithm – Contrastive Intrinsic Control (CIC) – which enables more effective exploration than prior unsupervised skill discovery algorithms by explicitly encouraging diverse behavior while distilling predictable behaviors into skills with a contrastive discriminator.
We showed that CIC is the ﬁrst competence-based approach to achieve leading performance on URLB.
We hope that this encourages further research in developing RL agents capable of generalization