Contrastive Learning as Goal-Conditioned Reinforcement Learning

In plain words

Teaching a robot arm to reach a target usually means the learner has to figure out, by itself, what the scene looks like and how its own actions change it. Standard reinforcement learning is fragile at this when the input is raw pictures, so people add extra training signals to help it along.

This paper shows that a different kind of learning, one that pulls matching pairs together and pushes random pairs apart, already contains what the robot needs. Each pair is a situation-and-action on one side and a state the robot later reached on the other. The learned match score then says how likely the robot is to reach that target if it acts this way, which is exactly the number a controller needs to pick its next move.

The resulting method is short to write down and needs no extra losses or image tricks. The authors report that it does well on robot and navigation tasks, including with pictures as input and in the offline setting.

TL;DR

• Contrastive learning over (state, action, future state) triples yields a goal-conditioned critic whose inner product is a Q-function up to a constant §4.2 eq. 6.
• The simplest instance, contrastive RL (NCE), is a binary-classification critic plus an actor that maximises it, with no target networks, augmentation or auxiliary loss §4.4 §4.5.
• Author-reported: on image-based goal reaching it beats DrQ, AE and CURL built on TD3+HER, and on the offline AntMaze benchmark it wins 5 of 6 tasks §5.2 tab. 1.

Problem

Deep RL is expected to learn its own representations end to end, but prior work finds this fragile, so methods add perception losses or data augmentation §1. The paper asks whether an RL algorithm can learn good representations without such auxiliary losses §1. It restricts the scope to goal-conditioned RL, where the agent is given a goal state and no hand-written reward §1 §3. Hindsight relabelling and temporal-difference learning are the prior routes for this problem §2. The paper argues that goal-conditioned RL can be treated as a data-driven problem rather than a reward-maximisation problem §2.

Method

Setup. The goal-conditioned reward at each step is the probability of reaching the goal at the next step, scaled by (1 − γ) eq. 1 §3. The goal- conditioned Q-function is the discounted sum of these rewards eq. 2 §3. The discounted state occupancy measure is a geometric mixture over future time steps eq. 3 §3. Sampling from it means drawing a time offset from a geometric distribution and reading the state visited at that offset §3.

Key identity. The Q-function equals the occupancy probability of the goal state eq. 5 §4.1.

Critic. The critic is a binary contrastive classifier. The positive pair is (state-action, future state drawn from the occupancy measure), and the negative pair is (state-action, future state drawn from the marginal) eq. 6 §4.2. The critic is parametrised as an inner product f(s, a, s_g) = φ(s, a)ᵀψ(s_g) §4.4. At the optimum, exp(f*) equals Q divided by the marginal density of the goal, so the constant can be ignored when choosing actions §4.2.

Actor. The policy maximises the critic at the commanded goal, with actions sampled from the policy and gradients taken through a reparameterisation eq. 7 §4.3. On image tasks an action-entropy term is added §4.3.

Algorithm. The complete method alternates critic fitting, actor update and data collection §4.4. The critic loss is the binary cross-entropy over the full batch matrix, with the diagonal labelled positive and the off-diagonal labelled negative alg. 1. Other batch elements act as negatives, so one goal representation per batch row is reused for all positives and negatives §4.4 alg. 1.

# Contrastive RL (NCE): critic loss  [alg. 1] [§4.4]
# Critic f(s, a, s_g) = phi(s, a)^T psi(s_g)  [§4.4]

def critic_loss(states, actions, future_states):
    # Encode each state-action pair and each goal (a future state)  [alg. 1]
    sa_repr = sa_encoder(states, actions)  # (B, repr_dim)
    g_repr = g_encoder(future_states)      # (B, repr_dim)
    # Score every pair against every goal in the batch  [alg. 1]
    logits = sa_repr @ g_repr.T            # (B, B)
    # Diagonal entries are positives, off-diagonal are negatives  [eq. 6]
    labels = identity(batch_size)          # (B, B)
    # Mean binary cross-entropy over all B*B entries  [eq. 4] [alg. 1]
    return sigmoid_cross_entropy(logits, labels)

# Contrastive RL (NCE): actor loss  [eq. 7] [§4.3]

def actor_loss(states, goals, entropy_coefficient):
    # Sample actions from the goal-conditioned policy  [§4.3]
    actions, log_probs = policy.sample(states, goals)  # (B, action_dim)
    # Score the sampled actions against their goals  [eq. 7] [alg. 1]
    sa_repr = sa_encoder(states, actions)
    g_repr = g_encoder(goals)
    scores = sum(sa_repr * g_repr, axis=1)  # (B,)
    # Maximise the score; entropy is used on image tasks only  [§4.3] [app. E]
    return mean(entropy_coefficient * log_probs - scores)

