Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Set offwidth and offheight properly before make MjrContext to support width > 640 or height > 480 in offscreen rendering #3044

Merged
merged 1 commit into from
Aug 24, 2022
Merged

Set offwidth and offheight properly before make MjrContext to support width > 640 or height > 480 in offscreen rendering #3044

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 13 additions & 17 deletions gym/envs/mujoco/mujoco_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
)
Expand Down Expand Up @@ -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,
):
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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):
Expand All @@ -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,
):
Expand Down Expand Up @@ -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":
Expand All @@ -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:
Expand All @@ -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']}"
Expand Down
27 changes: 10 additions & 17 deletions gym/envs/mujoco/mujoco_rendering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()

Expand Down