diff --git a/setup.py b/setup.py
index a71f3174..d9c2afb5 100644
--- a/setup.py
+++ b/setup.py
@@ -22,15 +22,16 @@
     long_description_content_type="text/markdown",
     install_requires=[
         "absl-py>=1.0.0",
-        "jax>=0.3.16",
-        "jaxlib>=0.3.15",  # necessary to build the doc atm
+        "jax>=0.4.16",
+        "jaxlib>=0.4.16",  # necessary to build the doc atm
         "jinja2<3.1.0",
-        "jumanji>=0.1.3",
-        "flax>=0.6, <0.6.2",
-        "brax>=0.0.15",
+        "jumanji>=0.3.1",
+        "flax>=0.7.4",
+        "chex>=0.1.83",
+        "brax>=0.9.2",
         "gym>=0.23.1",
         "numpy>=1.22.3",
-        "optax>=0.1, <0.1.5",
+        "optax>=0.1.7",
         "scikit-learn>=1.0.2",
         "scipy>=1.8.0",
     ],