FROM tensorflow/tensorflow:latest-gpu-jupyter

USER root

RUN apt-get install -y apt-utils

RUN apt-get update -y && \ 
    apt-get install -y xvfb && \
    apt-get install -y python-opengl 

# Optional, needed for some environments
RUN apt-get install -y cmake && \
    apt-get install -y zlib1g zlib1g-dev 

USER ${NB_USER}

RUN pip install \
        gym \
        pyvirtualdisplay 

# Needed for some environments
RUN pip install swig
RUN pip install box2d-py atari_py pystan
RUN pip install gym[accept-rom-license]
RUN pip install gym[atari] gym[all]
# stardard data sci packages
RUN pip install scikit-learn scipy statsmodels xgboost lightgbm scikit-image
RUN pip install opencv-python

#pytorch because ya i learned tensorflow has some opportunities for improvement
RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116