[Doc] Add Custom Options for VideoRecorder (#2259)

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_

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")
+```

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 ``<cur_dir>/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_<tag>_<step>.pt` file with torch.save),
             ``"memmap"`` (video saved as a `video_<tag>_<step>.memmap` file with :class:`~tensordict.MemoryMappedTensor`),
             ``"mp4"`` (video saved as a `video_<tag>_<step>.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

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 <ref_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 <gs_first_training>`!