Refactor Gym and DMControl wrappers (pytorch#194)

ai4co · Jun 15, 2022 · 97f4d48 · 97f4d48
1 parent 87eb981
commit 97f4d48
Show file tree

Hide file tree

Showing 21 changed files with 703 additions and 319 deletions.
diff --git a/.circleci/unittest/linux/scripts/environment.yml b/.circleci/unittest/linux/scripts/environment.yml
@@ -24,3 +24,4 @@ dependencies:
     - scipy
     - dm_control
     - mujoco_py
+    - pyrender
diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh
@@ -32,20 +32,17 @@ esac
 # submodules
 git submodule sync && git submodule update --init --recursive
 
-#printf "Installing PyTorch with %s\n" "${cudatoolkit}"
-#if [ "${os}" == "MacOSX" ]; then
-#    conda install -y -c "pytorch-${UPLOAD_CHANNEL}" "pytorch-${UPLOAD_CHANNEL}"::pytorch "${cudatoolkit}" pytest
-#else
-#    conda install -y -c "pytorch-${UPLOAD_CHANNEL}" "pytorch-${UPLOAD_CHANNEL}"::pytorch[build="*${version}*"] "${cudatoolkit}" pytest
-#fi
-
-#printf "Installing PyTorch with %s\n" "${CU_VERSION}"
+printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
-    pip install torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre
+    if [[ $OSTYPE == 'darwin'* ]]; then
+      pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    else
+      pip3 install torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre
+    fi
 else
-    conda install -y pytorch torchvision cudatoolkit=10.2 -c pytorch-nightly
+    pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu102
 fi
 
 printf "Installing functorch\n"
@@ -57,3 +54,18 @@ python -c "import functorch"
 
 printf "* Installing torchrl\n"
 python setup.py develop
+
+if [[ $OSTYPE == 'darwin'* ]]; then
+  PRIVATE_MUJOCO_GL=glfw
+else
+  conda install -y -c conda-forge mesa
+  conda install -y -c menpo osmesa
+  PRIVATE_MUJOCO_GL=osmesa
+fi
+
+conda env config vars set MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
+  DISPLAY=unix:0.0 \
+  MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1 \
+  LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin \
+  SDL_VIDEODRIVER=dummy \
+  MUJOCO_GL=$PRIVATE_MUJOCO_GL
diff --git a/.circleci/unittest/linux/scripts/run_test.sh b/.circleci/unittest/linux/scripts/run_test.sh
@@ -5,22 +5,27 @@ set -e
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
-if [[ $OSTYPE == 'darwin'* ]]; then
-  PRIVATE_MUJOCO_GL=glfw
-else
-  PRIVATE_MUJOCO_GL=glfw
-fi
-
 export PYTORCH_TEST_WITH_SLOW='1'
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
+
 root_dir="$(git rev-parse --show-toplevel)"
-export MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1
-export MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210
-export DISPLAY=unix:0.0
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/project/.mujoco/mujoco210/bin
+env_dir="${root_dir}/env"
+lib_dir="${env_dir}/lib"
+
+# solves ImportError: /lib64/libstdc++.so.6: version `GLIBCXX_3.4.21' not found
+export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$lib_dir
+#if [[ $OSTYPE == 'darwin'* ]]; then
+#  PRIVATE_MUJOCO_GL=glfw
+#else
+#  PRIVATE_MUJOCO_GL=egl
+#fi
+## pygame.error: No available video device
+#export SDL_VIDEODRIVER=dummy
+## solves AttributeError: 'NoneType' object has no attribute 'eglQueryString'
+#export MUJOCO_GL=$PRIVATE_MUJOCO_GL
 
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest test/smoke_test.py -v --durations 20
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest test/smoke_test_deps.py -v --durations 20
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest -v --durations 20
+pytest test/smoke_test.py -v --durations 20
+pytest test/smoke_test_deps.py -v --durations 20
+pytest -v --durations 20
diff --git a/.circleci/unittest/linux/scripts/setup_env.sh b/.circleci/unittest/linux/scripts/setup_env.sh
@@ -13,6 +13,7 @@ git config --global --add safe.directory '*'
 root_dir="$(git rev-parse --show-toplevel)"
 conda_dir="${root_dir}/conda"
 env_dir="${root_dir}/env"
+lib_dir="${env_dir}/lib"
 
 cd "${root_dir}"
 
@@ -39,7 +40,7 @@ conda activate "${env_dir}"
 
 # 3. Install mujoco
 printf "* Installing mujoco and related\n"
-mkdir $root_dir/.mujoco
+mkdir -p $root_dir/.mujoco
 cd $root_dir/.mujoco/
 wget https://github.com/deepmind/mujoco/releases/download/2.1.1/mujoco-2.1.1-linux-x86_64.tar.gz
 tar -xf mujoco-2.1.1-linux-x86_64.tar.gz
@@ -51,12 +52,10 @@ cd $this_dir
 printf "* Installing dependencies (except PyTorch)\n"
 echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
 cat "${this_dir}/environment.yml"
-conda env update --file "${this_dir}/environment.yml" --prune
 
-conda install -y -c conda-forge glfw
-conda install -y -c conda-forge glew
+#conda install -y -c conda-forge glfw
+#conda install -y -c conda-forge mesa
+#conda install -y -c menpo osmesa
+#conda install -y -c conda-forge glew
 
-#yum makecache
-#yum -y install glfw-devel
-#yum -y install libGLEW
-#yum -y install gcc-c++
+conda env update --file "${this_dir}/environment.yml" --prune
diff --git a/.circleci/unittest/linux_optdeps/scripts/run_test.sh b/.circleci/unittest/linux_optdeps/scripts/run_test.sh
@@ -15,4 +15,4 @@ export MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210
 export DISPLAY=unix:0.0
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/project/.mujoco/mujoco210/bin
 #MUJOCO_GL=glfw pytest --cov=torchrl --junitxml=test-results/junit.xml -v --durations 20
-MUJOCO_GL=glfw pytest -v --durations 20
+MUJOCO_GL=egl pytest -v --durations 20
diff --git a/.circleci/unittest/linux_stable/scripts/environment.yml b/.circleci/unittest/linux_stable/scripts/environment.yml
@@ -25,3 +25,4 @@ dependencies:
     - scipy
     - dm_control
     - mujoco_py
+    - pyrender
diff --git a/.circleci/unittest/linux_stable/scripts/install.sh b/.circleci/unittest/linux_stable/scripts/install.sh
@@ -36,9 +36,9 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     # conda install -y pytorch torchvision cpuonly -c pytorch-nightly
     # use pip to install pytorch as conda can frequently pick older release
-    pip install torch torchvision
+    pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
 else
-    conda install -y pytorch torchvision cudatoolkit=10.2 -c pytorch
+    pip3 install torch torchvision torchaudio
 fi
 
 printf "Installing functorch\n"
@@ -52,3 +52,18 @@ printf "g++ version: "
 gcc --version
 
 python setup.py install
+
+if [[ $OSTYPE == 'darwin'* ]]; then
+  PRIVATE_MUJOCO_GL=glfw
+else
+  conda install -y -c conda-forge mesa
+  conda install -y -c menpo osmesa
+  PRIVATE_MUJOCO_GL=osmesa
+fi
+
+conda env config vars set MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
+  DISPLAY=unix:0.0 \
+  MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1 \
+  LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin \
+  SDL_VIDEODRIVER=dummy \
+  MUJOCO_GL=$PRIVATE_MUJOCO_GL
diff --git a/.circleci/unittest/linux_stable/scripts/run_test.sh b/.circleci/unittest/linux_stable/scripts/run_test.sh
@@ -5,22 +5,27 @@ set -e
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
-if [[ $OSTYPE == 'darwin'* ]]; then
-  PRIVATE_MUJOCO_GL=glfw
-else
-  PRIVATE_MUJOCO_GL=glfw
-fi
-
 export PYTORCH_TEST_WITH_SLOW='1'
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
+
 root_dir="$(git rev-parse --show-toplevel)"
-export MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1
-export MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210
-export DISPLAY=unix:0.0
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/project/.mujoco/mujoco210/bin
+env_dir="${root_dir}/env"
+lib_dir="${env_dir}/lib"
+
+# solves ImportError: /lib64/libstdc++.so.6: version `GLIBCXX_3.4.21' not found
+export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$lib_dir
+#if [[ $OSTYPE == 'darwin'* ]]; then
+#  PRIVATE_MUJOCO_GL=glfw
+#else
+#  PRIVATE_MUJOCO_GL=egl
+#fi
+## pygame.error: No available video device
+#export SDL_VIDEODRIVER=dummy
+## solves AttributeError: 'NoneType' object has no attribute 'eglQueryString'
+#export MUJOCO_GL=$PRIVATE_MUJOCO_GL
 
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest test/smoke_test.py -v --durations 20
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest test/smoke_test_deps.py -v --durations 20
-MUJOCO_GL=$PRIVATE_MUJOCO_GL pytest -v --durations 20
+pytest test/smoke_test.py -v --durations 20
+pytest test/smoke_test_deps.py -v --durations 20
+pytest -v --durations 20
diff --git a/.circleci/unittest/linux_stable/scripts/setup_env.sh b/.circleci/unittest/linux_stable/scripts/setup_env.sh
@@ -13,6 +13,7 @@ git config --global --add safe.directory '*'
 root_dir="$(git rev-parse --show-toplevel)"
 conda_dir="${root_dir}/conda"
 env_dir="${root_dir}/env"
+lib_dir="${env_dir}/lib"
 
 cd "${root_dir}"
 
@@ -51,12 +52,10 @@ cd $this_dir
 printf "* Installing dependencies (except PyTorch)\n"
 echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
 cat "${this_dir}/environment.yml"
-conda env update --file "${this_dir}/environment.yml" --prune
 
-conda install -y -c conda-forge glfw
-conda install -y -c conda-forge glew
+#conda install -y -c conda-forge glfw
+#conda install -y -c conda-forge mesa
+##conda install -y -c menpo osmesa
+#conda install -y -c conda-forge glew
 
-#yum makecache
-#yum -y install glfw-devel
-#yum -y install libGLEW
-#yum -y install gcc-c++
+conda env update --file "${this_dir}/environment.yml" --prune
diff --git a/setup.cfg b/setup.cfg
@@ -15,6 +15,7 @@ per-file-ignores =
     __init__.py: F401, F403, F405
     ./hubconf.py: F401
     test/smoke_test.py: F401
+    test/smoke_test_deps.py: F401
     test_*.py: F841, E731, E266
 exclude = venv
 

diff --git a/test/smoke_test.py b/test/smoke_test.py
@@ -5,7 +5,7 @@ def test_imports():
         TensorDict,
         TensorSpec,
     )
