Goal-conditioned reinforcement learning

Overview

In goal-conditioned RL the agent receives a goal state as input and is scored on reaching it, so no reward has to be written by hand 2206.07568 §1 2206.07568 §3. Standard RL with raw image input is fragile, and methods usually add perception losses or augmentation to compensate 2206.07568 §1.

The contrastive variants share one idea: treat the goal-reaching value as a probability of visiting the goal in the future. Under the discounted occupancy measure, the goal-conditioned Q-function equals the occupancy probability of the goal 2206.07568 eq. 5 §4.1. Estimating it then becomes a contrastive problem: separate future states drawn from the agent's own trajectories from states drawn at random 2206.07568 §4.2. The invariant every contrastive variant keeps is that a learned score, up to a constant for a fixed goal, ranks actions in the same order as the Q-function, so an actor can maximise the score to pick actions 2206.07568 §4.2 §4.3.

The other cards keep the goal-conditioned objective but change what learns it. One trains a goal-conditioned world model and collects data from subgoal episodes 2411.02446 §3.1. Another plans over learned discrete landmarks for offline navigation 2411.07760 §4. A third leaves the policy alone and maps a new goal modality, audio, into the policy's existing goal space 2412.00949 §3.4.

Prior routes to the same problem are hindsight relabelling and temporal-difference learning 2206.07568 §2. The contrastive paper argues the problem can be handled as a data-driven one 2206.07568 §2.

Variants

Each variant below comes from one paper card. Results are author-reported, and each paper's own experiments are the only comparison for its variants.

Contrastive RL (NCE)

A binary contrastive critic with an inner-product form, f(s, a, g) = φ(s, a)ᵀψ(g) 2206.07568 §4.4. Positives are pairs of a state-action and a future state from the same trajectory. Negatives are the other goals in the batch 2206.07568 eq. 6 §4.2. The actor maximises the critic at the commanded goal 2206.07568 eq. 7 §4.3. The method uses no target networks, augmentation or auxiliary loss 2206.07568 §4.5. On image tasks built on TD3+HER, it beats DrQ, an autoencoder auxiliary loss and CURL on all tasks, author-reported 2206.07568 §5.2 fig. 4. On image tasks an action-entropy term is added to the actor 2206.07568 §4.3. Card: Contrastive Learning as Goal-Conditioned Reinforcement Learning.

Contrastive RL (CPC)

Replaces the binary loss with an InfoNCE softmax loss over each row, plus a regulariser on the log-sum-exp of each row, with coefficient 1e-2 2206.07568 app. C tab. 2. Author-reported: it beats the NCE form on three tasks 2206.07568 §5.3 fig. 5.

C-learning

The paper shows C-learning is a special case of the same contrastive family, learned with temporal-difference targets 2206.07568 §4.6. Author-reported: it beats NCE on three tasks and loses on others 2206.07568 §5.3 fig. 5.

Contrastive RL (NCE + C-learning)

Sums the NCE and C-learning objectives after a mixture-sampling step 2206.07568 app. D. Author-reported: it ranks consistently near the top across the design-variant tasks 2206.07568 §5.3 fig. 5.

Offline variant

Adds a goal-conditioned behavioural-cloning term to the actor, weighted by a coefficient λ. It learns several critics and takes their minimum 2206.07568 §5.5 tab. 3. Author-reported on D4RL AntMaze v2: it beats every listed baseline on 5 of 6 tasks, with a median gain of 15 points over goal-conditioned BC 2206.07568 §5.5 tab. 1. λ has no value in the paper text 2206.07568 §5.5; see Contradictions for the value the code uses.

Model-based: MUN (world model with subgoal episodes)

MUN keeps the goal-conditioned Dreamer framework (GC-Dreamer) and changes only how real data is collected 2411.02446 §3.1 app. D. Each iteration runs one episode between subgoal states drawn from the replay buffer, then one episode toward an environment goal, then trains the world model on both 2411.02446 alg. 1. Subgoals come from DAD, which runs farthest-point sampling over action vectors and keeps the states at the chosen steps 2411.02446 §3.2 alg. 2. The policy is trained only in imagination, with reward equal to minus a learned temporal distance 2411.02446 eq. 7. Author-reported: about 95% success on 3-block stacking, against about 60% for GC-Dreamer, MEGA-G and PEG-G within 2.5M steps 2411.02446 §4.3. On three navigation test sets it has the highest success rate of the methods compared 2411.02446 tab. 3 tab. 4 tab. 5. One run costs 36 to 70 GPU hours on an A100 2411.02446 tab. 1. Card: Learning World Models for Unconstrained Goal Navigation.

Offline hierarchical: QPHIL (landmark planning)

