jax
distrax
optax
flax
numpy
black
pre-commit
argparse
wandb
orbax-checkpoint
pygame
gymnax
chex
matplotlib
imageio
