From 78ce8b54d894d2fb496b7a19c948e1b2d90f9a47 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 29 Apr 2024 08:48:20 +0100 Subject: [PATCH 1/6] init --- .github/scripts/m1_script.sh | 2 +- .github/workflows/wheels.yml | 4 ++-- setup.py | 14 +++++++++----- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/.github/scripts/m1_script.sh b/.github/scripts/m1_script.sh index 6d2f194e3bc..6552d8e4622 100644 --- a/.github/scripts/m1_script.sh +++ b/.github/scripts/m1_script.sh @@ -1,5 +1,5 @@ #!/bin/bash -export BUILD_VERSION=0.4.0 +export TORCHRL_BUILD_VERSION=0.4.0 ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml index f46c20b8a7e..9b2e57db531 100644 --- a/.github/workflows/wheels.yml +++ b/.github/workflows/wheels.yml @@ -32,7 +32,7 @@ jobs: run: | export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" python3 -mpip install wheel - BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel + TORCHRL_BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel # NB: wheels have the linux_x86_64 tag so we rename to manylinux1 # find . -name 'dist/*whl' -exec bash -c ' mv $0 ${0/linux/manylinux1}' {} \; # pytorch/pytorch binaries are also manylinux_2_17 compliant but they @@ -72,7 +72,7 @@ jobs: shell: bash run: | python3 -mpip install wheel - BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel + TORCHRL_BUILD_VERSION=0.4.0 python3 setup.py bdist_wheel - name: Upload wheel for the test-wheel job uses: actions/upload-artifact@v2 with: diff --git a/setup.py b/setup.py index fee902d5486..0196cb4a8f4 100644 --- a/setup.py +++ b/setup.py @@ -32,8 +32,8 @@ def get_version(): version_txt = os.path.join(cwd, "version.txt") with open(version_txt, "r") as f: version = f.readline().strip() - if os.getenv("BUILD_VERSION"): - version = os.getenv("BUILD_VERSION") + if os.getenv("TORCHRL_BUILD_VERSION"): + version = os.getenv("TORCHRL_BUILD_VERSION") elif sha != "Unknown": version += "+" + sha[:7] return version @@ -68,11 +68,13 @@ def write_version_file(version): f.write("git_version = {}\n".format(repr(sha))) -def _get_pytorch_version(is_nightly): +def _get_pytorch_version(is_nightly, is_local): # if "PYTORCH_VERSION" in os.environ: # return f"torch=={os.environ['PYTORCH_VERSION']}" if is_nightly: return "torch>=2.4.0.dev" + elif is_local: + return "torch" return "torch>=2.3.0" @@ -178,10 +180,12 @@ def _main(argv): else: version = get_version() write_version_file(version) + TORCHRL_BUILD_VERSION = os.getenv("TORCHRL_BUILD_VERSION") logging.info("Building wheel {}-{}".format(package_name, version)) - logging.info(f"BUILD_VERSION is {os.getenv('BUILD_VERSION')}") + logging.info(f"TORCHRL_BUILD_VERSION is {TORCHRL_BUILD_VERSION}") - pytorch_package_dep = _get_pytorch_version(is_nightly) + is_local = TORCHRL_BUILD_VERSION is None + pytorch_package_dep = _get_pytorch_version(is_nightly, is_local) logging.info("-- PyTorch dependency:", pytorch_package_dep) # branch = _run_cmd(["git", "rev-parse", "--abbrev-ref", "HEAD"]) # tag = _run_cmd(["git", "describe", "--tags", "--exact-match", "@"]) From b4f4daf7ffe52cb7927529155cdb309d1b63d6bb Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Thu, 2 May 2024 14:17:17 +0100 Subject: [PATCH 2/6] [Doc] Fix links in doc (#2151) (cherry picked from commit 5b9cb440f7f507948b0431077c36706a99100b78) --- docs/source/reference/data.rst | 4 ++++ docs/source/reference/envs.rst | 4 ++-- docs/source/reference/trainers.rst | 2 ++ tutorials/sphinx-tutorials/coding_dqn.py | 12 ++++++------ tutorials/sphinx-tutorials/coding_ppo.py | 14 +++++++------- tutorials/sphinx-tutorials/multiagent_ppo.py | 4 ++-- tutorials/sphinx-tutorials/pendulum.py | 8 ++++---- 7 files changed, 27 insertions(+), 21 deletions(-) diff --git a/docs/source/reference/data.rst b/docs/source/reference/data.rst index efb1a755d0e..3c249bcfbaa 100644 --- a/docs/source/reference/data.rst +++ b/docs/source/reference/data.rst @@ -24,6 +24,8 @@ widely used replay buffers: Composable Replay Buffers ------------------------- +.. _ref_buffers: + We also give users the ability to compose a replay buffer. We provide a wide panel of solutions for replay buffer usage, including support for almost any data type; storage in memory, on device or on physical memory; @@ -796,6 +798,8 @@ such that they can be stacked together during sampling. TensorSpec ---------- +.. _ref_specs: + The `TensorSpec` parent class and subclasses define the basic properties of observations and actions in TorchRL, such as shape, device, dtype and domain. It is important that your environment specs match the input and output that it sends and receives, as diff --git a/docs/source/reference/envs.rst b/docs/source/reference/envs.rst index 5d4b6d0b7b5..3c5e69f221b 100644 --- a/docs/source/reference/envs.rst +++ b/docs/source/reference/envs.rst @@ -133,7 +133,7 @@ function. transform. -Our environment `tutorial `_ +Our environment :ref:`tutorial ` provides more information on how to design a custom environment from scratch. .. autosummary:: @@ -559,7 +559,7 @@ Transforms In most cases, the raw output of an environment must be treated before being passed to another object (such as a policy or a value operator). To do this, TorchRL provides a set of transforms that aim at reproducing the transform logic of `torch.distributions.Transform` and `torchvision.transforms`. -Our environment `tutorial `_ +Our environment :ref:`tutorial ` provides more information on how to design a custom transform. Transformed environments are build using the :class:`TransformedEnv` primitive. diff --git a/docs/source/reference/trainers.rst b/docs/source/reference/trainers.rst index e253ad7067e..2f0982257eb 100644 --- a/docs/source/reference/trainers.rst +++ b/docs/source/reference/trainers.rst @@ -3,6 +3,8 @@ torchrl.trainers package ======================== +.. _ref_trainers: + The trainer package provides utilities to write re-usable training scripts. The core idea is to use a trainer that implements a nested loop, where the outer loop runs the data collection steps and the inner loop the optimization steps. We believe this fits multiple RL training schemes, such as diff --git a/tutorials/sphinx-tutorials/coding_dqn.py b/tutorials/sphinx-tutorials/coding_dqn.py index 36434435e81..3b9d712736a 100644 --- a/tutorials/sphinx-tutorials/coding_dqn.py +++ b/tutorials/sphinx-tutorials/coding_dqn.py @@ -35,7 +35,7 @@ # - how to build an environment in TorchRL, including transforms (e.g. data # normalization, frame concatenation, resizing and turning to grayscale) # and parallel execution. Unlike what we did in the -# `DDPG tutorial `_, we +# :ref:`DDPG tutorial `, we # will normalize the pixels and not the state vector. # - how to design a :class:`~torchrl.modules.QValueActor` object, i.e. an actor # that estimates the action values and picks up the action with the highest @@ -46,7 +46,7 @@ # - and finally how to evaluate your model. # # **Prerequisites**: We encourage you to get familiar with torchrl through the -# `PPO tutorial `_ first. +# :ref:`PPO tutorial ` first. # # DQN # --- @@ -393,8 +393,8 @@ def get_replay_buffer(buffer_size, n_optim, batch_size): # Data collector # ~~~~~~~~~~~~~~ # -# As in `PPO `_ and -# `DDPG `_, we will be using +# As in :ref:`PPO ` and +# :ref:`DDPG `, we will be using # a data collector as a dataloader in the outer loop. # # We choose the following configuration: we will be running a series of @@ -691,7 +691,7 @@ def get_loss_module(actor, gamma): # In this case, a location must be explicitly passed (). This method gives # more control over the location of the hook but it also requires more # understanding of the Trainer mechanism. -# Check the `trainer documentation `_ +# Check the :ref:`trainer documentation ` # for a detailed description of the trainer hooks. # trainer.register_op("post_optim", target_net_updater.step) @@ -768,7 +768,7 @@ def print_csv_files_in_folder(folder_path): # - A prioritized replay buffer could also be used. This will give a # higher priority to samples that have the worst value accuracy. # Learn more on the -# `replay buffer section `_ +# :ref:`replay buffer section ` # of the documentation. # - A distributional loss (see :class:`~torchrl.objectives.DistributionalDQNLoss` # for more information). diff --git a/tutorials/sphinx-tutorials/coding_ppo.py b/tutorials/sphinx-tutorials/coding_ppo.py index 394af8e741b..eff6f31657b 100644 --- a/tutorials/sphinx-tutorials/coding_ppo.py +++ b/tutorials/sphinx-tutorials/coding_ppo.py @@ -26,12 +26,12 @@ We will cover six crucial components of TorchRL: -* `environments `__ -* `transforms `__ -* `models (policy and value function) `__ -* `loss modules `__ -* `data collectors `__ -* `replay buffers `__ +* :ref:`environments ` +* :ref:`transforms ` +* :ref:`models ` +* :ref:`loss modules ` +* :ref:`data collectors ` +* :ref:`replay buffers ` """ @@ -478,7 +478,7 @@ # Data collector # -------------- # -# TorchRL provides a set of `DataCollector classes `__. +# TorchRL provides a set of :ref:`DataCollector classes `. # Briefly, these classes execute three operations: reset an environment, # compute an action given the latest observation, execute a step in the environment, # and repeat the last two steps until the environment signals a stop (or reaches diff --git a/tutorials/sphinx-tutorials/multiagent_ppo.py b/tutorials/sphinx-tutorials/multiagent_ppo.py index 39f19242648..f19fb270908 100644 --- a/tutorials/sphinx-tutorials/multiagent_ppo.py +++ b/tutorials/sphinx-tutorials/multiagent_ppo.py @@ -183,7 +183,7 @@ # TorchRL API allows integrating various types of multi-agent environment flavours. # Some examples include environments with shared or individual agent rewards, done flags, and observations. # For more information on how the multi-agent environments API works in TorchRL, you can check out the dedicated -# `doc section `_. +# :ref:`doc section `. # # The VMAS simulator, in particular, models agents with individual rewards, info, observations, and actions, but # with a collective done flag. @@ -784,7 +784,7 @@ # # If you are interested in creating or wrapping your own multi-agent environments in TorchRL, # you can check out the dedicated -# `doc section `_. +# :ref:`doc section `. # # Finally, you can modify the parameters of this tutorial to try many other configurations and scenarios # to become a MARL master. diff --git a/tutorials/sphinx-tutorials/pendulum.py b/tutorials/sphinx-tutorials/pendulum.py index 5605af66a0c..4eda4ea8e91 100644 --- a/tutorials/sphinx-tutorials/pendulum.py +++ b/tutorials/sphinx-tutorials/pendulum.py @@ -33,9 +33,9 @@ In the process, we will touch three crucial components of TorchRL: -* `environments `__ -* `transforms `__ -* `models (policy and value function) `__ +* :ref:`environments ` +* :ref:`transforms ` +* :ref:`models ` """ @@ -389,7 +389,7 @@ def _reset(self, tensordict): # convenient shortcuts to the content of the output and input spec containers. # # TorchRL offers multiple :class:`~torchrl.data.TensorSpec` -# `subclasses `_ to +# :ref:`subclasses ` to # encode the environment's input and output characteristics. # # Specs shape From ad167371e59464ac7a59b74be9600ea736d8de44 Mon Sep 17 00:00:00 2001 From: N00bcak Date: Mon, 1 Jul 2024 23:02:07 +0800 Subject: [PATCH 3/6] proposed knowledge_base entry --- docs/source/reference/knowledge_base.rst | 2 + knowledge_base/VIDEO_CUSTOMISATION.md | 58 +++++++++++++++++++ torchrl/record/loggers/csv.py | 4 +- .../sphinx-tutorials/getting-started-4.py | 15 +++-- 4 files changed, 73 insertions(+), 6 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..78a57b5e924 100644 --- a/torchrl/record/loggers/csv.py +++ b/torchrl/record/loggers/csv.py @@ -86,7 +86,9 @@ 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) + torchvision.io.write_video( + filepath, vid_tensor, fps=self.video_fps, **kwargs + ) else: raise ValueError( f"Unknown video format {self.video_format}. Must be one of 'pt', 'memmap' or 'mp4'." 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 `! From 7c8c15604bda3ff445b67ed4c54cfbca43ffa2b8 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Tue, 2 Jul 2024 11:21:21 +0100 Subject: [PATCH 4/6] amend --- torchrl/record/loggers/csv.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/torchrl/record/loggers/csv.py b/torchrl/record/loggers/csv.py index 78a57b5e924..bb5bbab9df3 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 @@ -124,7 +136,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). @@ -165,12 +177,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.CVSLogger`, + 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 From 4e60774dedb6d5847c163c985b71928598123b46 Mon Sep 17 00:00:00 2001 From: N00bcak Date: Tue, 2 Jul 2024 18:50:59 +0800 Subject: [PATCH 5/6] correct typo --- torchrl/record/loggers/csv.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchrl/record/loggers/csv.py b/torchrl/record/loggers/csv.py index bb5bbab9df3..087d08a41dd 100644 --- a/torchrl/record/loggers/csv.py +++ b/torchrl/record/loggers/csv.py @@ -187,7 +187,7 @@ def log_video(self, name: str, video: Tensor, step: int = None, **kwargs) -> Non .. 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.CVSLogger`, + 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) From 48d8855b9051c5b19fb6481921ecdee2099ee746 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 3 Jul 2024 10:40:18 +0100 Subject: [PATCH 6/6] amend --- torchrl/record/loggers/csv.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/torchrl/record/loggers/csv.py b/torchrl/record/loggers/csv.py index 087d08a41dd..4de3de2b928 100644 --- a/torchrl/record/loggers/csv.py +++ b/torchrl/record/loggers/csv.py @@ -98,9 +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 - ) + 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'."