# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import importlib.util
import math
from collections.abc import Callable, Sequence
from copy import copy
from typing import Any
import numpy as np
import torch
from tensordict import NonTensorData, TensorDictBase
from tensordict.utils import NestedKey
from torchrl._utils import _can_be_pickled, _ends_with, logger as torchrl_logger
from torchrl.data.tensor_specs import NonTensor, TensorSpec, Unbounded
from torchrl.data.utils import CloudpickleWrapper
from torchrl.envs import EnvBase
from torchrl.envs.transforms import ObservationTransform, Transform
from torchrl.record.loggers import Logger
from torchrl.services.base import Service
_has_tv = importlib.util.find_spec("torchvision", None) is not None
_has_matplotlib = importlib.util.find_spec("matplotlib", None) is not None
def _make_video_grid(frames: torch.Tensor) -> torch.Tensor:
"""Tile ``[N, C, H, W]`` frames without requiring torchvision."""
nframes, channels, height, width = frames.shape
ncols = int(math.ceil(math.sqrt(nframes)))
nrows = int(math.ceil(nframes / ncols))
grid = frames.new_zeros(channels, nrows * height, ncols * width)
for index, frame in enumerate(frames):
row, col = divmod(index, ncols)
grid[
:, row * height : (row + 1) * height, col * width : (col + 1) * width
] = frame
return grid
[docs]
class VideoRecorder(ObservationTransform):
"""Video Recorder transform.
Will record a series of observations from an environment and write them
to a Logger object when needed.
Args:
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 :class:`~torchrl.record.loggers.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 :obj:`"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 :meth:`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
:meth:`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
:meth:`~torchrl.record.loggers.Logger.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 :class:`~torchrl.envs.GymEnv` or :class:`~torchrl.envs.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 :class:`~torchrl.envs.transforms.Compose`
which will dispatch the `dumps` to all its members.
>>> env.transform.dump()
.. note::
When recording a batched env (:class:`~torchrl.envs.SerialEnv` or
:class:`~torchrl.envs.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``!
"""
def __init__(
self,
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:
client = getattr(logger, "client", None)
if callable(client):
logger = client()
if max_frames is not None and max_frames <= 0:
raise ValueError(
f"max_frames must be a positive integer, got {max_frames}."
)
if in_keys is None:
in_keys = ["pixels"]
if out_keys is None:
out_keys = copy(in_keys)
super().__init__(in_keys=in_keys, out_keys=out_keys)
video_kwargs = {}
video_kwargs.update(kwargs)
if fps is not None:
video_kwargs["fps"] = fps
self.video_kwargs = video_kwargs
self.iter = 0
self.skip = skip
self.logger = logger
self.tag = tag
self.count = 0
self.center_crop = center_crop
self.make_grid = make_grid
self.max_frames = max_frames
self.dump_on_done = dump_on_done
if center_crop and not _has_tv:
raise ImportError(
"Could not load center_crop from torchvision. Make sure torchvision is installed."
)
self.obs = []
@property
def make_grid(self):
make_grid = self._make_grid
if make_grid is None:
if self.parent is not None:
self._make_grid = True
return True
self._make_grid = False
return False
return make_grid
@make_grid.setter
def make_grid(self, value):
self._make_grid = value
@property
def skip(self):
skip = self._skip
if skip is None:
parent = self.parent
if parent is None:
skip = 1
else:
batch_size = getattr(parent, "batch_size", ())
skip = 1 if len(batch_size) else 2
self._skip = skip
return skip
return skip
@skip.setter
def skip(self, value):
self._skip = value
def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
if isinstance(observation, NonTensorData):
observation_trsf = torch.tensor(observation.data)
else:
observation_trsf = observation
self.count += 1
if self.count % self.skip == 0:
if self.max_frames is not None and len(self.obs) >= self.max_frames:
return observation
if (
observation_trsf.ndim >= 3
and observation_trsf.shape[-3] in (1, 3)
and observation_trsf.shape[-2] > 3
and observation_trsf.shape[-1] > 3
):
# permute the channels to the last dim
observation_trsf = observation_trsf.permute(
*range(observation_trsf.ndim - 3), -2, -1, -3
)
# Handle grayscale (1-channel) by expanding to 3-channel for video
if observation_trsf.ndim >= 3 and observation_trsf.shape[-1] == 1:
observation_trsf = observation_trsf.expand(
*observation_trsf.shape[:-1], 3
)
if not (
observation_trsf.shape[-1] == 3 or observation_trsf.ndimension() == 2
):
raise RuntimeError(
f"Invalid observation shape, got: {observation.shape}"
)
observation_trsf = observation_trsf.clone()
if observation.ndimension() == 2:
observation_trsf = observation.unsqueeze(-3)
else:
if observation_trsf.shape[-1] != 3:
raise RuntimeError(
"observation_trsf is expected to have 3 dimensions, "
f"got {observation_trsf.ndimension()} instead"
)
trailing_dim = range(observation_trsf.ndimension() - 3)
observation_trsf = observation_trsf.permute(*trailing_dim, -1, -3, -2)
if self.center_crop:
if not _has_tv:
raise ImportError(
"Could not import torchvision, `center_crop` not available. "
"Make sure torchvision is installed in your environment."
)
from torchvision.transforms.functional import (
center_crop as center_crop_fn,
)
observation_trsf = center_crop_fn(
observation_trsf, [self.center_crop, self.center_crop]
)
if self.make_grid and observation_trsf.ndimension() >= 4:
obs_flat = observation_trsf.flatten(0, -4)
observation_trsf = _make_video_grid(obs_flat)
self.obs.append(observation_trsf.to("cpu", torch.uint8))
elif observation_trsf.ndimension() >= 4:
frames = observation_trsf.to("cpu", torch.uint8).flatten(0, -4)
if self.max_frames is not None:
frames = frames[: self.max_frames - len(self.obs)]
self.obs.extend(frames)
else:
self.obs.append(observation_trsf.to("cpu", torch.uint8))
return observation
def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
return self._call(tensordict)
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
next_tensordict = super()._step(tensordict, next_tensordict)
if self.dump_on_done and self._all_done(next_tensordict):
self.dump()
return next_tensordict
def _all_done(self, next_tensordict: TensorDictBase) -> bool:
"""Whether every ``"done"`` entry of the post-step data is ``True``."""
parent = self.parent
if parent is not None:
done_keys = [key for key in parent.done_keys if _ends_with(key, "done")]
else:
done_keys = ["done"]
dones = [next_tensordict.get(key, None) for key in done_keys]
dones = [done for done in dones if done is not None]
if not dones:
return False
return all(bool(done.all()) for done in dones)
def _check_batched_worker_compat(self) -> None:
torchrl_logger.warning(
"A VideoRecorder was found among the transforms of a "
"SerialEnv/ParallelEnv worker env. Worker-side transforms are not "
"reached by `dump` calls issued on the outer env or by collectors "
"and evaluators, so recorded frames accumulate until dumped "
"manually (consider `max_frames` or `dump_on_done` to bound "
"memory). Prefer attaching the recorder to the outer env, e.g. "
"TransformedEnv(ParallelEnv(N, make_env), VideoRecorder(...)), "
"which records all worker envs into a single grid video."
)
def to_animation(
self,
*,
title: str | None = None,
interval: int = 50,
repeat_delay: int = 1000,
clear: bool = False,
) -> Any:
"""Convert recorded frames to a Matplotlib animation.
This helper is intended for tutorials and notebooks where the recorded
frames should be rendered inline by Sphinx-Gallery or IPython instead of
being written through a logger. Frames are read from the same internal
buffer used by :meth:`dump`.
Args:
title: optional title for the rendered figure.
interval: delay between frames, in milliseconds.
repeat_delay: delay before repeating the animation, in milliseconds.
clear: if ``True``, clear the recorded frame buffer after creating
the animation.
Returns:
A :class:`matplotlib.animation.ArtistAnimation` built from the
frames currently stored by the recorder.
Examples:
>>> import torch
>>> from torchrl.record import VideoRecorder
>>> recorder = VideoRecorder(None, None)
>>> recorder._apply_transform(torch.zeros(3, 8, 8, dtype=torch.uint8))
>>> animation = recorder.to_animation() # doctest: +SKIP
"""
if not self.obs:
raise RuntimeError(
"VideoRecorder.to_animation() requires at least one recorded frame."
)
if not _has_matplotlib:
raise ImportError(
"VideoRecorder.to_animation() requires matplotlib to be installed."
)
import matplotlib.animation as mpl_animation
import matplotlib.pyplot as plt
fig, axis = plt.subplots()
axis.set_axis_off()
if title is not None:
axis.set_title(title)
artists = []
for frame in self.obs:
frame = frame.detach().cpu()
if frame.ndim == 3 and frame.shape[0] in (1, 3):
frame = frame.permute(1, 2, 0)
if frame.ndim == 3 and frame.shape[-1] == 1:
frame = frame.expand(*frame.shape[:-1], 3)
artists.append([axis.imshow(frame.numpy(), animated=True)])
out = mpl_animation.ArtistAnimation(
fig,
artists,
interval=interval,
blit=True,
repeat_delay=repeat_delay,
)
if clear:
self.obs.clear()
self.count = 0
return out
def dump(self, suffix: str | None = None, step: int | None = None) -> None:
"""Writes the video to the ``self.logger`` attribute.
Calling ``dump`` when no image has been stored in a no-op.
Args:
suffix (str, optional): a suffix for the video to be recorded.
step (int, optional): the step to log the video at. If not provided,
uses an internal counter that increments with each dump call.
"""
if self.obs:
obs = torch.stack(self.obs, 0).unsqueeze(0).cpu()
else:
obs = None
self.obs = []
if obs is not None:
if suffix is None:
tag = self.tag
else:
tag = "_".join([self.tag, suffix])
if self.logger is not None:
self.logger.log_video(
name=tag,
video=obs,
step=step if step is not None else self.iter,
**self.video_kwargs,
)
self.iter += 1
self.count = 0
self.obs = []
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
self._call(tensordict_reset)
return tensordict_reset
[docs]
class TensorDictRecorder(Transform):
"""TensorDict recorder.
When the 'dump' method is called, this class will save a stack of the tensordict resulting from :obj:`env.step(td)` in a
file with a prefix defined by the out_file_base argument.
Args:
out_file_base (str): a string defining the prefix of the file where the tensordict will be written.
skip_reset (bool): if ``True``, the first TensorDict of the list will be discarded (usually the tensordict
resulting from the call to :obj:`env.reset()`)
default: True
skip (int): frame interval for the saved tensordict.
default: 4
"""
def __init__(
self,
out_file_base: str,
skip_reset: bool = True,
skip: int = 4,
in_keys: Sequence[str] | None = None,
) -> None:
if in_keys is None:
in_keys = []
super().__init__(in_keys=in_keys)
self.iter = 0
self.out_file_base = out_file_base
self.td = []
self.skip_reset = skip_reset
self.skip = skip
self.count = 0
def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
self.count += 1
if self.count % self.skip == 0:
_td = next_tensordict
if self.in_keys:
_td = next_tensordict.select(*self.in_keys).to_tensordict()
self.td.append(_td)
return next_tensordict
def dump(self, suffix: str | None = None) -> None:
if suffix is None:
tag = self.tag
else:
tag = "_".join([self.tag, suffix])
td = self.td
if self.skip_reset:
td = td[1:]
torch.save(
torch.stack(td, 0).contiguous(),
f"{tag}_tensordict.t",
)
self.iter += 1
self.count = 0
del self.td
self.td = []
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
self._call(tensordict_reset)
return tensordict_reset
