in gym_pusht/envs/pusht.py [0:0]
def _get_img(self, screen, width, height, render_action=False):
img = np.transpose(np.array(pygame.surfarray.pixels3d(screen)), axes=(1, 0, 2))
img = cv2.resize(img, (width, height))
render_size = min(width, height)
if render_action and self._last_action is not None:
action = np.array(self._last_action)
coord = (action / 512 * [height, width]).astype(np.int32)
marker_size = int(8 / 96 * render_size)
thickness = int(1 / 96 * render_size)
cv2.drawMarker(
img,
coord,
color=(255, 0, 0),
markerType=cv2.MARKER_CROSS,
markerSize=marker_size,
thickness=thickness,
)
return img