# Goal sampling from one trajectory of length T  [§3] [eq. 3]

def sample_future_goal(states, t, discount):
    # Future offsets k >= 1 weighted by discount^k, truncated at episode end
    # The code excludes k = 0, which eq. 3 includes
    # [code: contrastive_rl/contrastive/builder.py#L136-L141]
    offsets = arange(1, T - t)                  # (T - t - 1,)
    weights = discount ** offsets
    k = categorical(weights)
    return states[t + offsets[k]]

Representations. Each encoder is an MLP with two hidden layers of 256 units and ReLU, followed by a linear output of size 64 app. E tab. 2 networks.py. The representation dimension is 64 tab. 2 config.py. On image tasks a CNN encodes each 64×64×3 observation, scaled to [0, 1], before the encoders app. E networks.py.

Offline variant. Adds a goal-conditioned behavioural-cloning term to the actor, with coefficient λ, and learns several critics and takes their minimum §5.5 tab. 3.

Variants. Contrastive RL (CPC) replaces the binary loss with an infoNCE softmax loss plus a regulariser on the log-sum-exp of each row, with coefficient 1e-2 app. C tab. 2. C-learning is a special case of this family that learns the classifier with temporal differences §4.6. Contrastive RL (NCE + C-learning) sums the two objectives after a mixture-sampling trick app. D.

Hyperparameters

Values from tab. 2 unless marked otherwise. Code values are from the official repository at code_commit.

Environments, as implemented in the code (the paper gives no episode lengths): fetch tasks 50 steps; sawyer tasks 150 steps (image drawer and window 50); point tasks 50 steps, or 100 for 11×11 mazes; AntMaze 700 (umaze) or 1000 steps env_utils.py. Success thresholds app. E: 5 cm for fetch and sawyer, 0.5 m for ant umaze (goal is the XY position), 2 m for point Spiral11x11.

Results

Author-reported. Five random seeds for online experiments, mean and standard deviation §5.1. The offline comparison uses D4RL AntMaze v2 tasks tab. 1.

Goal-conditioned RL against prior methods fig. 2. State-based: only the contrastive method solves sawyer bin, and on the two pushing tasks it is ahead of every prior method. The model-based baseline is strongest on ant umaze §5.1. Image-based: most methods make progress on fetch reach and point Spiral11x11. The contrastive method is best on the three harder tasks, though its success rate stays below 50% on sawyer push and sawyer bin, where the baselines reach none §5.1.

Against representation-learning add-ons fig. 4. On image tasks built on TD3+HER, contrastive RL (NCE) beats DrQ, an autoencoder auxiliary loss and CURL on all tasks, without augmentation or auxiliary losses §5.2.

Design variants fig. 5. Contrastive RL (CPC) beats NCE on three tasks. C-learning beats NCE on three tasks but loses on others. NCE + C-learning ranks consistently near the top §5.3.

Partial observability fig. 6. With a camera that moves with the hand, NCE reaches about 35% success on sawyer push, versus about 75% with a fixed camera §5.4.

Offline AntMaze tab. 1. Values are success rates (%), author-reported. The TD3+BC and IQL columns are numbers copied from prior work, and the paper marks them with an asterisk §5.5.

Author-reported claims: the method beats every baseline on 5 of 6 tasks; the median gain over GCBC is 15 points; on the large tasks it gains 7 to 9 points over IQL §5.5 tab. 1.

Speed §4.4. Author-reported on one TPUv2: 1100 batches per second for state tasks and 105 per second for image tasks, against 28 per second for the authors' DrQ implementation on the same hardware.

Other experiments app. F. A linear probe predicts shortest-path distance better from contrastive features than from TD3+HER or a random encoder (no numbers given in text). Contrastive learning beats a forward model when the goal has four or more dimensions. Robustness on image fetch push: success 78 ±5% originally, 73 ±10% with a recoloured object, 87 ±6% with a moved arm, and 0 ±0% with a recoloured table (100 trials in total). Pretrained representations give a small transfer benefit on some task pairs.

Paper vs code

Code: google-research/google-research at commit 3ef52111c1fabc301880818f5a547570e926a76b, directory contrastive_rl/, Apache-2.0. The authors' project page (https://ben-eysenbach.github.io/contrastive_rl), which the paper links §1 §4.4, links this directory as "Code". The code is a multi-directory monorepo; only contrastive_rl/ is used here.

