Merge pull request #84 from jeffreywu13579/requirements-update

Updated requirements

README.md

@@ -49,7 +49,7 @@ We fixed a major issue in the intervention action frame. See release [v0.1.1](ht
 
     - For GPU:
         ```bash
-        pip install --upgrade "jax[cuda12]==0.4.28" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+        pip install --upgrade "jax[cuda12_pip]==0.4.35" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
         ```
 
     - For TPU

serl_launcher/requirements.txt

@@ -1,15 +1,14 @@
 gym >= 0.26
-numpy==1.24.3
-jax>=0.4.13
-distrax==0.1.2
-flax==0.8.0
+numpy>=1.24.3
+flax>=0.8.0
+distrax>=0.1.2
 ml_collections >= 0.1.0
 tqdm >= 4.60.0
-chex==0.1.85
-optax==0.1.5
-orbax-checkpoint==0.5.10
+chex>=0.1.85
+optax>=0.1.5
+orbax-checkpoint>=0.5.10
 absl-py >= 0.12.0
-scipy <= 1.12.0
+scipy==1.11.4
 wandb >= 0.12.14
 tensorflow>=2.16.0
 tensorflow_probability>=0.24.0