diff --git a/README.md b/README.md index 45740fe..820c363 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,12 @@ This allows SB3 to maintain a stable and compact core, while still providing the Documentation is available online: [https://sb3-contrib.readthedocs.io/](https://sb3-contrib.readthedocs.io/) +## Stable-Baselines Jax (SBX) + +[Stable Baselines Jax (SBX)](https://github.com/araffin/sbx) is a proof of concept version of Stable-Baselines3 in Jax. + +It provides a minimal number of features compared to SB3 but can be much faster (up to 20x times!): https://twitter.com/araffin2/status/1590714558628253698 + ## Installation diff --git a/setup.py b/setup.py index dc7503c..3e26438 100644 --- a/setup.py +++ b/setup.py @@ -162,6 +162,7 @@ setup( "Changelog": "https://stable-baselines3.readthedocs.io/en/master/misc/changelog.html", "SB3-Contrib": "https://github.com/Stable-Baselines-Team/stable-baselines3-contrib", "RL-Zoo": "https://github.com/DLR-RM/rl-baselines3-zoo", + "SBX": "https://github.com/araffin/sbx", }, classifiers=[ "Programming Language :: Python :: 3",