gaspersic/Gymnasium
1
0
Fork 0
mirror of https://github.com/Farama-Foundation/Gymnasium.git synced 2025-08-01 22:11:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix#3043 (#3044)

Browse Source
This commit is contained in:
YouJiacheng
2022-08-25 00:20:24 +08:00
committed by GitHub
parent 6e2e921b5f
commit 44e9475cda
2 changed files with 23 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
gym/envs/mujoco/mujoco_env.py
View File

@@ -49,7 +49,9 @@ class BaseMujocoEnv(gym.Env):
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 @@ class BaseMujocoEnv(gym.Env):
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 @@ class BaseMujocoEnv(gym.Env):
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 @@ class MuJocoPyEnv(BaseMujocoEnv):
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 @@ class MujocoEnv(BaseMujocoEnv):
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 @@ class MujocoEnv(BaseMujocoEnv):
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 @@ class MujocoEnv(BaseMujocoEnv):
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 @@ class MujocoEnv(BaseMujocoEnv):
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 @@ class MujocoEnv(BaseMujocoEnv):
}:
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']}"

27
gym/envs/mujoco/mujoco_rendering.py
View File

@@ -44,6 +44,8 @@ class RenderContext:
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 @@ class RenderContext:
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 @@ class RenderContext:
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 @@ class RenderContext:
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()