QPHIL is an offline goal-reaching method for navigation. A VQ-VAE tokenizer maps each state to one of k landmark tokens, trained with reconstruction, commitment and temporal triplet losses 2411.07760 §4.1. A transformer plans a sequence of landmark tokens from the goal token and the compressed history 2411.07760 §4.2. A landmark policy, trained with IQL-style value learning on relabelled next-landmark goals, follows the plan, and a goal policy takes over for the final approach 2411.07760 §4.3 alg. 1. Author-reported: QPHIL has the highest success on both AntMaze-Ultra datasets in its table, 64.5% (play) and 70.3% (diverse) 2411.07760 tab. 1. On the new AntMaze-Extreme map it reaches up to 50.0% against 22.8% for HIQL without representation learning 2411.07760 fig. 5. Card: Navigation with QPHIL: Quantizing Planner for Hierarchical Implicit Q-Learning.

Audio as the goal (STEVE-Audio)

This is not a change to the RL update. A goal-conditioned Minecraft policy, STEVE-1, normally takes text or pictures as its goal. The authors train a contrastive audio-video model over frozen video and audio encoders, with only the two mapping networks trained 2412.00949 §3.1 §3.3. A conditional VAE then maps the audio embedding into the policy's visual goal space. The policy itself is not changed 2412.00949 §3.4. Author-reported: audio goals beat the visual goal on 4 of 6 item-collection tasks and the text goal on 5 of 6 2412.00949 §4.1 fig. 4. Audio is much weaker on the three "place" tasks 2412.00949 §4.2.1. Card: STEVE-Audio: Expanding the Goal Conditioning Modalities of Embodied Agents in Minecraft.

Proposed use in proof search (not evaluated)

Not a variant from the same family. The essay 2403.04571 §9 proposes goal-conditioned RL for proof search: the policy takes a goal as input, and intermediate lemmas are learned as subgoals. It names no algorithm, reward or prover, and reports no results 2403.04571 §9 §10. Card: Machine learning and information theory concepts towards an AI Mathematician.

When to use which

Trade-offs the cited sections support: the offline variant needs a behavioural prior and a λ choice, and the CPC and C-learning forms trade a different objective against NCE's simplicity 2206.07568 §4.6 app. C. MUN buys better world-model accuracy with a large compute budget 2411.02446 tab. 1 tab. 6. QPHIL's extra stages are the price of its planner 2411.07760 §4.

Implementation sketch

The NCE core is the same across tasks. The variants change the loss on the logits, and the offline variant adds a term to the actor loss. Pseudocode is in our own words. The sketch covers only the contrastive family; MUN and QPHIL need a world model and a landmark tokenizer respectively, and their cards give their own steps.

# Sample a goal: a future state from the same trajectory  [eq. 3] [§3]
def sample_future_goal(trajectory, t, gamma):
    # Offsets k >= 1 weighted by gamma^k, truncated at episode end.
    # Eq. 3 also allows k = 0 (the current state); the released code
    # does not  [code: contrastive_rl/contrastive/builder.py#L136-L141]
    offsets = arange(1, len(trajectory) - t)
    k = categorical(gamma ** offsets)
    return trajectory[t + offsets[k]]

# Critic loss for contrastive RL (NCE)  [§4.4] [alg. 1]
def critic_loss(states, actions, goals):
    # Encode each (state, action) pair and each goal  [§4.4]
    phi = sa_encoder(states, actions)  # (B, d)
    psi = g_encoder(goals)             # (B, d)
    # Score every pair in the batch by inner product  [§4.4]
    logits = phi @ psi.T               # (B, B)
    # Diagonal = positives, off-diagonal = negatives  [eq. 6]
    labels = eye(batch_size)
    # Mean binary cross-entropy over all B*B entries  [alg. 1]
    return sigmoid_bce(logits, labels).mean()

# Actor loss: maximise the critic at the commanded goal  [eq. 7] [§4.3]
def actor_loss(states, goals, entropy_weight):
    # Sample actions with reparameterised gradients  [§4.3]
    actions, log_probs = policy.sample(states, goals)
    scores = critic_score(states, actions, goals)  # (B,)
    # Entropy weight is 0 for state tasks and used on image tasks  [§4.3]
    return mean(entropy_weight * log_probs - scores)

Notes on where the variants differ:

• CPC swaps the critic loss for a softmax over each row, plus the log-sum-exp regulariser app. C.
• Offline adds a BC term to the actor loss, weighted λ, and takes the minimum over several critics §5.5 tab. 3.
• Image tasks feed a CNN on 64×64×3 observations, scaled to [0, 1], into the encoders app. E.

Contradictions

Within the contrastive paper, the paper and its released code differ in the places below. Each side is cited; the code claims come from google-research/google-research at commit 3ef52111c1fabc301880818f5a547570e926a76b, directory contrastive_rl/, Apache-2.0 2206.07568 §1.

