jax>=0.4.25
numpy>=1.26
flax
safetensors
chex
matplotlib
pillow
jaxtyping>=0.2.28
gymnax

[algs]
optax
distrax
flashbax>=0.1.0
wandb
hydra-core>=1.3.2
omegaconf>=2.3.0
tqdm>=4.66.0
imageio

[dev]
pytest
mpe2
ruff
pyright
pre-commit

[mabrax]
brax>=0.10.3
