jax<=0.9.0,>=0.4.37
distrax
optax
flax>=0.10.4
numpy
argparse
wandb
orbax-checkpoint
pygame
gymnax
chex
matplotlib
imageio
seaborn
hydra-core
safetensors
pygame_widgets==1.1.5
jaxgl
jax2d==1.0.1
moviepy==1.0.3