• Future-goal offset. Eq. 3 includes the current state (k = 0), with (1 − γ) normalisation 2206.07568 eq. 3 §3. The code samples offsets k ≥ 1 weighted by γ^k builder.py.
• Goals in the actor loss. The hyperparameter table lists random states 2206.07568 tab. 2. The default config uses a 0.5 mix, so half the batch uses each row's own future goal and half uses the previous row's config.py learning.py. Offline, the code uses future states, which matches 2206.07568 tab. 3 lp_contrastive.py.
• Target entropy on image tasks. The table says −dim(a) 2206.07568 tab. 2. The code default is 0.0 config.py, and the launcher leaves the entropy coefficient adaptive for image tasks lp_contrastive.py.
• Gradient steps per learner step. The table gives train-collect interval 16 for state and 64 for image 2206.07568 tab. 2. The code default is 64 config.py, and it sets 16 for image tasks when not run locally lp_contrastive.py. Whether the code's count matches the paper's definition was not verified.
• Negatives. Eq. 6 draws negatives from the marginal p(s_f) 2206.07568 eq. 6. The code uses the other rows of the batch, which are future states from the same sampler learning.py.
• Offline critics. Table 1 reports 5 critics 2206.07568 tab. 1. The code asserts exactly 2 learning.py. The 5-critic numbers have no code at this commit.
• Offline λ. The paper gives no value 2206.07568 §5.5. The code uses 0.05 lp_contrastive.py.

Across cards, no reported number contradicts another, but the results cannot be compared directly. The contrastive offline table reports D4RL AntMaze v2 tasks 2206.07568 tab. 1. The QPHIL card states no dataset version 2411.07760 app. C. MUN's navigation tables use start-goal pairs that its authors chose 2411.02446 §4.5. MUN's paper also says its DAD subgoal method is used on every task 2411.02446 §3.2. The released configs do not all use it; see the MUN card's Paper vs code section.

Pitfalls

• Batch size is part of the method. Negatives come from the batch, so the batch size of 256 is not only a speed setting 2206.07568 tab. 2 §4.4.
• Unbounded score. The critic is an unnormalised inner product, with no bound on its range 2206.07568 §3. The paper reports that normalisation and temperature hurt 2206.07568 app. G.
• Entropy term is for images only. State tasks use the critic score alone 2206.07568 §4.3.
• Proof assumptions. The policy-improvement result assumes tabular states and actions and an optimal critic, and the filter in the proof is not used in the experiments 2206.07568 §4.5 app. F. The convergence claim does not transfer directly to the neural-network setting used there.
• Offline baselines differ. TD3+BC and IQL numbers are copied from prior work, not re-run, and use -v0 tasks while this paper's runs use -v2 2206.07568 tab. 1.
• Missing details. The paper gives no episode lengths, no offline λ value, no offline seeds and no test-time action rule 2206.07568 §5 §5.5 app. E. The code sets test-time actions to the distribution mode networks.py.
• Essay, not method. The proof-search proposal has no algorithm, reward formula or evaluation 2403.04571 §6 §10. Do not implement it as specified; it is a direction, not a tested design.
• MUN's subgoal switching is learned, not measured. The paper pursues each subgoal until it is reached or the time budget runs out 2411.02446 alg. 1. The released code switches on the model's own distance estimate, so a wrong estimate changes which subgoals get visited. See the MUN card's Paper vs code section for this and for the released configs that skip DAD.
• MUN's world-model hyperparameters are not given. The paper says it uses the LEXA backbone's defaults and gives no values 2411.02446 app. E.3.
• QPHIL seeds and loss coefficients are ambiguous. The paper says 8 seeds but lists 7 values, and its VQ-VAE coefficients are not mapped to the loss symbols 2411.07760 §5.1 §8 app. B.2 §4.1. Plans are sampled at temperature 0.9, so results vary from run to run 2411.07760 app. B.2.
• QPHIL's large-maze claim does not hold for one row. The text says QPHIL outperforms all methods on larger maps, but on large-diverse HIQL with a learned representation scores higher 2411.07760 §5.2 tab. 1.
• No usable QPHIL code is cited. The paper says code will be released later 2411.07760 app. B.1. The repository found under the authors' name is not licensed for commercial use, so implement from the paper text. The card's Paper vs code section explains why.
• STEVE-Audio's training cannot be rerun from its release. The data is described in the paper 2412.00949 §3.2, but the repository does not include it or the starting checkpoint. The card's Pitfalls list the details.
• STEVE-Audio's conclusion overstates its table. The paper's conclusion says audio outperforms the text and visual goals, but its results section reports lower counts on some tasks 2412.00949 §5 §4.1.

Papers

• Contrastive Learning as Goal-Conditioned Reinforcement Learning: shows contrastive learning on action-labelled trajectories is a goal-conditioned RL method, and introduces contrastive RL with its NCE, CPC and offline forms.
• Learning World Models for Unconstrained Goal Navigation: trains a goal-conditioned world model on episodes between subgoals taken from the replay buffer, with subgoals picked by farthest-point sampling over actions.
• Navigation with QPHIL: Quantizing Planner for Hierarchical Implicit Q-Learning: plans over learned discrete landmarks with a transformer, and follows the plan with two IQL policies, for offline navigation.
• STEVE-Audio: Expanding the Goal Conditioning Modalities of Embodied Agents in Minecraft: maps audio into the goal space of a pretrained Minecraft policy without retraining the policy.
• Machine learning and information theory concepts towards an AI Mathematician: a position essay with no experiments. It proposes goal-conditioned RL, with lemmas as subgoals, as one route to proof search.