Fix#3043 (#3044)

openai · Aug 24, 2022 · 44e9475 · 44e9475
1 parent 6e2e921
commit 44e9475
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 34 deletions.
diff --git 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
@@ -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()