--pre

--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# === Core Dependencies ===
jax[cuda12]==0.6.2
flax
optax
chex
numpy
scipy

tensorflow
tensorflow-probability
distrax

mujoco==3.2.0
brax==0.12.0
gymnax
gym
craftax==1.5.0
