Corax: Core RL in JAX

Installation | Examples | Agents | Datasets

Corax is a library for reinforcement learning algorithms in JAX. It aims at providing modular, pure and functional components for RL algorithms that can be easily used in different training loops and accelerator configurations. Currently, we are exploring the design of a LearnerCore and ActorCore design that allows easy composition and scaling of RL algorithms. At the same time, Corax aims to provide strong baseline agents that can be forked and customized for future RL research.

Corax starts as a fork of the dm-acme library while aiming to provide a better experience for researchers working on Online/offline RL in JAX. Future development of Corax may diverge from the design in Acme.

Installation

You can install Corax with

pip install 'git+https://github.com/ethanluoyc/corax#egg=corax[tf,jax]'

To use Corax with GPU, you need to install JAX with GPU support. Follow the instructions here for how to install JAX with GPU support.

Note on TensorFlow dependency

The base corax package does not depend on specific deep learning frameworks. However, the JAX agent depends on TensorFlow for efficient data-processing.

We provide optional extras for installing tensorflow-cpu and compatible versions of TensorFlow Probability and Reverb.

However, should that be incompatible with your own dependency requirements, you can optionally specify these dependencies yourself and opt-out our extras. Check out the pyproject.toml for examples on how to specify compatible TensorFlow versions. Here is an example workflow of determining the compatible versions of TensorFlow, TensorFlow-Probability and Reverb. Assume that you will use tensorflow-cpu~=2.13.0, then

1. Looking at https://github.com/tensorflow/probability/releases, the tensorflow-probability version compatible with tensorflow~=2.13.0 is 0.21.0.
2. Looking at https://github.com/google-deepmind/reverb/tree/master#reverb-releases, the dm-reverb version compatible is 0.12.0.

Therefore, as an application developer, you should put the following in your requirements.txt

tensorflow-cpu~=2.13.0
tensorflow-probability~=0.21.0
dm-reverb~=0.12.0

If you use dm-launchpad. The workflow is similar, although as of 17 Oct, 2023 Launchpad does not provide an official build for tensorflow 2.13.0. We however, have an unofficial manylinux build for Python 3.9 and 3.10 available at https://github.com/ethanluoyc/launchpad/releases/tag/v0.6.0rc0. If you intend to use this version, you should include in your requirements.txt

# Use the exact link to the wheel file for your Python version
dm-launchpad @ https://github.com/ethanluoyc/launchpad/releases/download/v0.6.0rc0/dm_launchpad-0.6.0rc0-cp39-cp39-manylinux2014_x86_64.whl

We currently do not have build for tensorflow 2.14.0 due to https://github.com/google-deepmind/launchpad/issues/44.

Examples

Examples can be found in projects.

Development

git clone https://github.com/ethanluoyc/jax
cd corax
# Create a virtual environment with the method of your choice.
python3 -m venv .venv
source .venv/bin/activate
# Then run
pip install -e '.[tf,jax,test,dev]'
# Install pre-commit hooks if you intend to create PRs.
pre-commit install
# Install the baselines by running
pip install -r projects/baselines/requirements.txt -e projects/baselines

Agents

Corax includes high-quality implementation of many popular RL agents. These agents are meant to be forked and customized for future RL research.

The implementation has been used in numerous research projects and we intend to provide benchmark results for these agents in the future.

Corax currently implements the following agents JAX:

More agents, including those implemented in Magi may be added in the future. Contributions to include new agents are welcome!

Datasets

For online RL, Corax uses Reverb for online RL agents.

When working with offline RL, existing datasets provided by the community may come in different formats. It can be time-consuming to integrate existing algorithms with different datasets.

Therefore, for offline RL, Corax provides additional TFDS dataset builders that can build datasets stored in RLDS format. This allows easily running the same offline RL algorithm on offline RL datasets in a consistent manner. You may want to check out the list of datasets officially supported by the TFDS/RLDS team.

In addition to the official RLDS datasets, the following datasets can be built with Corax:

NOTE: Some of these datasets do not yet cover all splits provided by the original dataset. They will be added as the need arises.

Acknowledgements

We would like to thank the Acme authors who have provided a great starting point for Corax. Without them, Corax would not exist as a significant portion of the current code is forked from them. You should check out Acme if you are looking for more RL agent implementations.

We would like to thank the authors of the original papers for open-sourcing their code which has been a great help in our re-implementation.