Contrastive Learning as Goal-Conditioned Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
In reinforcement learning (RL), it is easier to solve a task if given a good representation.
While deep RL should automatically acquire such good representations,prior work often finds that learning representations in an end-to-end fashion isunstable and instead equip RL algorithms with additional representation learningparts (e.g., auxiliary losses, data augmentation).
How can we design RL algorithms that directly acquire good representations?
In this paper, instead of addingrepresentation learning parts to an existing RL algorithm, we show (contrastive)representation learning methods can be cast as RL algorithms in their own right.
To do this, we build upon prior work and apply contrastive representation learningto action-labeled trajectories, in such a way that the (inner product of) learned representations exactly corresponds to a goal-conditioned value function.
We use thisidea to reinterpret a prior RL method as performing contrastive learning, and thenuse the idea to propose a much simpler method that achieves similar performance.
Across a range of goal-conditioned RL tasks, we demonstrate that contrastive RLmethods achieve higher success rates than prior non-contrastive methods.
We alsoshow that contrastive RL outperforms prior methods on image-based tasks, withoutusing data augmentation or auxiliary objectives.
Introduction
Representation learning is an integral part of reinforcement learning (RL 2 ) algorithms.
While suchrepresentations might emerge from end-to-end training [7, 74, 113, 121], prior work has found itnecessary to equip RL algorithms with perception-specific loss functions [29, 40, 66, 84, 86, 95, 110,135] or data augmentations [64, 68, 110, 112], effectively decoupling the representation learningproblem from the reinforcement learning problem.
Given what prior work has shown about RLin the presence of function approximation and state aliasing [2, 130, 133], it is not surprising thatend-to-end learning of representations is fragile [64, 68]: an algorithm needs good representationsto drive the learning of the RL algorithm, but the RL algorithm needs to drive the learning of goodrepresentations.
So, can we design RL algorithms that do learn good representations, without theneed for auxiliary perception losses?
Rather than using a reinforcement learning algorithm to also solve a representation learning problem,we will use a representation learning algorithm to also solve certain types of reinforcement learningproblems, namely goal-conditioned RL.
Goal-conditioned RL is widely studied [6, 13, 20, 58,75 , 114], and intriguing from a representation learning perspective because it can be done in anentirely self-supervised manner, without manually-specified reward functions.
We will focus on
[Figure]
" state-
[Figure]
encoder, .
goal encoder,
[Figure]
contrastive (representation) learning methods, using observations from the same trajectory (as donein prior work [90, 103]) while also including actions as an additional input (See Fig. 1).
Intuitively,contrastive learning then resembles a goal-conditioned value function: nearby states have similarrepresentations and unreachable states have dissimilar representations.
We make this connectionprecise, showing that sampling positive pairs using the discounted state occupancy measure results inlearning representations whose inner product exactly corresponds to a value function.
In this paper, we show how contrastive representation learning can be used to perform goalconditioned RL.
We formally relate the learned representations to reward maximization, showingthat the inner product between representations corresponds to a value function.
This framework ofcontrastive RL generalizes prior methods, such as C-learning [27], and suggests new goal-conditioned
RL algorithms.
One new method achieves performance similar to prior methods but is simpler;another method consistently outperforms the prior methods.
On goal-conditioned RL tasks withimage observations, contrastive RL methods outperform prior methods that employ data augmentationand auxiliary objectives, and do so without data augmentation or auxiliary objectives.
Related Work
This paper will draw a connection between RL and contrastive representation learning, buildingupon a long line of contrastive learning methods in NLP and computer vision, and deep metriclearning [15, 49, 50, 50, 52, 72, 79, 81, 82, 89, 90, 102, 103, 107, 117, 124, 127].
Contrastive learningmethods learn representations such that similar (“positive”) examples have similar representationsand dissimilar (“negative”) examples have dissimilar representations.
3 While most methods generatethe “positive” examples via data augmentation, some methods generate similar examples usingdifferent camera viewpoints of the same scene [103, 117], or by sampling examples that occur closein time within time series data [4, 90, 103, 112].
Our analysis will focus on this latter strategy, as thedependence on time will allow us to draw a precise relationship with the time dependence in RL.
Deep RL algorithms promise to automatically learn good representations, in an end-to-end fashion.
However, prior work has found it challenging to uphold this promise [7, 74, 113, 121], promptingmany prior methods to employ separate objectives for representation learning and RL [29, 40, 66,84, 86, 95, 110, 112, 135].
Many prior methods choose a representation learning objectives thatreconstruct the input state [29, 43, 45, 46, 66, 86, 88, 136] while others use contrastive representationlearning methods [84, 90, 105, 110, 112].
Unlike these prior methods, we will not use a separaterepresentation learning objective, but instead use the same objective for both representation learningand reinforcement learning.
Some prior RL methods have also used contrastive learning to acquirereward functions [12, 18, 30, 35, 59, 63, 87, 128, 129, 140], often in imitation learning settings [34,
In contrast, we will use contrastive learning to directly acquire a value function, which (unlike areward function) can be used directly to take actions, without any additional RL.
This paper will focus on goal-conditioned RL problems, a problem prior work has approached usingtemporal difference learning [6, 27, 58, 75, 97, 100], conditional imitation learning [20, 37, 78, 99,114], model-based methods [21, 101], hierarchical RL [85], and planning-based methods [28, 88,99, 109].
The problems of automatically sampling goals and exploration [22, 32, 80, 93, 138] areorthogonal to this work.
Like prior work, we will parametrize the value function as an inner productbetween learned representations [31, 54, 100].
Unlike these prior methods, we will learn a valuefunction directly via contrastive learning, without using reward functions or TD learning.
Our analysis will be most similar to prior methods [9, 13, 27, 97] that view goal-conditioned RLas a data-driven problem, rather than as a reward-maximization problem.
Many of these methodsemploy hindsight relabeling [6, 24, 58, 73], wherein experience is relabeled with an outcome thatoccurred in the future.
Whereas hindsight relabeling is typically viewed as a trick to add on top ofan RL algorithm, this paper can roughly be interpreted as showing that the hindsight relabeling is astandalone RL algorithm.
Many goal-conditioned methods learn a value function that captures thesimilarity between two states [27, 58, 86, 120].
Such distance functions are structurally similar to thecritic function learned for contrastive learning, a connection we make precise in Sec. 4.
In fact, ouranalysis shows that C-learning [27] is already performing contrastive learning, and our experimentsshow that alternative contrastive RL methods can be much simpler and achieve higher performance.
Prior work has studied how representations related to reward functions using the framework ofuniversal value functions [10, 100] and successor features [8, 48, 76].
While these methods typicallyrequire additional supervision to drive representation learning (manually-specified reward functionsor features), our method is more similar to prior work that estimates the discounted state occupancymeasure as an inner product between learned representations [9, 126].
While these methods usetemporal difference learning, ours is akin to Monte Carlo learning.
While Monte Carlo learning isoften (but not always [21]) perceived as less sampling efficient, our experiments find that our approachcan be as sample efficient as TD methods.
Other prior work has focused on learning representationsthat can be used for planning [55, 77, 98, 99, 123].
Our method will learn representations usingan objective similar to prior work [99, 103], but makes the key observation that the representationalready encodes a value function: no additional planning or RL is necessary to choose actions.
Please see Appendix A for a discussion of how our work relates to unsupervised skill learning.
Preliminaries
Goal-conditioned reinforcement learning.
The goal-conditioned RL problem is defined by states 99 s t ∈ S , actions a t , an initial state distribution p 0 ( s ) , the dynamics p ( s t +1 | s t , a t ) , a distribution 100 over goals p g ( s g ) , and a reward function r g ( s, a ) for each goal.
This problem is equivalent to a 101 multi-task RL [5, 41, 116, 125, 134], where tasks correspond to reaching goals states.
Followingprior work [9, 13, 27, 97], we define the reward as the probability (density) of reaching the goal atthe next time step:
[Equation]
This reward function is appealing because it avoids the need for a human user to specify a distance 105 metric (unlike, e.g., [6]).
Even though our method will not estimate the reward function, we will stilluse the reward function for analysis.
For a goal-conditioned policy π ( a | s, s g ) , we use π ( τ | s t ) todenote the probability of sampling an infinite-length trajectory τ =
( s 0 , a 0 , s 1 , a 1 , · · · ) .
We defined 108 the expected reward objective and Q-function as
[Equation]
Intuitively, this objective corresponds to sampling a goal s g and then optimizing the policy to go to 110 that goal and stay there.
Finally, we define the discounted state occupancy measure as [51, 137] [Equation] [Equation]
[Equation]
where p πt ( s ) is the probability density over states that policy π visits after t steps.
Sampling from thediscounted state occupancy measure is easy: first sample a time offset from a geometric distribution( t ∼ G EOM (1 − γ ) ), and then look at what state the policy visits after exactly t steps.
We will uses t + to denote states sampled from the discounted state occupancy measure.
Because our method willcombine experience collected from multiple policies, we also define the average stationary distributionas p π ( ·|· ) ( s t + = s | s, a ) ≜ (cid:82) p π ( ·|· ,s g ) ( s t
+ = s | s, a ) p
π ( s g | s, a ) ds g , where p π ( s g | s, a ) is theprobability of the commanded goal given the current state-action pair.
This stationary distribution isequivalent to that of the policy π ( a | s ) ≜ (cid:82) π ( a | s, s g )
p π ( s g | s ) ds g
Contrastive representation learning.
Contrastive representation learning methods [15, 42, 49, 50, 120 57, 72, 79, 81, 82, 117, 119, 124] take as input pairs of positive and negative examples, and learnrepresentations so that positive pairs have similar representations and negative pairs have dissimilar 122 representations.
We use ( u, v ) to denote an input pair (e.g., u is an image and v is an augmentedversion of that image).
Positive examples are sampled from a joint distribution p ( u, v ) while negative 124 examples are sampled from the product of marginal distributions, p ( u ) p ( v ) .
We will use an objectivebased on binary classification [72, 81, 82, 89].
Let f ( u, v ) =
ϕ ( u ) T ψ ( v ) be the similarity betweenthe representations of u and v .
We will call f the critic function 5 and note that its range is ( −∞ , ∞ ) .
We will use NCE-binary [79] objective (also known as InfoMAX [50]):
[Equation]
[Table]
[Table]
[Table]
u )
[Table]
log
[Table]
, v + ) (cid:124)
[Table]
[Table]
+ ) ) + log(1 − σ ( f ( u, v − ) (cid:124) [Table] u ) T
[Table]
− ) ))
[Table]
Contrastive Learning as an RL Algorithm
This section shows how to use contrastive representation to directly perform goal-conditioned RL.
The key idea (Lemma 4.1) is that contrastive learning estimates the Q-function for a certain policyand reward function.
To prove this result, we relate the Q-function to the state occupancy measure(Sec. 4.1) and then relate the optimal critic function to the state occupancy measure (Sec. 4.2).
This result allows us to propose a new algorithm for goal-conditioned RL based on contrastivelearning.
Unlike prior work, this algorithm is not adding contrastive learning on top of an existing RLalgorithm.
This framework generalizes C-learning [27], offering a cogent explanation for its goodperformance while also suggesting new methods that are simpler and can achieve higher performance.
4.1 Relating the Q-function to probabilities
This section sets the stage for the main results of this section by providing a probabilistic perspectivegoal-conditioned RL.
The objective expected reward objective and associated Q-function in (Eq. 5)can equivalently be expressed as the probability (density) of reaching a goal in the future:
Proposition 1 (rewards → probabilities) .
The Q-function for the goal-conditioned reward functionr g (Eq. 1) is equivalent to the probability of state s g under the discounted state occupancy measure:
[Equation]
The proof is in Appendix B. Translating rewards into probabilities not only makes it easier to analyzethe goal-conditioned problem, but also means that any method for estimating probabilities (e.g.,contrastive learning) can be turned into a method for estimating this Q-function.
4.2 Contrastive Learning Estimates a Q-Function
We will use contrastive learning to learn a value function by carefully choosing the inputs u and v .
The first input, u , will correspond to a state-action pair, u = ( s t , a t ) ∼ p ( s, a ) .
In practice, thesepairs are sampled from the replay buffer.
Including the actions in the input is important because it will 151 allow us to determine which actions to take to reach a desired future state.
The second variable, v , is afuture state, v = s f .
For the “positive” training pairs, the future state is sampled from the discountedstate occupancy measure, s f ∼ p π ( ·|· ) ( s t + | s t , a t ) .
The “negative” training pairs, we sample a future state from a random state-action pair: s f ∼ p ( s t + ) ≜ (cid:82) p π ( ·|· ) ( s t + | s, a ) p ( s, a ) dsda .
With theseinputs, the contrastive learning objective (Eq. 4) can be written as
[Equation]
where L ( s, a, s + f , s − f ) ≜ log σ ( f ( s, a, s + f ) (cid:124) (cid:123)(cid:122) (cid:125)
ϕ
( s,a ) T ψ ( s + f ) )
+ log(1 − σ ( f ( s, a, s − f ) (cid:124) (cid:123)(cid:122) (cid:125) ϕ
( s
,a ) T ψ ( s − f ) )) .
Intuitively, the critic function f ( u = ( s t , a t ) , v = s f ) now tells us the correlation between the currentstate-action pair and future outcomes, analogous to a Q-function.
We therefore can use the criticfunction in the same way as actor-critic RL algorithms [62], figuring which actions lead to the desiredoutcome.
Because the Bayes-optimal critic function is a function of the state occupancy measure [79],f ∗ ( s, a, s g ) = log (cid:16)
[Equation]
π [Equation]
[Equation]
s t
+ = s g
| s,a ) p ( s g ) (cid:17) , it can be used to express the Q-function:
Lemma 4.1.
The critic function that optimizes Eq.
6 is a Q-function for the goal-conditioned rewardfunction (Eq. 1), up to a multiplicative constant
[Equation]
The critic function can be viewed as an unnormalized density model, where p ( s g ) is the partitionfunction.
Much of the appeal of contrastive learning is it avoids estimating the partition function [42],which can be challenging; in the RL setting, it will turn out that this constant can be ignored whenselecting actions.
Our experiments show that learning a normalized density model works well whens g is low-dimensional, but struggles to solve higher-dimensional tasks.
This lemma relates the critic function to Q π ( ·|· ) s f ( s, a ) , not Q π ( ·|· ,s f ) s f ( s, a ) .
The underlying reason isthat the critic function combines together experience collected when commanding different goals.
Prior goal-conditioned behavioral cloning methods [20, 37, 78, 114] perform similar sharing, but donot analyze the relationship between the learned policies and Q functions.
Sec.
4.5 shows that thiscritic can be used as the basis for a convergent RL algorithm, under some assumptions.
4.3 Learning the Goal-Conditioned Policy
The learned critic function not only tells us the likelihood of future states, but also tells us how 175 different actions change the likelihood of a state occurring in the future.
Thus, to learn a policy forreaching a goal state, we choose the actions that make that state most likely to occur in the future:
[Equation]
The approximation above reflects errors in learning the optimal critic, and will allow us to prove thatthis policy loss corresponds to policy improvement in Sec.
4.5, under some assumptions.
In practice, we parametrize the goal-conditioned policy as a neural network that takes as input thestate and goal and outputs a distribution over actions.
The actor loss (Eq. 7) is computed by samplingstates and random goals from the replay buffer, sampling actions from the policy, and then takinggradients on the policy using a reparametrization gradient.
On tasks with image observations, we addan action entropy term to the policy objective.
4.4 A Complete Goal-Conditioned RL Algorithm
The complete algorithm alternates between fitting the critic function using contrastive learning,updating the policy using Eq. 7, and collecting more data.
Alg.
1 provides a JAX [11] implementationof the actor and critic losses.
Note that the critic is parameterized as an inner product between a representation of the state-action pair, and a representation of the goal state:
f
( s, a, s g ) =
ϕ ( s, a ) T ψ ( s g ) .
This parametrization allows for efficiency computation, as we can compute the goal representationsjust once, and use them both in the positive pairs and the negative pairs.
While this is commonpractice in representation learning, it is not exploited by most goal-conditioned RL algorithms.
Werefer to this method as contrastive RL (NCE) .
In Appendix C, we derive a variant of this method( contrastive RL (CPC) ) that uses the infoNCE bound on mutual information.
Contrastive RL (NCE) is an on-policy algorithm because it only estimates the Q-function for thepolicy that collected the data.
However, in practice we take as many gradient steps on each transition from , eye from binary cross entropy def critic_loss(states, actions, future_states): sa_repr = sa_encoder(states, actions) # (batch_dim, repr_dim) g_repr = g_encoder(future_states)
# (batch_dim, repr_dim) logits = einsum('ik,jk->ij', sa_repr, g_repr) # outer product: <sa_repr[i], g_repr[j]> return sigmoid_binary_cross_entropy(logits= logits, labels=eye(batch_size))
def actor_loss(states, future_states): actions =
policy.sample(states, goal=future_states) # (batch_size, action_dim)
sa_repr = sa_encoder(states, actions) # (batch_dim, repr_dim) g_repr
= g_encoder(goals) # (batch_dim, repr_dim) logits =
einsum('ik,ik->i', sa_repr, g_repr) # inner product: <sa_repr[i], g_repr[i]> return -1.0 * logits as standard off-policy RL algorithms [36, 44].
Please see Appendix E for full implementation details.
We will also release an efficient implementation based on ACME [53] and JAX [11].
On a single
TPUv2, training proceeds at 1 , 100 batchessec for state-based tasks and 105 batchessec for image-based tasks;for comparison, our implementation of DrQ on the same hardware setup runs at 28 batchessec ( 3 . 9 ×slower).
6 Architectures and hyperparameters are described in Appendix E. {+
+}
4.5 Convergence Guarantees
In general, providing convergence guarantees for methods that perform relabeling is challenging.
Mostprior work offers no such guarantees [6, 20, 21] or guarantees under only restrictive assumptions [37,114].
Our convergence analysis will also make assumptions, though our experiments show that themethod continues to perform well when those assumptions are violated.
Assumption 1. Assume that each goal conditioned policy does not visit different goal states:
[Equation]
This assumption will only have to hold for the commanded goals, and is weaker than assuming thatthe policies visit disjoint states. Assumption 2.
Further, assume that these goal-conditioned policies induce an average policyπ ( a | s ) that has a uniform distribution over subset of actions: π ( a | s ) = U NIF ( A ( s ) ⊆ A ) .
Lemma 4.2 (Policy improvement) . Let p g
( s g ) be a distribution over a discrete set of goals and letπ ( a | s, s g ) be a goal-conditioned policy.
Assume that this goal distribution and goal-conditioned 213 policy satisfy Assumptions 1 and 2, and that the critic is Bayes-optimal.
Then the goal-conditionedpolicy obtained after one interation of contrastive RL, π ′ ( a | s, s g ) , achieves higher rewards than the 215 initial goal-conditioned policy:
[Equation]
217 In summary, we have shown that applying contrastive learning to a particular choice of inputsresults in an RL algorithm, one that learns a Q-function and (under some assumptions)
convergesto the reward-maximizing policy.
Contrastive RL (NCE) is simple: it does not require multiple
Q-values [36], target Q networks [83], data augmentation [64, 68],
4.6 C-learning as Contrastive Learning
C-learning [27] is a special case of contrastive RL: it learns a critic function to distinguish future goals from random goals.
Compared with contrastive RL (NCE), C-learning learns the classifier using temporal difference learning
8 Viewing C-learning as a special case of contrastive RL suggests [Figure] that contrastive RL algorithms might be implemented in a variety of different ways, each with relativemerits.
For example, contrastive RL (NCE) is much simpler than C-learning and tends to performa bit better.
Appendix D introduces another member of the contrastive RL family ( contrastive RL(NCE + C-learning) )
that tends to yield the best performance .
Experiments
Our experiments use goal-conditioned RL problems to compare contrastive RL algorithms to priornon-contrastive methods, including those that use data augmentation and auxiliary objectives.
Appendices E, G, and H contain additional experiments, visualizations, and failed experiments.
5.1 Comparing to prior goal-conditioned RL methods
Baselines.
We compare to three baselines.
“HER” [75] is a goal-conditioned RL method thatuses hindsight relabeling [6] with a high-performance actor-critic algorithm (TD3).
This baselineis representative of a large class of prior work that uses hindsight relabeling [6, 71, 96, 100].
Likecontrastive RL, this baseline does not assume access to a reward function.
The second baseline isgoal-conditioned behavioral cloning (“GCBC”)
[14, 20, 23, 37, 78, 91, 111, 114], which trains apolicy to reach goal s g by performing behavioral cloning on trajectories that reach state s g .
GCBCis a simple method that achieves excellent results [14, 23] and has the same inputs as our method( ( s, a, s f ) triplets).
The third baseline is a model-based approach that fits a density model to the futurestate distribution p π ( ·|· ) ( s t + | s, a ) and trains a goal-conditioned policy to maximize the probabilityof the commanded goal.
This baseline is similar successor representations [19] and prior multi-stepmodels [21, 56].
Both contrastive RL (Alg. 1) and this model-based approach encode the future statedistribution, but the output dimension of this model-based method depends on the state dimension.
We therefore expect this approach to excel in low-dimensional settings but struggle on image-basedtasks.
Where possible, we use the same architectures and hyperparameters for all methods.
We willinclude additional representation learning baselines when studying representations in the subsequentsection.
Tasks.
We compare to a suite of goalconditioned tasks, mostly taken from priorwork.
Four standard manipulation tasks includefetch reach and fetch push from Plappertet al.
[92] and sawyer push and sawyer binfrom Yu et al.
We evaluate these tasksboth with state-based observations and (unlikemost prior work) image-based observations.
The sawyer bin task poses an exploration challenge,as the agent must learn to pick up an object from one bin and place it at a goal location in anotherbin; the agent does not receive any reward shaping or demonstrations. We include two navigation [Figure] cc e ss r a t e fetch reach
[Figure] sawyer push [Figure] fetch push  41e sawyer bin contrastive RL (NCE)
TD3+HER
TD
[Figure]
HER
[Figure]
DrQ
TD
[Figure]
HER
[Figure]
AE
TD3+HER
[Figure]
CURL tasks: point Spiral11x11 is a 2D maze task with image observations and ant umaze [33] isa 111-dimensional locomotion task that presents a challenging low-level control problem.
Wherepossible, we use the same initial state distribution, goal distribution, observations, and definition ofsuccess as prior work.
Goals have the same dimension as the states, with one exception: on the antumaze task, we used the global XY position as the goal.
We illustrate three of the tasks to the right.
The agent does not have access to any ground truth reward function in any of these tasks.
We report results in Fig.
2, using five random seeds for each experiment and plotting the mean andstandard deviation across those random seeds.
One the state-based tasks (Fig. 2a), most methods solvethe easiest task ( fetch reach ) while only our method solves the most challenging task ( sawyerbin ).
Our method also outperforms all prior methods on the two pushing tasks.
The model-basedbaseline performs best on the ant umaze task, likely because learning a model is relatively easywhen the goal is lower-dimensional goal (just the XY location).
On the image-based tasks (Fig.
2b),most methods make progress on the two easiest tasks ( fetch reach and point Spiral11x11 );our method outperforms the baselines on the three more challenging tasks.
Of particular note is thesuccess on sawyer push and sawyer bin : while the success rate of our method remains below50%, no baselines make any progress on learning these tasks.
Taken together, these results suggestthat contrastive RL (NCE) is a competitive goal-conditioned RL algorithm.
5.2 Comparing to prior representation learning methods
We hypothesize that contrastive RL may automatically learn good representations.
To test thishypothesis, we compare contrastive RL (NCE) to techniques proposed by prior work for representationlearning.
These include data augmentation [64, 68, 131] (“DrQ”) and auxiliary objectives based onan autoencoder [29, 86, 88, 132] (“AE”) and a contrastive learning objective (“CURL”) that generatespositive examples using data augmentation, similar to prior work [84, 110, 112].
Because prior workhas demonstrated these techniques in combination with actor-critic RL algorithms, we will use thesetechniques in combination with the actor-critic baseline from the previous section (“TD3 + HER”).
While contrastive RL (NCE) resembles a contrastive representation learning method, it does notinclude any data augmentation or auxiliary representation learning objectives.
We show results in Fig.
4, with error bars again showing the mean and standard deviation across5 random seeds.
While adding the autoencoder improves the baseline on the fetch reach andadding DrQ improves the baseline on the sawyer push , contrastive RL (NCE) outperforms theprior methods on all tasks.
Unlike these methods, contrastive RL does not use auxiliary objectivesor additional domain knowledge in the form of image-appropriate data augmentations.
Theseexperiments do not show that representation learning is never useful, and do not show that contrastive
RL cannot be improved with additional representation learning machinery.
Rather, they show thatdesigning RL algorithms that structurally resemble contrastive representation learning yields biggerimprovements than simply adding representation learning tricks on top of existing RL algorithms.
5.3 Probing the dimensions of contrastive RL
Up to now, we have focused on the specific instantiation of contrastive RL spelled out in Alg. 1.
However, there is a whole family of RL algorithms with contrastive characteristics.
C-learning is acontrastive RL algorithm that uses temporal difference learning (Sec. 4.6).
Contrastive RL (CPC) is [Figure] [Figure] a variant of Alg.
1 based on the infoNCE objective [90] that we derive in Appendix C Contrastive
RL (NCE + C-learning) is variant that combines C-learning with Alg.
D (see Appendix D.).
Theaim of these experiments are to study whether generalizing C-learning to a family of contrastive RLalgorithms was useful: do the simpler methods achieve similar performance, and do other methodsachieve better performance?
We present results in Fig.
5, again plotting the mean and standard deviation across five randomseeds.
Contrastive RL (CPC) outperforms contrastive RL (NCE) on three, suggesting that swappingone mutual information estimator for another can sometimes improve performance, though bothestimators can be effective.
C-learning outperforms contrastive RL (NCE) on three tasks but performsworse on other tasks.
Contrastive RL (NCE + C-learning) consistently ranks among the best methods.
These experiments demonstrate that the prior contrastive RL method, C-learning [27], achieves goodresults on most tasks; generalizing C-learning to a family of contrastive RL algorithms resulting innew algorithms that achieve higher performance and can be much simpler. 5.4 Linear regression with the learned features [Figure]
To study the learned representations in isolation we takethe state-action representations ϕ
( s, a )
trained on theimage-based point NineRooms task, and run a linearprobe [3, 47] experiment to see whether the representations have learned to encode task-relevant information (theshortest path distance to the goal). As shown in Fig. 6,contrastive RL (NCE) learns representations that achievelower test error than those learned by TD3+HER and bya random CNN encoder.  Conclusion
In this paper, we showed how contrastive representation learning can be used for goal-conditioned RL.
This connection not only lets us re-interpret a prior RL method as performing contrastive learning,but also suggests a family of contrastive RL methods, which includes simpler algorithms, as wellas algorithms that attain better overall performance.
While this paper might be construed to implythat RL is more or less important than representation learning [67, 70, 106, 108], we have a differenttakeaway: that it may be enough to build RL algorithms that look like representation learning.
One limitation of this work is that it looks only at the goal-conditioned RL problems.
How thesemethods might be applied to arbitrary RL problems remains an open problem, though we note thatrecent algorithms for this setting [26] already bear a resemblance to contrastive RL.
Whether the richset of ideas from contrastive learning might be used to construct even better RL algorithms likewiseremains an open question.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] The main claims are that (1) contrastive learning can be
used to learn a Q-function (Proof in Appendix B) and that (2) contrastive RL methodscan outperform non-contrastive RL algorithms on goal-conditioned RL tasks (results
in Fig. 2).
(b) Did you describe the limitations of your work?
[TODO]
(c) Did you discuss any potential negative societal impacts of your work?
[No] While RL
broadly might be used for applications with both positive and negative outcomes, ouralgorithmic contributions are not tied to any particular application.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] E
g., Assumptions 1 and
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B3.
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)?
[No] We have
included all experimental details in Appendix E; code will be released upon acceptance.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix E.
(c) Did you report error bars (e.g., with respect to the random seed after running ex- periments multiple times)?
[Yes] All figures show 5 random seeds, witht error barscorresponding to the mean and standard deviation across these seeds.
(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] Sec. 4.4 describes the training
speed on one TPUv2.
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?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(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]
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 amountspent on participant compensation?
[N/A]