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
Primary area: reinforcement learning
Keywords: Deep Reinforcement Learning, GPU-accelerated Physics Simulators, Contrastive Learning, Unsupervised Reinforcement Learning

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)

Decision: accept (spotlight)
Reviewer scores: 8, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730405552951

thread_type:
official_review

<thread_object>
creator: reviewer_3dlk

<content>
summary:
The authors introduce JaxGCRL, a benchmark and framework for evaluating goal-conditioned RL algorithms based on contrastive learning.
They re-implement a number of goal-conditioned tasks from prior literature and evaluate their implementation on it. 

They 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.

strengths:
The paper has several significant strengths.

- Although not novel, JAX implementations are to be commended. They improve research iteration speed significantly.
- The paper is very well written. The authors communicate their results clearly and unambiguously.
- The authors evaluate using the inter-quartile mean and bootstrapped confidence intervals. This is more sound than using learning curves etc.
- 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. 

- 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.
- 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.

I also have a number of more minor points:
- 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? 
- 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.
- 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. 
- A number of the papers listed in the appendix have incorrect citations or are missing authors.
- Line 1032 in the appendix contains a typo (lenght -> length)

questions:
See weaknesses.

[1] Brax documentation. https://github.com/google/brax?tab=readme-ov-file#one-api-four-pipelines

[2] Gymnax: A JAX-based Reinforcement Learning Library. Robert Lange. https://github.com/RobertTLange/gymnax

[3] Lu, Chris, et al. "Discovered policy optimisation." Advances in Neural Information Processing Systems 35 (2022): 16455-16468. https://github.com/luchris429/purejaxrl
</content>
</thread_object>

</thread>
<thread>
creation date:
1730496677670

thread_type:
official_review

<thread_object>
creator: reviewer_fgpt

<content>
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.
In additional, it provides a stable baseline algorithm for the goal-conditioned reinforcement learning problems that's benchmarked in the 8 diverse continuous environments.

strengths:
* The JaxGCRL codebase is significantly faster than the original codebase. 
* The proposed baseline consistently outperform the counterpart in all 8 environments, demonstrating the stableness from simple to complex environments.
* 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.

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.
* In Fig 1. Are the num_actors same for the JaxGCRL and CRL?
* How do you define if the agent is within goal's proximity?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730566651269

thread_type:
official_review

<thread_object>
creator: reviewer_j3bj

<content>
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.

strengths:
- Adding even more environments and algorithms to the JAX ecosystem is great, especially for goal-conditioned RL which is lacking in this space.
- The proof-of-concept experiments demonstrate what this library can allow, namely, more thorough investigation of design decisions in goal-conditioned RL
- 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.

questions:
- What is your support plan going forward with JaxGCRL, are you planning on adding new environments or algorithms?
- 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?
- What about other non-contrastive GCRL algorithms? Are you planning on adding support for those?
	- Relatedly, how easy would it be for someone else to implement a new GCRL algorithm to fit within your framework?
	- 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?
- 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?
- Is there an easy way to transfer a JaxGCRL agent to existing environments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730618173949

thread_type:
official_review

<thread_object>
creator: reviewer_mvly

<content>
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.

strengths:
- The library is well-motivated, as speeding up RL leads to better experimentation
- The authors implement many environments and energy functions
- 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.

- There is no documentation, it is unclear:
    - Which approaches are implemented
    - How to use these approaches
    - How to add new models
    - The structure of the codebase
- There are no unit tests, so the correctness of the code (and the ability to maintain the code as time goes on) is unclear
- The train script is solely for the authors, relying on a pre-existing conda environment
- There are no tutorials beyond a single bash command that runs a parameter sweep
- The library relies on wandb, and does not seem to run without a wandb account
- 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

questions:
- Figure 4 typo: "though DPO policies remain at the goal for a shorter"
</content>
</thread_object>

</thread>
<thread>
creation date:
1731870555638

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their time reviewing our manuscript. The insights provided by the reviewers have allowed us to increase further the quality and readability of our manuscript and JaxGCRL. So far, we have made the changes listed below in response to reviewers' questions and suggestions. 

**Improving the JaxGCRL user experience with proper documentation.** We would like to thank reviewer MvLy for their helpful suggestions on the essential components of a modern RL library. In the coming days, we will share updates on the changes to JaxGCRL to enhance its usability, starting with a documentation (MkDocs Material) update today:

- We added a list of implemented environments.
- We added a concise example of using different environments and methods.
- We described the structure of the codebase.
- We added information on how to run experiments without Weights&Biases by setting a run flag
Reviewers can see these new changes on the anonymised repo: https://anonymous.4open.science/r/JaxGCRL-2316/README.md.

**Readability of our manuscript** - following comments from all the reviewers, we have made several changes to our manuscript, which we believe further improve its clarity:

- In Section 2.2 (Accelerating Deep Reinforcement Learning), we clarified JAX's importance for GPU-accelerated environments.
- In Section 5.1 (Experimental Setup), we added the definition of the goal proximity criterion and information about Brax implementation dependency.
- In Section 5.2 (JaxGCRL Benchmark Results), we added new baselines: TD3, TD3+HER and PPO.
- In Appendix (A.6), we have included additional results with scaled-up CRL in a data-rich setting (see below: Extended results on 300M steps for big architecture).
- We also corrected all the typos and missing citations pointed out by reviewers.

These new changes are visible in the updated manuscript: https://anonymous.4open.science/r/paperJaxGCRL-EFEB/Accelerating_Goal_Conditioned_RL_Algorithms_and_Research_rebuttal.pdf. We have highlighted the changes with orange colour, both text and modified figures.

**Extended results on 300M steps for big architecture** - Reviewer 3DLK pointed out that we haven’t evaluated the highest performing architecture (network with 1024 layer width and depth of 4 and layer norm) on every environment while training for 300M steps as in Section 5.5.

We find that these changes can increase the fraction of trials where the agent reaches the goal at least once, but they do not enable the agent to stabilise around the goal (e.g., on 5 tasks, the best agent spends less than 50% of an episode at 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 to recover when it pushes the ball too far away. The key takeaway is that there is still room for substantial improvement (e.g., can any RL method spend 80% of time steps at goal on Humanoid task), underscoring one dimension on which these tasks are challenging.

Links to figures from this experiment:
- [Time near goal](https://anonymous.4open.science/r/paperJaxGCRL-EFEB/time_near_goal.png)
- [Success rate](https://anonymous.4open.science/r/paperJaxGCRL-EFEB/success_rate.png)

We have also added a new task, Ant Hard Maze. No method reaches the goal in more than 50% of episodes on this new task. We also want to highlight that modifying proposed environments to make them more difficult is straightforward and needs just modifying a single method (`random_target`), which defines the distribution from which goals are sampled.

We believe that these changes increase the quality of JaxGCRL and our manuscript, and again, we are grateful to the reviewers for their suggestions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a new benchmark for self-supervised goal-conditioned RL (GCRL) which is optimized for running on a GPU and can thus yield high-throughout experiments. The authors benchmark a contrastive RL method against other popular RL algorithms and demonstrate strong gains.

The main weaknesses raised during the review process were around differentiation with Brax, and the difficulty of the provided environments. The authors seem to have addressed all the reviewer concerns.

The reviewers are unanimous in recommending acceptance for this work, and after going through the paper and discussion, I agree. 
A minor comment is that there is a broken reference in line 33.