def render(self, mode: str = "rgb", *args, **kwargs) -> np.ndarray:
assert mode in ["rgb",
"depth"], "only rgb and depth rendering is implemented"
if mode == "rgb":
return tile_images(self.env.current_frames)
elif mode == "depth":
return tile_images(self.env.current_depths)
def render(
self, mode: str = "human", *args, **kwargs
) -> Union[np.ndarray, None, List[np.ndarray]]:
"""Render observations from all Tasks in a tiled image or list of
images."""
images = self.command(
commands=RENDER_COMMAND,
data_list=[(args, {"mode": "rgb", **kwargs})] * self.num_unpaused_tasks,
)
if mode == "raw_rgb_list":
return images
tile = tile_images(images)
if mode == "human":
import cv2
cv2.imshow("vectask", tile[:, :, ::-1])
cv2.waitKey(1)
return None
elif mode == "rgb_array":
return tile
else:
raise NotImplementedError
def make_vid(self, images):
max_length = max([len(ep) for ep in images])
if max_length == 0:
return None
valid_im = None
for ep in images:
if len(ep) > 0:
valid_im = ep[0]
break
frames = []
for it in range(max_length):
current_images = []
for ep in images:
if it < len(ep):
current_images.append(ep[it])
else:
if it == 0:
current_images.append(np.zeros_like(valid_im))
else:
gray = ep[-1].copy()
gray[:, :, 0] = gray[:, :, 2] = gray[:, :, 1]
current_images.append(gray)
frames.append(tile_images(current_images))
return process_video(
frames, self.max_clip_length, self.max_video_length, fps=self.fps
)
def render(
self, mode: str = "human", *args, **kwargs
) -> Union[np.ndarray, None, List[np.ndarray]]:
"""Render observations from all Tasks in a tiled image or a list of
images."""
images = [
g.send((RENDER_COMMAND, (args, {"mode": "rgb", **kwargs})))
for g in self._vector_task_generators
]
if mode == "raw_rgb_list":
return images
for index, _ in reversed(self._paused):
images.insert(index, np.zeros_like(images[0]))
tile = tile_images(images)
if mode == "human":
import cv2
cv2.imshow("vectask", tile[:, :, ::-1])
cv2.waitKey(1)
return None
elif mode == "rgb_array":
return tile
else:
raise NotImplementedError
