in mujoco_py/mjrenderpool.py [0:0]
def render(self, width, height, states=None, camera_name=None,
depth=False, randomize=False, copy=True):
"""
Renders the simulations in batch. If no states are provided,
the max_batch_size will be used.
Args:
- width (int): width of image to render.
- height (int): height of image to render.
- states (list): list of MjSimStates; updates the states before
rendering. Batch size will be number of states supplied.
- camera_name (str): name of camera to render from.
- depth (bool): if True, also return depth.
- randomize (bool): calls modder.rand_all() before rendering.
- copy (bool): return a copy rather than a reference
Returns:
- rgbs: NxHxWx3 numpy array of N images in batch of width W
and height H.
- depth: NxHxW numpy array of N images in batch of width W
and height H. Only returned if depth=True.
"""
if self._closed:
raise RuntimeError("The pool has been closed.")
if (width * height) > self._max_image_size:
raise ValueError(
"Requested image larger than maximum image size. Create "
"a new RenderPool with a larger maximum image size.")
if states is None:
batch_size = self._max_batch_size
states = [None] * batch_size
else:
batch_size = len(states)
if batch_size > self._max_batch_size:
raise ValueError(
"Requested batch size larger than max batch size. Create "
"a new RenderPool with a larger max batch size.")
self.pool.starmap(
MjRenderPool._worker_render,
[(i, state, width, height, camera_name, randomize)
for i, state in enumerate(states)])
rgbs = self._shared_rgbs_array[:width * height * 3 * batch_size]
rgbs = rgbs.reshape(batch_size, height, width, 3)
if copy:
rgbs = rgbs.copy()
if depth:
depths = self._shared_depths_array[:width * height * batch_size]
depths = depths.reshape(batch_size, height, width).copy()
if copy:
depths = depths.copy()
return rgbs, depths
else:
return rgbs