From 472ff8edb815070c405da913dcbe64c1a06e0e7d Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@ensta.org>
Date: Fri, 23 Jun 2023 12:38:31 +0200
Subject: [PATCH] [ci skip] Update README with SBX

---
 README.md | 6 ++++++
 setup.py  | 1 +
 2 files changed, 7 insertions(+)

diff --git a/README.md b/README.md
index 45740fedc..820c3634c 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 dc7503c76..3e2643866 100644
--- a/setup.py
+++ b/setup.py
@@ -162,6 +162,7 @@
         "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",