reference/generated/torchrl.record.VideoRecorder

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

VideoRecorder

torchrl.record.VideoRecorder(logger: Logger | Service | None, tag: str | None, in_keys: Sequence[NestedKey] | None = None, skip: int | None = None, center_crop: int | None = None, make_grid: bool | None = None, out_keys: Sequence[NestedKey] | None = None, fps: int | None = None, max_frames: int | None = None, dump_on_done: bool = False, **kwargs) → None

Video Recorder transform.

Will record a series of observations from an environment and write them to a Logger object when needed.

Parameters:

• logger (Logger or Service) – a logger or logger-service owner where the video should be written. To save the video under a memmap tensor or an mp4 file, use the CSVLogger class.

• tag (str) – the video tag in the logger.

• in_keys (Sequence of NestedKey, optional) – keys to be read to produce the video. Default is "pixels".

• skip (int) – frame interval in the output video. Defaults to 1 for vector environments and standalone use, and 2 for a single parent environment.

• center_crop (int, optional) – value of square center crop.

• make_grid (bool, optional) – if True, a grid is created assuming that a tensor of shape [B x W x H x 3] is provided, with B being the batch size. Default is True if the transform has a parent environment, and False if not.

• out_keys (sequence of NestedKey, optional) – destination keys. Defaults to in_keys if not provided.

• fps (int, optional) – Frames per second of the output video. Defaults to the logger predefined fps, and overrides it if provided.

• max_frames (int, optional) – maximum number of frames held in the recorder buffer. Once the buffer is full, new frames are discarded until dump() empties it. Use this to bound memory when dumps are infrequent (or could be missed altogether). Unbounded by default.

• dump_on_done (bool, optional) – if True, the recorder calls dump() on its own whenever a step ends the episode, i.e. whenever all the "done" entries of the root tensordict are True (for a batched env, this means every sub-env is done at the same step). Defaults to False.

• **kwargs (Dict[str, Any], optional) – additional keyword arguments for log_video().

Examples

The following example shows how to save a rollout under a video. First a few imports:

>>> from torchrl.record import VideoRecorder
>>> from torchrl.record.loggers.csv import CSVLogger
>>> from torchrl.envs import TransformedEnv, DMControlEnv

The video format is chosen in the logger. Wandb and tensorboard will take care of that on their own, CSV accepts various video formats.

>>> logger = CSVLogger(exp_name="cheetah", log_dir="cheetah_videos", video_format="mp4")

Some envs (eg, Atari games) natively return images, some require the user to ask for them. Check GymEnv or DMControlEnv to see how to render images in these contexts.

>>> base_env = DMControlEnv("cheetah", "run", from_pixels=True)
>>> env = TransformedEnv(base_env, VideoRecorder(logger=logger, tag="run_video"))
>>> env.rollout(100)

All transforms have a dump function, mostly a no-op except for VideoRecorder, and Compose which will dispatch the dumps to all its members.

>>> env.transform.dump()

Note

When recording a batched env (SerialEnv or ParallelEnv), attach the recorder to the outer env, e.g. TransformedEnv(ParallelEnv(N, make_env), VideoRecorder(...)), so that the batch is tiled into a single grid video (make_grid=True). Recorders living inside the worker envs of a batched env are not reached by dump calls issued on the outer env or by collectors and evaluators, and would accumulate frames indefinitely.

The transform can also be used within a dataset to save the video collected. Unlike in the environment case, images will come in a batch. The skip argument will enable to save the images only at specific intervals.

>>> from torchrl.data.datasets import OpenXExperienceReplay
>>> from torchrl.envs import Compose
>>> from torchrl.record import VideoRecorder, CSVLogger
>>> # Create a logger that saves videos as mp4 using 24 frames per sec
>>> logger = CSVLogger("./dump", video_format="mp4", video_fps=24)
>>> # We use the VideoRecorder transform to save register the images coming from the batch.
>>> #  Setting the fps to 12 overrides the one set in the logger, not doing so keeps it unchanged.
>>> t = VideoRecorder(logger=logger, tag="pixels", in_keys=[("next", "observation", "image")], fps=12)
>>> # Each batch of data will have 10 consecutive videos of 200 frames each (maximum, since strict_length=False)
>>> dataset = OpenXExperienceReplay("cmu_stretch", batch_size=2000, slice_len=200,
...             download=True, strict_length=False,
...             transform=t)
>>> # Get a batch of data and visualize it
>>> for data in dataset:
...     t.dump()
...     break

Our video is available under ./cheetah_videos/cheetah/videos/run_video_0.mp4!