Accelerating Goal-Conditioned RL Algorithms and Research

2024/08/20 by Michał Bortkiewicz, Bortkiewicz, Michał, Vivek Myers +10 · 8 citations
Computer Science · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Fuzzy Logic and Control Systems #Machine Learning (cs.LG)

paper · pdf · doi:10.48550/arxiv.2408.11052

openalex publication_date 2024/08/20 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Self-supervision has the potential to transform reinforcement learning (RL), paralleling the breakthroughs it has enabled in other areas of machine learning. While self-supervised learning in other domains aims to find patterns in a fixed dataset, self-supervised goal-conditioned reinforcement learning (GCRL) agents discover new behaviors by learning from the goals achieved during unstructured interaction with the environment. However, these methods have failed to see similar success, both due to a lack of data from slow environment simulations as well as a lack of stable algorithms. We take a step toward addressing both of these issues by releasing a high-performance codebase and benchmark (JaxGCRL) for self-supervised GCRL, enabling researchers to train agents for millions of environment steps in minutes on a single GPU. By utilizing GPU-accelerated replay buffers, environments, and a stable contrastive RL algorithm, we reduce training time by up to 22×. Additionally, we assess key design choices in contrastive RL, identifying those that most effectively stabilize and enhance training performance. With this approach, we provide a foundation for future research in self-supervised GCRL, enabling researchers to quickly iterate on new ideas and evaluate them in diverse and challenging environments. Website + Code: https://github.com/MichalBortkiewicz/JaxGCRL

Cited by

• Dense and Diverse Goal Coverage in Multi Goal Reinforcement Learning
• Consistent Zero-Shot Imitation with Contrastive Goal Inference
• Automaton Constrained Q-Learning
• Offline Goal-conditioned Reinforcement Learning with Quasimetric Representations
• Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Exploration
• NiceWebRL: a Python library for human subject experiments with reinforcement learning environments
• Contrastive Representations for Temporal Reasoning
• Identifying nonequilibrium degrees of freedom in high-dimensional stochastic systems

Related

• 1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities
• OGBench: Benchmarking Offline Goal-Conditioned RL
• Rethinking Goal-conditioned Supervised Learning and Its Connection to Offline RL
• Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching from Offline Data
• Demystifying the Mechanisms Behind Emergent Exploration in Goal-conditioned RL