diff --git a/qdax/__init__.py b/qdax/__init__.py
index d31c31ea..788da1fb 100644
--- a/qdax/__init__.py
+++ b/qdax/__init__.py
@@ -1 +1 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"
diff --git a/setup.py b/setup.py
index a71f3174..0065bf18 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",
     ],
@@ -45,7 +46,7 @@
         "License :: OSI Approved :: MIT License",
         "Operating System :: POSIX :: Linux",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
 )