Equivariant Goal Conditioned Contrastive Reinforcement Learning
This work addresses sample efficiency and generalization challenges in robotic manipulation for reinforcement learning practitioners, representing an incremental improvement by applying equivariant constraints to an existing contrastive framework.
The paper tackled the problem of improving sample efficiency and spatial generalization in goal-conditioned robotic manipulation tasks by proposing Equivariant Contrastive Reinforcement Learning (ECRL), which leverages symmetries through a rotation-invariant critic and rotation-equivariant actor, consistently outperforming baselines in simulated tasks.
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.