From e65d351213ed4f381d629012a33baf92470ac595 Mon Sep 17 00:00:00 2001 From: YouJiacheng <1503679330@qq.com> Date: Wed, 24 Aug 2022 23:58:19 +0800 Subject: [PATCH] Fix#3043 --- gym/envs/mujoco/mujoco_env.py | 30 +++++++++++++---------------- gym/envs/mujoco/mujoco_rendering.py | 27 ++++++++++---------------- 2 files changed, 23 insertions(+), 34 deletions(-) diff --git a/gym/envs/mujoco/mujoco_env.py b/gym/envs/mujoco/mujoco_env.py index 3a9abca35ba..1b7408d09cd 100644 --- a/gym/envs/mujoco/mujoco_env.py +++ b/gym/envs/mujoco/mujoco_env.py @@ -49,7 +49,9 @@ def __init__( if not path.exists(self.fullpath): raise OSError(f"File {self.fullpath} does not exist") - self._initialize_simulation() + self.width = width + self.height = height + self._initialize_simulation() # may use width and height self.init_qpos = self.data.qpos.ravel().copy() self.init_qvel = self.data.qvel.ravel().copy() @@ -76,8 +78,6 @@ def __init__( self.render_mode = render_mode render_frame = partial( self._render, - width=width, - height=height, camera_name=camera_name, camera_id=camera_id, ) @@ -126,8 +126,6 @@ def _step_mujoco_simulation(self, ctrl, n_frames): def _render( self, mode: str = "human", - width: int = DEFAULT_SIZE, - height: int = DEFAULT_SIZE, camera_id: Optional[int] = None, camera_name: Optional[str] = None, ): @@ -249,11 +247,10 @@ def _step_mujoco_simulation(self, ctrl, n_frames): def _render( self, mode: str = "human", - width: int = DEFAULT_SIZE, - height: int = DEFAULT_SIZE, camera_id: Optional[int] = None, camera_name: Optional[str] = None, ): + width, height = self.width, self.height assert mode in self.metadata["render_modes"] if mode in { "rgb_array", @@ -350,6 +347,9 @@ def __init__( def _initialize_simulation(self): self.model = mujoco.MjModel.from_xml_path(self.fullpath) + # MjrContext will copy model.vis.global_.off* to con.off* + self.model.vis.global_.offwidth = self.width + self.model.vis.global_.offheight = self.height self.data = mujoco.MjData(self.model) def _reset_simulation(self): @@ -376,8 +376,6 @@ def _step_mujoco_simulation(self, ctrl, n_frames): def _render( self, mode: str = "human", - width: int = DEFAULT_SIZE, - height: int = DEFAULT_SIZE, camera_id: Optional[int] = None, camera_name: Optional[str] = None, ): @@ -406,16 +404,16 @@ def _render( camera_name, ) - self._get_viewer(mode).render(width, height, camera_id=camera_id) + self._get_viewer(mode).render(camera_id=camera_id) if mode in {"rgb_array", "single_rgb_array"}: - data = self._get_viewer(mode).read_pixels(width, height, depth=False) + data = self._get_viewer(mode).read_pixels(depth=False) # original image is upside-down, so flip it return data[::-1, :, :] elif mode in {"depth_array", "single_depth_array"}: - self._get_viewer(mode).render(width, height) + self._get_viewer(mode).render() # Extract depth part of the read_pixels() tuple - data = self._get_viewer(mode).read_pixels(width, height, depth=True)[1] + data = self._get_viewer(mode).read_pixels(depth=True)[1] # original image is upside-down, so flip it return data[::-1, :] elif mode == "human": @@ -427,7 +425,7 @@ def close(self): super().close() def _get_viewer( - self, mode, width=DEFAULT_SIZE, height=DEFAULT_SIZE + self, mode ) -> Union["gym.envs.mujoco.Viewer", "gym.envs.mujoco.RenderContextOffscreen"]: self.viewer = self._viewers.get(mode) if self.viewer is None: @@ -443,9 +441,7 @@ def _get_viewer( }: from gym.envs.mujoco import RenderContextOffscreen - self.viewer = RenderContextOffscreen( - width, height, self.model, self.data - ) + self.viewer = RenderContextOffscreen(self.model, self.data) else: raise AttributeError( f"Unexpected mode: {mode}, expected modes: {self.metadata['render_modes']}" diff --git a/gym/envs/mujoco/mujoco_rendering.py b/gym/envs/mujoco/mujoco_rendering.py index 5521200480f..9102ab122fe 100644 --- a/gym/envs/mujoco/mujoco_rendering.py +++ b/gym/envs/mujoco/mujoco_rendering.py @@ -44,6 +44,8 @@ def __init__(self, model, data, offscreen=True): self.model = model self.data = data self.offscreen = offscreen + self.offwidth = model.vis.global_.offwidth + self.offheight = model.vis.global_.offheight max_geom = 1000 mujoco.mj_forward(self.model, self.data) @@ -70,22 +72,10 @@ def _set_mujoco_buffers(self): if self.con.currentBuffer != mujoco.mjtFramebuffer.mjFB_WINDOW: raise RuntimeError("Window rendering not supported") - def update_offscreen_size(self, width, height): - if width != self.con.offWidth or height != self.con.offHeight: - self.model.vis.global_.offwidth = width - self.model.vis.global_.offheight = height - self.con.free() - self._set_mujoco_buffers() - - def render(self, width, height, camera_id=None, segmentation=False): + def render(self, camera_id=None, segmentation=False): + width, height = self.offwidth, self.offheight rect = mujoco.MjrRect(left=0, bottom=0, width=width, height=height) - # Sometimes buffers are too small. - if width > self.con.offWidth or height > self.con.offHeight: - new_width = max(width, self.model.vis.global_.offwidth) - new_height = max(height, self.model.vis.global_.offheight) - self.update_offscreen_size(new_width, new_height) - if camera_id is not None: if camera_id == -1: self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE @@ -126,7 +116,8 @@ def render(self, width, height, camera_id=None, segmentation=False): self.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = 0 self.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = 0 - def read_pixels(self, width, height, depth=True, segmentation=False): + def read_pixels(self, depth=True, segmentation=False): + width, height = self.offwidth, self.offheight rect = mujoco.MjrRect(left=0, bottom=0, width=width, height=height) rgb_arr = np.zeros(3 * rect.width * rect.height, dtype=np.uint8) @@ -231,8 +222,10 @@ def close(self): class RenderContextOffscreen(RenderContext): """Offscreen rendering class with opengl context.""" - def __init__(self, width, height, model, data): - + def __init__(self, model, data): + # We must make GLContext before MjrContext + width = model.vis.global_.offwidth + height = model.vis.global_.offheight self._get_opengl_backend(width, height) self.opengl_context.make_current()