{"id":43876361,"url":"https://github.com/ethanluoyc/corax","last_synced_at":"2026-02-06T14:44:39.606Z","repository":{"id":198559864,"uuid":"699444586","full_name":"ethanluoyc/corax","owner":"ethanluoyc","description":"Corax: Core RL in JAX","archived":false,"fork":false,"pushed_at":"2024-02-22T14:17:32.000Z","size":459,"stargazers_count":29,"open_issues_count":3,"forks_count":0,"subscribers_count":3,"default_branch":"main","last_synced_at":"2024-04-16T07:35:43.788Z","etag":null,"topics":["jax","machine-learning","reinforcement-learning"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ethanluoyc.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null}},"created_at":"2023-10-02T16:41:39.000Z","updated_at":"2024-04-04T11:14:32.000Z","dependencies_parsed_at":"2024-02-20T15:36:37.913Z","dependency_job_id":null,"html_url":"https://github.com/ethanluoyc/corax","commit_stats":{"total_commits":33,"total_committers":2,"mean_commits":16.5,"dds":"0.030303030303030276","last_synced_commit":"cf507542085de0129698ad6e3ced5f26974fbfe7"},"previous_names":["ethanluoyc/corax"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/ethanluoyc/corax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ethanluoyc%2Fcorax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ethanluoyc%2Fcorax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ethanluoyc%2Fcorax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ethanluoyc%2Fcorax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ethanluoyc","download_url":"https://codeload.github.com/ethanluoyc/corax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ethanluoyc%2Fcorax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":29165068,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-02-06T14:37:12.680Z","status":"ssl_error","status_checked_at":"2026-02-06T14:36:22.973Z","response_time":59,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","machine-learning","reinforcement-learning"],"created_at":"2026-02-06T14:44:38.768Z","updated_at":"2026-02-06T14:44:39.591Z","avatar_url":"https://github.com/ethanluoyc.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Corax: Core RL in JAX\n[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)\n[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)\n[![test](https://github.com/ethanluoyc/corax/actions/workflows/test.yml/badge.svg)](https://github.com/ethanluoyc/corax/actions/workflows/test.yml)\n\n**[Installation](#installation)** |\n**[Examples](#examples)** |\n**[Agents](#agents)** |\n**[Datasets](#datasets)**\n\nCorax is a library for reinforcement learning algorithms in JAX. It aims at providing\nmodular, pure and functional components for RL algorithms that can be easily used in\ndifferent training loops and accelerator configurations. Currently, we are exploring the\ndesign of a LearnerCore and ActorCore design that allows easy composition and scaling of\nRL algorithms. At the same time, Corax aims to provide strong baseline agents that can\nbe forked and customized for future RL research.\n\nCorax starts as a fork of the\n[dm-acme](https://github.com/google-deepmind/acme/tree/master) library while aiming to\nprovide a better experience for researchers working on Online/offline RL in JAX. Future\ndevelopment of Corax may diverge from the design in Acme.\n\n## Installation\nYou can install Corax with\n```\npip install 'git+https://github.com/ethanluoyc/corax#egg=corax[tf,jax]'\n```\n\nTo use Corax with GPU, you need to install JAX with GPU support. Follow the instructions\n[here](https://jax.readthedocs.io/en/latest/installation.html) for how to install JAX\nwith GPU support.\n\n### Note on TensorFlow dependency\nThe base corax package does not depend on specific deep\nlearning frameworks. However, the JAX agent depends on TensorFlow for efficient\ndata-processing.\n\nWe provide optional extras for installing `tensorflow-cpu` and\ncompatible versions of [TensorFlow\nProbability](https://github.com/tensorflow/probability) and\n[Reverb](https://github.com/google-deepmind/reverb/tree/master).\n\nHowever, should that be\nincompatible with your own dependency requirements, you can optionally specify these\ndependencies yourself and opt-out our extras. Check out the\n[pyproject.toml](./pyproject.toml) for examples on how to specify compatible TensorFlow\nversions.\nHere 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\n1. Looking at https://github.com/tensorflow/probability/releases,\nthe tensorflow-probability version compatible with `tensorflow~=2.13.0` is `0.21.0`.\n2. Looking at https://github.com/google-deepmind/reverb/tree/master#reverb-releases, the\n`dm-reverb` version compatible is `0.12.0`.\n\nTherefore, as an application developer, you should put the following in your `requirements.txt`\n```bash\ntensorflow-cpu~=2.13.0\ntensorflow-probability~=0.21.0\ndm-reverb~=0.12.0\n```\n\nIf you use `dm-launchpad`. The workflow is similar, although as of 17 Oct, 2023 Launchpad\ndoes not provide an official build for `tensorflow 2.13.0`. We however, have an unofficial\nmanylinux build for Python 3.9 and 3.10 available at\nhttps://github.com/ethanluoyc/launchpad/releases/tag/v0.6.0rc0. If you intend to use\nthis version, you should include in your `requirements.txt`\n```bash\n# Use the exact link to the wheel file for your Python version\ndm-launchpad @ https://github.com/ethanluoyc/launchpad/releases/download/v0.6.0rc0/dm_launchpad-0.6.0rc0-cp39-cp39-manylinux2014_x86_64.whl\n```\nWe currently do not have build for tensorflow 2.14.0 due to\nhttps://github.com/google-deepmind/launchpad/issues/44.\n\n## Examples\nExamples can be found in [projects](projects/).\n\n## Development\n\n```bash\ngit clone https://github.com/ethanluoyc/jax\ncd corax\n# Create a virtual environment with the method of your choice.\npython3 -m venv .venv\nsource .venv/bin/activate\n# Then run\npip install -e '.[tf,jax,test,dev]'\n# Install pre-commit hooks if you intend to create PRs.\npre-commit install\n# Install the baselines by running\npip install -r projects/baselines/requirements.txt -e projects/baselines\n```\n\n## Agents\nCorax includes high-quality implementation of many popular RL agents. These agents are\nmeant to be forked and customized for future RL research.\n\nThe implementation has been used in numerous research projects and we intend to provide\nbenchmark results for these agents in the future.\n\nCorax currently implements the following agents JAX:\n\n| Agent                | Paper                    | Code                                                           |\n|----------------------|--------------------------|----------------------------------------------------------------|\n| CalQL                | [Nakamoto et al., 2023]  | [calql](corax/agents/jax/calql/)                               |\n| CQL                  | [Kumar et al., 2020]     | [calql](corax/agents/jax/calql/)                               |\n| IQL                  | [Kostrikov et al., 2021] | [iql](corax/agents/jax/iql/)                                   |\n| RLPD                 | [Ball et al., 2023]      | [redq](corax/agents/jax/redq/)                                 |\n| Decision Transformer | [Chen et al., 2021a]     | [decision_transformer](corax/agents/jax/decision_transformer/) |\n| DrQ-v2(-BC)          | [Yarats et al., 2021]    | [drq_v2](corax/agents/jax/drq_v2/)                             |\n| ORIL                 | [Zolna et al., 2020]     | [oril](corax/agents/jax/oril/)                                 |\n| OTR                  | [Luo et al., 2023]       | [otr](corax/agents/jax/otr/)                                   |\n| REDQ                 | [Chen et al., 2021b]     | [redq](corax/agents/jax/redq/)                                 |\n| TD3                  | [Fujimoto et al., 2018]  | [td3](corax/agents/jax/td3/)                                   |\n| TD3-BC               | [Fujimoto et al., 2021]  | [td3](corax/agents/jax/td3/)                                   |\n| TD-MPC               | [Hansen et al., 2021]    | [tdmpc](corax/agents/jax/tdmpc/)                               |\n\nMore agents, including those implemented in\n[Magi](https://github.com/ethanluoyc/magi/tree/main/magi) may be added in the future.\nContributions to include new agents are welcome!\n\n## Datasets\nFor online RL, Corax uses [Reverb](https://github.com/google-deepmind/reverb/) for online RL agents.\n\nWhen working with offline RL, existing datasets provided by the community may come in\ndifferent formats. It can be time-consuming to integrate existing algorithms with\ndifferent datasets.\n\nTherefore, for offline RL, Corax provides additional\n[TFDS](https://github.com/tensorflow/datasets/tree/master) dataset builders that can\nbuild datasets stored in [RLDS](https://github.com/google-research/rlds) format. This\nallows easily running the same offline RL algorithm on offline RL datasets in a\nconsistent manner. You may want to check out the\n[list](https://github.com/google-research/rlds/tree/main#available-datasets) of datasets\nofficially supported by the TFDS/RLDS team.\n\nIn addition to the official RLDS datasets, the following datasets can be built with\nCorax:\n\n| Dataset         | Paper                    | Code                                                |\n|-----------------|--------------------------|-----------------------------------------------------|\n| V-D4RL          | [Lu et al., 2023]        | [vd4rl](corax/datasets/tfds/vd4rl/)                 |\n| Watch and Match | [Haldar et al., 2022]    | [rot](corax/datasets/tfds/rot/)                     |\n| ExoRL           | [Yarats et al., 2022]    | [exorl](corax/datasets/tfds/exorl/)                 |\n| GWIL            | [Fickinger et al., 2022] | [gwil](corax/datasets/tfds/gwil/)                   |\n| Adroit Binary   | [Nair et al., 2022]      | [adroit_binary](corax/datasets/tfds/adroit_binary/) |\n\nNOTE: Some of these datasets do not yet cover all splits provided by the original\ndataset. They will be added as the need arises.\n\n## Acknowledgements\nWe would like to thank the [Acme](https://github.com/google-deepmind/acme) authors who\nhave provided a great starting point for Corax. Without them, Corax would not exist as a\nsignificant portion of the current code is forked from them. You should check out Acme\nif you are looking for more RL agent implementations.\n\nWe would like to thank the authors of the original papers for open-sourcing their code\nwhich has been a great help in our re-implementation.\n\n\u003c!-- Agents --\u003e\n[Nakamoto et al., 2023]: https://arxiv.org/abs/2303.05479\n[Chen et al., 2021a]: https://arxiv.org/abs/2106.01345\n[Yarats et al., 2021]: https://arxiv.org/abs/2107.09645\n[Zolna et al., 2020]: https://arxiv.org/pdf/2011.13885.pdf\n[Kostrikov et al., 2021]: https://openreview.net/forum?id=68n2s9ZJWF8\n[Chen et al., 2021b]: https://arxiv.org/abs/2101.05982\n[Ball et al., 2023]: https://arxiv.org/abs/2302.02948\n[Fujimoto et al., 2018]: https://arxiv.org/abs/1802.09477\n[Fujimoto et al., 2021]: https://arxiv.org/abs/2106.06860.pdf\n[Hansen et al., 2021]: https://arxiv.org/abs/2203.04955\n[Kumar et al., 2020]: https://arxiv.org/abs/2006.04779\n[Luo et al., 2023]: https://arxiv.org/abs/2303.13971\n\n\u003c!-- Papers --\u003e\n[Lu et al., 2023]: https://arxiv.org/abs/1806.06920\n[Haldar et al., 2022]: https://openreview.net/forum?id=ZUtgUA0Fuwd\n[Yarats et al., 2022]: https://arxiv.org/abs/2201.13425\n[Fickinger et al., 2022]: https://arxiv.org/abs/2110.03684\n[Nair et al., 2022]: https://arxiv.org/abs/2006.09359\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fethanluoyc%2Fcorax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fethanluoyc%2Fcorax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fethanluoyc%2Fcorax/lists"}