"""File for rendering and recording of vector-based environments."""
from __future__ import annotations
import gc
import os
from collections.abc import Callable, Sequence
from copy import deepcopy
from typing import TYPE_CHECKING, Any
import numpy as np
import gymnasium as gym
from gymnasium import error, logger
from gymnasium.core import ActType, ObsType, RenderFrame
from gymnasium.error import DependencyNotInstalled
from gymnasium.logger import warn
from gymnasium.vector import VectorEnv, VectorWrapper
from gymnasium.vector.vector_env import ArrayType
if TYPE_CHECKING:
import pygame
[docs]
class HumanRendering(
VectorWrapper[ObsType, ActType, ObsType, ActType, ArrayType],
gym.utils.RecordConstructorArgs,
):
"""Adds support for Human-based Rendering for Vector-based environments."""
ACCEPTED_RENDER_MODES = [
"rgb_array",
"rgb_array_list",
"depth_array",
"depth_array_list",
]
screen_size: tuple[int, int] | None
scaled_subenv_size: tuple[int, int] | None
num_rows: int | None
num_cols: int | None
window: pygame.Surface | None
clock: pygame.time.Clock | None
metadata: dict[str, Any]
def __init__(
self,
env: VectorEnv[ObsType, ActType, ArrayType],
screen_size: tuple[int, int] | None = None,
):
"""Constructor for Human Rendering of Vector-based environments.
Args:
env: The vector environment
screen_size: The rendering screen size otherwise the environment sub-env render size is used
"""
VectorWrapper.__init__(self, env)
gym.utils.RecordConstructorArgs.__init__(self, screen_size=screen_size)
self.screen_size = screen_size
self.scaled_subenv_size, self.num_rows, self.num_cols = None, None, None
self.window = None # Has to be initialized before asserts, as self.window is used in auto close
self.clock = None
if self.env.render_mode not in self.ACCEPTED_RENDER_MODES:
raise ValueError(
f"Expected env.render_mode to be one of {self.ACCEPTED_RENDER_MODES} but got '{self.env.render_mode}'"
)
if "render_fps" not in self.env.metadata:
raise ValueError(
"The base environment must specify 'render_fps' to be used with the HumanRendering wrapper"
)
if "human" not in self.metadata["render_modes"]:
self.metadata = deepcopy(self.env.metadata)
self.metadata["render_modes"].append("human")
@property
def render_mode(self) -> str:
"""Always returns ``'human'``."""
return "human"
def step(
self, actions: ActType
) -> tuple[ObsType, ArrayType, ArrayType, ArrayType, dict[str, Any]]:
"""Perform a step in the base environment and render a frame to the screen."""
result = super().step(actions)
self._render_frame()
return result
def reset(
self,
*,
seed: int | None = None,
options: dict[str, Any] | None = None,
) -> tuple[ObsType, dict[str, Any]]:
"""Reset the base environment and render a frame to the screen."""
result = super().reset(seed=seed, options=options)
self._render_frame()
return result
def _render_frame(self):
"""Fetch the last frame from the base environment and render it to the screen."""
try:
import pygame
except ImportError as e:
raise DependencyNotInstalled(
"pygame is not installed, run `pip install gymnasium[classic-control]`"
) from e
assert self.env.render_mode is not None
if self.env.render_mode.endswith("_last"):
subenv_renders = self.env.render()
assert isinstance(subenv_renders, list)
subenv_renders = subenv_renders[-1]
else:
subenv_renders = self.env.render()
assert subenv_renders is not None
assert len(subenv_renders) == self.num_envs
assert all(isinstance(render, np.ndarray) for render in subenv_renders), (
f"Expected `env.render()` to return a numpy array, actually returned {[type(render) for render in subenv_renders]}"
)
subenv_renders = np.array(subenv_renders, dtype=np.uint8)
subenv_renders = np.transpose(subenv_renders, axes=(0, 2, 1, 3))
# shape = (num envs, width, height, channels)
if self.screen_size is None:
self.screen_size = subenv_renders.shape[1:3]
if self.scaled_subenv_size is None:
subenv_size = subenv_renders.shape[1:3]
width_ratio = subenv_size[0] / self.screen_size[0]
height_ratio = subenv_size[1] / self.screen_size[1]
num_rows, num_cols = 1, 1
while num_rows * num_cols < self.num_envs:
row_ratio = num_rows * height_ratio
col_ratio = num_cols * width_ratio
if row_ratio == col_ratio:
num_rows, num_cols = num_rows + 1, num_cols + 1
elif row_ratio > col_ratio:
num_cols += 1
else:
num_rows += 1
scaling_factor = min(
self.screen_size[0] / (num_cols * subenv_size[0]),
self.screen_size[1] / (num_rows * subenv_size[1]),
)
assert (
num_cols * subenv_size[0] * scaling_factor == self.screen_size[0]
) or (num_rows * subenv_size[1] * scaling_factor == self.screen_size[1])
self.num_rows: int = num_rows
self.num_cols: int = num_cols
self.scaled_subenv_size = (
int(subenv_size[0] * scaling_factor),
int(subenv_size[1] * scaling_factor),
)
assert self.num_rows * self.num_cols >= self.num_envs
assert self.scaled_subenv_size[0] * self.num_cols <= self.screen_size[0]
assert self.scaled_subenv_size[1] * self.num_rows <= self.screen_size[1]
# print(f'{self.num_envs=}, {self.num_rows=}, {self.num_cols=}, {self.screen_size=}, {self.scaled_subenv_size=}')
try:
import cv2
except ImportError as e:
raise DependencyNotInstalled(
'opencv (cv2) is not installed, run `pip install "gymnasium[other]"`'
) from e
merged_rgb_array = np.zeros(self.screen_size + (3,), dtype=np.uint8)
assert self.num_cols is not None
assert self.num_rows is not None
cols, rows = np.meshgrid(np.arange(self.num_cols), np.arange(self.num_rows))
for i, col, row in zip(range(self.num_envs), cols.flatten(), rows.flatten()): # noqa: B905
scaled_render = cv2.resize(subenv_renders[i], self.scaled_subenv_size[::-1])
x = col * self.scaled_subenv_size[0]
y = row * self.scaled_subenv_size[1]
merged_rgb_array[
x : x + self.scaled_subenv_size[0],
y : y + self.scaled_subenv_size[1],
] = scaled_render
if self.window is None:
pygame.display.init()
self.window = pygame.display.set_mode(self.screen_size)
if self.clock is None:
self.clock = pygame.time.Clock()
surf = pygame.surfarray.make_surface(merged_rgb_array)
self.window.blit(surf, (0, 0))
pygame.event.pump()
self.clock.tick(self.metadata["render_fps"])
pygame.display.flip()
def close(self, **kwargs: Any) -> None:
"""Close the rendering window."""
if self.window is not None:
import pygame
pygame.display.quit()
pygame.quit()
super().close(**kwargs)
[docs]
class RecordVideo(
VectorWrapper[ObsType, ActType, ObsType, ActType, np.ndarray],
gym.utils.RecordConstructorArgs,
):
"""Adds support for video recording for Vector-based environments.
The class is the same as gymnasium.wrappers.rendering.RecordVideo, but
expects multiple frames when rendering the environment (one for each
environment of the VectorEnv). Frames are concatenated into one frame such
that its aspect ratio is as close as possible to the desired one.
Examples - Run 5 environments for 200 timesteps, and save the video every 5 episodes:
>>> import os
>>> import gymnasium as gym
>>> from gymnasium.wrappers.vector import RecordVideo
>>> envs = gym.make_vec("CartPole-v1", num_envs=5, render_mode="rgb_array")
>>> envs = RecordVideo(
... envs,
... video_folder="save_videos_5envs",
... video_aspect_ratio=(1,1),
... episode_trigger=lambda t: t % 5 == 0,
... )
>>> _ = envs.reset(seed=123)
>>> _ = envs.action_space.seed(123)
>>> for i in range(200):
... actions = envs.action_space.sample()
... _ = envs.step(actions)
>>> envs.close()
>>> len(os.listdir("save_videos_5envs"))
2
"""
autoreset_mode: gym.vector.AutoresetMode
has_autoreset: bool
episode_trigger: Callable[[int], bool] | None
step_trigger: Callable[[int], bool] | None
gc_trigger: Callable[[int], bool] | None
disable_logger: bool
record_first_only: bool
video_aspect_ratio: tuple[int, int]
frame_cols: int
frame_rows: int
video_folder: str
frames_per_sec: int
name_prefix: str
_video_name: str | None
video_length: float
recording: bool
recorded_frames: list[np.ndarray]
render_history: list[np.ndarray]
step_id: int
episode_id: int
def __init__(
self,
env: VectorEnv[ObsType, ActType, np.ndarray],
video_folder: str,
video_aspect_ratio: tuple[int, int] = (1, 1),
record_first_only: bool = False,
episode_trigger: Callable[[int], bool] | None = None,
step_trigger: Callable[[int], bool] | None = None,
video_length: int = 0,
name_prefix: str = "rl-video",
fps: int | None = None,
disable_logger: bool = True,
gc_trigger: Callable[[int], bool] | None = lambda episode: True,
):
"""Wrapper records videos of environment rollouts.
Args:
env: The environment that will be wrapped
video_folder (str): The folder where the recordings will be stored
video_aspect_ratio (tuple): the desired aspect ratio of the video concatenating all environments. For example, (1, 1) means an
aspect ratio of 1:1, while (16, 9) means 16:9.
record_first_only (bool): if True, only the first environment is recorded.
episode_trigger: Function that accepts an integer and returns ``True`` iff a recording should be started at this episode
step_trigger: Function that accepts an integer and returns ``True`` iff a recording should be started at this step
video_length (int): The length of recorded episodes. If 0, entire episodes are recorded.
Otherwise, snippets of the specified length are captured
name_prefix (str): Will be prepended to the filename of the recordings
fps (int): The frame per second in the video. Provides a custom video fps for environment, if ``None`` then
the environment metadata ``render_fps`` key is used if it exists, otherwise a default value of 30 is used.
disable_logger (bool): Whether to disable moviepy logger or not, default it is disabled
gc_trigger: Function that accepts an integer and returns ``True`` iff garbage collection should be performed after this episode
Note:
For vector environments that use same-step autoreset (see https://farama.org/Vector-Autoreset-Mode for more details)
then the final frame of the episode will not be included in the video.
"""
VectorWrapper.__init__(self, env)
gym.utils.RecordConstructorArgs.__init__(
self,
video_folder=video_folder,
episode_trigger=episode_trigger,
step_trigger=step_trigger,
video_length=video_length,
name_prefix=name_prefix,
disable_logger=disable_logger,
)
if env.render_mode in {None, "human", "ansi"}:
raise ValueError(
f"Render mode is {env.render_mode}, which is incompatible with RecordVideo. "
"Initialize your environment with a render_mode that returns an image, such as rgb_array."
)
if episode_trigger is None and step_trigger is None:
from gymnasium.utils.save_video import capped_cubic_video_schedule
episode_trigger = capped_cubic_video_schedule
autoreset_mode = env.metadata.get("autoreset_mode", None)
if autoreset_mode is not None:
self.autoreset_mode = autoreset_mode
else:
warn(
f"{env} metadata doesn't specify its autoreset mode ({env.metadata!r}), therefore, defaulting to next step."
)
self.autoreset_mode = gym.vector.AutoresetMode.NEXT_STEP
if self.autoreset_mode == gym.vector.AutoresetMode.SAME_STEP:
logger.warn(
"Vector environment's autoreset mode is same-step (https://farama.org/Vector-Autoreset-Mode). Recorded episodes will not contain the last frame of the episode."
)
self.has_autoreset = False
self.episode_trigger = episode_trigger
self.step_trigger = step_trigger
self.disable_logger = disable_logger
self.gc_trigger = gc_trigger
self.record_first_only = record_first_only
self.video_aspect_ratio = video_aspect_ratio
self.frame_cols = -1
self.frame_rows = -1
self.video_folder = os.path.abspath(video_folder)
if os.path.isdir(self.video_folder):
logger.warn(
f"Overwriting existing videos at {self.video_folder} folder "
f"(try specifying a different `video_folder` for the `RecordVideo` wrapper if this is not desired)"
)
os.makedirs(self.video_folder, exist_ok=True)
if fps is None:
fps = self.metadata.get("render_fps", 30)
self.frames_per_sec = fps
self.name_prefix = name_prefix
self._video_name = None
self.video_length = video_length if video_length != 0 else float("inf")
self.recording = False
self.recorded_frames = []
self.render_history = []
self.step_id = -1
self.episode_id = -1
try:
import moviepy # noqa: F401
except ImportError as e:
raise error.DependencyNotInstalled(
'MoviePy is not installed, run `pip install "gymnasium[other]"`'
) from e
def _get_concat_frame_shape(self, n_frames, h, w):
"""Finds the right shape to concatenate frames from all environments into one frame."""
target_video_aspect_ratio = (
self.video_aspect_ratio[0] / self.video_aspect_ratio[1]
)
best_rows, best_cols = 1, n_frames
min_diff = np.inf
for rows in range(1, int(n_frames**0.5) + 1):
if n_frames % rows == 0:
cols = n_frames // rows
total_height = rows * h
total_width = cols * w
aspect = total_width / total_height
diff = abs(aspect - target_video_aspect_ratio)
if diff < min_diff:
min_diff = diff
best_rows, best_cols = rows, cols
self.frame_rows = best_rows
self.frame_cols = best_cols
def _concat_frames(self, frames: np.typing.ArrayLike) -> np.ndarray:
"""Concatenates a list of frames into one large frame."""
frames = np.array(frames)
n_frames, h, w, c = frames.shape
grid = np.zeros(
(self.frame_rows * h, self.frame_cols * w, c), dtype=frames.dtype
)
for idx in range(n_frames):
r = idx // self.frame_cols
c_ = idx % self.frame_cols
grid[r * h : (r + 1) * h, c_ * w : (c_ + 1) * w] = frames[idx]
return grid
def _capture_frame(self) -> None:
assert self.recording, "Cannot capture a frame, recording wasn't started."
envs_frame = self.env.render()
assert isinstance(envs_frame, Sequence), type(envs_frame)
assert len(envs_frame) == self.num_envs
if self.record_first_only:
envs_frame = [envs_frame[0]]
if self.frame_cols == -1 or self.frame_rows == -1:
n_frames = len(envs_frame)
h, w, c = envs_frame[0].shape
self._get_concat_frame_shape(n_frames, h, w)
concatenated_envs_frame = self._concat_frames(envs_frame)
self.recorded_frames.append(concatenated_envs_frame)
def reset(
self, *, seed: int | None = None, options: dict[str, Any] | None = None
) -> tuple[ObsType, dict[str, Any]]:
"""Reset the environment and eventually starts a new recording."""
if options is None or "reset_mask" not in options or options["reset_mask"][0]:
self.episode_id += 1
if self.recording and self.video_length == float("inf"):
self.stop_recording()
if self.episode_trigger and self.episode_trigger(self.episode_id):
self.start_recording(f"{self.name_prefix}-episode-{self.episode_id}")
obs, info = super().reset(seed=seed, options=options)
if self.recording:
self._capture_frame()
if len(self.recorded_frames) > self.video_length:
self.stop_recording()
self.has_autoreset = False
return obs, info
def step(
self, actions: ActType
) -> tuple[ObsType, np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
"""Steps through the environment using action, recording observations if :attr:`self.recording`."""
obs, rewards, terminations, truncations, info = self.env.step(actions)
self.step_id += 1
if self.autoreset_mode == gym.vector.AutoresetMode.NEXT_STEP:
if self.has_autoreset:
self.episode_id += 1
if self.recording and self.video_length == float("inf"):
self.stop_recording()
if self.episode_trigger and self.episode_trigger(self.episode_id):
self.start_recording(
f"{self.name_prefix}-episode-{self.episode_id}"
)
self.has_autoreset = terminations[0] or truncations[0]
elif self.autoreset_mode == gym.vector.AutoresetMode.SAME_STEP and (
terminations[0] or truncations[0]
):
self.episode_id += 1
if self.recording and self.video_length == float("inf"):
self.stop_recording()
if self.episode_trigger and self.episode_trigger(self.episode_id):
self.start_recording(f"{self.name_prefix}-episode-{self.episode_id}")
if self.step_trigger and self.step_trigger(self.step_id):
self.start_recording(f"{self.name_prefix}-step-{self.step_id}")
if self.recording:
self._capture_frame()
if len(self.recorded_frames) > self.video_length:
self.stop_recording()
return obs, rewards, terminations, truncations, info
def render(self) -> RenderFrame | list[RenderFrame]:
"""Compute the render frames as specified by render_mode attribute during initialization of the environment."""
render_out = super().render()
if self.recording and isinstance(render_out, list):
self.recorded_frames += render_out
if len(self.render_history) > 0:
tmp_history = self.render_history
self.render_history = []
return tmp_history + render_out
else:
return render_out
def close(self, **kwargs: Any) -> None:
"""Closes the wrapper then the video recorder."""
super().close(**kwargs)
if self.recording:
self.stop_recording()
def start_recording(self, video_name: str) -> None:
"""Start a new recording. If it is already recording, stops the current recording before starting the new one."""
if self.recording:
self.stop_recording()
self.recording = True
self._video_name = video_name
def stop_recording(self) -> None:
"""Stop current recording and saves the video."""
assert self.recording, "stop_recording was called, but no recording was started"
if len(self.recorded_frames) == 0:
logger.warn("Ignored saving a video as there were zero frames to save.")
else:
try:
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
except ImportError as e:
raise error.DependencyNotInstalled(
'MoviePy is not installed, run `pip install "gymnasium[other]"`'
) from e
clip = ImageSequenceClip(self.recorded_frames, fps=self.frames_per_sec)
moviepy_logger = None if self.disable_logger else "bar"
path = os.path.join(self.video_folder, f"{self._video_name}.mp4")
clip.write_videofile(path, logger=moviepy_logger)
del clip
del self.recorded_frames
self.recorded_frames = []
self.recording = False
self._video_name = None
if self.gc_trigger and self.gc_trigger(self.episode_id):
gc.collect()
def __del__(self) -> None:
"""Warn the user in case last video wasn't saved."""
# `__init__` can raise before `recorded_frames` is assigned, and `__del__`
# still runs on the half-built wrapper. Reading the attribute directly
# raised an AttributeError that Python reports as "Exception ignored in",
# on top of the error that actually stopped construction.
if getattr(self, "recorded_frames", None):
logger.warn("Unable to save last video! Did you call close()?")
