Equivariant Goal Conditioned Contrastive Reinforcement Learning

TL;DR

This paper introduces Equivariant CRL, a method that incorporates symmetry constraints into contrastive reinforcement learning to improve sample efficiency and generalization in goal-conditioned robotic manipulation tasks.

Contribution

The paper proposes a novel equivariant CRL framework that leverages task symmetries, formalizes goal-conditioned group-invariant MDPs, and demonstrates improved performance in simulation and offline settings.

Findings

Outperforms baselines in simulated tasks

Enhances sample efficiency and spatial generalization

Effective in both state-based and image-based settings

Abstract

Contrastive Reinforcement Learning (CRL) provides a promising framework for extracting useful structured representations from unlabeled interactions. By pulling together state-action pairs and their corresponding future states, while pushing apart negative pairs, CRL enables learning nontrivial policies without manually designed rewards. In this work, we propose Equivariant CRL (ECRL), which further structures the latent space using equivariant constraints. By leveraging inherent symmetries in goal-conditioned manipulation tasks, our method improves both sample efficiency and spatial generalization. Specifically, we formally define Goal-Conditioned Group-Invariant MDPs to characterize rotation-symmetric robotic manipulation tasks, and build on this by introducing a novel rotation-invariant critic representation paired with a rotation-equivariant actor for Contrastive RL. Our approach consistently outperforms strong baselines across a range of simulated tasks in both state-based and image-based settings. Finally, we extend our method to the offline RL setting, demonstrating its effectiveness across multiple tasks.