From c4db149883c687162d87678e3e293a21d2fda83d Mon Sep 17 00:00:00 2001 From: kurtamohler Date: Thu, 4 Jul 2024 02:27:51 -0700 Subject: [PATCH 1/5] Fix "Run in Colab" and "Download Notebook" links in tutorials (#2268) --- docs/source/_static/js/theme.js | 2 +- docs/source/_static/js/torchrl_theme.js | 2 +- docs/source/_templates/layout.html | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/_static/js/theme.js b/docs/source/_static/js/theme.js index 2b23c856d49..5c7c27f82b1 100644 --- a/docs/source/_static/js/theme.js +++ b/docs/source/_static/js/theme.js @@ -946,7 +946,7 @@ if (downloadNote.length >= 1) { tutorialUrlArray[0] = tutorialUrlArray[0] + "/sphinx-tutorials" var githubLink = "https://github.com/pytorch/rl/blob/main/" + tutorialUrlArray.join("/") + ".py", - notebookLink = $(".reference.download")[1].href, + notebookLink = $(".sphx-glr-download-jupyter").find(".download.reference")[0].href, notebookDownloadPath = notebookLink.split('_downloads')[1], colabLink = "https://colab.research.google.com/github/pytorch/rl/blob/gh-pages/main/_downloads" + notebookDownloadPath; diff --git a/docs/source/_static/js/torchrl_theme.js b/docs/source/_static/js/torchrl_theme.js index 05e1a5b7a63..dc4588cdfa7 100644 --- a/docs/source/_static/js/torchrl_theme.js +++ b/docs/source/_static/js/torchrl_theme.js @@ -945,7 +945,7 @@ if (downloadNote.length >= 1) { var tutorialUrlArray = $("#tutorial-type").text().split('/'); var githubLink = "https://github.com/pytorch/rl/tree/tutorial_py_dup/sphinx-tutorials/" + tutorialUrlArray[tutorialUrlArray.length - 1] + ".py", - notebookLink = $(".reference.download")[1].href, + notebookLink = $(".sphx-glr-download-jupyter").find(".download.reference")[0].href, notebookDownloadPath = notebookLink.split('_downloads')[1], colabLink = "https://colab.research.google.com/github/pytorch/rl/blob/gh-pages/_downloads" + notebookDownloadPath; diff --git a/docs/source/_templates/layout.html b/docs/source/_templates/layout.html index 543edfac202..7babfdd1ab0 100644 --- a/docs/source/_templates/layout.html +++ b/docs/source/_templates/layout.html @@ -61,7 +61,7 @@ if (downloadNote.length >= 1) { var tutorialUrl = $("#tutorial-type").text(); var githubLink = "https://github.com/pytorch/rl/blob/main/tutorials/sphinx-tutorials/" + tutorialUrl + ".py", - notebookLink = $(".reference.download")[1].href, + notebookLink = $(".sphx-glr-download-jupyter").find(".download.reference")[0].href, notebookDownloadPath = notebookLink.split('_downloads')[1], colabLink = "https://colab.research.google.com/github/pytorch/rl/blob/gh-pages/main/_downloads" + notebookDownloadPath; From 55f0a52ac7025c459dd71ec180c9c7d1c7576f9c Mon Sep 17 00:00:00 2001 From: N00bcak Date: Fri, 5 Jul 2024 04:19:20 +0800 Subject: [PATCH 2/5] [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 `! From 65d43b9c7be5287afc026565fa7f98404f4cbfc5 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 8 Jul 2024 10:43:09 +0100 Subject: [PATCH 3/5] [CI] Fix wheels (#2274) --- .github/workflows/build-wheels-linux.yml | 47 +++++ .github/workflows/build-wheels-windows.yml | 48 +++++ .github/workflows/wheels.yml | 193 --------------------- 3 files changed, 95 insertions(+), 193 deletions(-) create mode 100644 .github/workflows/build-wheels-linux.yml create mode 100644 .github/workflows/build-wheels-windows.yml delete mode 100644 .github/workflows/wheels.yml diff --git a/.github/workflows/build-wheels-linux.yml b/.github/workflows/build-wheels-linux.yml new file mode 100644 index 00000000000..5171a7c3e2a --- /dev/null +++ b/.github/workflows/build-wheels-linux.yml @@ -0,0 +1,47 @@ +name: Build Linux Wheels + +on: + pull_request: + push: + branches: + - nightly + - main + - release/* + tags: + # NOTE: Binary build pipelines should only get triggered on release candidate builds + # Release candidate tags look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + generate-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: linux + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build: + needs: generate-matrix + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/rl + smoke-test-script: test/smoke_test.py + package-name: torchrl + name: pytorch/rl + uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + package-name: ${{ matrix.package-name }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml new file mode 100644 index 00000000000..683f2a93f69 --- /dev/null +++ b/.github/workflows/build-wheels-windows.yml @@ -0,0 +1,48 @@ +name: Build Windows Wheels + +on: + pull_request: + push: + branches: + - nightly + - main + - release/* + tags: + # NOTE: Binary build pipelines should only get triggered on release candidate builds + # Release candidate tags look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + generate-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: windows + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build: + needs: generate-matrix + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/rl + post-script: "python packaging/wheel/relocate.py" + smoke-test-script: test/smoke_test.py + package-name: torchrl + name: pytorch/rl + uses: pytorch/test-infra/.github/workflows/build_wheels_windows.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + package-name: ${{ matrix.package-name }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml deleted file mode 100644 index 9b2e57db531..00000000000 --- a/.github/workflows/wheels.yml +++ /dev/null @@ -1,193 +0,0 @@ -name: Wheels -on: - pull_request: - types: [opened, synchronize, reopened] - push: - branches: - - release/* - -concurrency: - # Documentation suggests ${{ github.head_ref }}, but that's only available on pull_request/pull_request_target triggers, so using ${{ github.ref }}. - # On master, we want all builds to complete even if merging happens faster to make it easier to discover at which point something broke. - group: ${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && format('ci-master-{0}', github.sha) || format('ci-{0}', github.ref) }} - cancel-in-progress: true - -jobs: - - build-wheel-linux: - runs-on: ubuntu-20.04 - strategy: - matrix: - python_version: [["3.8", "cp38-cp38"], ["3.9", "cp39-cp39"], ["3.10", "cp310-cp310"], ["3.11", "cp311-cp311"]] - cuda_support: [["", "--index-url https://download.pytorch.org/whl/cpu", "\"['cpu', '11.3', '11.6']\"", "cpu"]] - container: pytorch/manylinux-${{ matrix.cuda_support[3] }} - steps: - - name: Checkout torchrl - uses: actions/checkout@v2 - - name: Install PyTorch RC - run: | - export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" - python3 -mpip install torch ${{ matrix.cuda_support[1] }} - - name: Build wheel - run: | - export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" - python3 -mpip install 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 - # pretend that they're manylinux1 compliant so we do the same. - - name: Show auditwheel output; confirm 2-17 - run: | - python3 -mpip install auditwheel - auditwheel show dist/* - - name: Upload wheel for the test-wheel job - uses: actions/upload-artifact@v2 - with: - name: torchrl-linux-${{ matrix.python_version[0] }}.whl - path: dist/torchrl-*.whl - - name: Upload wheel for download - uses: actions/upload-artifact@v2 - with: - name: torchrl-batch.whl - path: dist/*.whl - - build-wheel-windows: - runs-on: windows-latest - strategy: - matrix: - python_version: [["3.8", "3.8"], ["3.9", "3.9"], ["3.10", "3.10.3"], ["3.11", "3.11"]] - steps: - - name: Setup Python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python_version[1] }} - - name: Checkout torchrl - uses: actions/checkout@v2 - - name: Install PyTorch RC - shell: bash - run: | - python3 -mpip install torch --index-url https://download.pytorch.org/whl/cpu - - name: Build wheel - shell: bash - run: | - python3 -mpip install 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: - name: torchrl-win-${{ matrix.python_version[0] }}.whl - path: dist/torchrl-*.whl - - name: Upload wheel for download - uses: actions/upload-artifact@v2 - with: - name: torchrl-batch.whl - path: dist/*.whl - - - test-wheel: - needs: [build-wheel-linux] - strategy: - matrix: - os: [["linux", "ubuntu-20.04"]] - python_version: [ "3.8", "3.9", "3.10", "3.11" ] - runs-on: ${{ matrix.os[1] }} - steps: - - name: Setup Python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python_version }} - architecture: x64 - - name: Checkout torchrl - uses: actions/checkout@v2 - - name: Install PyTorch RC - run: | - python3 -mpip install torch torchvision --index-url https://download.pytorch.org/whl/cpu - - name: Upgrade pip - run: | - python3 -mpip install --upgrade pip - - name: Install tensordict - run: | - python3 -mpip install git+https://github.com/pytorch/tensordict.git - - name: Install test dependencies - run: | - python3 -mpip install numpy pytest pytest-cov codecov unittest-xml-reporting pillow>=4.1.1 scipy av networkx expecttest pyyaml - - name: Download built wheels - uses: actions/download-artifact@v2 - with: - name: torchrl-${{ matrix.os[0] }}-${{ matrix.python_version }}.whl - path: /tmp/wheels - - name: Install built wheels - run: | - python3 -mpip install /tmp/wheels/* - - name: Log version string - run: | - # Avoid ambiguity of "import torchrl" by deleting the source files. - rm -rf torchrl/ - python -c "import torchrl; print(torchrl.__version__)" - - name: Run tests - run: | - set -e - export IN_CI=1 - mkdir test-reports - python -m torch.utils.collect_env - python -c "import torchrl; print(torchrl.__version__)" - EXIT_STATUS=0 - pytest test/smoke_test.py -v --durations 200 - exit $EXIT_STATUS - - test-wheel-windows: - needs: build-wheel-windows - strategy: - matrix: - python_version: [ "3.8", "3.9", "3.10", "3.11" ] - runs-on: windows-latest - steps: - - name: Setup Python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python_version }} - - name: Checkout torchrl - uses: actions/checkout@v2 - - name: Install PyTorch RC - shell: bash - run: | - python3 -mpip install torch torchvision --index-url https://download.pytorch.org/whl/cpu - - name: Upgrade pip - shell: bash - run: | - python3 -mpip install --upgrade pip - - name: Install tensordict - shell: bash - run: | - python3 -mpip install git+https://github.com/pytorch/tensordict.git - - name: Install test dependencies - shell: bash - run: | - python3 -mpip install numpy pytest pytest-cov codecov unittest-xml-reporting pillow>=4.1.1 scipy av networkx expecttest pyyaml - - name: Download built wheels - uses: actions/download-artifact@v2 - with: - name: torchrl-win-${{ matrix.python_version }}.whl - path: wheels - - name: Install built wheels - shell: bash - run: | - python3 -mpip install wheels/* - - name: Log version string - shell: bash - run: | - # Avoid ambiguity of "import torchrl" by deleting the source files. - rm -rf torchrl/ - python -c "import torchrl; print(torchrl.__version__)" - - name: Run tests - shell: bash - run: | - set -e - export IN_CI=1 - mkdir test-reports - python -m torch.utils.collect_env - python -c "import torchrl; print(torchrl.__version__)" - EXIT_STATUS=0 - pytest test/smoke_test.py -v --durations 200 - exit $EXIT_STATUS From 95b6a1689f84d5ca5ab3e844f24e576d09043582 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 8 Jul 2024 16:09:34 +0100 Subject: [PATCH 4/5] [BugFix] Fix tensordict private imports (#2275) --- test/test_cost.py | 2 +- torchrl/envs/batched_envs.py | 2 +- torchrl/envs/transforms/transforms.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_cost.py b/test/test_cost.py index 76fc4e651f4..36a81a6906d 100644 --- a/test/test_cost.py +++ b/test/test_cost.py @@ -12,7 +12,7 @@ from dataclasses import asdict, dataclass from packaging import version as pack_version -from tensordict._tensordict import unravel_keys +from tensordict._C import unravel_keys from tensordict.nn import ( InteractionType, diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index 7f462782757..dab0c52cdf7 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -27,7 +27,7 @@ TensorDict, TensorDictBase, ) -from tensordict._tensordict import unravel_key +from tensordict._C import unravel_key from torch import multiprocessing as mp from torchrl._utils import ( _check_for_faulty_process, diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py index bec76c603e6..eb9cdce923d 100644 --- a/torchrl/envs/transforms/transforms.py +++ b/torchrl/envs/transforms/transforms.py @@ -39,7 +39,7 @@ unravel_key, unravel_key_list, ) -from tensordict._tensordict import _unravel_key_to_tuple +from tensordict._C import _unravel_key_to_tuple from tensordict.nn import dispatch, TensorDictModuleBase from tensordict.utils import expand_as_right, expand_right, NestedKey from torch import nn, Tensor From a151923d5e1eab68a8256fa02d0946a0f57bd5d0 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 8 Jul 2024 16:14:30 +0100 Subject: [PATCH 5/5] Revert "[BugFix] Fix tensordict private imports" (#2276) --- test/test_cost.py | 2 +- torchrl/envs/batched_envs.py | 2 +- torchrl/envs/transforms/transforms.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_cost.py b/test/test_cost.py index 36a81a6906d..76fc4e651f4 100644 --- a/test/test_cost.py +++ b/test/test_cost.py @@ -12,7 +12,7 @@ from dataclasses import asdict, dataclass from packaging import version as pack_version -from tensordict._C import unravel_keys +from tensordict._tensordict import unravel_keys from tensordict.nn import ( InteractionType, diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index dab0c52cdf7..7f462782757 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -27,7 +27,7 @@ TensorDict, TensorDictBase, ) -from tensordict._C import unravel_key +from tensordict._tensordict import unravel_key from torch import multiprocessing as mp from torchrl._utils import ( _check_for_faulty_process, diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py index eb9cdce923d..bec76c603e6 100644 --- a/torchrl/envs/transforms/transforms.py +++ b/torchrl/envs/transforms/transforms.py @@ -39,7 +39,7 @@ unravel_key, unravel_key_list, ) -from tensordict._C import _unravel_key_to_tuple +from tensordict._tensordict import _unravel_key_to_tuple from tensordict.nn import dispatch, TensorDictModuleBase from tensordict.utils import expand_as_right, expand_right, NestedKey from torch import nn, Tensor