Add Renderer class to Python native bindings.
PiperOrigin-RevId: 492994751 Change-Id: Iebbe5980668a15b1cd62c0076cc131255462b2eb
This commit is contained in:
Kyle Bayes
committed by
Copybara-Service
Copybara-Service
parent
341578450f
commit
ef695bb8a8
@@ -0,0 +1,230 @@
|
||||
# Copyright 2022 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Defines a renderer class for the MuJoCo Python native bindings."""
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
from mujoco import _enums
|
||||
from mujoco import _functions
|
||||
from mujoco import _render
|
||||
from mujoco import _structs
|
||||
from mujoco import gl_context
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Renderer:
|
||||
"""Renders MuJoCo scenes."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: _structs.MjModel,
|
||||
height: int = 240,
|
||||
width: int = 320,
|
||||
max_geom: int = 10000
|
||||
) -> None:
|
||||
"""Initializes a new `Renderer`.
|
||||
|
||||
Args:
|
||||
model: an MjModel instance.
|
||||
height: image height in pixels.
|
||||
width: image width in pixels.
|
||||
max_geom: Optional integer specifying the maximum number of geoms that can
|
||||
be rendered in the same scene. If None this will be chosen automatically
|
||||
based on the estimated maximum number of renderable geoms in the model.
|
||||
Raises:
|
||||
ValueError: If `camera_id` is outside the valid range, or if `width` or
|
||||
`height` exceed the dimensions of MuJoCo's offscreen framebuffer.
|
||||
"""
|
||||
buffer_width = model.vis.global_.offwidth
|
||||
buffer_height = model.vis.global_.offheight
|
||||
if width > buffer_width:
|
||||
raise ValueError(f"""
|
||||
Image width {width} > framebuffer width {buffer_width}. Either reduce the image
|
||||
width or specify a larger offscreen framebuffer in the model XML using the
|
||||
clause:
|
||||
<visual>
|
||||
<global offwidth="my_width"/>
|
||||
</visual>""".lstrip())
|
||||
|
||||
if height > buffer_height:
|
||||
raise ValueError(f"""
|
||||
Image height {height} > framebuffer height {buffer_height}. Either reduce the
|
||||
image height or specify a larger offscreen framebuffer in the model XML using
|
||||
the clause:
|
||||
<visual>
|
||||
<global offheight="my_height"/>
|
||||
</visual>""".lstrip())
|
||||
|
||||
self._width = width
|
||||
self._height = height
|
||||
self._model = model
|
||||
|
||||
self._scene = _structs.MjvScene(model=model, maxgeom=max_geom)
|
||||
self._scene_option = _structs.MjvOption()
|
||||
|
||||
self._rect = _render.MjrRect(0, 0, self._width, self._height)
|
||||
|
||||
# Internal buffers.
|
||||
self._rgb_buffer = np.empty((self._height, self._width, 3), dtype=np.uint8)
|
||||
self._depth_buffer = np.empty((self._height, self._width), dtype=np.float32)
|
||||
|
||||
# Create render contexts.
|
||||
self._gl_context = gl_context.GLContext(width, height)
|
||||
self._gl_context.make_current()
|
||||
self._mjr_context = _render.MjrContext(
|
||||
model, _enums.mjtFontScale.mjFONTSCALE_150
|
||||
)
|
||||
_render.mjr_setBuffer(
|
||||
_enums.mjtFramebuffer.mjFB_OFFSCREEN, self._mjr_context
|
||||
)
|
||||
|
||||
# Default render flags.
|
||||
self._depth_rendering = False
|
||||
self._segmentation_rendering = False
|
||||
|
||||
@property
|
||||
def model(self):
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def scene(self) -> _structs.MjvScene:
|
||||
return self._scene
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self._height
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self._width
|
||||
|
||||
def enable_depth_rendering(self):
|
||||
self._segmentation_rendering = False
|
||||
self._depth_rendering = True
|
||||
|
||||
def disable_depth_rendering(self):
|
||||
self._depth_rendering = False
|
||||
|
||||
def enable_segmentation_rendering(self):
|
||||
self._segmentation_rendering = True
|
||||
self._depth_rendering = False
|
||||
|
||||
def disable_segmentation_rendering(self):
|
||||
self._segmentation_rendering = False
|
||||
|
||||
def render(self) -> np.ndarray:
|
||||
"""Renders the scene as a numpy array of pixel values.
|
||||
|
||||
Returns:
|
||||
A numpy array of pixels with dimensions (H, W, 3). The array will be
|
||||
mutated by future calls to `render`.
|
||||
"""
|
||||
original_flags = self._scene.flags.copy()
|
||||
|
||||
if self._segmentation_rendering:
|
||||
self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True
|
||||
self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True
|
||||
|
||||
self._gl_context.make_current()
|
||||
|
||||
# Render scene and read contents of RGB and depth buffers.
|
||||
_render.mjr_render(self._rect, self._scene, self._mjr_context)
|
||||
_render.mjr_readPixels(self._rgb_buffer, self._depth_buffer, self._rect,
|
||||
self._mjr_context)
|
||||
|
||||
if self._depth_rendering:
|
||||
# Get the distances to the near and far clipping planes.
|
||||
extent = self._model.stat.extent
|
||||
near = self._model.vis.map.znear * extent
|
||||
far = self._model.vis.map.zfar * extent
|
||||
|
||||
# Convert from [0 1] to depth in meters, see links below:
|
||||
# http://stackoverflow.com/a/6657284/1461210
|
||||
# https://www.khronos.org/opengl/wiki/Depth_Buffer_Precision
|
||||
pixels = near / (1 - self._depth_buffer * (1 - near / far))
|
||||
|
||||
elif self._segmentation_rendering:
|
||||
# Convert 3-channel uint8 to 1-channel uint32.
|
||||
image3 = self._rgb_buffer.astype(np.uint32)
|
||||
segimage = (image3[:, :, 0] +
|
||||
image3[:, :, 1] * (2**8) +
|
||||
image3[:, :, 2] * (2**16))
|
||||
# Remap segid to 2-channel (object ID, object type) pair.
|
||||
# Seg ID 0 is background -- will be remapped to (-1, -1).
|
||||
ngeoms = self._scene.ngeom
|
||||
segid2output = np.full((ngeoms + 1, 2), fill_value=-1,
|
||||
dtype=np.int32) # Seg id cannot be > ngeom + 1.
|
||||
visible_geoms = [g for g in self._scene.geoms[:ngeoms] if g.segid != -1]
|
||||
visible_segids = np.array([g.segid + 1 for g in visible_geoms], np.int32)
|
||||
visible_objid = np.array([g.objid for g in visible_geoms], np.int32)
|
||||
visible_objtype = np.array([g.objtype for g in visible_geoms], np.int32)
|
||||
segid2output[visible_segids, 0] = visible_objid
|
||||
segid2output[visible_segids, 1] = visible_objtype
|
||||
pixels = segid2output[segimage]
|
||||
|
||||
# Reset scene flags.
|
||||
np.copyto(self._scene.flags, original_flags)
|
||||
else:
|
||||
pixels = self._rgb_buffer
|
||||
return np.flipud(pixels)
|
||||
|
||||
def update_scene(
|
||||
self,
|
||||
data: _structs.MjData,
|
||||
camera: Union[int, str, _structs.MjvCamera] = -1,
|
||||
scene_option: Optional[_structs.MjvOption] = None
|
||||
):
|
||||
"""Updates geometry used for rendering.
|
||||
|
||||
Args:
|
||||
data: An instance of `MjData`.
|
||||
camera: An instance of `MjvCamera`, a string or an integer
|
||||
scene_option: A custom `MjvOption` instance to use to render
|
||||
the scene instead of the default.
|
||||
"""
|
||||
if not isinstance(camera, _structs.MjvCamera):
|
||||
camera_id = camera
|
||||
if isinstance(camera_id, str):
|
||||
camera_id = _functions.mj_name2id(self._model,
|
||||
_enums.mjtObj.mjOBJ_CAMERA, camera_id)
|
||||
if camera_id < -1:
|
||||
raise ValueError('camera_id cannot be smaller than -1.')
|
||||
if camera_id >= self._model.ncam:
|
||||
raise ValueError(
|
||||
f'model has {self._model.ncam} fixed cameras. '
|
||||
f'camera_id={camera_id} is invalid.'
|
||||
)
|
||||
|
||||
# Render camera.
|
||||
camera = _structs.MjvCamera()
|
||||
camera.fixedcamid = camera_id
|
||||
|
||||
# Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
|
||||
# camera explicitly defined in the model.
|
||||
if camera_id == -1:
|
||||
camera.type = _enums.mjtCamera.mjCAMERA_FREE
|
||||
_functions.mjv_defaultFreeCamera(self._model, camera)
|
||||
else:
|
||||
camera.type = _enums.mjtCamera.mjCAMERA_FIXED
|
||||
|
||||
scene_option = scene_option or self._scene_option
|
||||
_functions.mjv_updateScene(
|
||||
self._model,
|
||||
data,
|
||||
scene_option,
|
||||
None,
|
||||
camera, _enums.mjtCatBit.mjCAT_ALL,
|
||||
self._scene,
|
||||
)
|
||||
Reference in New Issue
Block a user