chex
distrax
flashbax
flax
gymnax
jax
pyyaml
tensorboard
tensorboardX
tyro
wandb

[utils]
numpy
optax
orbax-checkpoint
pandas
pytest
tqdm

[viz]
imageio
imageio[ffmpeg]
ipython
ipywidgets
jinja2
matplotlib
pillow
pygame
scipy
seaborn
