FROM tensorflow/tensorflow:2.5.0-gpu

RUN apt-get update && apt-get install -y git

RUN apt-get install -y wget rar unrar ffmpeg libsm6 libxext6 libpython3.6 && rm -rf /var/lib/apt/lists/* 

RUN pip install --upgrade pip setuptools wheel
RUN pip install --upgrade "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

RUN pip install atari-py==0.2.9
RUN wget http://www.atarimania.com/roms/Roms.rar
RUN unrar e Roms.rar
RUN unzip ROMS.zip
RUN python -m atari_py.import_roms /ROMS

RUN pip install numpy nltk click tqdm tables joblib opencv-python dataclasses 

RUN pip install absl-py==0.12.0 bsuite==0.3.5 
RUN pip install dm-control==0.0.364896371 dm-env==1.4 dm-haiku==0.0.4
RUN pip install dm-sonnet==2.0.0 dm-tree==0.1.6 optax==0.0.8 
RUN pip install pillow==8.2.0 pytype==2021.6.17 pytest-xdist==2.3.0 
 
RUN pip install rlax
RUN pip install dopamine-rl 
RUN pip install tensorflow_datasets tensorflow-probability

RUN pip install git+https://github.com/deepmind/acme.git#egg=dm-acme

ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu"
RUN pip install dm-reverb[tensorflow]

RUN pip install scikit-learn