Refactor Mujoco Rendering and bump to Mujoco 2.3.0 (#112)

gymnasium/envs/mujoco/ant_v4.py

@@ -240,7 +240,12 @@ def __init__(
         )
 
         MujocoEnv.__init__(
-            self, xml_file, 5, observation_space=observation_space, **kwargs
+            self,
+            xml_file,
+            5,
+            observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
+            **kwargs
         )
 
     @property
@@ -348,11 +353,3 @@ def reset_model(self):
         observation = self._get_obs()
 
         return observation
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        for key, value in DEFAULT_CAMERA_CONFIG.items():
-            if isinstance(value, np.ndarray):
-                getattr(self.viewer.cam, key)[:] = value
-            else:
-                setattr(self.viewer.cam, key, value)

gymnasium/envs/mujoco/half_cheetah_v4.py

@@ -180,7 +180,12 @@ def __init__(
             )
 
         MujocoEnv.__init__(
-            self, "half_cheetah.xml", 5, observation_space=observation_space, **kwargs
+            self,
+            "half_cheetah.xml",
+            5,
+            observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
+            **kwargs
         )
 
     def control_cost(self, action):
@@ -237,11 +242,3 @@ def reset_model(self):
 
         observation = self._get_obs()
         return observation
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        for key, value in DEFAULT_CAMERA_CONFIG.items():
-            if isinstance(value, np.ndarray):
-                getattr(self.viewer.cam, key)[:] = value
-            else:
-                setattr(self.viewer.cam, key, value)

gymnasium/envs/mujoco/hopper_v4.py

@@ -203,7 +203,12 @@ def __init__(
             )
 
         MujocoEnv.__init__(
-            self, "hopper.xml", 4, observation_space=observation_space, **kwargs
+            self,
+            "hopper.xml",
+            4,
+            observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
+            **kwargs
         )
 
     @property

gymnasium/envs/mujoco/humanoid_v4.py

@@ -268,7 +268,12 @@ def __init__(
             )
 
         MujocoEnv.__init__(
-            self, "humanoid.xml", 5, observation_space=observation_space, **kwargs
+            self,
+            "humanoid.xml",
+            5,
+            observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
+            **kwargs
         )
 
     @property
@@ -366,11 +371,3 @@ def reset_model(self):
 
         observation = self._get_obs()
         return observation
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        for key, value in DEFAULT_CAMERA_CONFIG.items():
-            if isinstance(value, np.ndarray):
-                getattr(self.viewer.cam, key)[:] = value
-            else:
-                setattr(self.viewer.cam, key, value)

gymnasium/envs/mujoco/humanoidstandup_v4.py

@@ -4,6 +4,13 @@
 from gymnasium.envs.mujoco import MujocoEnv
 from gymnasium.spaces import Box
 
+DEFAULT_CAMERA_CONFIG = {
+    "trackbodyid": 1,
+    "distance": 4.0,
+    "lookat": np.array((0.0, 0.0, 0.8925)),
+    "elevation": -20.0,
+}
+
 
 class HumanoidStandupEnv(MujocoEnv, utils.EzPickle):
     """
@@ -201,6 +208,7 @@ def __init__(self, **kwargs):
             "humanoidstandup.xml",
             5,
             observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
             **kwargs
         )
         utils.EzPickle.__init__(self, **kwargs)
@@ -255,10 +263,3 @@ def reset_model(self):
             ),
         )
         return self._get_obs()
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        self.viewer.cam.trackbodyid = 1
-        self.viewer.cam.distance = self.model.stat.extent * 1.0
-        self.viewer.cam.lookat[2] = 0.8925
-        self.viewer.cam.elevation = -20

gymnasium/envs/mujoco/inverted_double_pendulum_v4.py

@@ -4,6 +4,12 @@
 from gymnasium.envs.mujoco import MujocoEnv
 from gymnasium.spaces import Box
 
+DEFAULT_CAMERA_CONFIG = {
+    "trackbodyid": 0,
+    "distance": 4.1225,
+    "lookat": np.array((0.0, 0.0, 0.12250000000000005)),
+}
+
 
 class InvertedDoublePendulumEnv(MujocoEnv, utils.EzPickle):
     """
@@ -133,6 +139,7 @@ def __init__(self, **kwargs):
             "inverted_double_pendulum.xml",
             5,
             observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
             **kwargs
         )
         utils.EzPickle.__init__(self, **kwargs)
