C ONTRASTIVE V ALUE L EARNING : I MPLICIT M ODELS FOR S IMPLE O FFLINE RL
A BSTRACT
Model-based reinforcement learning (RL) methods are appealing in the ofﬂine setting because they allow an agent to reason about the consequences of actions without interacting with the environment.
Prior methods learn a 1-step dynamics model, which predicts the next state given the current state and action.
These models do not immediately tell the agent which actions to take, but must be integrated into a larger RL framework.
Can we model the environment dynamics in a different way, such that the learned model does directly indicate the value of each action?
In this paper, we propose Contrastive Value Learning (CVL), which learns an implicit, multi-step model of the environment dynamics.
This model can be learned without access to reward functions, but nonetheless can be used to directly estimate the value of each action, without requiring any TD learning.
Because this model represents the multi-step transitions implicitly, it avoids having to predict high-dimensional observations and thus scales to high-dimensional tasks.
Our experiments demonstrate that CVL outperforms prior ofﬂine RL methods on complex continuous control benchmarks.
I NTRODUCTION
While the ofﬂine RL setting is relevant to many real-world applications where the ability for online data collection is limited, it often requires RL algorithms to ﬁnd policies that are not well-supported by the training data.
Instead of learning via trial-and-error, ofﬂine RL algorithms must leverage logged historical data to learn about the outcome of different actions, potentially by capturing environment dynamics as a proxy signal.
Many prior approaches for this ofﬂine RL setting have been proposed, whether in model-free (Wu et al., 2019; Fujimoto et al., 2019; Kumar et al., 2020) or model-based (Kidambi et al., 2020; Yu et al., 2021) settings.
Our focus will be on those that address this prediction problem head-on: by learning a predictive model of the environment which can be used in conjunction with most model-free algorithms.
In this paper, we propose to learn a different type of model for ofﬂine RL, a model which (1) will not require predicting high-dimensional observations and (2) can be directly used to estimate Q-values without requiring either model-based rollouts or model-free temporal difference learning.
Precisely, we will learn an implicit model of the discounted state occupancy measure, which answers the question “where will the agent be in the time-averaged future?”
We will learn this implicit model via contrastive learning, treating it as a classiﬁer rather than a generative model of observations.
Once learned, we predict the likelihood of reaching every reward-labeled state.
By weighting these predictions by the corresponding rewards, we form an unbiased estimate of the Q-function.
Whereas methods like Q-learning estimate the Q-function of a state “backing up” reward values, our approach goes in the opposite direction, “propagating forward” predictions about where the agent will go.
We name our proposed algorithm Contrastive Value Learning(CVL).
CVL is a simple algorithm for ofﬂine RL which learns the future state occupancy measure using contrastive learning and re-weights it with the future reward samples to construct a quantity proportional to the true value function.
Because CVL represents multi-step transitions implicitly, it avoids having to predict high-dimensional observations and thus scales to high-dimensional tasks.
Using the same algorithm, we can handle settings where reward-free data is provided, which cannot be directly handled by classical ofﬂine
RL methods such as FQI (Munos, 2003) or BCQ (Fujimoto et al., 2019).
We compare our proposed method to competitive ofﬂine RL baselines, notably CQL (Kumar et al., 2020) and CQL+UDS (Yu et al., 2022) on an ofﬂine version of the multi-task Metaworld benchmark (Yu et al., 2020), and ﬁnd that CVL greatly outperforms the baseline approaches as measured by the rliable library (Agarwal et al., 2021b).
Additional experiments on image-based tasks from this same benchmark show that our approach scales to high-dimension tasks more seamlessly than the baselines.
We also conduct a series of ablation experiments highlighting critical components of our method.
R ELATED WORKS
Prior work has given rise to multiple ofﬂine RL algorithms, which often rely on behavior regularization in order to be well-supported by the training data.
The key idea of ofﬂine RL methods is to balance interpolation and extrapolation errors, while ensuring proper diversity of out-of-dataset actions.
Popular ofﬂine RL algorithms such as BCQ and CQL rely on a behavior regularization loss (Wu et al., 2019) as a way to control the extrapolation error.
This regularization term ensures that the learned policy is well-supported by the data, i.e. does not stray too far away from the logging policy.
The major issue with current ofﬂine RL algorithms is that they fail to fully capture the entire distribution over state-action pairs present in the training data.
To directly learn a value function using policy or value iteration, one needs to have information about the transition model in the form of sequences of state-action pairs, as well as the reward emitted by this transition.
However, in some real-world scenarios, the reward might only be available for a small subset of data.
For instance, in the case of recommending products available in an online catalog to the user, the true long-term reward (user buys the product) is only available for users who have browsed the item list for long enough and have purchased a given item.
It is possible to decompose the value function into reward-dependent and reward-free parts, as was done by (Barreto et al., 2016) through the successor representation framework (Dayan, 1993).
More recent approaches (Janner et al., 2020;
Eysenbach et al., 2020; 2022) use a generative model to learn the occupancy measure over future states for each state-action pair in the dataset; its expectation corresponds to the successor representation.
However, learning an explicit multi-step model such as (Janner et al., 2020) can be unstable due to the bootstrapping term in the temporal difference loss.
Similarly to model-based approaches, our method will learn a reward-free representation of the world, but will do so without having to predict high-dimensional observations and without having to do costly autoregressive rollouts.
Thus, while our critic is trained without requiring rewards, it is much more similar to a value function than a standard 1-step model.
Learning a conditional probability distribution over a highly complex space can be a challenging task, which is why it is often easier to instead approximate it using a density ratio speciﬁed by an inner product in a much lower-dimensional latent space.
To learn an occupancy measure over future states without passing via the temporal difference route, one can use noise-contrastive estimation (NCE, Gutmann and Hyv¨arinen, 2010; Oord et al., 2018) to approximate the corresponding log ratio of densities as an implicit function.
Contrastive learning was originally proposed as an alternative to classical maximum likelihood estimation, but has since then seen successes in static self-supervised learning (He et al., 2020; Chen et al., 2020).
In reinforcement learning, NCE was shown to improve the robustness of state representations to exogenous noise (Srinivas et al., 2020;
Mazoure et al., 2020; Agarwal et al., 2021a) and, more recently, to be an efﬁcient replacement for traditional goal-conditioned methods (Eysenbach et al., 2022).
P RELIMINARIES
Reinforcement learning We assume a Markov decision process M deﬁned by the tuple M = h S , S 0 , A , T , r,  i , where S is a state space, S 0 ✓ S is the set of starting states, A is an action space,
T = P [ ·| s t , a t ] : S ⇥ A !
( S ) is a one-step transition function 1 , r : S ⇥ A !
[ r min , r max ] is a reward function and  2 [0 , 1) is a discount factor.
The system starts in one of the initial states s 0 2 S 0 .
At every timestep t = 1 , 2 , 3 , .. , the policy ⇡ : S !
( A ) , samples an action a t ⇠ ⇡ ( ·| o t ) .
The environment transitions into a next state s t +1 ⇠ T ( ·| s t , a t ) and emits a reward r t = r
( s t , a t ) .
The aim is to learn a Markovian policy ⇡ ( a | s ) that maximizes the discounted sum of returns over an episode of length H :
[Equation]
where P ⇡ t : t + K denotes the joint distribution of { s t + k , a t + k } Kk =1 obtained by executing ⇡ in the en- vironment for K timesteps starting at timestep t .
We assume that the ofﬂine RL algorithm cannot interact with the environment, but instead must learn from an ofﬂine dataset of logged trajectories { ( s 0 , a 0 , s 1 , a 1 , · · · ) } .
Value-based RL algorithms maximize cumulative episodic rewards by estimating the state-action value function under a policy ⇡
[Equation]
where P ⇡ t : H ( s | s t , a t ) = (1   )
P
[Equation]
as deﬁned in Janner et al.
This decomposition of the value function has already been used in previous works based on the successor representation (Dayan, 1993; Barreto et al., 2016) and, more recently,  -models (Janner et al., 2020).
We will use it to efﬁciently learn an implicit density ratio proportional to the state occupancy measure using contrastive learning.
Noise-contrastive estimation Noise-contrastive estimation (NCE, Gutmann and Hyv¨arinen, 2010) spans a broad class of learning algorithms, at the core of which is negative sampling (Mikolov et al., 2013), i.e., learning an implicit metric space from positive and negative examples.
Given reference samples, samples from a positive distribution (i.e., high similarity with reference points) and samples from a negative distribution (i.e., low similarity with reference points), contrastive learning methods learn an embedding where positive examples are located closer to the reference points than negative examples.
One of the most well-known and commonly used NCE objectives is
InfoNCE (Oord et al., 2018), which solves
[Equation]
over some hypothesis class  : {  : X ! Z} for input space X , latent space Z , x ⇠ P ( X ) , y ⇠ P positives ( X ) and y ⇠ P negatives ( X ) .
Contrastive learning has been widely studied in the static unsupervised/ supervised learning settings (Hjelm et al., 2018;
Chen et al., 2020; He et al., 2020), as well as in reinforcement learning (Kim et al., 2018; Mazoure et al., 2020) for learning state representations with desirable properties such as alignment and uniformity (Wang and Isola, 2020).
Solving Equation (3) for (  ⇤ , ⇤ ) yields a critic f : X ⇥ Y !
R which decomposes as f ⇤ ( x, y ) =
⇤ ( x ) >
⇤ ( y ) and, at optimality 2 , captures the log-ratio of P positives ( X ) and P negatives ( X ) :
[Equation]
Implicit dynamics models via NCE.
Various prior works (Du and Mordatch, 2019; Mazoure et al., 2020; Nachum and Yang, 2021) have studied the use of NCE to approximate a single-step dynamics model, where triplets ( s t , a t , s t +1 ) have higher similarity than ( s t , a t , s t 0
= t +1 ) , effectively deﬁning positive and negative distributions over trajectory data.
More recently, contrastive goalconditioned RL (Eysenbach et al., 2022) used InfoNCE to condition the critic on goal states sampled from the replay buffer.
These methods use asymetric encoders, using  ( s t , a t ) and ( s t +  t ) , where positive samples of s t +  t are sampled from the discounted state occupancy measure for t  0 .
The conditional probability distribution of future states given the current state-action pair can be efﬁciently estimated using an implicit model trained via contrastive learning over positive and negative feature distributions, as shown in Equation (5).
[Equation]
Minimizing ` InfoNCE over trajectory data yields a critic which, at optimality, approximates the future discounted state occupancy measure up to a multiplicative term as per Equation (4),
[Equation]
Intuitively, f ⇤ approximates a H -step dynamics model which has an implicit dependence on policy ⇡ that collected the training data.
Ordinarily, training state-space models is hard when the dimensions are large, e.g. image-based domains.
However, by using contrastive learning, we can learn this model without having to require it predict high-dimensional observations, as similarity is evaluated in a lower-dimensional latent space (observe that in Equation (5) the inner product is computed in Z , whose dimension we control, instead of X , which is speciﬁed externally).
An apparent limitation of the approach is that the probability of future states s t +  t is recovered only up to a constant.
However, it turns out that we can still use this model to get accurate estimates of the Q-values, as is described in the next section.
E STIMATING AND M AXIMIZING R ETURNS VIA C ONTRASTIVE L EARNING
In this section, we show how NCE can be used to learn a quantity proportional to a value function, and how the later can be used in a policy iteration scheme.
4. E STIMATING Q- VALUES USING THE C ONTRASTIVE M ODEL
As shown in Equation (2), the Q-function at ( s t , a t ) can be thought of as evaluating the reward function at states sampled from the discounted occupancy measure P ⇡ t : H ( s t , a t ) .
That is, to estimate a quantity akin to Q ⇡ , we can ﬁrst estimate the occupancy measure and take a weighted average of rewards over future states using the probabilities from the log-density ratio learned by the contrastive model.
Precisely, Equation (2) corresponds to using an importance-weighted estimator, where an optimal critic that minimizes Equation (5) approximates the density ratio from Equation (6).
The critic itself can be trained using the occupancy measure formulation speciﬁed in ?? over all state-action pairs in a given episode.
However, ?? needs to be re-adjusted to account for ﬁnitehorizon truncation of the geometric mass function presented in Deﬁnition 1.
Deﬁnition 1 (Truncated distribution) Let X be a random variable with distribution function F X .
Y is a called the truncated distribution of X with support [ m, M ] s.t. 0 < m < M if
[Equation]
We denote the special case of the truncated geometric distribution as TruncGeom ( p, m, M ) .
The contrastive objective to train the critic to approximate the discounted occupancy measure over a dataset D is then
[Equation]
It is possible that multiple optimal critics exist s.t.
the multiplicative proportionality constant depends on the action.
To avoid this, we adopt a similar approach as Eysenbach et al.
and introduce a regularization term over the partition function, making the critic training objective be
[Equation]
Now, suppose we found an optimal critic f .
Combining ??
with Deﬁnition 1, we obtain the following form of the Q-function for an optimal critic f which minimizes Equation (5):
[Equation]
Here, the offset  t is a random variable sampled from TruncGeom (1   , t, H ) where H is the horizon of the MDP 3 .
We can also show that Q ( s, a ) < Q ( s,
[Equation]
Q [Equation]
for all s 2 S and a, a 0 2 A , which makes the contrastive Q-values suitable for policy evaluation.
However, we do not, in general, expect Q NCE to recover the optimal Q function, as the recovered
Q-values are on-policy with respect to ⇡ .
4. E FFICIENT E STIMATION USING R ANDOM F OURIER F EATURES
A major issue with using Q NCE out-of-the-box is that it is computationally expensive, requiring evaluation of the inner product  ( s t , a t ) >
( s t +  t ) with a large number of future states and hence multiple forward passes through .
The underlying cause of this computational overhead is the RBF kernel term e
[Equation] +  t ) .
If we instead used a linear kernel, the constant term  ( s t , a t ) would be factored out, and we could separately keep track of reward-weighted future expected features.
This would (1) reduce the computational complexity of N actor updates over D from O ( |D| N ) to
O ( |D| + N ) and (2) reduce the variance of the representation if averaging features of future states using exponential moving average.
It turns out that the RBF kernel can be approximately linearized by using random Fourier features (Rahimi and Recht, 2007; Nachum and Yang, 2021).
Lemma 1 Let x , y 2 R d be unit vectors, and let F W , b ( x ) = » 2 ed cos( W x + b ) where W ⇠
Normal (0 , I ) and b ⇠ Uniform (0 , 2 ⇡ ) .
Then, E [ F W , b ( x ) > F W , b ( y )] =
e x > y .
Lemma 1 is a straightforward modiﬁcation of the result from Rahimi and Recht (2007) and allows us to reduce the RBF kernel to an expectation over d -dimensional random feature vectors:
[Equation]
The advantage of using the RFF approximation is that it allows us to split the exponential term inside the expectation and separately keep track of the policy-dependent, reward-weighted future state probability term, while the state-action dependence term is learned online.
Speciﬁcally, we keep track of ⇠ ( ⇡ ) via an exponential-moving average during the entire duration of training 4 .
4. L EARNING THE P OLICY
Once the policy evaluation phase completes and we have an estimate Q NCE , we optimize a policy to maximize the returns predicted by this Q-value.
We can decode the policy by minimizing its
Kullback-Leibler divergence to the Boltzmann Q-value distribution (see Haarnoja et al.
(2018)), which can be efﬁciently done by minimizing the following objective:
[Equation]
Note that in discrete action spaces, minimizing Equation (12) leads to a soft version of the greedy policy decoding ⇡ greedy ( s )
=
arg max a 2 A Q NCE ( s, a ) for s 2 S .
In practice, we approximate the KL term in Equation (12) using N a Monte-Carlo action samples.
Decoding ⇡ in such a way can lead to sampling out-of-distribution actions in regions where the Qfunction might be inaccurate due to poor dataset coverage.
To mitigate this issue, we follow prior work (Cobbe et al., 2021;
Zhao et al., 2021; Schwarzer et al., 2021) and add a behavior cloning term which prevents the new policy from straying too far away from the data:
[Equation]
(13) for entropy estimator H ( ⇡ ( s ))
=
[Equation] ⇠ ⇡ ( s ) [log ⇡ ( a | s )] .
We add this extra loss to ` Policy to learn a policy ⇡ which prioritizes high Q-values that are well-supported by the ofﬂine dataset D .
Thus, the ﬁnal policy optimization objective becomes
[Equation]
The policy found by minimizing ` Policy has, on average, non-decreasing returns, as per Lemma 2.
Lemma 2 (Contrastive policy improvement) Let µ be a policy and let
[Equation] NCE [Equation] (  , )] .
If
[Equation]
then Q ⇡ ( s, a )  Q µ ( s, a ) for all ( s, a ) 2 D µ .
The proof of Lemma 2 is located in Section 6.2.
Speciﬁcally, Lemma 2 tells us that using CVL as a surrogate Q-function corresponds to one step of conservative policy improvement, where ⇡ satisﬁes soft constraints of Equation (12) and small E D µ [ D KL ( ⇡ ( s ) || µ ( s ))] via the BC term.
4. P RACTICAL I MPLEMENTATION
We now present our complete method, which can be viewed as an actor-critic method for ofﬂine RL.
We learn the critic via contrastive learning (Equation (9)) and learn the policy via Equation (14).
We will interleave these steps in most of our experiments, but experiments in Section 5.1 show that the critic can be pretrained e.g. in the presence of unlabeled data from related tasks.
We summarize the method in Algorithm 1.
4. I NTERPRETATIONS AND C ONNECTIONS WITH P RIOR WORK
The main distinction between Contrastive Value Learning and prior works consists speciﬁcally in representing the Q-values in a two-step decomposition: the Q-value is represented as an occupancy measure weighted by the reward signal; the occupancy measure itself is represented using a powerful likelihood-based model parameterized using an implicit function.
Decoupling the learning of the occupancy measure from reward maximization allows, among others, for efﬁcient pretraining strategies on unlabeled data, i.e. trajectory data without reward information, and can be used to learn provably optimal state representations for any reward function (Touati and Ollivier, 2021).
While CVL is similar in spirit to the successor representation (Dayan, 1993; Barreto et al., 2016), the occupancy measure learned by CVL is much richer than that of SR, as it captures the entire distribution over future states instead of only the ﬁrst moment.
Another method,  -models (Janner et al., 2020), is closely related to CVL, but uses a surrogate single-step TD objective to learn the occupancy measure, similarly to C-learning (Eysenbach et al., 2020).
Input : Dataset D ⇠ µ , ,  networks, temperature parameter ⌧ , exponential moving average parameter  1 for epoch
j = 1 , 2 , .., J do for minibatch B ⇠ D do /* Update density ratio estimator using Equation (9) */
Update  ( j +1) , ( j +1) using r  , ` Critic (  ( j ) , ( j ) ) ; / * Estimate the contrastive Q-function */ Q ( s, a ) Equation (11) if using RFF, otherwise Equation (10); /
* Decode policy from Q-function using Equation (14) */
Update
[Figure]
✓
[Figure]
r ✓ { ` Policy ( ✓ ) } ; /* Update future state encoder using EMA */
[Figure]
future
[Figure]
features weighted by rewards */
[Figure]
E XPERIMENTS
Our experiments aim to answer three questions.
First, we study how CVL compares with baseline approaches on a large benchmark of state-based tasks.
Our second set of experiments look at imagebased tasks, testing the hypothesis that CVL scales to these tasks more effectively than the baselines.
We conclude with ablation experiments.
Our main point of comparison will be a high-performing ofﬂine RL method, CQL (Kumar et al., 2020).
While CVL learns an implicit model, that model is structurally much more similar to value-based RL methods than model-based methods, motivating our comparison to a value-based baseline (CQL).
We will also include behavioral cloning as a baseline.
Metaworld.
We ﬁrst test our approach on the complex MetaWorld benchmark (Yu et al., 2020), which consists of 50 robotic manipulation tasks such as open a door, pick up an object, reach a certain area of the table, executed by a robotic arm (see Figure 1 (left)).
This domain is an ideal testbed for CVL, as it allows for both full-state and image-based experiments, has a dense and informative reward function thus decoupling the problem of representation learning from exploration, and is challenging for model-free methods which leaves room for improvement.
While the original
MetaWorld domain has been used to evaluate online RL agents, we create an ad hoc dataset suitable for ofﬂine learning.
To do so, we train Soft Actor-Critic (Haarnoja et al., 2018) from full states on each of the 50 tasks separately for 500k frames, and save the resulting replay buffer, which forms the training dataset.
As shown in Figure 1 (right), CVL manages to considerably improve upon strong baselines such as behavior cloning, CQL and CQL with UDS (Yu et al., 2022) 5 .
We report the results on all tasks of the MetaWorld suite over 5 random seeds, according to the aggregation methodology proposed by Agarwal et al.
(2021b).
Per-environment scores are available in Table 4.
Image-based experiments Our working hypothesis is that contrastive formulation of the value function acts in itself as a pre-training mechanism via the prism of representation learning.
For this reason, we conduct further experiments on 4 image-based tasks from the MetaWorld suite (similarly to full-states, the dataset was obtained from the SAC replay buffer trained on rendered images).
Pretraining the reward-free occupancy measure on related tasks allows CVL to outperform baseline approaches and even CVL trained tabula rasa .
[Table]
Results presented in Table 1 show that CVL is also able to learn meaningful Q-values and achieve good empirical performance on hard image-based tasks.
5. A BLATION EXPERIMENTS .
When is pretraining the model useful?
In theory, the model can be pretrained on the data from other tasks, however, we do not always expect this to help (e.g., when the pretraining tasks are very different).
We ran an experiment to test this capability.
The results, shown in Fig. 2, show that pretraining sometimes speeds up learning.
In particular, we observe that pretraining is effective when the pretraining tasks are similar to the target task and contain a diverse set of state-action pairs.
How reliable is the Q NCE approximation?
Given that contrastive Q-values are proportional to the true Q-function, a natural question to ask is how good is Q NCE at capturing the topology of Q ?
First, we conduct an ablation demonstrating how linearizing the RBF kernel via random Fourier features provides a performance gain on the ofﬂine MetaWorld tasks Figure 3.
Speciﬁcally, we hypothesize that this is due to the reduced variance of the RFF Q-value estimator which keeps track of future reward-weighted state features using a rolling average.
Next, we qualitatively assess the similarity between contrastive and true Q-values on the continuous
Mountain Car environment (Moore, 1990) by ﬁrst pre-training SAC online on the task and then ﬁtting CVL to the data from SAC’s replay buffer.
Figure 4 (left) shows the contrastive Q-values on a log-scale, evaluated on trajectories from the SAC replay; for comparison, we also show the Q-values learned by online SAC in Figure 4 (right).
Note that the value function learned by CVL conserves the same topology as the true value function, up to a multiplicative rescaling.
D ISCUSSION
This paper presented an RL algorithm that learns a contrastive model of the world, and uses that model to obtain Q values by estimating the likelihood of visiting future states.
Our experiments demonstrate that this approach can effectively solve a large number of ofﬂine RL tasks, includ- ing from image-based observations.
Our pretraining results hinted that CVL can be pretrained on datasets from other tasks, and we are excited to pretrain our model on datasets of increasing size.
Limitations.
One limitation of our approach is that it corresponds to a single step of policy improvement.
This limitation might be lifted by training the contrastive model using a temporal difference update for the contrastive model (Eysenbach et al., 2020; Blier et al., 2021).
A second limitation is that the RFF approximation can be poor when the feature dimension is small.
We tried to train the contrastive model using non-exponentiated features (akin to HaoChen et al.
, but failed to achieve satisfactory results.
Figuring out how to effectively train these spectral (i e., non-exponentiated) models remains an important question