• Future-goal distribution. Eq. 3 includes the current time step (t = 0), with (1 − γ) normalisation eq. 3 §3. The code samples strictly later states, offsets k ≥ 1, weighted by γ^k and truncated at the episode end builder.py.
• Goals in the actor loss. The paper's hyperparameter table says random states tab. 2. The default config uses 0.5, so half the actor batch uses the row's own sampled future goal and half uses the goal of the previous row config.py learning.py. Offline the code uses future states lp_contrastive.py, matching tab. 3.
• Target entropy on image tasks. The paper says −dim(a) tab. 2. The code default is 0.0 config.py. The launcher leaves the entropy coefficient unset for image tasks, so it is adaptive lp_contrastive.py. A helper that computes −dim(a) exists config.py, but the launcher does not call it, and its use was not found elsewhere in the code read.
• Gradient steps per learner step. The table says train-collect interval 16 for state and 64 for image tab. 2. The code default is 64 config.py, and it sets 16 for image tasks when not run locally lp_contrastive.py. The two are reversed. Whether the code's per-step count maps to the paper's "N collected, N updates" definition tab. 2 was not verified.
• Offline critic count. Table 1 reports 5 critics tab. 1. The code asserts exactly 2 learning.py and stacks only 2 in the critic network networks.py. The 5-critic numbers have no code at this commit.
• Offline λ. The paper writes λ symbolically and gives no value §5.5. The code uses 0.05 lp_contrastive.py. It mixes terms as bc_coef × BC loss + (1 − bc_coef) × actor loss learning.py.
• Offline rate limiter. The code sets samples per insert to 1e6 with a very large tolerance, which the paper's Table 3 does not list lp_contrastive.py.
• Image encoder. The paper names a CNN from prior work app. E. The code uses an Atari-style torso on 64×64×3 inputs networks.py. It applies the same torso instance to state and goal within one call networks.py. The policy applies a separate torso to the concatenated state and goal networks.py.
• Negatives. Eq. 6 draws negatives from the marginal p(s_f) §4.2. The code uses the other rows of the batch, which are themselves future states from the same sampler learning.py.
• Filtering. The convergence proof uses a filter with threshold ε §4.5. The paper does not use it in the experiments §4.5 app. F. No filter step was found in contrastive/.
• Representation normalisation. Off by default config.py, as the paper's failed-experiments list reports app. G.
• Optimiser detail. Adam eps = 1e-7 is in code builder.py and not in the paper.
• Test-time action. The evaluation policy uses the distribution mode networks.py. The paper does not say how actions are chosen at test time.
• Random initial actions. The paper collects 10,000 random transitions first tab. 2. The code also has a flag for a uniform-random initial actor, enabled in the launcher config.py. The paper does not mention the actor part.
• Target network. The code keeps a target critic updated by EMA for all modes learning.py, but the NCE loss does not read it learning.py. The paper says NCE needs no target networks §4.5.

Pitfalls

• Positive-goal offset. Eq. 3 includes the current state. The code never picks offset 0, so the current state is never a positive builder.py. Follow the code if you want the released behaviour, and treat eq. 3 as the theory.
• Batch size matters. Negatives are the other rows of the batch, so the 256-element batch is part of the method, not just a speed setting tab. 2 §4.4.
• Scale of the critic. The score is an unnormalised inner product of 64-d vectors, so no bound is enforced on its range §3 networks.py. The paper reports that normalisation and temperature hurt app. G.
• Entropy only on image tasks. On state tasks there is no entropy term, so the actor is the critic score alone §4.3 lp_contrastive.py.
• Adaptive alpha starts at 1. The log-alpha is initialised to 0, with Adam at 3e-4 learning.py. The paper gives no initial value.
• Data reuse. The method is on-policy in the sense that the critic estimates the Q-function of the policy that collected the data, yet it takes many gradient steps per transition §4.4. Samples per insert is 256 tab. 2.
• Proof assumptions. The policy-improvement result assumes tabular states and actions and a Bayes-optimal critic §4.5. Experiments use neural networks, and the filter is off, so the convergence claim does not transfer directly.
• Offline baselines. TD3+BC and IQL numbers are copied from prior work rather than re-run, and TD baselines use -v0 tasks while this paper's runs use -v2 tab. 1.
• Costs. The method needs a full trajectory replay buffer with per-trajectory sampling app. E item 2. Throughput figures are for a TPUv2 §4.4.
• Missing values. No episode lengths, no offline λ value, no seeds for offline runs, and no test-time action rule in the paper §5 §5.5 app. E.

Lineage

• Builds on: C-learning, which this paper shows is a special case of contrastive RL §4.6. Hindsight relabelling and goal-conditioned behavioural cloning provide the data and baseline setup §2 §5.1.
• Compared against: HER, goal-conditioned BC, a model-based density baseline, DrQ, autoencoder auxiliary losses, CURL, DT, BC, TD3+BC and IQL §5.1 §5.2 §5.5 tab. 1.
• Superseded by: not stated in the paper.