Hotfix: dependencies in requirements.txt and setup.py are out of sync (…

…#165)

* update imports in setup.py
* changed required version of python from 3.8 to 3.9
* change QDax version from 0.2.3 to 0.2.4

qdax/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"

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",
     ],
 )