Contrastive Learning as Goal-Conditioned Reinforcement Learning

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 is unstable and instead equip RL algorithms with additional representation learning parts (e.g., auxiliary losses, data augmentation). How can we design RL algorithms that directly acquire good representations? In this paper, instead of adding representation 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 learning to action-labeled trajectories, in such a way that the (inner product of) learned representations exactly corresponds to a goal-conditioned value function. We use this idea to reinterpret a prior RL method as performing contrastive learning, and then use the idea to propose a much simpler method that achieves similar performance. Across a range of goal-conditioned RL tasks, we demonstrate that contrastive RL methods achieve higher success rates than prior non-contrastive methods, including in the offline RL setting. We also show that contrastive RL outperforms prior methods on image-based tasks, without using data augmentation or auxiliary objectives.

• The paper introduces a framework where contrastive learning directly represents goal-conditioned value functions by linking learned inner products to reward maximization.
• The method simplifies conventional RL by eliminating the need for multiple Q-values, elaborate data augmentation, or auxiliary objectives, achieving superior results on diverse tasks.
• Strong theoretical and empirical evidence, including convergence guarantees and successful transfer across tasks, supports the practicality and robust performance of the proposed approach.

Contrastive Learning as Goal-Conditioned Reinforcement Learning

This paper introduces a novel perspective on goal-conditioned reinforcement learning (RL) by framing contrastive representation learning as an RL algorithm itself. Instead of augmenting existing RL algorithms with representation learning techniques like auxiliary losses or data augmentation, the authors demonstrate that contrastive learning can directly acquire effective representations corresponding to goal-conditioned value functions. The paper's key contribution lies in formally connecting contrastive learning with reward maximization, highlighting that the inner product between learned representations corresponds to a value function. This framework generalizes prior methods like C-learning and suggests new, simpler, and higher-performing goal-conditioned RL algorithms.

Theoretical Foundation

The paper grounds its approach in a rigorous theoretical framework, beginning with the definition of the goal-conditioned RL problem, characterized by states $s_t \in$, actions $a_t$, an initial state distribution $p_0(s)$, dynamics $p(s_{t+1} \mid s_t, a_t)$, a distribution over goals $p_g(s_g)$, and a reward function $r_g(s, a)$ for each goal. The reward is defined as the probability density of reaching the goal at the next time step: $$r_g(s_t, a_t) \triangleq (1 - \gamma) p(s_{t+1} = s_g \mid s_t, a_t)$$.

Proposition 1 establishes the equivalence between the Q-function and the probability of reaching a goal state $s_g$ under the discounted state occupancy measure: $$Q_{s_g}^\pi(s, a) = p^{\pi(\cdot \mid \cdot, s_g)}(s_{t+} = s_g \mid s, a)$$.

The core contribution is Lemma 1, which demonstrates that the critic function $f^*(s, a, s_f)$ that optimizes the contrastive learning objective (Eq. 4) is a Q-function for the goal-conditioned reward function (Eq. 1), up to a constant factor: $$\exp(f^*(s, a, s_f)) = \frac{1}{p(s_f)} \cdot Q_{s_f}^{\pi(\cdot \mid \cdot)}(s, a)$$.

Method Implementation

The contrastive RL algorithm involves alternating between fitting the critic function using contrastive learning and updating the policy using the actor loss (Eq. 5): $$\max_{\pi(a \mid s, s_g)} E_{\pi(a \mid s, s_g)p(s)p(s_g)}\left[f(s, a, s_f = s_g) \right] \approx E_{\pi(a \mid s, s_g)p(s)p(s_g)}\left[\log Q_{s_g}^{\pi(\cdot \mid \cdot)}(s, a) - \log p(s_g) \right]$$.

Figure 1: Reinforcement learning via contrastive learning. Our method uses contrastive learning to acquire representations of state-action pairs ($\phi(s, a)$) and future states ($\psi(s_f)$), so that the representations of future states are closer than the representations of random states. We prove that learned representation corresponds to a value function for a certain reward function. To select actions for reaching goal $s_g$, the policy chooses the action where $\phi(s, a)$ is closest to $\psi(s_g)$.

The critic is parameterized as an inner product between representations of state-action pairs and goal states: $f(s, a, s_g) = \phi(s, a)^T\psi(s_g)$. The paper provides a JAX implementation of the actor and critic losses (Algorithm 1). Contrastive RL (NCE) is presented as a simple algorithm that does not require multiple Q-values, target Q networks, data augmentation, or auxiliary objectives. A variant, contrastive RL (CPC), is derived using the infoNCE bound on mutual information.

Convergence Analysis

Lemma 2 provides a convergence guarantee for contrastive RL under the assumption of tabular states and actions, a Bayes-optimal critic, and an additional filtering step that excludes training examples where the probability of the trajectory under the commanded goal differs significantly from that under the actually reached goal. The result shows that performing contrastive RL on a static dataset results in one step of approximate policy improvement. Re-collecting data and iteratively applying contrastive RL leads to approximate policy improvement.

Experimental Evaluation

The paper evaluates contrastive RL algorithms on a suite of goal-conditioned tasks, including fetch reach, fetch push, sawyer push, sawyer bin, ant umaze, and point Spiral11x11. The results demonstrate that contrastive RL (NCE) outperforms prior methods on most tasks, including those with image observations. The method also demonstrates competitive performance in offline RL, outperforming baselines on five out of six D4RL AntMaze tasks. Ablation studies probe the design decisions of contrastive RL, demonstrating the benefits of the NCE objective and the importance of sampling random goals for the actor loss. Further experiments demonstrate the transferability of learned representations across tasks and the robustness of contrastive RL to environment perturbations.

Figure 2: Representation learning for image-based tasks. While adding data augmentation and auxiliary representation objectives can boost the performance of the TD3+HER baseline, replacing the underlying goal-conditioned RL algorithm with one that resembles contrastive representation learning (i.e., ours) yields a larger increase in success rates. Baselines: DrQ augments images and averages the Q-values across 4 augmentations; auto encoder (AE) adds an auxiliary reconstruction loss; CURL applies RL on top of representations learned via augmentation-based contrastive learning.

Implications and Future Directions

The paper demonstrates that contrastive representation learning can serve as a foundation for goal-conditioned RL, offering a fresh perspective on representation learning in RL. The proposed framework generalizes prior methods and suggests new algorithms that are simpler and more effective. The finding that RL algorithms can be constructed to resemble representation learning has broad implications for the design of future RL algorithms. The authors note that applying these methods to arbitrary RL problems remains an open question, while acknowledging that recent algorithms for this setting already bear a resemblance to contrastive RL. Future research could explore the rich set of ideas from contrastive learning to construct even better RL algorithms.

Conclusion

This paper makes a significant contribution to the field of reinforcement learning by bridging the gap between contrastive representation learning and goal-conditioned RL. The theoretical analysis, algorithmic innovations, and comprehensive experimental results provide a strong foundation for future research in this area. By demonstrating that RL algorithms can be designed to resemble representation learning, the authors open new avenues for developing more effective and efficient RL agents.