Spaces:
Sleeping
Sleeping
| """Utility functions to save rendering videos.""" | |
| import os | |
| from typing import Callable, Optional | |
| import gym | |
| from gym import logger | |
| try: | |
| from moviepy.video.io.ImageSequenceClip import ImageSequenceClip | |
| except ImportError: | |
| raise gym.error.DependencyNotInstalled( | |
| "MoviePy is not installed, run `pip install moviepy`" | |
| ) | |
| def capped_cubic_video_schedule(episode_id: int) -> bool: | |
| """The default episode trigger. | |
| This function will trigger recordings at the episode indices 0, 1, 4, 8, 27, ..., :math:`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 | |
| def save_video( | |
| frames: list, | |
| video_folder: str, | |
| episode_trigger: Callable[[int], bool] = None, | |
| step_trigger: Callable[[int], bool] = None, | |
| video_length: Optional[int] = None, | |
| name_prefix: str = "rl-video", | |
| episode_index: int = 0, | |
| step_starting_index: int = 0, | |
| **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. | |
| **kwargs: The kwargs that will be passed to moviepy's ImageSequenceClip. | |
| You need to specify either fps or duration. | |
| Example: | |
| >>> import gym | |
| >>> from gym.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): | |
| ... action = env.action_space.sample() | |
| ... _, _, done, _ = env.step(action) | |
| ... if done: | |
| ... save_video( | |
| ... env.render(), | |
| ... "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") | |
| 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") | |