Add flashbax support
This gives us a JAX experience buffer for Soft Actor-Critic.
Add jaxtyping in too so we can type things better.
Change-Id: I4c9f0071f9e912dcab0c883da8d9d7990ed06c46
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/tools/python/requirements.txt b/tools/python/requirements.txt
index e7f23c3..2b806f5 100644
--- a/tools/python/requirements.txt
+++ b/tools/python/requirements.txt
@@ -25,9 +25,13 @@
bokeh
tabulate
+flask
+
casadi>=3.6.6
+# ML libraries
jax[cuda12]
+jaxtyping
optax
flax
clu
@@ -38,4 +42,5 @@
tensorflow
tensorflow_datasets
-flask
+# Experience buffer for reinforcement learning
+flashbax