docker_requirements.txt

# This list of dependencies was partially generated with `pip freeze` for
# the Dockerfile. If using this list in a different context, ignore the
# transitive dependencies.

# DeepMind libraries.
chex==0.1.3
distrax==0.1.2
dm-env==1.5
dm-haiku==0.0.6
dm-tree==0.1.7
optax==0.1.2
rlax==0.1.2

# JAX libraries with GPU support.
jax==0.3.10
jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.10%2Bcuda11.cudnn805-cp39-none-manylinux2014_x86_64.whl

# ALE and Gym libraries.
atari-py==0.2.9
gym==0.19.0

# Other explicit dependencies.
absl-py==1.1.0
numpy==1.21.5
Pillow==10.0.0
python-snappy==0.6.1

# Transitive dependencies.
cloudpickle==1.6.0
decorator==5.1.1
flatbuffers==2.0
gast==0.5.3
jmp==0.0.2
opt-einsum==3.3.0
scipy==1.8.1
six==1.16.0
tabulate==0.8.10
tensorflow-probability==0.17.0
toolz==0.11.2
typing_extensions==4.2.0