# requirements are alligned with nvcr.io/nvidia/jax:23.10-py3 image
jax==0.4.33
jaxlib==0.4.33
flax==0.7.4
chex==0.1.84
optax==0.1.7
dotmap==1.3.30
evosax==0.1.5
distrax==0.1.5
brax==0.10.3
mujoco==3.1.3
gymnax==0.0.6
safetensors==0.4.2
flashbax==0.1.0
# less sensitive libs
wandb
pytest
pygame
numpy>=1.26.1
hydra-core>=1.3.2
omegaconf>=2.3.0
matplotlib>=3.8.3
pillow>=10.2.0
pettingzoo>=1.24.3
tqdm>=4.66.0
scipy<=1.12