From 55f0a52ac7025c459dd71ec180c9c7d1c7576f9c Mon Sep 17 00:00:00 2001 From: N00bcak Date: Fri, 5 Jul 2024 04:19:20 +0800 Subject: [PATCH] [Doc] Add Custom Options for VideoRecorder (#2259) --- docs/source/reference/knowledge_base.rst | 2 + knowledge_base/VIDEO_CUSTOMISATION.md | 58 +++++++++++++++++++ torchrl/record/loggers/csv.py | 25 +++++++- .../sphinx-tutorials/getting-started-4.py | 15 +++-- 4 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 knowledge_base/VIDEO_CUSTOMISATION.md diff --git a/docs/source/reference/knowledge_base.rst b/docs/source/reference/knowledge_base.rst index 2815c4c5732..5f18842d172 100644 --- a/docs/source/reference/knowledge_base.rst +++ b/docs/source/reference/knowledge_base.rst @@ -1,6 +1,8 @@ Knowledge Base ============== +.. _ref_knowledge_base: + .. include:: ../../../knowledge_base/README.md :start-line: 1 :parser: myst_parser.sphinx_ diff --git a/knowledge_base/VIDEO_CUSTOMISATION.md b/knowledge_base/VIDEO_CUSTOMISATION.md new file mode 100644 index 00000000000..956110d89aa --- /dev/null +++ b/knowledge_base/VIDEO_CUSTOMISATION.md @@ -0,0 +1,58 @@ +# Customising Video Renders + +## Tweaking Video Rendering Settings +TorchRL relies heavily on the [torchvision.io](https://pytorch.org/vision/main/io.html) +and [PyAV](https://github.com/PyAV-Org/PyAV) modules for its video logging +capabilities. Though these libraries are quite convenient and powerful, it is +not easy to access the variety of knobs and settings at your disposal. + +This guide hopes to clarify what appear to be the general principles behind +customising video rendering, and show you how you can manually adjust your +rollouts' rendering settings to your liking. + +## General Principles +Ultimately, [torchvision.io](https://pytorch.org/vision/main/io.html) and +[PyAV](https://github.com/PyAV-Org/PyAV) make calls to [FFmpeg](https://ffmpeg.org/) +libraries in order to render videos. + +In other words: + +- Whatever can be fed into [FFmpeg](https://ffmpeg.org/), we can also feed +into TorchRL's `Loggers`. +- For any custom settings we wish to use, we must reference them from +[FFmpeg's documentation](https://trac.ffmpeg.org/) + +## Video Rendering Customization Example + +Suppose the following snippet gave us extremely blurry videos, even though +we provided it clear, frame-by-frame images to stitch together: +```python +from torchrl.envs import GymEnv, TransformedEnv +from torchrl.record import CSVLogger, VideoRecorder + +logger = CSVLogger(exp_name="my_exp") +env = GymEnv("CartPole-v1", from_pixels=True, pixels_only=False) + +recorder = VideoRecorder(logger, tag="my_video") +record_env = TransformedEnv(env, recorder) +rollout = record_env.rollout(max_steps=3) +recorder.dump() +``` + +Since TorchRL's default video codec is [H264](https://trac.ffmpeg.org/wiki/Encode/H.264), +the settings that we must change should be in there. + +For the purposes of this example, let us choose a +[Constant Rate Factor (CRF)](https://trac.ffmpeg.org/wiki/Encode/H.264#crf) of +`17` and a [preset](https://trac.ffmpeg.org/wiki/Encode/H.264#Preset) of `slow`, +as advised by the documentation. + +We can improve the video quality by appending all our desired settings +(as keyword arguments) to `recorder` like so: +```python +# The arguments' types don't appear to matter too much, as long as they are +# appropriate for Python. +# For example, this would work as well: +# logger = CSVLogger(exp_name="my_exp", crf=17, preset="slow") +logger = CSVLogger(exp_name="my_exp", crf="17", preset="slow") +``` diff --git a/torchrl/record/loggers/csv.py b/torchrl/record/loggers/csv.py index 9e179699bd5..4de3de2b928 100644 --- a/torchrl/record/loggers/csv.py +++ b/torchrl/record/loggers/csv.py @@ -51,6 +51,18 @@ def add_scalar(self, name: str, value: float, global_step: Optional[int] = None) fd.flush() def add_video(self, tag, vid_tensor, global_step: Optional[int] = None, **kwargs): + """Writes a video on a file on disk. + + The video format can be one of + - `"pt"`: uses :func:`~torch.save` to save the video tensor); + - `"memmap"`: saved the file as memory-mapped array (reading this file will require + the dtype and shape to be known at read time); + - `"mp4"`: saves the file as an `.mp4` file using torchvision :func:`~torchvision.io.write_video` + API. Any ``kwargs`` passed to ``add_video`` will be transmitted to ``write_video``. + These include ``preset``, ``crf`` and others. + See ffmpeg's doc (https://trac.ffmpeg.org/wiki/Encode/H.264) for some more information of the video format options. + + """ if global_step is None: global_step = self.videos_counter[tag] self.videos_counter[tag] += 1 @@ -86,7 +98,8 @@ def add_video(self, tag, vid_tensor, global_step: Optional[int] = None, **kwargs vid_tensor = vid_tensor.flatten(0, vid_tensor.ndim - 4) vid_tensor = vid_tensor.permute((0, 2, 3, 1)) vid_tensor = vid_tensor.expand(*vid_tensor.shape[:-1], 3) - torchvision.io.write_video(filepath, vid_tensor, fps=self.video_fps) + kwargs.setdefault("fps", self.video_fps) + torchvision.io.write_video(filepath, vid_tensor, **kwargs) else: raise ValueError( f"Unknown video format {self.video_format}. Must be one of 'pt', 'memmap' or 'mp4'." @@ -122,7 +135,7 @@ class CSVLogger(Logger): exp_name (str): The name of the experiment. log_dir (str or Path, optional): where the experiment should be saved. Defaults to ``/csv_logs``. - video_format (str, optional): how videos should be saved. Must be one of + video_format (str, optional): how videos should be saved when calling :meth:`~torchrl.record.loggers.csv.CSVExperiment.add_video`. Must be one of ``"pt"`` (video saved as a `video__.pt` file with torch.save), ``"memmap"`` (video saved as a `video__.memmap` file with :class:`~tensordict.MemoryMappedTensor`), ``"mp4"`` (video saved as a `video__.mp4` file, requires torchvision to be installed). @@ -163,12 +176,18 @@ def log_scalar(self, name: str, value: float, step: int = None) -> None: self.experiment.add_scalar(name, value, global_step=step) def log_video(self, name: str, video: Tensor, step: int = None, **kwargs) -> None: - """Log videos inputs to a .pt file. + """Log videos inputs to a .pt (or other format) file. Args: name (str): The name of the video. video (Tensor): The video to be logged. step (int, optional): The step at which the video is logged. Defaults to None. + **kwargs: other kwargs passed to the underlying video logger. + + .. note:: If the video format is `mp4`, many more arguments can be passed to the :meth:`~torchvision.io.write_video` + function. + For more information on video logging with :class:`~torchrl.record.loggers.csv.CSVLogger`, + see the :meth:`~torchrl.record.loggers.csv.CSVExperiment.add_video` documentation. """ # check for correct format of the video tensor ((N), T, C, H, W) # check that the color channel (C) is either 1 or 3 diff --git a/tutorials/sphinx-tutorials/getting-started-4.py b/tutorials/sphinx-tutorials/getting-started-4.py index 195738abca9..32731661a7b 100644 --- a/tutorials/sphinx-tutorials/getting-started-4.py +++ b/tutorials/sphinx-tutorials/getting-started-4.py @@ -40,7 +40,7 @@ # # Usually, building a logger requires # at least an experiment name and possibly a logging directory and other -# hyperapameters. +# hyperparameters. # from torchrl.record import CSVLogger @@ -66,9 +66,12 @@ # # Let's first see how we can create a Gym environment that outputs images # alongside its observations. :class:`~torchrl.envs.GymEnv` accept two keywords -# for this purpose: ``from_pixels=True`` will make the env ``step`` function +# for this purpose: +# - ``from_pixels=True`` will make the env ``step`` function # write a ``"pixels"`` entry containing the images corresponding to your -# observations, and the ``pixels_only=False`` will indicate that you want the +# observations, and +# +# - ``pixels_only=False`` will indicate that you want the # observations to be returned as well. # @@ -94,8 +97,8 @@ ##################################### # When running this environment, all the ``"pixels"`` entries will be saved in -# a local buffer and dumped in a video on demand (it is important that you -# call this method when appropriate): +# a local buffer (i.e. RAM) and dumped in a video on demand (to prevent excessive +# RAM usage, you are advised to call this method whenever appropriate!): rollout = record_env.rollout(max_steps=3) # Uncomment this line to save the video on disk: @@ -105,6 +108,8 @@ # In this specific case, the video format can be chosen when instantiating # the CSVLogger. # +# (If you want to customise how your video is recorded, have a look at :ref:`our knowledge base `.) +# # This is all we wanted to cover in the getting started tutorial. # You should now be ready to code your # :ref:`first training loop with TorchRL `!