Spaces:
Running
Running
File size: 4,191 Bytes
375a1cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
"""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")
|