@@ -169,10 +176,3 @@ def reset_model(self):
             self.init_qvel + self.np_random.standard_normal(self.model.nv) * 0.1,
         )
         return self._get_obs()
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        v = self.viewer
-        v.cam.trackbodyid = 0
-        v.cam.distance = self.model.stat.extent * 0.5
-        v.cam.lookat[2] = 0.12250000000000005  # v.model.stat.center[2]

gymnasium/envs/mujoco/inverted_pendulum_v4.py

@@ -4,6 +4,11 @@
 from gymnasium.envs.mujoco import MujocoEnv
 from gymnasium.spaces import Box
 
+DEFAULT_CAMERA_CONFIG = {
+    "trackbodyid": 0,
+    "distance": 2.04,
+}
+
 
 class InvertedPendulumEnv(MujocoEnv, utils.EzPickle):
     """
@@ -104,6 +109,7 @@ def __init__(self, **kwargs):
             "inverted_pendulum.xml",
             2,
             observation_space=observation_space,
+            default_camera_config=DEFAULT_CAMERA_CONFIG,
             **kwargs
         )
 
@@ -128,9 +134,3 @@ def reset_model(self):
 
     def _get_obs(self):
         return np.concatenate([self.data.qpos, self.data.qvel]).ravel()
-
-    def viewer_setup(self):
-        assert self.viewer is not None
-        v = self.viewer
-        v.cam.trackbodyid = 0
-        v.cam.distance = self.model.stat.extent

gymnasium/envs/mujoco/mujoco_env.py

@@ -52,12 +52,9 @@ def __init__(
 
         self.init_qpos = self.data.qpos.ravel().copy()
         self.init_qvel = self.data.qvel.ravel().copy()
-        self._viewers = {}
 
         self.frame_skip = frame_skip
 
-        self.viewer = None
-
         assert self.metadata["render_modes"] == [
             "human",
             "rgb_array",
@@ -90,12 +87,6 @@ def reset_model(self):
         """
         raise NotImplementedError
 
