optax==0.2.3
jax<0.5
gymnasium==0.29
equinox
jaxtyping
haliax
gymnax