cleaning up jax bs

2026-05-31 08:33:36 +00:00 · 2026-03-08 19:15:58 +01:00
parent 73246d7dd8
commit 4c658a93a7
27 changed files with 173 additions and 3146 deletions
--- a/docker/trainer.requirements.txt
+++ b/docker/trainer.requirements.txt
@@ -5,9 +5,3 @@ gymnasium>=0.29.0
 stable-baselines3>=2.2.0
 tensorboard>=2.15.0
 wandb>=0.17.0
-tensorflow-probability==0.24.0
-flax==0.10.7
-optax==0.2.7
-distrax==0.1.5
-orbax-checkpoint==0.11.32
-chex==0.1.90