Fix mujoco_py only installation (#3072)

YouJiacheng · web-flow · commit 05d52cfacaf9 · 2022-09-12T08:19:22.000-04:00
Only import `mujoco_rendering` when `mujoco` installed.
Otherwise the `try import` code in `gym.envs.mujoco.mujoco_env` is useless.
diff --git a/gym/envs/mujoco/__init__.py b/gym/envs/mujoco/__init__.py
@@ -1,11 +1,5 @@
 from gym.envs.mujoco.mujoco_env import MujocoEnv, MuJocoPyEnv  # isort:skip
-from gym.envs.mujoco.mujoco_rendering import (  # isort:skip
-    RenderContextOffscreen,
-    Viewer,
-)
 
-# ^^^^^ so that user gets the correct error
-# message if mujoco is not installed correctly
 from gym.envs.mujoco.ant import AntEnv
 from gym.envs.mujoco.half_cheetah import HalfCheetahEnv
 from gym.envs.mujoco.hopper import HopperEnv
diff --git a/gym/envs/mujoco/mujoco_env.py b/gym/envs/mujoco/mujoco_env.py
@@ -7,20 +7,19 @@
 from gym import error, logger, spaces
 from gym.spaces import Space
 
-MUJOCO_PY_NOT_INSTALLED = False
-MUJOCO_NOT_INSTALLED = False
-
 try:
     import mujoco_py
 except ImportError as e:
     MUJOCO_PY_IMPORT_ERROR = e
-    MUJOCO_PY_NOT_INSTALLED = True
+else:
+    MUJOCO_PY_IMPORT_ERROR = None
 
 try:
     import mujoco
 except ImportError as e:
     MUJOCO_IMPORT_ERROR = e
-    MUJOCO_NOT_INSTALLED = True
+else:
+    MUJOCO_IMPORT_ERROR = None
 
 
 DEFAULT_SIZE = 480
@@ -183,7 +182,7 @@ def __init__(
         camera_id: Optional[int] = None,
         camera_name: Optional[str] = None,
     ):
-        if MUJOCO_PY_NOT_INSTALLED:
+        if MUJOCO_PY_IMPORT_ERROR is not None:
             raise error.DependencyNotInstalled(
                 f"{MUJOCO_PY_IMPORT_ERROR}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)"
             )
@@ -305,7 +304,7 @@ def __init__(
         camera_id: Optional[int] = None,
         camera_name: Optional[str] = None,
     ):
-        if MUJOCO_NOT_INSTALLED:
+        if MUJOCO_IMPORT_ERROR is not None:
             raise error.DependencyNotInstalled(
                 f"{MUJOCO_IMPORT_ERROR}. (HINT: you need to install mujoco)"
             )
@@ -395,15 +394,18 @@ def close(self):
 
     def _get_viewer(
         self, mode
-    ) -> Union["gym.envs.mujoco.Viewer", "gym.envs.mujoco.RenderContextOffscreen"]:
+    ) -> Union[
+        "gym.envs.mujoco.mujoco_rendering.Viewer",
+        "gym.envs.mujoco.mujoco_rendering.RenderContextOffscreen",
+    ]:
         self.viewer = self._viewers.get(mode)
         if self.viewer is None:
             if mode == "human":
-                from gym.envs.mujoco import Viewer
+                from gym.envs.mujoco.mujoco_rendering import Viewer
 
                 self.viewer = Viewer(self.model, self.data)
             elif mode in {"rgb_array", "depth_array"}:
-                from gym.envs.mujoco import RenderContextOffscreen
+                from gym.envs.mujoco.mujoco_rendering import RenderContextOffscreen
 
                 self.viewer = RenderContextOffscreen(self.model, self.data)
             else:
diff --git a/tests/envs/utils.py b/tests/envs/utils.py
@@ -4,7 +4,7 @@
 import numpy as np
 
 import gym
-from gym import logger
+from gym import error, logger
 from gym.envs.registration import EnvSpec
 
 
@@ -17,7 +17,7 @@ def try_make_env(env_spec: EnvSpec) -> Optional[gym.Env]:
     if "gym.envs." in env_spec.entry_point:
         try:
             return env_spec.make(disable_env_checker=True).unwrapped
-        except ImportError as e:
+        except (ImportError, error.DependencyNotInstalled) as e:
             logger.warn(f"Not testing {env_spec.id} due to error: {e}")
     return None
 
