Source code for gymnasium.utils.save_video

"""Utility functions to save rendering videos."""
from __future__ import annotations

import os
from typing import Callable

import gymnasium as gym
from gymnasium import logger


try:
    from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
except ImportError as e:
    raise gym.error.DependencyNotInstalled(
        'moviepy is not installed, run `pip install "gymnasium[other]"`'
    ) from e


[docs] def capped_cubic_video_schedule(episode_id: int) -> bool: r"""The default episode trigger. This function will trigger recordings at the episode indices :math:`\{0, 1, 4, 8, 27, ..., k^3, ..., 729, 1000, 2000, 3000, ...\}` Args: episode_id: The episode number Returns: If to apply a video schedule number """ if episode_id < 1000: return int(round(episode_id ** (1.0 / 3))) ** 3 == episode_id else: return episode_id % 1000 == 0
[docs] def save_video( frames: list, video_folder: str, episode_trigger: Callable[[int], bool] = None, step_trigger: Callable[[int], bool] = None, video_length: int | None = None, name_prefix: str = "rl-video", episode_index: int = 0, step_starting_index: int = 0, save_logger: str | None = None, **kwargs, ): """Save videos from rendering frames. This function extract video from a list of render frame episodes. Args: frames (List[RenderFrame]): A list of frames to compose the video. video_folder (str): The folder where the recordings will be stored 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 it isn't specified, the entire episode is recorded. Otherwise, snippets of the specified length are captured. name_prefix (str): Will be prepended to the filename of the recordings. episode_index (int): The index of the current episode. step_starting_index (int): The step index of the first frame. save_logger: If to log the video saving progress, helpful for long videos that take a while, use "bar" to enable. **kwargs: The kwargs that will be passed to moviepy's ImageSequenceClip. You need to specify either fps or duration. Example: >>> import gymnasium as gym >>> from gymnasium.utils.save_video import save_video >>> env = gym.make("FrozenLake-v1", render_mode="rgb_array_list") >>> _ = env.reset() >>> step_starting_index = 0 >>> episode_index = 0 >>> for step_index in range(199): # doctest: +SKIP ... action = env.action_space.sample() ... _, _, terminated, truncated, _ = env.step(action) ... ... if terminated or truncated: ... save_video( ... frames=env.render(), ... video_folder="videos", ... fps=env.metadata["render_fps"], ... step_starting_index=step_starting_index, ... episode_index=episode_index ... ) ... step_starting_index = step_index + 1 ... episode_index += 1 ... env.reset() >>> env.close() """ if not isinstance(frames, list): logger.error(f"Expected a list of frames, got a {type(frames)} instead.") if episode_trigger is None and step_trigger is None: episode_trigger = capped_cubic_video_schedule video_folder = os.path.abspath(video_folder) os.makedirs(video_folder, exist_ok=True) path_prefix = f"{video_folder}/{name_prefix}" if episode_trigger is not None and episode_trigger(episode_index): clip = ImageSequenceClip(frames[:video_length], **kwargs) clip.write_videofile( f"{path_prefix}-episode-{episode_index}.mp4", logger=save_logger ) if step_trigger is not None: # skip the first frame since it comes from reset for step_index, frame_index in enumerate( range(1, len(frames)), start=step_starting_index ): if step_trigger(step_index): end_index = ( frame_index + video_length if video_length is not None else None ) clip = ImageSequenceClip(frames[frame_index:end_index], **kwargs) clip.write_videofile( f"{path_prefix}-step-{step_index}.mp4", logger=save_logger )