-    def viewer_setup(self):
-        """
-        This method is called when the viewer is initialized.
-        Optionally implement this method, if you need to tinker with camera position and so forth.
-        """
-
     def _initialize_simulation(self):
         """
         Initialize MuJoCo simulation data structures mjModel and mjData.
@@ -159,9 +150,8 @@ def do_simulation(self, ctrl, n_frames):
         self._step_mujoco_simulation(ctrl, n_frames)
 
     def close(self):
-        if self.viewer is not None:
-            self.viewer = None
-            self._viewers = {}
+        """Close all processes like rendering contexts"""
+        raise NotImplementedError
 
     def get_body_com(self, body_name):
         """Return the cartesian position of a body frame"""
@@ -197,6 +187,9 @@ def __init__(
             "you are trying to precisely replicate previous works)."
         )
 
+        self.viewer = None
+        self._viewers = {}
+
         super().__init__(
             model_path,
             frame_skip,
@@ -223,6 +216,9 @@ def set_state(self, qpos, qvel):
         self.sim.set_state(state)
         self.sim.forward()
 
+    def get_body_com(self, body_name):
+        return self.data.get_body_xpos(body_name)
+
     def _step_mujoco_simulation(self, ctrl, n_frames):
         self.sim.data.ctrl[:] = ctrl
 
@@ -297,8 +293,17 @@ def _get_viewer(
 
         return self.viewer
 
-    def get_body_com(self, body_name):
-        return self.data.get_body_xpos(body_name)
+    def close(self):
+        if self.viewer is not None:
+            self.viewer = None
+            self._viewers = {}
+
+    def viewer_setup(self):
+        """
+        This method is called when the viewer is initialized.
+        Optionally implement this method, if you need to tinker with camera position and so forth.
+        """
+        raise NotImplementedError
 
 
 class MujocoEnv(BaseMujocoEnv):
@@ -314,11 +319,13 @@ def __init__(
         height: int = DEFAULT_SIZE,
         camera_id: Optional[int] = None,
         camera_name: Optional[str] = None,
+        default_camera_config: Optional[dict] = None,
     ):
         if MUJOCO_IMPORT_ERROR is not None:
             raise error.DependencyNotInstalled(
                 f"{MUJOCO_IMPORT_ERROR}. (HINT: you need to install mujoco)"
             )
+
         super().__init__(
             model_path,
             frame_skip,
@@ -330,6 +337,12 @@ def __init__(
             camera_name,
         )
 
+        from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
+
+        self.mujoco_renderer = MujocoRenderer(
+            self.model, self.data, default_camera_config
+        )
+
     def _initialize_simulation(self):
         self.model = mujoco.MjModel.from_xml_path(self.fullpath)
         # MjrContext will copy model.vis.global_.off* to con.off*
@@ -359,85 +372,13 @@ def _step_mujoco_simulation(self, ctrl, n_frames):
         mujoco.mj_rnePostConstraint(self.model, self.data)
 
     def render(self):
-        if self.render_mode is None:
-            assert self.spec is not None
-            gym.logger.warn(
-                "You are calling render method without specifying any render mode. "
-                "You can specify the render_mode at initialization, "
-                f'e.g. gym.make("{self.spec.id}", render_mode="rgb_array")'
-            )
-            return
-
-        if self.render_mode in {
-            "rgb_array",
-            "depth_array",
-        }:
-            camera_id = self.camera_id
-            camera_name = self.camera_name
-
-            if camera_id is not None and camera_name is not None:
-                raise ValueError(
-                    "Both `camera_id` and `camera_name` cannot be"
-                    " specified at the same time."
-                )
-
-            no_camera_specified = camera_name is None and camera_id is None
-            if no_camera_specified:
-                camera_name = "track"
-
-            if camera_id is None:
-                camera_id = mujoco.mj_name2id(
-                    self.model,
-                    mujoco.mjtObj.mjOBJ_CAMERA,
-                    camera_name,
-                )
-
-                self._get_viewer(self.render_mode).render(camera_id=camera_id)
-
-        if self.render_mode == "rgb_array":
-            data = self._get_viewer(self.render_mode).read_pixels(depth=False)
-            # original image is upside-down, so flip it
-            return data[::-1, :, :]
-        elif self.render_mode == "depth_array":
-            self._get_viewer(self.render_mode).render()
-            # Extract depth part of the read_pixels() tuple
-            data = self._get_viewer(self.render_mode).read_pixels(depth=True)[1]
-            # original image is upside-down, so flip it
-            return data[::-1, :]
-        elif self.render_mode == "human":
-            self._get_viewer(self.render_mode).render()
+        return self.mujoco_renderer.render(
+            self.render_mode, self.camera_id, self.camera_name
+        )
 
     def close(self):
-        if self.viewer is not None:
-            self.viewer.close()
-        super().close()
-
-    def _get_viewer(
-        self, mode
-    ) -> 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 gymnasium.envs.mujoco.mujoco_rendering import Viewer
-
-                self.viewer = Viewer(self.model, self.data)
-            elif mode in {"rgb_array", "depth_array"}:
-                from gymnasium.envs.mujoco.mujoco_rendering import (
-                    RenderContextOffscreen,
-                )
-
-                self.viewer = RenderContextOffscreen(self.model, self.data)
-            else:
-                raise AttributeError(
-                    f"Unexpected mode: {mode}, expected modes: {self.metadata['render_modes']}"
-                )
-
-            self.viewer_setup()
-            self._viewers[mode] = self.viewer
-        return self.viewer
+        if self.mujoco_renderer is not None:
+            self.mujoco_renderer.close()
 
     def get_body_com(self, body_name):
         return self.data.body(body_name).xpos