jax[cuda12]
hydra-core
rootutils
flax
optax
distrax
shutup
scikit-learn
scikit-image
ml_collections
rich
gymnasium
matplotlib
ogbench
jaxtyping
shapely
moviepy==1.0.3
imageio
wandb==0.18.0