From a2e17dfa54ac30ebe0e27a125892fd93319e6412 Mon Sep 17 00:00:00 2001 From: Rodrigo de Lazcano Date: Fri, 16 Dec 2022 23:39:24 -0500 Subject: [PATCH] Update to refactored mujoco renderer in Gymnasium v0.27.0 (#60) * update mujoco renderer gymnasium v27 * remove rgb_array_list from metadata * skip render check pytest * pre-commit black --- gymnasium_robotics/envs/fetch_env.py | 29 +++++++++---- gymnasium_robotics/envs/hand_env.py | 34 +++++++++------- gymnasium_robotics/envs/robot_env.py | 61 ++++++++++------------------ tests/test_envs.py | 4 +- 4 files changed, 66 insertions(+), 62 deletions(-) diff --git a/gymnasium_robotics/envs/fetch_env.py b/gymnasium_robotics/envs/fetch_env.py index 3c4e3307..97b29650 100644 --- a/gymnasium_robotics/envs/fetch_env.py +++ b/gymnasium_robotics/envs/fetch_env.py @@ -5,6 +5,13 @@ from gymnasium_robotics.envs.robot_env import MujocoPyRobotEnv, MujocoRobotEnv from gymnasium_robotics.utils import rotations +DEFAULT_CAMERA_CONFIG = { + "distance": 2.5, + "azimuth": 132.0, + "elevation": -14.0, + "lookat": np.array([1.3, 0.75, 0.55]), +} + def goal_distance(goal_a, goal_b): assert goal_a.shape == goal_b.shape @@ -143,14 +150,6 @@ def get_gripper_xpos(self): raise NotImplementedError - def _viewer_setup(self): - lookat = self.get_gripper_xpos() - for idx, value in enumerate(lookat): - self.viewer.cam.lookat[idx] = value - self.viewer.cam.distance = 2.5 - self.viewer.cam.azimuth = 132.0 - self.viewer.cam.elevation = -14.0 - def _sample_goal(self): if self.has_object: goal = self.initial_gripper_xpos[:3] + self.np_random.uniform( @@ -238,6 +237,17 @@ def _render_callback(self): self.sim.model.site_pos[site_id] = self.goal - sites_offset[0] self.sim.forward() + def _viewer_setup(self): + lookat = self.get_gripper_xpos() + for idx, value in enumerate(lookat): + self.viewer.cam.lookat[idx] = value + 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) + def _reset_sim(self): self.sim.set_state(self.initial_state) @@ -279,6 +289,9 @@ def _env_setup(self, initial_qpos): class MujocoFetchEnv(get_base_fetch_env(MujocoRobotEnv)): + def __init__(self, default_camera_config: dict = DEFAULT_CAMERA_CONFIG, **kwargs): + super().__init__(default_camera_config=default_camera_config, **kwargs) + def _step_callback(self): if self.block_gripper: self._utils.set_joint_qpos( diff --git a/gymnasium_robotics/envs/hand_env.py b/gymnasium_robotics/envs/hand_env.py index ba529502..ce1241da 100644 --- a/gymnasium_robotics/envs/hand_env.py +++ b/gymnasium_robotics/envs/hand_env.py @@ -4,6 +4,13 @@ from gymnasium_robotics.envs.robot_env import MujocoPyRobotEnv, MujocoRobotEnv +DEFAULT_CAMERA_CONFIG = { + "distance": 0.5, + "azimuth": 55.0, + "elevation": -25.0, + "lookat": np.array([1, 0.96, 0.14]), +} + def get_base_hand_env( RobotEnvClass: Union[MujocoPyRobotEnv, MujocoRobotEnv] @@ -21,24 +28,19 @@ def __init__(self, relative_control, **kwargs): # RobotEnv methods # ---------------------------- - def _get_palm_xpos(self): - raise NotImplementedError def _set_action(self, action): assert action.shape == (20,) - def _viewer_setup(self): - lookat = self._get_palm_xpos() - for idx, value in enumerate(lookat): - self.viewer.cam.lookat[idx] = value - self.viewer.cam.distance = 0.5 - self.viewer.cam.azimuth = 55.0 - self.viewer.cam.elevation = -25.0 - return BaseHandEnv class MujocoHandEnv(get_base_hand_env(MujocoRobotEnv)): + def __init__( + self, default_camera_config: dict = DEFAULT_CAMERA_CONFIG, **kwargs + ) -> None: + super().__init__(default_camera_config=default_camera_config, **kwargs) + def _set_action(self, action): super()._set_action(action) ctrlrange = self.model.actuator_ctrlrange @@ -60,10 +62,6 @@ def _set_action(self, action): self.data.ctrl[:] = actuation_center + action * actuation_range self.data.ctrl[:] = np.clip(self.data.ctrl, ctrlrange[:, 0], ctrlrange[:, 1]) - def _get_palm_xpos(self): - body_id = self._model_names.body_name2id["robot0:palm"] - return self.data.xpos[body_id] - class MujocoPyHandEnv(get_base_hand_env(MujocoPyRobotEnv)): def _set_action(self, action): @@ -92,3 +90,11 @@ def _set_action(self, action): def _get_palm_xpos(self): body_id = self.sim.model.body_name2id("robot0:palm") return self.sim.data.body_xpos[body_id] + + def _viewer_setup(self): + lookat = self._get_palm_xpos() + for idx, value in enumerate(lookat): + self.viewer.cam.lookat[idx] = value + self.viewer.cam.distance = 0.5 + self.viewer.cam.azimuth = 55.0 + self.viewer.cam.elevation = -25.0 diff --git a/gymnasium_robotics/envs/robot_env.py b/gymnasium_robotics/envs/robot_env.py index 00465224..c706f6dc 100644 --- a/gymnasium_robotics/envs/robot_env.py +++ b/gymnasium_robotics/envs/robot_env.py @@ -2,7 +2,6 @@ import os from typing import Optional, Union -import gymnasium as gym import numpy as np from gymnasium import error, logger, spaces @@ -36,7 +35,6 @@ class BaseRobotEnv(GoalEnv): "render_modes": [ "human", "rgb_array", - "rgb_array_list", ], "render_fps": 25, } @@ -202,12 +200,6 @@ def _env_setup(self, initial_qpos): """ pass - def _viewer_setup(self): - """Initial configuration of the viewer. Can be used to set the camera position, - for example. - """ - pass - def _render_callback(self): """A custom callback that is called before rendering. Can be used to implement custom visualizations. @@ -222,7 +214,7 @@ def _step_callback(self): class MujocoRobotEnv(BaseRobotEnv): - def __init__(self, **kwargs): + def __init__(self, default_camera_config: Optional[dict] = None, **kwargs): if MUJOCO_IMPORT_ERROR is not None: raise error.DependencyNotInstalled( f"{MUJOCO_IMPORT_ERROR}. (HINT: you need to install mujoco)" @@ -233,6 +225,12 @@ def __init__(self, **kwargs): super().__init__(**kwargs) + 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 = self._mujoco.MjModel.from_xml_path(self.fullpath) self.data = self._mujoco.MjData(self.model) @@ -258,35 +256,11 @@ def _reset_sim(self): def render(self): self._render_callback() - if self.render_mode == "rgb_array": - self._get_viewer(self.render_mode).render() - 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 == "human": - self._get_viewer(self.render_mode).render() - - def _get_viewer( - self, mode - ) -> Union["gym.envs.mujoco.Viewer", "gym.envs.mujoco.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", - "rgb_array_list", - }: - from gymnasium.envs.mujoco.mujoco_rendering import ( - RenderContextOffscreen, - ) + return self.mujoco_renderer.render(self.render_mode) - self.viewer = RenderContextOffscreen(model=self.model, data=self.data) - self._viewer_setup() - self._viewers[mode] = self.viewer - return self.viewer + def close(self): + if self.mujoco_renderer is not None: + self.mujoco_renderer.close() @property def dt(self): @@ -334,7 +308,6 @@ def render(self): self._render_callback() if self.render_mode in { "rgb_array", - "rgb_array_list", }: self._get_viewer(self.render_mode).render(width, height) # window size used for old mujoco-py: @@ -346,6 +319,11 @@ def render(self): elif self.render_mode == "human": self._get_viewer(self.render_mode).render() + def close(self): + if self.viewer is not None: + self.viewer = None + self._viewers = {} + def _get_viewer( self, mode ) -> Union["mujoco_py.MjViewer", "mujoco_py.MjRenderContextOffscreen"]: @@ -356,7 +334,6 @@ def _get_viewer( elif mode in { "rgb_array", - "rgb_array_list", }: self.viewer = self._mujoco_py.MjRenderContextOffscreen(self.sim, -1) self._viewer_setup() @@ -369,3 +346,9 @@ def dt(self): def _mujoco_step(self, action): self.sim.step() + + def _viewer_setup(self): + """Initial configuration of the viewer. Can be used to set the camera position, + for example. + """ + pass diff --git a/tests/test_envs.py b/tests/test_envs.py index 5495977e..b569d15e 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -29,10 +29,12 @@ def test_env(spec): # Capture warnings env = spec.make(disable_env_checker=True).unwrapped + warnings.simplefilter("always") # Test if env adheres to Gym API with warnings.catch_warnings(record=True) as w: - check_env(env) + check_env(env, skip_render_check=True) + env.close() for warning in w: if warning.message.args[0] not in CHECK_ENV_IGNORE_WARNINGS: raise Error(f"Unexpected warning: {warning.message}")