-    from torchrl.envs import GymLikeEnv
     from torchrl.envs import TransformedEnv, Transform
+    from torchrl.envs.gym_like import GymLikeEnv
     from torchrl.modules import TensorDictModule
     from torchrl.objectives.costs.common import _LossModule
diff --git a/test/smoke_test_deps.py b/test/smoke_test_deps.py
@@ -2,9 +2,17 @@
 
 from torch.utils.tensorboard import SummaryWriter
 from torchrl.envs import DMControlEnv, GymEnv
+from torchrl.envs.libs.dm_control import _has_dmc
+from torchrl.envs.libs.gym import _has_gym
 
 
 def test_dm_control():
+    import dm_control
+    import dm_env
+    from dm_control import suite
+    from dm_control.suite.wrappers import pixels
+
+    assert _has_dmc
     env = DMControlEnv("cheetah", "run")
     env.reset()
 
@@ -16,12 +24,10 @@ def test_dm_control():
 
 
 def test_gym():
-    env = GymEnv("ALE/Pong-v5")
-    env.reset()
+    import gym
 
-
-def test_gym_pixels():
-    env = GymEnv("ALE/Pong-v5", from_pixels=True)
+    assert _has_gym
+    env = GymEnv("ALE/Pong-v5")
     env.reset()
 
 

diff --git a/test/test_gym.py b/test/test_gym.py