Accelerating Goal-Conditioned Reinforcement Learning Algorithms and Research
Authors: Michał Bortkiewicz
Władysław Pałucki
Vivek Myers
Tadeusz Dziarmaga
Tomasz Arczewski
Łukasz Kuciński
Benjamin Eysenbach
Venue: ICLR 2025 Spotlight (2025)
Keywords: Deep Reinforcement Learning
GPU-accelerated Physics Simulators
Contrastive Learning
Unsupervised Reinforcement Learning
Primary area: reinforcement learning
TL;DR: This paper presents JaxGCRL, a high-performance codebase and benchmark designed for self-supervised goal-conditioned reinforcement learning, offering faster training and promoting more efficient RL research.

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\times$. 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. Code: [https://anonymous.4open.science/r/JaxGCRL-2316/README.md](https://anonymous.4open.science/r/JaxGCRL-2316/README.md)

Published as a conference paper at ICLR 2025
ACCELERATING GOAL-CONDITIONED REINFORCE- MENT LEARNING ALGORITHMS AND RESEARCH
Michal Bortkiewicz'
Whadyslaw Palucki2
Vivek Myers
Tadeusz Dziarmaga'
Tomasz Arczewski
Lukasz Kucinski2:5.0
Benjamin Eysenbach"
Warsaw University of Technology University of Warsaw UC Berkeley
'Jagiellonian University 'Polish Academy of Sciences
"DDEA NCBR PPrinceton University
michalborricclw..ccsggg....... com
Baadek.pallukk0@gaail.
ABSTRACT
Self-supervision has the potential to transform reinforcement learning (RL), paral- leling the breakthroughs it has enabled in other areas of machine learning. While self-vupe learning in other domains aims to find patterns in fixed dataset, self-supervised goal- coal-conditioneds reinforcement learning (GCRL) agents discover new behaviors by learning from the goals achieved during unstructured interaction with environment. However, these methods have failed to see similar success, both due to lack of data from slow environment simulations well lack of sta- ble algorithms. We take 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 single GPU. By utilizing GPU-accelerated replay buffers, environments, and sta- ble contrastive RL algorithm, we reduce training time by up to 22x. Additionally, we assess key design choices in contrastive RL, identifying those that most effec- tively stabilize and enhance training performance. With this approach, we provide foundation for future research in self-supervised GCRL, enabling researchers to quickly iterate on new ideas and evaluate them in diverse and challenging environments. Code: https: hete:://ttthnne /ggtttb...../. htt./Miichaaaartkkcca JaxGGLLL
INTRODUCTION
Self-supervised learning has significantly influenced machine learning over the last decade, transform- ing how research is done in domains such as natural language processing and computer vision (Chen al., 2020b; Dosovitskiy et al., 2020; Vaswani et al., 2017). In the context of reinforcement learning (RL), most self-supervised prior methods apply the same recipe that has been successful other domains: learning representations or models from a large, fixed dataset (Hoffmann et al., 2022; Sardana et al., 2024; Muennighoff et al., 2023). However, the RL setting also enables a fundamentally different type of self-supervised learning: rather than learning from fixed dataset (as done in NLP and computer vision), self- self-supevvised reinforcement leamer can collect its own dataset. Thus, rather than learning a representation of a dataset. the self-supervised reinforcement learner acquires representation of environment or of behaviors and optimal policies therein. Self-supervised reinforcement learners may address many of the challenges that stymie today's foundation's 's models: reasoning about the consequences of actions (Rajani et al., 2019; Kwon et al., 2023) (i.c., counterfac- tuals (Bhargava and Ng, 2022; Jin et al., 2023)) and long horizon planning (Bhargava and Ng, 2022; Du et al., 2023; Guan et al., 2023).
In this paper, we study self-supervised RL in an online setting: an agent interacts in an environment without reward learn representations, which later used to quickly solve downstream tasks. We focus on goal-condn goal-conditioned reinforcement learning (GCRL) algorithms, which aim to use these unsupervised interactions to learn policies for achieving various goals essential capability for

---

Published as a conference paper at ICLR 2025
Fast Training
Parallel Environments
Vectorized Replay Buffers
Contrastive Self- Supervision
Juo03I, tesert Usite ieqummeatiik
OM envisteps
How
orDroore P. poak pastomaass
122
0.5M env stogs
(Ct0PP))
(X1024)
aditrani d
couuttyy preciiss strind anf
Extract Gook-Beaching Pollcy: rie g) ang min (6((, siaal
>20 times faster training!
Figure I: JaxGCRL is fast. It learns real-reaching policies for Ant in 10 minutes on 1 GPU. This paper releases a GCRL benchmark and baseline algorithms that enable research and experiments to be done in minutes.
multipurpose robots. Prior work has proposed several algorithms for self-supervised RL (Eysenbach al., 2022; Zheng et al., 2024; Myers et al., 2023; 2024a), including algorithms that focus on learning goals (Erraqabi et al., 2022). However, these methods were limited by small datasets and infrequent online interactions with the environment, which has prevented them from exploring the potential emergent properties of self-supervised reinforcement learning on large-scale data.
The main goal of this paper is introduce JaxGCRL: extremely fast GPU-accelerated codebase and benchmark for effective self-supervised GCRL research. For instance, an experiment with 10 million environment steps lasts only around 10 minutes on a single GPU, which is up to 22x faster than in the original contrastive RL codebase (Eysenbach et al., 2022) (Fig. 1). This speed allows researchers to "interactively" debug their algorithms, changing pieces and getting results in near real-time for multiple random seeds on single GPU without the hustle of distributed training. Consequently, JaxGCRL eliminates the barriers to entry to state-of-the-aa GCRL research, making it more accessible to under-resourced institutions.
To achieve this training and performance improvement, we combine insights from self-supervised RL with recent advances in GPU-accelerated simulation. The first key ingredient is recent work on GPU-accelerated simulators, both for physics (Freeman et al., 2021; Thibault al., 2024; Liang et al., 2018) and other tasks (Matthews et al., 2024; Dalton et al., 2020; Fischer et al., 2009; Bonnet et al., 2024; Rutherford et al., 2024) that enable users to collect data up to 1000 faster (Freeman et al., 2021) than prior methods based on CPUs. The second key ingredient is highly stable algorithm build upon recent work on contrastive RL (CRL) (Eysenbach et al., 2022; Zheng et al., 2024), which uses temporal contrastive learning learn a value function. The third key ingredient is suite of tasks for evaluating self-supervised RL agents which is not only blazing fast but also stress tests the exploration and long-horizon reasoning capabilities of RL policies.
The contributions of this work are as follows:
JaxGCRL codebase: blazingly fast J'T-compiled training pipeline for GCRL experiment. JaxGCRL benchmark: we introduce suite of 8 GPU-accelerated state-based environments that help to accurately assess GCRL algorithm capabilities and limitations.
Extensive empirical analysis: we evaluate important CRL design choices, focusing on key algo- rithm components, architecture scaling, and training in data-rich settings.
2 RELATED WORK
We build upon recent advances in GCRL, self-supervisedd RL, and hardware-ace accelleraee physics simulators, showing that CRL enables fast and reliable training across a diverse suite of environments.
2.1 GOAL-CCNDITTONEE REINFORCEMENT LEARNING
GCRL is special case of the general multi-task RL setting, in which the potential tasks are defined by goal states that the agent is trying reach in an environment (Kaelbling, 1993; Ghosh et al., 2019;
2

---

Published as a conference paper at ICLR 2025
Nair al., 2018). Achieving any goal is appealing for generalist agents, it allows for diverse behaviors without needing specific reward functions for each task (each state defines a task when seen as a goal) (Schaul et al., 2015). GCRL techniques have manipulation (Andrychowicz et al., 2017; Eysenbach et al., 2021; Ghosh et al., 2021; Ding et al., 2022; Walke et al., 2023) and navigation (Shah et al., 2023; Levine and Shah, 2023; Manderson et al., 2020; Hoang et al., 2021). Recent work has shown that representations of goals can be imbued with additional structure enable capabilities such as language grounding (Myers et al., 2023; Ma et al., 2023), compositionality (Liu et al., 2023; Myers et al., 2024b; Wang et al., 2023), and planning (Park al., 2023; Eysenbach et al., 2024). We show how GCRL techniques based on contrastive learning (Eysenbach et al., 2022) can be scaled with GPU acceleration to enable fast and stable training.
success in domains such as robotic
2.2 ACCELERATING DEEP REINFORCEMENT LEARNING
Deep RL has only recently become practical for many tasks, in part due to improvements in hardware support for these algorithms. Distributed training has enabled RL algorithms to scale across hundreds of GPUs (Mnih et al., 2016; Espeholt et al., 2018; 2020; Hoffman et al., 2022). To resolve the bottleneck of environment interaction with CPU-bound environments, various GPU-accelerated environments have been proposed (Matthews al., 2024; Freeman al., 10 Bonnet et al., 2024; Rutherford et al., 2024; Liang et al., 2018; Dalton et al., 2020; Makoviychuk et al., 2021; Lange, 2022; Lu et al., 2022). Most of these works rely on JAX (Bradbury al., 2018: 8; Heek et al., 2023; Hennigan et al., 2020), which enables JIT compilation, operator fusion and other components necessary for efficient vectorized code execution. These features significantly accelerate data collection by supporting rollouts in hundreds of parallelized environments. We build on these advances to scale self-supervised RL to data-rich settings.
2.3 SELF-SUPERVISED RL
Self-supervised training has enabled key breakthroughs in language modeling and computer vision (Sermanet et al., 2017; Zhu et 2020; Devlin et al., 2019; He et al., 2022; Mikolov et al., 2013). In the context of RL. by the term "selffsuperriiee".. we mean techniques that can be learned through interaction with an environment without reward signal. Perhaps the most successful form of self-supervision has been in multi-agent games that can be rapidly simulated, such as Go and Chess, where self-play has enabled the creation of superhuman agents (Silver et al., 2016; 2017; Zha et al., 2021). When learning goal-reaching agents, another basic form of self-supervision is to relabel trajectories successful demonstrations of the goal that was reached, even if it differs from the original commanded goal (Kaelbling. 1993; Venkattaramanujam et al., 2020). This technique has seen recent adoption "hindsight experience replay" for various deep RL algorithms (Andrychowicz et al., 2017; 7; Ghosh et al., 2021; Chebotar et al., 2021; Rauber et al., 2021; ; Eysenbach et al., 2022).
Another perspective self-supervision is intrinsic motivation, broadly defined as when an agent computes its own reward signal (Barto, 2013). Intrinsic motivation methods include curiosity (Barto, 2013; Bellemare et al., 2016; Baumli et al., 2021), surprise minimization (Berseth et al., 2019; Rhinehart et al., 20 and empowerment (Klyubin et al., 2005; de Abril and Kanai, 2018; Choi al., 2021; Myers et al., 2024a). Closely related are skill discovery methods, which aim construct intrinsic rewards for diverse collections of behavior (Gregor et al., 2016; Eysenbach et al., 2019; Sharma et al., 2020; Kim et al., 2021; Park et al., 2021: Laskin et al., 2022; Park et al., 2024). Self-supervised RL. methods have been difficult to scale due to the need for many environment interactions (Franke et al., 2021; Mnih et al., 2015). This work addresses that challenge by offering a fast and scalable contrastive RL algorithm on benchmark of diverse tasks.
2.4 RL BENCHMARKS
The RL community has recently started to pay greater attention to how RL research is conducted, reported, and evaluated (Henderson et al., 2018; Jordan et al., 2020; Agarwal et al., 2022). A key issue is a lack of reliable and efficient benchmarks: it becomes hard to rigorously compare novel methods when the number of trials needed to see statistically significant results across diverse settings ranges in the thousands of training hours (Jordan al., 2024). Some benchmarks that have adoption include OpenAl sym/Gymnasium (Brockman et al., 2022; Towers et al., 2024), DeepMind Control Suite (Tassa et al., 2018), and D4RL (Fu et al., 2021). More recently, hardware-acceleran
3

---

Published as a conference paper at ICLR 2025
versions of some of these benchmarks have been proposed (Gu et al., 2021; Matthews al., 2024; Koyamada et al., 2023; Makoviychuk et al., 2021; Nikulin et al., 2023). However, the RL community still lacks advanced benchmarks for goal-conditioned methods. We address this gap with JaxGCRL, significantly lowering the GCRL evaluation cost and thereby enabling impactful RL research.
3 PRELIMINARIES
In this section, we introduce notation and preliminary definitions for goal-conditioned RL and the contrastive RL. method, which serves as the foundation for this work.
In the goal-conditioned reinforcement learning setting. an agent interacts with a controlled Markov process (CMP) M (S, p. P., 7) to reach arbitrary goals (Kaelbling, 1993; Andrychowiez et al., 2017; Blier et al., 2021). At any time / the agent will observe a state and select a corresponding action de E A. The dynamics this interaction are defined by the distribution p($++1 St, at), with an initial distribution (5())) over the state the start of trajectory, for 8t E S and a, A.
For any goal g E S, we cast optimal goal-reaching as a problem of inference (Borsa et al., 2019; Barreto et al., 2022; Blier et al., 2021; Eysenbach et al., 2022): given the current state and desired goal, what is the most likely action that will bring us toward that goal? As we will see, this is equivalent solving the Markov decision process (MDP) Mg obtained by augmenting M with the goal-conditioned reward ry (st, a+) - (1 7) $ = g St, at). Formally, goal-conditioned policy π (a s, g) receives both the current observation of the environment as well as a goal S.
We denote the k-step oction-conditioned policy distribution ( So: ag) as the distribution of states k steps in the future given the initial state 8p and action dg under We define the discounted
state visitation distribution as př, (s+ 8, s,a) _ (1 -0 7ip (s+ 8, s,a). which we interpret
as the distribution of the state T steps in the future for T ~ Geom 7). This last expression
is precisely the Q-function of the policy πt. ',9) for the reward rg: Q7 a) _ pa 8, a) (see
Appendix 1. For a given distribution over goals g ~ Pg. we can now write the overall objective as
max Byo(so [g Su, ao)] <(b--)
3.1 CONTRASTIVE CRITIC LEARNING
CRL is an actor-critic method that aims to solve Eq. (1). The critic is represented as a state-action- goal value function [s, a, g), which provides the likelihood of future states and how various actions influence the likelihood of future states. This function satisfies:
[s, g) x (g 8, a, = (8,
Therefore, we can treat f(s,,
approximation of the Q-function and use it to train the actor. This approach builds on previous research that frames learning of this critic as a classification problem (Eysenbach et al., 2022; 2021; Zheng et al., 2023; 2024; Farebrother et al., 2024; Myers et al., 2024b). Training is performed on batches of (s, a, g) to classify whether or not is the future state corresponding to the trajectory starting from (s, a). Thus, in each sample from batch (s; Qs gi) E B the goal 9; is sampled from the future of the trajectory containing (s),,.).
The family of CRL algorithms consists of the following components: (a) the state-action pair and goal state representations, (68, a) and (gg)), respectively; (b) the critic, which is defined as an energy function (,. a, measuring some form of similarity between Φ(s, a) and ((g); and (c) contrastive loss function, which is function of the matrix containing the critic values {ff.. (si, a, 3 over the elements of the batch B. The base contrastive loss we study will be the infoNCE objective (Sohn, 2016), modified to use symmetrized (Radford et al., 2021) critic parameterized with fz-distances (Eysenbach et al., 2024). The final objective for the critic can thus be expressed as:
Σ
|8|
elovieion))
efo (onwaen S ]
min Eg 0.
log
log k
1=1
5
i=1
4

---

Published as a conference paper at ICLR 2025
Figure 2: JaxGCRL benchmark: New suite of GPU-acoelerated environments for studying GCRL. In this setting, the agent does not receive any rewards demonstrations, making some of these tasks excellent testbed for studying exploration and oong-hoizoon reasoning. Our accompanying implementation GCRL algorithms trains with more than environment steps per second on a single GPU, enabling rapid experimentation.
where
fo. (s, = |o(s, a) - (2g)ll2.
This loss contrasts each positive sample with the batch of negative samples. Other losses, which tested in Section 5.3, are further discussed in Appendix A.2.
3.2 POLICY LEARNING
We use a DtPG-style policy extraction loss to learn a goal-conditioned policy, Ta(als, g). by optimizing the critic fa.s (Lillicrap et al., 2016):
max E plsssaatllll.e (s, a'))
0
Each batch is formed by sampling (s, a) pairs uniformly and then sampling goals from the states that occur after in a trajectory.
4 JAXGCRL: A NEW BENCHMARK AND IMPLEMENTATION
JaxGCRL is an efficient tool for developing and evaluating new GCRL algorithms. It shifts the bottleneck from compute to implementation time, allowing researchers test new ideas, like contrastive objectives, within minutes a significant improvement over traditional RL workflows.
4.1 JAXGCRL SPEEDUP ON A SINGLE GPU
We compare the proposed fully JIT-compiled implementation of CRL in JaxGCRL to the original implementation from Eysenbach al. (2022). In particular, Fig. 1 shows the performance in Ant environment along with the experiment's wall-clock time. The speedup in this configuration is 22-fold, with the new implementation reaching training speed of over 16500 environment steps per second with 1:16 update to data (UTD) ratio. For results with other UTD, refer to Section 5.6. Complete speedup results are provided in Appendix A. Importantly, BRAX physies simulator differs from the original MuJoCo, so performance numbers here may vary slightly from prior work.
4.2 JAXGCRL ENVIRONMENTS IN THE BENCHMARK
To evaluate the performance of GCRL methods, we propose JaxGCRL benchmark consisting of diverse continuous control environments. These environments range from simple, ideal for quick checks, to complex, requiring long-term planning and exploration for success. The following list provides brief description of each environment, with the technical details summarized in Table 1:

---

Published as a conference paper at ICLR 2025
Reacher (Brockman et al., 2022). A 2D manipulation task involves positioning the end part of 2-segment robotic arm a specific location sampled uniformly from a disk around an agent.
Half Cheetah (Wawrzyíski, 2009). In this 2d task, a 2-legged agent has get goal that is sampled from one of the possible places, one on each side of the starting position.
Pusher (Brockman al., 2022). This is 3d robotic task, that consists robotic arm and movable circular object resting on the ground. The goal is to use the arm to push the object into the goal position. With each reset, both position of the goal and movable object are selected randomly.
Ant (Schulman et al., 2016). A re-implementation of the MuloCo Ant with quadruped robot that needs to walk to the goal randomly sampled from circle centred the starting position.
Ant Maze (Fu et al., 2021). This environment uses the same quadruped model as the previous Ant, but the agent must navigate a maze to reach target. We prepared 3 different mazes varying size and difficulty. In each maze, the goals are sampled from set of listed possible positions.
Ant Soccer TTunaauuuuaakon et al., 2020). this environment, the Ant has to push a spherical object into goal position that is sampled uniformly from circle around the starting position. The position of the movable sphere is randomized on the line between an agent and goal.
Ant Push (Fu et al., 2021). To reach goal, the Ant has to push a movable box the way. If the box is pushed in the wrong direction, the task becomes unsolvable. Succeeding requires exploration, understanding block dynamics, and how it changes the layout of the maze.
Humanoid (Tassa et al., 2012). Re-implementing the Mujoco task involves navigating a complex humanoid-like robot to walk towards goal sampled from a disk centred at the starting position.
In Fig. 3, we report the baseline results for the proposed benchmark. It is worth noting that for most environments, experiments involving 50M environment steps can be completed in less than an hour.
4.3 CONTRASTIVE RL DESIGN CHOICES
JaxGCRL streamlines and accelerates the evaluation of new GCRL algorithms, enabling quick assess- ment of key CRI. design choices:
Energy functions. Mcasuring the similarity between samples can be achieved various ways, resulting in potentially different agent behaviors. Our analysis in following sections include cosine similarity (Chen et al., 2020a), dot product (Radford et al., 2021), and negative L1 and Lg dis- tance (Hu et al., 2023), detailed list can be found in Appendix A.2. Even though there is consensus on the choice energy functions for temporal representations, recent works showed that they should abide by quasimetric properties (Wang al., 2023; Myers al., 2024b).
Contrastive losses. Beside InfoNCE-type losses (van den Oord et al., 2019; Eysenbach et al., 2022), we evaluate FlatNCE-like losses (Chen al., 2021), and Monte Carlo version of unsupervised loss (Touati and Ollivier, 2021). Additionally, we test novel objectives inspired by preference optimization for large language models (Calandriello et al., 2024). Specifically, we evaluate DPO, IPO, and SPPO, which increase the scores of positive samples and reduce the scores of negative ones. full list of contrastive objectives can be found in Appendix A.2.
Architecture scaling. Scaling neural network architectures to improve performance is common practice in other areas of deep learning, but it remains relatively underexplored in RL models, as noted in Nauman al. (2024a;b). Recently, Zheng et al. (2024) showed that CRL might benefit from deeper and wider architectures with Layer Normalization (Ba et al., 2016) for offline CRL in pixel-based environments; want to examine whether this also holds for online state-based settings.

---

Published as a conference paper at ICLR 2025
5 EXAMPLES OF FAST EXPERIMENTS POSSIBLE WITH THE NEW BENCHMARK
The goal of our experiments is twofold: (1) to establish a baseline for the proposed JaxGCRL environments, and (2) evaluate CRL performance in relation key design choices. In Section 5.1, we define setup that is used for most of the experiments unless explicitly stated otherwise. First in Section 5.2, we report baseline results on JaxGCRL. Second, in Sections 5.3 and 5.4, we try to understand the influence of design choices on CRL learning performance. Third, in Section 5.5, we asses those design choices in data-rich setting with 300M environment steps. Lastly, in Section 5.6, we explore the relation between performance and UTD.
5.1 EXPERIMENTAL SETUP
Our experiments use JaxGCRL suite of simulated environments described in Section 4.2. We evaluate algorithms in an online setting for 50M environment steps. We compare CRL with Soft Actor- Critic (SAC) (Haarnoja et al., 2018), SAC with Hindsight Experience Replay (HER) (Andrychowicz et al., 2017), TD3 (Fujimoto et al., 2018), TD3+HER, and PPO (Schulman et al., 2017). For algorithms with HER, we use final strategy relabelling, i.c. relabeling goals with states achieved at the end of the trajectory. In the majority of experiments, we use CRL with L2 energy function, symmetric InfoNCE objective, and tuneable entropy coefficient for all methods. See Appendix B for details. We use two performance metrics: success rate and time near goal. Success rate measures whether the agent reached the goal least once during the episode, while time near goal indicates how long the agent stayed close to it. We use a sparse reward for all baselines, with r = 1 when the agent is in goal proximity and 0 otherwise. We define goal-reaching as achieving proximity below the goal distance threshold defined in Table 1. The implementations of PPO and SAC are sourced from the Brax repository (Freeman et al., 2021), while TD3, HER, and CRL are partially based on Brax.
5.2 JAXGCRL BENCHMARK RESULTS
We establish baseline results for JaxGCRL with all algorithms in Fig. 3. Clearly, CRL achieves the highest performance across tested methods, resulting in non-trivial policies even in the hardest tasks in terms of high-ddmensional state and action spaces (Humanoid) and exploration (Pusher and Ant Push). However, the performance in these challenging environments is low, indicating room for improvement for future contrastive RL methods. As expected, HER can improve performance for both TD3 and SAC. In contrast, PPO performs poorly across all tasks, likely due to the challenges posed by the sparse reward setting. Additional experiments on JaxGCRL environments with design choices discussed in the following sections can be found in Appendix A.4.
Rescher
Pusher Hard
Hummaoid
Seecoms rate
CRL
Success rale
Success rale
0.50
0.9
Success
0.2
0.5
0.25
0.66
0.00
0.0
0.0
10
20
30
40
10
60
0
10
40
0
10
20
30
Time minates
minutes
Time
milautes
Time in. minutes
Ant U-Maze
Big, Maze
Aut Seccer
Aut Push
0.21
Suoves tale
0.4
0.2
Surees rale
Success rate
(3.2
Successs
0.1
0.2
0.0
0.0
10
220
30
40
50
01
10
201
30
10
30
62)
70
0
10
200
30
10 20
30
40
50
600
70
Time in miiuiee
Time
Time minnuees
Time minauees
Figure 3: Bascline results in JaxGCRL benchmark. Success rates of all the baseline algorithms for 50M environment steps for every JaxGCRL environment. CRL outperforms other baselines in most of the environments. The training speed is a function of the environment complexity, method complexity. and physics backend; see Appendix A.4. Specifically, due to differences in how each method works, the speed varies greatly in the same environments; this can be best seen with the PPO method being significantly faster than others due to it not using replay buffer, which frees up GPU memory for more parallel environment simulations. Results are reported as the interquartile mean (IQM) along with its standard error, based 10 seeds.
7

---

Published as a conference paper at ICLR 2025
IQM
IQM
Symmetric InfoNCE
Symmetric InfoNCE
FlatNCE
FlatNCE
InfoNCE Backward
InfoNCE Backward
InfoNCE
InfoNCE
Binary
Binary
FlatNCE Backward
FlatNCE Backward
Forward-Backward
Forward-Backyee
SPPO
SPPO
DPO
DPO
IPO
IPO
0.2
0.3
40
80
120
Success rate
Time near goal
Figure 4: InfoNCE-based loss functions perform best. The critic loss functions that achieve the highest success rates are based on InfoNCE and DPO. However, DPO policies tend to stay at the goal for shorter duration. IQMs averaged over 10 seeds and plotted with one standard error.
5.3 CONTRASTIVE ORJECTIVES AND ENERGY FUNCTIONS COMPARISON
Contrastive objective and energy function are the two main components of contrastive methods, serv- ing as the primary drivers their final performance. We evaluate CRL with 10 different contrastive objectives, as defined in Section 4.3 across five environments: Ant Soccer, Ant, Ant Big Maze, Ant U-Maze, Ant Push, and Pusher, and report aggregated performance. For the energy function, we use L2, as it consistently resulted in the highest performance for the CRL method, especially regarding time near goal, see Appendix A.2. Additionally, we apply logsumexp regularization with coefficient of 0.1 to each objective. This auxiliary objective is essential, as without it, the performance of InfoNCE deteriorates significantly (Eysenbach et al., 2022). The analysis presented in Fig. reveals that the originally proposed NCE-binary objective, along with forward-backe SPPO, are the least effective objectives among those evaluated. However, for other InfoNCE-derived objectives, it is difficult to determine the best one, their performance is similar. Interestingly, CRL seems fairly robust to the choice of contrastive objective.
IPO, and
5.4 SCALING THE ARCHITECTURE
In this section, explore how increasing the size of actor and critic networks, in terms of both depth and width, influences CRL performance. We evaluate the aggregated performance during the final 10M steps of 50 million-step training process across three environments: Ant, Ant Soccer, and Ant U-Maze. We use the 1.2, Symmetric InfoNCE and logsumexp regularizer coefficient 0.1 for all architecture sizes.
IOM
Depth Deptn
Depth Depth Depth
Layar Norm Depth 4 #1024
No Layer Nom
Depth #1024
0.4
0.5
0.6
120
180
Depth 4 #512 LN
Depth #512
Success rate
Time near goal
Depth 4 #256 LN
Dopth #256
Depth 3 #1024
Depth #1024
(a) Neurons per layer 256
Depth #256
Depth #256
Depth
IOM
Depth Depth
IQM
Depth 2 #1024 LN
Depth 2 #512 LN
Depth #256 LN
Depth #512
Depth #1102
Depth #256
Depti
1.00
0.4
0.5
0.6 0.7
120
180
240
P(X > Y)
Success rate
Time near goal
Success Rate
Time Neur Goal
(b) Neurons per layer 512
330
IQM
Depth
Depth Depth
220
Depth
120
160
420
6x1024
Success rate
Time near goal
2.0
(c) Neurons per layer = 1024
Mumber of environmnnttt atoza (n
Number anatonmentt milliona
Figure Scaling the critic and actor networks. Increasing the width and depth generally enbances performance, but performance levels off for deeper architectures width of 1024. Aggregated metries, seeds per configuration.
Figure Layer normalization enables stable per- formance improvement. Using Layer Normaliza- tion (LN) in the largest architecture allows for contin- ued learning even after reaching the saturation point a standard large architecture.
8

---

Published as a conference paper at ICLR 2025
We present the results of this scaling experiment in Fig. 5 and observe that increasing both the width and depth tends to increase performance. However, performance does not further improve when increasing the depth for width 1024 neurons. Our next experiment studies whether layer normalization can stabilize the performance of these biggest networks (width of 1024 neurons, depth of 4). Indeed, the results in Fig. 6 show that adding layer normalization before every activation allows better scaling properties, especially for bigger networks.
5.5 SCALING THE DATA
We evaluate the benefits CRL gains from training in a data-rich setting. In particular, we report performance for large architectures (studied in Section 5.4) with different combinations of energy functions and contrastive objectives for 300M environment steps in Fig. 7. We observe that the L2 energy function with InfoNCE objective configuration outperforms all others by substantial margin, leading to a higher success rate and time near the goal across three locomotion tasks. Interestingly. the dot product energy function performs best in the object manipulation task (Ant Soccer). This indicates that only a subset of wide array of design choices performs well when scaling CRL. Additionally, there is still room for improvement in scaling CRL with data, as the success rate in Ant Soccer and Ant Big Maze remains around 40%. For additional experiments, refer to Appendix A.6
5.6 GRADIENT UPDATES TO DATA RATIO
JaxGCRL enables efficient exccution of extensive experiments. Leveraging this capability, explore the effect the model's update frequency in CRL by evaluating a range of UTD ratios. In particular, examine ratios (1:1, 1:8, 1:16, 1:24, 1:32, 1:49, 1:48) in five environments: Ant, Ant Soccer, Ant U-Maze, Pusher Hard, and Ant Push. Interestingly, we only observe a significant increase in performance for Pusher Hard with higher number of updates, while in other environments, it leads to decreased or similar performance. With a UTD ratio of 1:16, our code is 22x faster than prior implementations, and with a lower frequency of gradient updates, can be even faster.
0.4
0.6
Ant Soccer
Ant
Ant U Maze
Pushor
Ant Push
Success
Symeent eesiCE
Li levite
Kimber oteeriioneer ooss Fooe)
Samber of ofw.rroeeeeeeettte wew.emeeeeee kon
Success ralu
0.8
Blg Mace
Art Soccar
Ant
Ant U Maze
Pushor
Ant Push
0.5
0.2 12
0.2
20
30
100
150
80
160
200
400
60
120
Time near goal
Kimber Kunber mtereooreerr madc (h Booe!
Surber dE wereerrreee di miliinn
Figure 8: More gradients & better performance. Success rate (top) and Time near goal (bottom) for different UTD ratios. Increasing the UTD ratio for CRL increases performance only for Pusher.
Figure 7: JaxGCRL allows researchers to study en- ergy functions and critic losses over hundreds of millions of steps. Among tested configurations, L. with InfoNCE objective performs best locomotion environments when data is abundant.
Key takeaways from empirical experiments:
Experiments with 10M steps can be completed in minutes, while those with billions of environment steps can be done in a few hours using JaxGCRL on single GPU.
CRL is the only method that can learn effectively in all proposed environments without needing a high UTD ratio in most cases. It benefits greatly from using large architectures, especially when Layer Normalization applied.
Different combinations of energy and contrastive functions lead to different outcomes: some primarily improve success rates, while others extend the time spent near the goal.

---

Published as a conference paper at ICLR 2025
Taken together, these experiments not only provide guidance good design decisions for self- supervised RL, but also highlight how our fast codebase and benchmark can enable researchers to quickly iterate ideas and hyperparameters.
6 CONCLUSION
In this paper, we introduce JaxGCRL, a very fast benchmark and codebase for coal-conditioned RL. The speed of the new benchmark enables us to rapidly study design choices for state-based CRL, including the network architectures and contrastive losses. We expect that self-supervised RI. methods will open the door to entirely new learning algorithms with broad capabilities that go beyond the capabilities of today's foundational models. The key step towards this goal is accelerating and democratising self-supervised RL research that any lab carry it out regardless of its computing capabilities. Open-sourcing the proposed codebase with easy-to-implement self-supervised RL methods is important step in this direction.
Limitations. The GCRL paradigm complicates the process of defining goals that are not easily expressed single state, making it infeasible for some applications. Additionally, our benchmark environments and methods assume full observability and that goals are being sampled from known goal distribution during training rollouts. Future work should relax these assumptions to make self-supervised RI. agents useful in more practical settings. We also only investigate online GCRL settings.
Reproducibility Statement. All experiments can be replicated using the provided publicly avail- JaXGGLL This repository
able JaxGCRL code https:/ /ttpp::/giiiii thub. com/Mi includes comprehensive instructions for setting up the environment, running the experiments, and evaluating the results, making it straightforward reproduce the findings.
Acknowledgments. This research was substantially supported by the National Science Centre, Poland (grant no. / and the Warsaw University of Technology through the Excellence Initiative: Research University (IDUB) program. We gratefully acknowledge the Polish high-performance computing infrastructure, PCSS PI.Cloud, for providing computational resources and support under grant no. p10334-01, and PLGrid (HPC Center: ACK Cyfronet AGH), for providing resources and support under grant no. PLG/20240110040. We would also like to recognize the founding of the DoD NDSEG fellowship for Vivek Myers. This work was partially conducted using Princeton Research Computing resources Princeton University, consortium led by the Princeton Institute for Computational Science and Engineering (PICSciE) and Research Computing.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G. Bellemare. Deep Reinforcement Learning at the Edge of the Statistical Precipice. arXiv:2108. 132224,08..3364, 2022.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight Experience Replay. Neural Information Processing Systems, volume 30, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. International Conference on Advances in Artificial Intelligence, 2016.
André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor Features for Transfer in Reinforcement Learning. Curran Associates, wec.aaxivv1600.053311 2022.
Andrew G. Barto. Intrinsic Motivation and Reinforcement Learning. Gianluca Baldassarre and Marco Mirolli, editors, Intrinsically Motivaled Learning in Natural and Artificial Systems, pp.
17-47. Springer, 2013.
Kate Baumli, David Warde-Farley, Steven Hansen, and Volodymyr Mnih. Relative Variational Intrinsic Control, AAAI Conference Artificial Intelligence, volume 35, pp. 0732,6740, 2021.
10

---

Published as a conference paper at ICLR 2025
Mare G. Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying Count-Based Exploration and Intrinsic Motivation. Neural Information Processing Systems, 2016.
Glen Berseth, Daniel Geng, Coline Devin, Nicholas Rhinehart, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine. Smirl: Surprise Minimizing Reinforcement Learning in Unstable Environments. International Conference on Learning Representations, 2019.
Prajjwal Bhargava and Vincent Ng. Commonsense Knowledge Reasoning and Generation With Pre- ree-Trained Language Models: A Survey. 2022.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning Successor States and Goal-Dependent
Values: A Mathematical Viewpoint. arXiv:2101.077
2021.
Clément Bonnet, Daniel Luo, Donal Byrne, Shikha Surana, Sasha Abramowitz, Paul Duckworth, Vincent Coyette, Laurence I. Midgley, et al. Jumanji: A Diverse Suite of Scalable Reinforcement Learning Environments in JAX. International Conference on Learning Representations, 2024.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado van Hasselt, David Silver, and Tom Schaul. Universal Successor Features Approximators. International Conference on Learning Representations, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary. Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable Transformations of Python+Numpy Programs. 2018.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai Gym. J. Open Source Softw, 2022.
Daniele Calandriello, Daniel Guo, Remi Munos, Mark Rowland, Yunhao Tang, Bernardo Avila Pires, Pierre Harvey Richemond, Charline Le Lan, et al. Human Alignment of Large Language Models Through Online Preference Optimisation. arXiv:2403.0866
2024.
Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian. Chelsea Finn, et al. Actionable Models: Unsupervised Offline Reinforcement Learning of Robotic Skills. International Conference on Machine Learning, 2021.
Junya Chen, Zhe Gan, Xuan Li, Qing Guo, Liqun Chen, Shuyang Gao, Tagyoung Chung, Yi Xu, et al. Simpler, Faster, Stronger: Breaking the Log-K K Curse on Contrastive Learners With Flatnce. arXiv:2107.0113 1100000.1155 1552 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A Simple Framework for Contrastive Learning of Visual Representations. arXiv:2002.05709) 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A Simple Framework for Contrastive Learning of Visual Representations. International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607, 2020b.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Varia- tional Empowerment as Representation Learning for Goal-Conditioned Reinforcement Learning. International Conference on Machine Learning, pp. 1953-1963, 2021.
Steven Dalton, luri Frosio, and Michael Garland. Accelerating Reinforcement Learning Through GPU Atari Emulation. Neural Information Processing Systems, volume 33, pp. 19773-19782, 2020.
Ildefons Magrans de Abril and Ryota Kanai. A Unified Strategy for Implementing Curiosity and Empowerment Driven Reinforcement Learning. arXiv: 1.X.v:1806606660 1806006505 2018
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding Conference of the North American Chapter of the Association for Computational Linguistics, 2019.
Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing Goal-Conditioned Reinforcement Learning With Variational Causal Reasoning. Neural Information Processing Systems, 2022.
A. Dosovitskiy, L. Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, M. Dehghani. Matthias Minderer, G. Heigold, S. Gelly, Jakob Uszkoreit, and N. Houlsby, An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale. International Conference on Learning Representations, 2020.
Yuqing Du, Eliza Kosoy, Alyssa Dayan, Maria Rufova, Pieter Abbeel, and Alison Gopnik. What Can Al Learn From Human Exploration? Intrisically-Mootiaaees Humans and Agents in Open- World

---

Published as a conference paper at ICLR 2025
Exploration. NeurlPS 2023 Workshop: Information-Theorer Principles in Cognitive Systems, 2023.
Akram Erraqabi, Marlos C. Machado, Mingde Zhao, Sainbayar Sukhbaatar, Alessandro Lazaric, Ludovic Denoyer, and Yoshua Bengio. Temporal AbstraactionsAAugmennn Temporally Contrastive Learning: An Alternative to the Laplacian in RL. arXiv:2203. 1369, 2022.
Lasse Espeholt, Raphaël Marinier, Piotr Stanezyk, Ke Wang, and Marein Michalski. SEED RL:
Scalable and Efficient Re With Accelerated Central Inference. International Conference on
Learning Representations, 2020.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: Scalable Distributed Deep-RL. With Importance Weighted Actor-Leamer Architectures. Interna- tional Conference on Machine Learning, 2018
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity Is All You Need: Learning Skills Without a Reward Function. International Conference on Learning Representations, 2019.
Benjamin Eysenbach. Vivek Myers, Ruslan Salakhutdinov, and Sergey Levine. Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference. Neural Information Processing Systems, 2024.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-Learning: Learning to Achieve Goals via Recursive Classification. International Conference on Learning Representations, 2021. Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R. Salakhutdinov. Contrastive Learning as Coal-Coond Reinforcement Learning. Neural Information Processing Systems, volume 35, pp. 35603-35620, 2022.
Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, and Aleksandra Faust. Stop Regressing: Training Value Functions via Classification for Scalable Deep RL. International Conference on Machine Learning, arXiv:2403.03950, 2024.
Leonardo G. Fischer, Renato Silveira, and Luciana Nedel. Gpu Accelerated Path-Planning for Multi- Agents in Virtual Environments. VIII Brazilian Symposium on Games and Digital Entertainment, pp. 101-110, 2009.
Jörg K. H. Franke, Gregor Köhler, André Biedenkapp, and Frank Hutter. Sample- Automated Deep Reinforcement Learning. International Conference on Learning Representations, 2021. C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem.
Brax- Differentiable Physics Engine for Large Scale Rigid Body Simulation. NeurIPS Datasets and Benchmarks, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. MxXiv:2004.02214 2021.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing Function Approximation Error in Actor-Critic Methods. Jennifer G. Dy and Andreas Krause, editors, International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1582-1591, 2018
Dibya Ghosh, Abhishek Gupta, and Sergey Levine. Learning Actionable Representations With Gond-Cond Policies. International Conference on Learning Representations, 2019.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to Reach Goals via Iterated Supervised Learning. International Conference on Learning Representations, 2021.
Karol Gregor, Danilo Jimenez. Rezende, and Daan Wierstra. Variational Intrinsic Control. Interna-
tional Conference on Learning Representattinss, 2016.
Shixiang Shane Gu, Manfred Diaz, Daniel C. Freeman, Hiroki Furuta, Seyed Kamyar Seyed Ghasemipour, Anton Raichuk, Byron David, Erik Frey, Erwin Coumans, and Olivier Bachem. Braxlines: Fast and Interactive Toolkit for RL-Drivee Behavior Engineering Beyond Reward Maximization. arXiv:21 10.04686, 2021.
Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging Pre- Trained Large Language Models to Construct and Utilize World Models for Model-Bass Task Planning. Neural Information Processing Systems, volume 36, pp. 79081-79094, 2023.
12

---

Published as a conference paper at ICLR 2025
Tuomas Haarnoja, Aurick Zhou, P. Abbeel, and S. Levine. Soft Actor-Critics Off-Policy Maximum Entropy Deep Reinforcement Learning With a Stochastic Actor. International Conference on Machine Learning, 2018.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked Autoencoders Are Scalable Vision Learners. IEEE/CVF Conference Computer Vision and Puttern Recognition, 2022.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A Neural Network Library and Ecosystem for JAX. 2023.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep Reinforcement Learning That Matters. AAAI Conference on Artificial Intelligence, 2018.
Tom Hennigan, Trevor Cai, Tamara Norman, Lena Martens, and Igor Babuschkin. Haiku: Sonnet for JAX. 2020.
Christopher Hoang, Sungryull Sohn, Jongwook Choi, Wilka Carvalho, and Honglak Lee. Succes- sor Feature Landmarks for Long-Horizon Coal-Conditioned Reinforcement Learning. Neural Information Processing Systems, volume 34, pp. 26963- 26975, 2021.
Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Danila Sinopalnikov, Piotr Stańczyk, Sabela Ramos, al. Acme: Research Framework for Distributed Reinforcement Learning. Stti:20000009979 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, et al. Training Compute-Optima Large
Language Models. arXiv:2203.15555 2022.
Tianyang Hu, Zhili Liu, Fengwei Zhou, Wenjia Wang, and Weiran Huang. Your Contrastive
Learning Is Secretly Doing Stochastic Neighbor Embedding. International Conference on Learning Representations, 2023.
Zhijing Jin, Yuen Chen, Felix Leeb, Luigi Gresele, Ojasv Kamal, L. Y. U. Zhiheng, Kevin Blin, Fernando Gonzalez Adauto, Max Kleiman- Kleiman-Weiner, and Mrinmaya Sachan. Cladder: Assessing Causal Reasoning in Language Models. Neural Information Processing Systems, 2023.
Scott Jordan, Yash Chandak, Daniel Cohen, Mengxue Zhang, and Philip Thomas. Evaluating the Performance of Reinforcement Learning Algorithms. International Conference on Machine Learning, pp. 4962-4973, 2020.
Scott M. Jordan, Adam White, Bruno Castro da Silva, Martha White, and Philip Thomas. Position: Benchmarking Is Limited in Reinforcement Learning Research. arXiv:2406.1622 2024. Lestie Pack Kaelbling. Learning to Achieve Goals. International Joint Conference on Artificial Intelligence, 1993.
Jaekyeom Kim, Seohong Park, and Gunhee Kim. Unsupervised Skill Discovery With Bottleneck
Option Learning. International Conference on Machine Learning, 2021.
Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. Empowerment: Universal Agent-Centric Measure of Control. IEEE Congress Evolutionary Computation, volume 1, pp. 128-135, 2005.
Sotetsu Koyamada, Shinri Okano, Soichiro Nishimori, Yu Murata, Keigo Habara, Haruka Kita, and
Shin Ishii. Pgx: Hardware-AAccelenan Parallel Game Simulators for Reinforcement Learning.
Neural Information Processing Systems, 36:45716-45744 2023.
Minae Kwon, Hengyuan Hu, Vivek Myers, Siddharth Karamcheti, A. Dragan, and Dorsa Sadigh. To- ward Grounded Commonsense Reasoning. International Conference Robotics and Automation, 2023.
Robert Tjarko Lange. Gymnax: A JAX-Based Reinforcement Learning Environment Library. 2022. Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel, CIC: Contrastive Intrinsic Control for Unsupervised Skill Discovery. arXiv:2202.00166 2022. Sergey Levine and Dhruv Shah. Learning Robotic Navigation From Experience: Principles, Methods, and Recent Results. Philosophical Transactions of the Royal Society B: Biological Sciences, 378(1 (1881866)200214447 2023.
Jacky Liang, Viktor Makoviychuk, Ankur Handa, Nuttapong Chentanez, Miles Macklin, and Dieter Fox. Gpu-Accelerated Robotic Simulation for Distributed Reinforcement Learning. Conference Robot Learning, pp. 270-282, 2018.
13

---

Published as a conference paper at ICLR 2025
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous Control With Deep Reinforcement Learning. Inter- national Conference on Learning Representations, 2016.
Bo Liu, Yihao Feng, Qiang Liu, and Peter Stone. Metric Residual Network for Sample Efficient Goal-Conditooned Reinforcement Learning. AAAI Conference on Artificial Intelligence, volume 37, pp. $799-8806, 2023.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered Policy Optimisation. Neural Information Processing Systems, 35:16455-16465 2022.
Yecheng Jason Ma, William Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-Image Representations and Rewards for Robotic Control. International Conference Machine Learning, 2023.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac Gym: High Performance GPU-Based Physics Simulation for Robot Learning. NeurlPS Datasets and Benchmarks, 2021.
Travis Manderson, Juan Camilo Gamboa Higuera, Stefan Wapnick, Jean- -aan-Francois Tremblay, Flo- rian Shkurti, David Meger, and Gregory Dudek. Vision-Based Coal-Connitioned Policies for Underwater Navigation in the Presence of Obstacles. Robotics Science and Systems, 2020.
Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Thomas Jackson, Samuel Coward, and Jakob Nicolaus Foerster. Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning. International Conference on Machine Learning, 2024.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg Corrado, and Jeff Dean. Distributed Representations of Words and Phrases and Their Compositionality. Neural Information Processing Systems, volume 26, 2013.
Volodymyr Mnih, Adrià Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning. International Conference on Machine Learning, 2016.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, et al. Human-Level Control Through Deep Reinforcement Learning. Nature, 5 (1887540):529-555 533, 2015.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A. Raffel. Scaling Data-Constraines Language Models. Neural Information Processing Systems, volume 36, pp. 11358-50376,, 2023.
Vivek Myers, Evan Ellis, Sergey Levine, Benjamin Eysenbach, and Anca Dragan. Learning to Assist Humans Without Inferring Rewards. Neural Information Processing Systems, 2024a.
Vivek Myers, Andre Wang He, Kuan Fang, Homer Rich Walke, Philippe Hansen-Estruch, Ching-An Cheng, Mihai Jalobeanu, Andrey Kolobov, Anca Dragan, and Sergey Levine. Goal Representations for Instruction Following: A Semi-SSuerriied Language Interface to Control. Conference on Robot Learning, pp. 3894-3999999 3908, 2023.
Vivek Myers, Chongyi Zheng, Anca Dragan, Sergey Levine, and Benjamin Eysenbach. Learning Temporal Distances: Contrastive Successor Features Can Provide Metric Structure for Decision- Making. International Conference on Machine Learning, arXiv:2406. 1 1098,, 2024b.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual Reinforcement Learning With Imagined Goals. Neural Information Processing Systems, volume 31, 2018.
Michal Nauman, Michal Bortkiewicz, Piotr Mitoš, Tomasz Trzcinski, Mateusz Ostaszewski, and Marek Cygan. Overestimation, Overfitting, and Plasticity in Actor-Critic: The Bitter Lesson of Reinforcement Learning. International Conference on Machine Learning, arXiv:2403.00514 2024a.
Michal Nauman, Mateusz Ostaszewski, Krzysztof Jankowski, Piotr Mitoš, and Marek Cygan. Big- ger, Regularized, Optimistic: Scaling for Compute and Sample-Fificient Continuous Control. arXiv:2405. 111.2400..16188, 2024b.
Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Xland-Mon Minigrid Scalable Reit-oReinorremment Learning Environments in JAX.
12244 2023.
14

---

Published as a conference paper at ICLR 2025
Seohong Park, Jongwook Choi, Jackyeom Kim, Honglak Lee, and Gunhee Kim. Lipschitz- Constrained Unsupervised Skill Discovery. International Conference on Learning Representations, 2021.
Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. HIQL: Offline Goal-
Conditioned RL With Latent States as Actions. Neural Information Processing Systems, 2023.
Seohong Park, Oleh Rybkin, and Sergey Levine. METRA: Scalable Unsupervised RL With Metric-
Aware Abstraction. International Conference on Learning Representations, 2024.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning Transferable Visual Models From Natural Language Supervision. International Confer- ence on Machine Learning, arXiv:2103.000 00000000000.....0000000000000000 2021.
Nazneen Fatema Rajani, Bryan McCann, Caiming Xiong, and Richard Socher. Explain Yourself! Leveraging Language Models for Commonsense Reasoning. Association for Computational Linguistics, pp. 4932-4942, 2019.
Paulo Rauber, Avinash Ummadisingu, Filipe Mutz, and Juergen Schmidhuber. Hindsight Policy Gradients. Neural Comput, 2021.
Nicholas Rhinehart, Jenny Wang, Glen Berseth, John Co-Reyes, Danijar Hafner, Chelsea Finn, and Sergey Levine. Information Is Power: Intrinsic Control via Information Capture. Neural Information Processing Systems, volume 34, pp. 10745-10758, 2021.
Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Garðar Ing- varsson, Timon Willi, Akbir Khan, Christian Schroeder de Witt, and Alexandra Souly. Jaxmarl: Multi-Agent RL Environments in JAX. Agent Learning in Open-Endedness Workshop, 2024.
Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. Beyond Chinchilla-Optimal Accounting for Inference in Language Model Scaling Laws. International Conference on Machine Learning, 2024.
Tom Schaul, Dan Horgan, Karol Gregor, and David Silver. Universal Value Function Approximators.
Nd International Conference on Machine Learning, pp. 1312-1320, 2015.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-Dimenssooal
Continuous Control Using Generalized Advantage Estimation. International Conference on Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. arXiv: LeXXiv1707.000 17077706777.77 2017.
Pierre Sermanet, Corey Lynch. Jasmine Hsu, and Sergey Levine. Time-Contrastive Networks: Self- Supervised Learning From Multi- View Observation. JEEE Conference on Computer Vision and Pattern Recognition Workshops, pp.
2017.
Dhruv Shah, Ajay Sridhar, Nitish Dashora, Kyle Stachowicz, Kevin Black, Noriaki Hirose, and Sergey Levine. Vint: A Foundation Model for Visual Navigation. Conference Robot Learning, 2023.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-Aware Unsupervised Discovery of Skills. International Conference on Learning Representations, 2020.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, et al. Mastering the Game of Go With Deep Neural Networks and Tree Search. Nature, 529(7587):484 489, 2016.
David Silver. Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, et al. Mastering Chess and Shogi by Self-Play With a General Reinforcement Learning Algorithm. 112.0:1712.01815, 2017.
Kihyuk Sohn. Improved Deep Metric Learning With Multi-Class N-Pair Loss Objective. Neural Information Processing Systems, volume 29, 2016.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancy, Timothy Lillicrap, and Martin Riedmiller. Deepmind Control Suite. arXiv:1801. 100001800..0069,, 2018.
Yuval Tassa, Tom Erez, and Emanuel Todorov. Synthesis and Stabilization of Complex Behaviors Through Online Trajectory Optimization. JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 4906-4913, 2012.
15

---

Published as a conference paper at ICLR 2025
William Thibault, William Melek, and Katja Mombaur. Learning Velocity- Based Humanoid Loco- motion: Massively Parallel Learning With Brax and MJX. arXiv:2407.055 148, 2024.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A Physies Engine for Model-Based Control. JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012.
Ahmed Touati and Yann Ollivier. Learning One Representation to Optimize All Rewards. Neural Information Processing Systems, volume 34, pp. 13-23, 2021.
Mark Towers, Ariel K Kwiatkowski, Jordan Terry, John U. Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, et al. Gymnasium: A Standard Interface for Reinforcement Learning Environments. mxXiv2407.170032. 2024.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. DM_Control, Software and Tasks for Continuous Control. Software Impacts, 6:100022, 2020.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation Learning With Contrastive Predictive Coding. trXiv:1807.000 .07.003748, 2019.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. Neural Information Processing
Systems, 2017.
Srinivas Venkattaramanujam, Eric Crawford, Thang Doan, and Doina Precup. Self Supervuue Learning of Distance Functions for Goal-Conditioned Reinforcement Learning. arXiv:1907 .2999, 2020.
Homer Rich Walke, Kevin Black, Tony Z. Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen- Estruch, Andre Wang He, Vivek Myers, et al. Bridgedata V2: A Dataset for Robot Learning at Scale. Conference on Robot Learning, pp. 1723-1736, 2023,
Tongzhou Wang. Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal Goal-Reaching Rein- forcement Learning via Quasimetric Learning. International Conference on Machine Learning, pp.
36411-36430, 2023.
Pawel Wawzzyíski. A Cat-Like Robot Real-Time Learning to Run. Mikko Kolehmainen, Pekka Toivanen, and Bartlomiej Beliczynski, editors, Adaptive and Natural Computing Algorithms, pp- 380-390, 2009.
Daochen Zha, Jingru Xie, Wenye Ma, Sheng Zhang, Xiangru Lian, Xia Hu, and Ji Liu. Douzero: Mastering Doudizhu With Self-Play Deep Reinforcement Learning. International Conference on Machine Learning, pp. 12333-12344, 2021.
Chongyi Zheng, Benjamin Eysenbach, Homer Walke, Patrick Yin, Kuan Fang, Ruslan Salakhut- dinov, and Sergey Levine. Stabilizing Contrastive RL: Techniques for Offline Goal Reaching. International Conference on Learning Representations, 2024.
Chongyi Zheng, Ruslan Salakhutdinov, and Benjamin Eysenbach. Contrastive Difference Predictive Coding. International Conference on Learning Representations, 2023.
Yizhe Zhu, Martin Renqiang Min, Asim Kadav, and Hans Peter Graf. S3vae: Self- Self-Supervised Sequential VAE for Representation Disentanglement and Data Generation. JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6538-6547, 2020.
A ADDITIONAL RESULTS
A.1 SPEEDUP COMPARISON ACROSS VARIOUS NUMBERS OF PARALLEL ENVIRONMENTS
We present an extended version of the plot from Fig. 1, with additional experiments, as depicted on the left side of Fig. 9. Each experiment was run for 10M environment steps, and for the original repository, we varied the number of parallel actors for data collection, testing configurations with 4, 8. 16, and 32 actors, each running on separate CPU threads. Each configuration was tested with three different random seeds, and we present the results along with the corresponding standard deviations. The novel repository used 1024 actors for data collection. We used NVIDIA V100 GPU for this experiment.
A notable observation from these experiments is the variation in success rates associated with different numbers of parallel actors. We hypothesize that this discrepancy arises due to the increased diversity
16

---

Published as a conference paper at ICLR 2025
of data supplied the replay buffer as the number of independent parallel environments increases, leading to more varied experiences for each policy update. We conducted similar experiments using method with varying numbers of parallel environments to further investigate. The results presented on the right side of Fig. 9. This observation, while interesting, is beyond the scope of the current work and is proposed
area for further investigation.
1:16 (env steps: grad staps)
1:16 (env stops: grad steps)
1.0
Contrantive RI aclm)
JaXGRR y addnngg
0.9
Contraction aclmm)
0.9
JodiO addens)
0.8
Contraxive (22 actrra)
0.8
(32 actons
0.7
JaxOCRL (1024 actors)
0.7
JaxXCCL, (1024 actora)
Succes
6.1
0.0
minutee)
missteal)
Figure 9: Speedup in ant environment for the 1:16 ratio of SGD steps environment steps for different numbers of parallel actors.
A.2
ENERGY FUNCTIONS AND CONTRASTIVE OBJECTIVES
The full list evaluated energy functions:
(0,0,os(s,,,9)) (ó(s, a),
So (8, (8, ag (((s,)),(g));
fo.VLL a,
-Il -
fo. (s, s,
-l (g))l22.
fe. wro egee (8, (,
-Iló(s,a) (0))
(8)
The full list of tested contrastive objectives:
Σ log
[B]
LIntONCE-wwdB::C vb)
(9)
-=1 (lo.
K
>ie 1=1
|B]
LISGNNCE
bwd budl (B: =
log
(10)
2==1 (lo.o(s;a3a))
LInfoNCE-sym (B v)=
LIntONCE-fud (B; Φ. +) + LIntONCE-bw v3),
(11)
Σ B
[B] log
j=1
(12)
detach
B c/
[
-0
CFINNCEboddB:0 = Σ
|B|
log
j=1
(13)
imll
detach
LrB(B: Σ [88]
(efow((taaaa)) + 2(
Σ
(14)
[88]
[s|
LDrO(B; Φ, (6)
log o [Jon as, gi) So,
(80,
(15)
(Br 6. Σ [8]
18
[f (Si, a.) (s, 1]2
(16)
j=1
17

---

Published as a conference paper at ICLR 2025
IQM
IQM
Cos Dot L2 w/o sqrt L2
Cos Dot
w/o sqrt L2
L1
L1
0.24
0.32
60
90
120
Success rate
Time near goal
Energy functions influence CRL performance metrics in multiple ways. Success rate (left) and time near goal (right) results for 5 energy functions. IQMs indicate better performance of p-norms and dot product energy functions cosine similarity for CRL. Interestingly. 1.2 results in much higher time near goal than other energy functions. Results averaged over five seeds and plotted with one standard error.
Figure 10:
a.,;;) [fo, (4) 10 112++ (f,) +1]² CSPPO (B:0, - Σ 18 E [f
(17)
where we have highlighted the indices corresponding to positive and negative samples for clarity.
The last three of those losses, that is DPO, IPO. and SPPO were inspired by the structure of losses the Preference Optimization domain. Unlike in the losses from InfoNCE family, here the samples are compared in pairs.
The DPO loss simply drives the difference between scores of positive and negative samples to be larger, and doesn't regularize those scores in any other way.
The IPO loss can be seen as a restriction DPO, where scores are regularized to always have difference of one between positive and negative sample pairs.
The SPPO loss restricts this even further, and regularizes the scores to be equal to one for positive samples and negative one for negative samples.
A.3 ENERGY FUNCTIONS RESULTS
The performance of contrastive learning is sensitive energy function choice (Sohn, 2016). This section aims to understand how different energy functions impact CRL performance. In particular, we evaluate five energy functions: LI, L2, L2 w/o sqrt. dot product and cosine. For every energy function, we use symmetric InfoNCE as a contrastive objective, with a 0.1 logsumexp penalty coefficient.
In Fig. 10, we report the performance of every energy function for four different environments and pusher. We find that p-norms and dot-product significantly outperform cosine similarity. Additionally, removing the root square from the L2 norm (L2 w/o sqrt) results in performance degradation, especially regarding time near goal. This modification makes the energy function no longer abide by the triangle inequality, which, as pointed out by (Myers et al., 2024b), desirable for temporal contrastive features. Results per environment reported in Fig. 11.
Fig. 11 shows the results of different energy functions per environment, success rate and time near goal. Clearly, no single energy function performs well in all the tested environments, for instance, LI and L2, which perform well in Ant environments, work poorly in Pusher. In addition, we observed high variability in every configuration performance, indicated by relatively wide standard errors.
Figure 11: Success rate (top) and time near goal (bottom) results in different energy functions and environments. Best performing energy function varies across environments.
18

---

Published as a conference paper at ICLR 2025
Reocher
Pusher Har
Hemazdld
Agt
Sucope sate
0.0
S&C HER
Stocces
6.5
Su
staps Ball)
Envinoneent stops millionn)
alag milliuas)
stages ltmos))
Aut U-Mase
Aut Big Museo
Aut Socuor
Ant. Posk
Sucods rore
Sucvees rar
6.5
0.5
Succen
stops milloone)
stops millione)
stage millioms)
strgp moilbone))
Figure 12: Baseline results JaxGCRL benchmark. Success rates for each our benchmark environments
using bigger architecture. Reported as 10 seeds and standard error.
A.4 BENCHMARK PERFORMANCE
In this section, report additional results for more advanced architectures the benchmark environ- ment. In particular, we report the results for architecture with 4 hidden layers of size 1024 and Layer Normalization in Fig. 12. Unsurprisingly, the performance is significantly better in most environ- ments, particularly for Humanoid. This suggests that a larger architecture is needed to effectively manage the nigh-dimensional state and action spaces involved in this environment.
A.5 HINDSIGHT EXPERIENCE REPLAY DETAILS
In HER, we relabel, on average, 50% of goals with states achieved at the end of the rollout. The humanoid environment was trained on goals sampled from distance in the range [1.0, 5.0) meters and evaluated on goals sampled from a distance 5.0 meters. All other environments were trained and evaluated on identical environments.
We observe poor performance for SAC+HER in the Pusher environment as result of HER generating a "successful" experience. which is trivial. In particular, the goal in the pusher environment is the desired location of the puck, which is different from its initial position, so that agent should push the puck to that position. Relabeling the goal to the final location of the puck reached at the end of the episode often results in just changing the goal to the puck's s initial position. This happens because, during the early stages of training, the random policy usually doesn't 't interact with the puck.
A.6 SCALED-UP CRL ARCHITECTURE IN DATA-RICH SETTING
In Figures 13 and 14, report success rates and time near goal for scaled-up CRL agent with architecture consisting of layers with 1024 neurons per layer and layer norm. We find that these architectures increase the fraction of trials where the agent reaches the goal least once, but they do not enable the agent to stabilise around the goal (e.g., on 7 tasks, the best agent spends less than 50% of an episode the goal).
When visualizing the rollouts, we observe that the Humanoid agent falls immediately after reaching the goal state, and the Ant Soccer agent struggles recover when it pushes the ball too far away. The Humanoid merely "flings" itself toward the goal, while the optimal policy would involve running to the goal and remaining there. This inability stabilize around the goal suggests that the agent not effectively optimizing the actor's s objective, pointing to potential area for further research.
B TECHNICAL DETAILS
JaxGCRL is fast implementation of state-based self-supervised reinforcement learning algorithms and new benchmark of GPU-accelerated environments. Our implementation leverages the power of GPU-accelerated simulators (BRAX and MuJoCo MJX) (Freeman et al., 2021; Todorov al., 2012)
19

---

Published as a conference paper at ICLR 2025
20 30
300
100 300
Figure 13: CRL with big architecture success rates in data-rich setting.
12.0
10/0
1800
1800
1800
Figure 14: CRL with big architecture time near goal data-rich setting.
to reduce the time required for data collection and training, allowing researchers to run extensive experiments in a fraction of the time previously needed. The bottleneck of former self-supervised RL implementations
twofold. Firstly, data collection was executed many CPU threads, single thread was often used for a single actor. This reduced the number of possible parallel workers, as only high-compute servers could run hundreds of parallel actors. Secondly, the necessity of data migration between CPU (data collection) and GPU (training) posed additional overhead. These two problems are mitigated by fully JIT-compiled algorithm implementation and execution of all environments and replay buffer operations directly on the GPU.
Notably, our implementation uses only one CPU thread and has low RAM usage as all the operations. including those on the replay buffer, are computed on GPU. It's important to note that the BRAX physics simulator is not exactly the same as the original MuloCo simulator, so the performance numbers reported here are slightly different from those in prior work. All methods and baselines we report are run on the same BRAX simulator.
B. ENVIRONMENT DETAILS
In each of the environments, there are number of parameters that can change the learning process. non-exhaustive list of such details for each environment is presented below.
B.2 BENCHMARK DETAILS
Our experiments JaxGCRL suite of simulated environments described in Section 4.2. We evaluate algorithms in an online setting, with a UTD ratio 1:16 for CRL, TD3, TD3+HER, SAC. SAC+HER, and 1 5 for PPO. We PPO, for which we use a discount factor of 0.97. For every environment, we sample evaluation goals from the same distribution as training ones and use replay buffer of size 10M for CRL, TD3, TD3+HER, SAC, and SAC+HER. We use 1024 parallel environments for all methods except for PPO, where use 4096 parallel environments to collect data. All experiments are conducted for 50 million environment steps.
batch size of 256 and discount factor of 0.99 for all methods except
in 3 dimensions
20

---

Published as a conference paper at ICLR 2025
Table Environments details
Environment
Goal distance
Termination
Brax pipeline
Goal sampling
Reacher
0.05
No
Spring
Disc with radius sampled from [0.0, 0.2] Fixed goal location
Half-Cheetah Pusher Easy
0.5
No
MJX Generalized
0.1
No
coordinate sampled from 2 -0.25), coordinate
sampled from [-0.2, 0.2)
Pusher Hard
0.1
No
Generalized
x coordinate sampled from [-0.65, 0.35]. y coordinate sampled from [-0.55, 0.45]
Humanoid
0.5
Yes
Spring
Disc with radius sampled from [1.0, 5.0]
Ant
0.5
Yes
Spring
Circle with radius 10.0
Ant Maze
0.5
Yes
Spring
Maze specific Circle with radius 5.0
Ant Soccer
0.5
Yes
Spring
Ant Push
0.5
Yes
MJX
Two possible locations with
uniform noise added
Table 2: Hyperparameters
Hyperparameter
Value
num timesteps
50,000,000
max max_replay_size
10,000
min sin_replay_o size
1,000
episode_ @pistde_length discounting
1,000
0.99
num_envs batch size
1024 (512 for humanoid)
256
multiplier Laceeer__um___ num_ aum_sdd_steps
action cepoonteeeettt_
unroll_lengts
62
policy_lr
6e-4
critic_lr
3e-4
contrastive_loss tosistittiooosssuuuiiiioooo function
symmetric infonce
energy_function
L2
logsumexp.penaltty
0.1
hidden layers (for both encoders and actor)
[256,256]
representation dimension
64
B.3 BENCHMARK PARAMETERS
The parameters used for benchmarking experiments can be found in Table 2.
min Sain_repla____ ssizesizeei is parameter that controls how many transitions per environment should be gathered to prefill the replay buffer. max_ree tera_repla______ size is parameter that controls how many transitions maximally stored in replay buffer per environment.
C RANDOM GOALS
Our loss in Eq. (3) differs from the original CRL algorithm (Eysenbach et al., 2022) by sampling goals from the same trajectories as states during policy extraction, rather than random goals from the replay buffer. Mathematically,
generalize Eq. (3) to account for either of these strategies
21

---

Published as a conference paper at ICLR 2025
by adding hyperparameter to the loss controlling the degree of random goal sampling during training.
(1 -1-o8 a) Bep.earppoolsossssllllll (s, 9)]
+ (a (s, a'
The hyperparameter controls the rate of counterfactual goal learning, where the policy is updated based on the critic's evaluation of goals that did not actually occur in the trajectory. We find that taking, 0 (i.e., random goal sampling) leads to better performance, and suggest using the policy loss in Eq. (3) for training contrastive RL. methods.
D PROOFS
D.1 Q-FUNCTION IS PROBABILITY
This proof follows closely the one presented in Eysenbach al. (2022).
We want relate the Q-function discounted state visitation distribution: (s, pg 8, The Q-function is usually defined in terms of rewards:
27 [s. a) A (-
Σ (sr az)
sama, a0=a
(18)
We will define rewards conditioned with goal as:
(s(80) a) - (1 (1--)) (p(s0 = g) >p(s = so, aol):
t = 0
(19)
(11 7)+p(s++111111 =9 [18,aa))
t> 0.
Lastly, we define discounted state visitation distribution:
(T) - (1 - (o)
(20)
For t >> O, the term PF (g) is probability of reaching the goal g at timestep t with policy conditioned on and thus:
pf (g) (m((1)) [pr (t-1,a_11)
= E [p(s = 8 Gt-)]]
On the second line, we have used the Markov property. We can now substitute this into Eq. (20):
p7(0)= = (1 - 7) Σ (p)
1=
= -11 (1- + (1 -7))
' [p(s: = - (1-11
1=1
= (1 - (g) + (1 7
Σ
7
1+1 Ex((10) [p(s(+1 = ser sc)
1=0
En(-1o)
(1 (1-7)p(s0 = g) + (1 - 7)
Σ
7 1+ p(sc+1 = g at)
1=0
Ex((11))
(1-7) (p(s0 = 9) + xp(s1 = g so, ao))
Σ 7d (1-7))p(54+ = se, at)
(50.000))
22

---

Published as a conference paper at ICLR 2025
o
(rrg(5e; an)
=
(-1o))
Σ
Thus a set state-action pair (s, a), we have:
o0
(g s, = (1)))
Σ
80 ( az)
= (g
d==2
l==
which relates the Q-function to discounted state visitation distribution.
FAILED EXPERIMENTS
Weight Decay: Prior work (Nauman et al., 2024b) indicated that regularizing critic weights might improve learning stability. We did not observe significant upgrades in the CRL setup, perhaps due to much lower ratio of updates per environment step. We tested this only for small architectures.
2. Random Goals: In prior implementation (Eysenbach et al., 2022) using random goals in actor loss resulted in higher performance. We did observe that in our online setting.
F PSEUDOCODE
Pseudocode for the contrastive learning algorithms studied is presented in Algorithm
Algorithm Contrastive Reinforcement Learning
Contrastive loss Leriee energy function f 2: Initialize Ø, t), and an empty replay buffer D
Input:
3: repeat
4: in parallel over environments
5:
Observe state and sample an action a T (8,g)
6:
Execute in the environment
7:
Observe next state s' and done signal to indicate whether terminal Append (s, s') current trajectory for this environment is terminal then
8:
9:
10:
Reset environment state and sample new goal Store current trajectory for this environment in DD Start new trajectory for this environment
11:
13:
for j L.... num updates do
Randomly sample (with discount) batch B from of state-action pairs and goals from their future
Update critie:
(ø, +)) +- (Ø, 1b a d. [Criii v) + 8 Loogsmmcxp (B, o, b)]
16:
Update policy:
V [CAcuOr (B; 0, v)
17: until convergence
23


Reviews:
[{"review_id": "8XAr169EeI", "reviewer": "Reviewer_MvLy", "summary": "The authors propose a new library for goal conditioned reinforcement learning (GCRL). Unlike prior work, their method runs end to end on the GPU, making training faster. They implement 8 environments in JAX, as well a few algorithms and various objectives. Then, they evaluate existing methods across a number of axes, investigating replay ratios, model sizes, and energy functions.", "questions": "- Figure 4 typo: \"though DPO policies remain at the goal for a shorter\"", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "- The library is well-motivated, as speeding up RL leads to better experimentation\n- The authors implement many environments and energy functions\n- The scale, energy function, update-to-data ratio experiments are interesting and useful for future work on GCRL", "weaknesses": "The library appears like a \"one-and-done\" sort of thing that will not be maintained after publication. In RL there is already a large graveyard of abandoned RL projects that no longer run and provide no value to the community. Given this fact, I can only review the current state of the library. In its current state, I think the library needs a bit more work before publication. Please see https://docs.cleanrl.dev for an example of what I think a modern RL library should look like.\n\n- There is no documentation, it is unclear:\n    - Which approaches are implemented\n    - How to use these approaches\n    - How to add new models\n    - The structure of the codebase\n- There are no unit tests, so the correctness of the code (and the ability to maintain the code as time goes on) is unclear\n- The train script is solely for the authors, relying on a pre-existing conda environment\n- There are no tutorials beyond a single bash command that runs a parameter sweep\n- The library relies on wandb, and does not seem to run without a wandb account\n- As far as I understand, the authors only implement 3 algorithms, and I would like to see more than three baselines so that we can do proper comparisons", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zkHMtKBkWw", "reviewer": "Reviewer_j3bj", "summary": "This paper introduces JaxGCRL, a codebase that contains environments and a scalable goal-conditioned RL algorithm, all implemented in JAX. This allows researchers to train GC agents much faster than before, making these experiments more accessible. This work also analyses several design decisions of contrastive RL algorithms, enabled by a fast simulator & algorithm implementation.", "questions": "- What is your support plan going forward with JaxGCRL, are you planning on adding new environments or algorithms?\n- It seems like JaxGCRL is very much focused on brax-type environments, is there a part of goal conditioned RL research that potentially focuses rather on discrete action environments that you are leaving out?\n- What about other non-contrastive GCRL algorithms? Are you planning on adding support for those?\n\t- Relatedly, how easy would it be for someone else to implement a new GCRL algorithm to fit within your framework?\n\t- And how easy is it to add another goal conditioned environment, based on an existing JAX environment? For instance, minigrid or craftax or xland minigrid, etc?\n- In the maze, for instance, can you dynamically, within the JIT, change the maze layout, or does it have to be statically known at compile time?\n- Is there an easy way to transfer a JaxGCRL agent to existing environments?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Adding even more environments and algorithms to the JAX ecosystem is great, especially for goal-conditioned RL which is lacking in this space.\n- The proof-of-concept experiments demonstrate what this library can allow, namely, more thorough investigation of design decisions in goal-conditioned RL\n- The writing and motivation is quite clear.", "weaknesses": "- I can't see any major weaknesses, apart from the limited number of environments, although 8 is pretty respectable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KV6ZAxzi1r", "reviewer": "Reviewer_fgPt", "summary": "The paper provides a JIT-complied codebase with vectorized environments that can speed up the training and iterating new ideas on goal-conditioned reinforcement learning problems.\nIn additional, it provides a stable baseline algorithm for the goal-conditioned reinforcement learning problems that's benchmarked in the 8 diverse continuous environments.", "questions": "* In Sec 5.3, why is the contrastive objective only evaluated on part of the 8 environments? Similar question in sec 5.6 for examining different UTD ratios.\n* In Fig 1. Are the num_actors same for the JaxGCRL and CRL?\n* How do you define if the agent is within goal's proximity?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* The JaxGCRL codebase is significantly faster than the original codebase. \n* The proposed baseline consistently outperform the counterpart in all 8 environments, demonstrating the stableness from simple to complex environments.\n* The performance of different design choice is extensively tested and the result metric is easy to interpret.", "weaknesses": "The paper mentions it leverages the power of GPU-accelerated simulators, but by comparing against the brax training code under https://github.com/google/brax/tree/main/brax/training, there are some similarities for the training code as well, and it's not mentioned in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ouc9GxNfS1", "reviewer": "Reviewer_3DLK", "summary": "The authors introduce JaxGCRL, a benchmark and framework for evaluating goal-conditioned RL algorithms based on contrastive learning.\nThey re-implement a number of goal-conditioned tasks from prior literature and evaluate their implementation on it. \n\nThey then evaluate the effect of different losses, more samples, and larger networks on their implementation. They demonstrate that their Jax-based implementation is significantly faster than previous libraries, accelerating future research.", "questions": "See weaknesses.\n\n[1] Brax documentation. https://github.com/google/brax?tab=readme-ov-file#one-api-four-pipelines\n\n[2] Gymnax: A JAX-based Reinforcement Learning Library. Robert Lange. https://github.com/RobertTLange/gymnax\n\n[3] Lu, Chris, et al. \"Discovered policy optimisation.\" Advances in Neural Information Processing Systems 35 (2022): 16455-16468. https://github.com/luchris429/purejaxrl", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper has several significant strengths.\n\n- Although not novel, JAX implementations are to be commended. They improve research iteration speed significantly.\n- The paper is very well written. The authors communicate their results clearly and unambiguously.\n- The authors evaluate using the inter-quartile mean and bootstrapped confidence intervals. This is more sound than using learning curves etc.\n- The authors provide a number of ablations and experiments that explain the performance of their implementation.", "weaknesses": "However, I have a number of issues with this paper, which is why I recommend rejection. \n\n- The authors claim that their setting is challenging, but do not effectively demonstrate that this is the case. The authors demonstrate that by using a bigger network (1024 layer width and depth of 4) and layer norm, the performance significantly improves. They also run experiments where they train for significantly more interactions. However, as best I can tell (it is not always clear which network is used in which experiment), the authors never run their biggest, highest performing network for 300M steps on all the tasks. The authors do not pitch their work as focussing on sample efficiency, and therefore I am not sure why their evaluation framework should be compelling if the tasks can be solved by scaling up networks and using more samples. If the authors can provide a demonstration that this does not satisfactorily solve their benchmark, **I will raise my score**. However, without this demonstration, I do not believe that the experimental insights and JAX implementation are enough to warrant acceptance.\n- I am confused about the experiments concerning the update-to-date ratio (UTD). Given a fixed step budget, doing fewer or more updates is a pure trade-off. You can do fewer, less noisy updates, or do more, noisier updates. This occurs all over RL, for example when choosing the number of parallel environments to use in PPO. I am not sure why a high or low number of updates would be beneficial, or this quantity would be interesting to examine.\n\nI also have a number of more minor points:\n- The authors claim that they cannot directly compare brax and mujoco because brax uses a different physics engine, but the MuJoCo physics engine has been available in brax for a while now [1] -- what exactly is the issue here? \n- The discussion of related work on jax-based environments is missing some work. Gymnax [2] and PureJaxRL [3], both were important landmarks in the use of and benefits of JAX in RL and warrant inclusion.\n- The authors should probably rephrase line 117, which begins with \"In addition to reducing CPU-GPU data-transfer overhead...\". While implementing an environment in JAX *does* do this, there are also significant other factors such as JIT compilation and the resulting operator fusion and the ability to use more vectorised environments than a typical CPU + pytorch approach that lead to the significant speedups. \n- A number of the papers listed in the appendix have incorrect citations or are missing authors.\n- Line 1032 in the appendix contains a typo (lenght -> length)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)