From a98c120e50d8bdf5e3983a444ad6c8b5be0d1753 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Tue, 26 May 2026 18:48:22 -0700 Subject: [PATCH 01/11] [Feature] Add chunked value estimator calls --- test/objectives/test_values.py | 110 +++++++++++++++- torchrl/objectives/value/advantages.py | 121 +++++++++++++++++- .../trainers/algorithms/configs/objectives.py | 4 + 3 files changed, 227 insertions(+), 8 deletions(-) diff --git a/test/objectives/test_values.py b/test/objectives/test_values.py index c5092d926c4..9f2909890a9 100644 --- a/test/objectives/test_values.py +++ b/test/objectives/test_values.py @@ -69,7 +69,18 @@ class TestValues: ], ) @pytest.mark.parametrize("shifted", [False, True]) - def test_value_chunk_size_matches_unchunked(self, estimator_cls, kwargs, shifted): + @pytest.mark.parametrize("deactivate_vmap", [False, True]) + @pytest.mark.parametrize( + "chunk_kwargs", + [ + {"value_chunk_size": 3}, + {"num_chunks": 3}, + {"num_chunk": 3}, + ], + ) + def test_chunked_value_calls_match_unchunked( + self, estimator_cls, kwargs, shifted, deactivate_vmap, chunk_kwargs + ): torch.manual_seed(0) value_net = TensorDictModule( nn.Linear(3, 1), @@ -95,12 +106,14 @@ def test_value_chunk_size_matches_unchunked(self, estimator_cls, kwargs, shifted **kwargs, value_network=value_net, shifted=shifted, + deactivate_vmap=deactivate_vmap, ) chunked = estimator_cls( **kwargs, value_network=value_net, shifted=shifted, - value_chunk_size=3, + deactivate_vmap=deactivate_vmap, + **chunk_kwargs, ) expected = unchunked(td.clone()) @@ -108,6 +121,99 @@ def test_value_chunk_size_matches_unchunked(self, estimator_cls, kwargs, shifted torch.testing.assert_close(actual["advantage"], expected["advantage"]) torch.testing.assert_close(actual["value_target"], expected["value_target"]) + def test_value_chunk_options_are_mutually_exclusive(self): + value_net = TensorDictModule( + nn.Linear(3, 1), + in_keys=["obs"], + out_keys=["state_value"], + ) + with pytest.raises(ValueError, match="value_chunk_size and num_chunks"): + GAE( + gamma=0.9, + lmbda=0.95, + value_network=value_net, + value_chunk_size=3, + num_chunks=2, + ) + with pytest.raises(ValueError, match="num_chunks and num_chunk"): + GAE( + gamma=0.9, + lmbda=0.95, + value_network=value_net, + num_chunks=2, + num_chunk=3, + ) + + @pytest.mark.parametrize("vectorized", [False, True]) + @pytest.mark.parametrize("deactivate_vmap", [False, True]) + @pytest.mark.parametrize("method", ["forward", "value_estimate"]) + @pytest.mark.parametrize( + "chunk_kwargs", + [ + {"value_chunk_size": 3}, + {"num_chunks": 3}, + {"num_chunk": 3}, + ], + ) + def test_gae_chunked_functional_calls_match_unchunked( + self, vectorized, deactivate_vmap, method, chunk_kwargs + ): + torch.manual_seed(0) + value_net = TensorDictModule( + nn.Linear(3, 1), + in_keys=["obs"], + out_keys=["state_value"], + ) + td = TensorDict( + { + "obs": torch.randn(4, 5, 3), + "next": { + "obs": torch.randn(4, 5, 3), + "reward": torch.randn(4, 5, 1), + "done": torch.zeros(4, 5, 1, dtype=torch.bool), + "terminated": torch.zeros(4, 5, 1, dtype=torch.bool), + }, + }, + [4, 5], + ) + td["next", "done"][:, -1] = True + td["next", "terminated"][:, -1] = True + params = TensorDict.from_module(value_net) + target_params = params.clone(False) + + unchunked = GAE( + gamma=0.9, + lmbda=0.95, + value_network=value_net, + vectorized=vectorized, + deactivate_vmap=deactivate_vmap, + ) + chunked = GAE( + gamma=0.9, + lmbda=0.95, + value_network=value_net, + vectorized=vectorized, + deactivate_vmap=deactivate_vmap, + **chunk_kwargs, + ) + if method == "forward": + expected = unchunked( + td.clone(), params=params, target_params=target_params + ) + actual = chunked(td.clone(), params=params, target_params=target_params) + torch.testing.assert_close(actual["advantage"], expected["advantage"]) + torch.testing.assert_close( + actual["value_target"], expected["value_target"] + ) + else: + expected = unchunked.value_estimate( + td.clone(), params=params, target_params=target_params + ) + actual = chunked.value_estimate( + td.clone(), params=params, target_params=target_params + ) + torch.testing.assert_close(actual, expected) + @pytest.mark.parametrize( "estimator_cls,kwargs", [ diff --git a/torchrl/objectives/value/advantages.py b/torchrl/objectives/value/advantages.py index 07f3a285719..99eb567216f 100644 --- a/torchrl/objectives/value/advantages.py +++ b/torchrl/objectives/value/advantages.py @@ -238,6 +238,8 @@ def __init__( device: torch.device | None = None, deactivate_vmap: bool = False, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__() @@ -249,7 +251,16 @@ def __init__( self._tensor_keys = None self.differentiable = differentiable self.deactivate_vmap = deactivate_vmap - self.value_chunk_size = value_chunk_size + self.value_chunk_size = self._check_positive_int( + value_chunk_size, "value_chunk_size" + ) + self.num_chunks = self._resolve_num_chunks(num_chunks, num_chunk) + if self.value_chunk_size is not None and self.num_chunks is not None: + raise ValueError( + "value_chunk_size and num_chunks cannot both be set. " + "Use value_chunk_size to specify a chunk size or num_chunks " + "to specify a number of chunks." + ) if shifted_budget < 1: raise ValueError(f"shifted_budget must be >= 1, got {shifted_budget}.") self.shifted_budget = shifted_budget @@ -271,6 +282,43 @@ def __init__( "Setting 'value_key' via constructor is deprecated, use .set_keys(value_key='some_key') instead.", ) + @staticmethod + def _check_positive_int(value: int | None, name: str) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer or None, got {type(value)}.") + if value < 1: + raise ValueError(f"{name} must be >= 1, got {value}.") + return value + + @classmethod + def _resolve_num_chunks( + cls, num_chunks: int | None, num_chunk: int | None + ) -> int | None: + num_chunks = cls._check_positive_int(num_chunks, "num_chunks") + num_chunk = cls._check_positive_int(num_chunk, "num_chunk") + if num_chunk is None: + return num_chunks + if num_chunks is not None and num_chunks != num_chunk: + raise ValueError( + "num_chunks and num_chunk cannot both be set to different values." + ) + return num_chunk + + def _split_value_net_input( + self, data: TensorDictBase + ) -> tuple[TensorDictBase, ...]: + if self.num_chunks is not None: + num_chunks = min(self.num_chunks, data.shape[0]) + if num_chunks == 1: + return (data,) + return tuple(data.chunk(num_chunks, dim=0)) + chunk_size = self.value_chunk_size + if chunk_size is None or data.numel() <= chunk_size: + return (data,) + return tuple(data.split(chunk_size, dim=0)) + @property def tensor_keys(self) -> _AcceptedKeys: if self._tensor_keys is None: @@ -675,11 +723,11 @@ def _call_value_nets( data = self._sanitize_next_obs_nan(data, in_keys) def _call_value_net(data_in: TensorDictBase) -> torch.Tensor: - chunk_size = self.value_chunk_size - if chunk_size is None or data_in.numel() <= chunk_size: + chunks = self._split_value_net_input(data_in) + if len(chunks) == 1: return value_net(data_in).get(value_key) values = [] - for chunk in data_in.split(chunk_size, dim=0): + for chunk in chunks: values.append(value_net(chunk).get(value_key)) return torch.cat(values, dim=0) @@ -713,10 +761,31 @@ def _call_value_net(data_in: TensorDictBase) -> torch.Tensor: elif ( params is None and next_params is None - and self.value_chunk_size is not None + and (self.value_chunk_size is not None or self.num_chunks is not None) ): value = _call_value_net(data_root) value_ = _call_value_net(data_next) + elif params is not None and ( + self.value_chunk_size is not None or self.num_chunks is not None + ): + params_stack = torch.stack([params, next_params], 0).contiguous() + values = [] + next_values = [] + for root_chunk, next_chunk in zip( + self._split_value_net_input(data_root), + self._split_value_net_input(data_next), + ): + data_out = _vmap_func( + value_net, + (0, 0), + randomness=vmap_randomness, + pseudo_vmap=self.deactivate_vmap, + )(torch.stack([root_chunk, next_chunk], 0), params_stack) + value_est = data_out.get(value_key) + values.append(value_est[0]) + next_values.append(value_est[1]) + value = torch.cat(values, dim=0) + value_ = torch.cat(next_values, dim=0) elif params is not None: params_stack = torch.stack([params, next_params], 0).contiguous() data_out = _vmap_func( @@ -731,7 +800,7 @@ def _call_value_net(data_in: TensorDictBase) -> torch.Tensor: data_out = _pseudo_vmap(value_net, (0,), randomness=vmap_randomness)( data_in ) - if self.value_chunk_size is None or params is not None: + if self.value_chunk_size is None and self.num_chunks is None: value_est = data_out.get(value_key) value, value_ = value_est[0], value_est[1] data.set(value_key, value) @@ -867,6 +936,10 @@ class TD0Estimator(ValueEstimatorBase): value_chunk_size (int, optional): if set, splits value-network calls into chunks of this many elements along the leading dimension. Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -889,6 +962,8 @@ def __init__( device: torch.device | None = None, deactivate_vmap: bool = False, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__( @@ -902,6 +977,8 @@ def __init__( device=device, deactivate_vmap=deactivate_vmap, value_chunk_size=value_chunk_size, + num_chunks=num_chunks, + num_chunk=num_chunk, shifted_budget=shifted_budget, ) self.register_buffer("gamma", torch.tensor(gamma, device=self._device)) @@ -1136,6 +1213,10 @@ class TD1Estimator(ValueEstimatorBase): value_chunk_size (int, optional): if set, splits value-network calls into chunks of this many elements along the leading dimension. Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1159,6 +1240,8 @@ def __init__( time_dim: int | None = None, deactivate_vmap: bool = False, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__( @@ -1172,6 +1255,8 @@ def __init__( device=device, deactivate_vmap=deactivate_vmap, value_chunk_size=value_chunk_size, + num_chunks=num_chunks, + num_chunk=num_chunk, shifted_budget=shifted_budget, ) self.register_buffer("gamma", torch.tensor(gamma, device=self._device)) @@ -1417,6 +1502,10 @@ class TDLambdaEstimator(ValueEstimatorBase): value_chunk_size (int, optional): if set, splits value-network calls into chunks of this many elements along the leading dimension. Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1442,6 +1531,8 @@ def __init__( time_dim: int | None = None, deactivate_vmap: bool = False, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__( @@ -1455,6 +1546,8 @@ def __init__( device=device, deactivate_vmap=deactivate_vmap, value_chunk_size=value_chunk_size, + num_chunks=num_chunks, + num_chunk=num_chunk, shifted_budget=shifted_budget, ) self.register_buffer("gamma", torch.tensor(gamma, device=self._device)) @@ -1738,6 +1831,10 @@ class GAE(ValueEstimatorBase): value_chunk_size (int, optional): if set, splits value-network calls into chunks of this many elements along the leading dimension. Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1788,6 +1885,8 @@ def __init__( auto_reset_env: bool = False, deactivate_vmap: bool = False, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__( @@ -1801,6 +1900,8 @@ def __init__( device=device, deactivate_vmap=deactivate_vmap, value_chunk_size=value_chunk_size, + num_chunks=num_chunks, + num_chunk=num_chunk, shifted_budget=shifted_budget, ) self.register_buffer( @@ -2323,6 +2424,10 @@ class VTrace(ValueEstimatorBase): value_chunk_size (int, optional): if set, splits value-network calls into chunks of this many elements along the leading dimension. Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -2356,6 +2461,8 @@ def __init__( device: torch.device | None = None, time_dim: int | None = None, value_chunk_size: int | None = None, + num_chunks: int | None = None, + num_chunk: int | None = None, shifted_budget: int = 1, ): super().__init__( @@ -2368,6 +2475,8 @@ def __init__( skip_existing=skip_existing, device=device, value_chunk_size=value_chunk_size, + num_chunks=num_chunks, + num_chunk=num_chunk, shifted_budget=shifted_budget, ) if not isinstance(gamma, torch.Tensor): diff --git a/torchrl/trainers/algorithms/configs/objectives.py b/torchrl/trainers/algorithms/configs/objectives.py index d02e13b9f21..9177836ccad 100644 --- a/torchrl/trainers/algorithms/configs/objectives.py +++ b/torchrl/trainers/algorithms/configs/objectives.py @@ -278,6 +278,10 @@ class GAEConfig(LossConfig): time_dim: int | None = None auto_reset_env: bool = False deactivate_vmap: bool = False + value_chunk_size: int | None = None + num_chunks: int | None = None + num_chunk: int | None = None + shifted_budget: int = 1 _target_: str = "torchrl.objectives.value.GAE" _partial_: bool = False From 84e833334815fbfc9b133b52897a259babe97228 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 07:57:25 -0700 Subject: [PATCH 02/11] linter --- test/objectives/test_values.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/test/objectives/test_values.py b/test/objectives/test_values.py index 9f2909890a9..41baec97ce6 100644 --- a/test/objectives/test_values.py +++ b/test/objectives/test_values.py @@ -197,14 +197,10 @@ def test_gae_chunked_functional_calls_match_unchunked( **chunk_kwargs, ) if method == "forward": - expected = unchunked( - td.clone(), params=params, target_params=target_params - ) + expected = unchunked(td.clone(), params=params, target_params=target_params) actual = chunked(td.clone(), params=params, target_params=target_params) torch.testing.assert_close(actual["advantage"], expected["advantage"]) - torch.testing.assert_close( - actual["value_target"], expected["value_target"] - ) + torch.testing.assert_close(actual["value_target"], expected["value_target"]) else: expected = unchunked.value_estimate( td.clone(), params=params, target_params=target_params From 0db0c6f78791b6187c6e838157ef2d2f23a8ddd7 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 08:57:13 -0700 Subject: [PATCH 03/11] init --- test/objectives/test_values.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/objectives/test_values.py b/test/objectives/test_values.py index 41baec97ce6..9bdd2aecbd8 100644 --- a/test/objectives/test_values.py +++ b/test/objectives/test_values.py @@ -81,6 +81,8 @@ class TestValues: def test_chunked_value_calls_match_unchunked( self, estimator_cls, kwargs, shifted, deactivate_vmap, chunk_kwargs ): + if deactivate_vmap and _TORCH_VERSION < version.parse("2.7"): + pytest.skip("_pseudo_vmap is not supported for torch<2.7") torch.manual_seed(0) value_net = TensorDictModule( nn.Linear(3, 1), @@ -158,6 +160,8 @@ def test_value_chunk_options_are_mutually_exclusive(self): def test_gae_chunked_functional_calls_match_unchunked( self, vectorized, deactivate_vmap, method, chunk_kwargs ): + if deactivate_vmap and _TORCH_VERSION < version.parse("2.7"): + pytest.skip("_pseudo_vmap is not supported for torch<2.7") torch.manual_seed(0) value_net = TensorDictModule( nn.Linear(3, 1), From 1335f9adb062d9d04f4d17b77b507648d908e935 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 09:45:30 -0700 Subject: [PATCH 04/11] [CI] Fix olddeps and systemic workflow failures --- .github/scripts/pre-build-script-win.sh | 3 +- .github/scripts/pre-build-script.sh | 11 +++---- .github/scripts/td_script.sh | 3 +- .github/scripts/version_script.bat | 6 +++- .../linux_libs/scripts_gym/run_all.sh | 4 +-- .../linux_libs/scripts_habitat/install.sh | 12 ++++---- .../linux_libs/scripts_habitat/run_all.sh | 4 +-- .../scripts_meltingpot/environment.yml | 1 - .../linux_libs/scripts_meltingpot/install.sh | 7 +++-- .../linux_libs/scripts_minari/install.sh | 2 +- .../scripts_robohive/install_and_run_test.sh | 7 +++-- .../linux_libs/scripts_vmas/environment.yml | 2 -- .../linux_libs/scripts_vmas/install.sh | 9 ++++-- .../linux_olddeps/scripts_gym_0_13/install.sh | 8 +++-- .../unittest/linux_optdeps/scripts/run_all.sh | 7 +++-- .github/unittest/llm/scripts_llm/install.sh | 7 +++-- .github/unittest/llm/scripts_llm/run_test.sh | 2 +- .../unittest/llm/scripts_sglang/install.sh | 13 ++++++++- .github/workflows/benchmarks.yml | 5 ++-- .github/workflows/benchmarks_pr.yml | 7 +++-- .../workflows/build-wheels-aarch64-linux.yml | 1 + .github/workflows/build-wheels-linux.yml | 1 + .github/workflows/build-wheels-m1.yml | 1 + .github/workflows/build-wheels-windows.yml | 1 + .github/workflows/docs.yml | 8 +++-- .github/workflows/nightly_build.yml | 11 +++++-- .github/workflows/test-linux-habitat.yml | 2 +- docs/source/reference/objectives_common.rst | 1 + packaging/build_wheels.sh | 7 ++--- packaging/prepare_nightly_build.py | 8 ++--- packaging/verify_nightly_version.py | 4 ++- setup.py | 4 ++- test/objectives/test_loss_module.py | 12 +++++++- test/objectives/test_values.py | 17 +++++++++++ test/test_collectors.py | 9 ++++-- torchrl/objectives/value/advantages.py | 29 ++++++++++++++++++- 36 files changed, 164 insertions(+), 72 deletions(-) diff --git a/.github/scripts/pre-build-script-win.sh b/.github/scripts/pre-build-script-win.sh index 8bd2f61b28d..69405c8d73e 100644 --- a/.github/scripts/pre-build-script-win.sh +++ b/.github/scripts/pre-build-script-win.sh @@ -2,6 +2,7 @@ pip install --upgrade setuptools packaging ${CONDA_RUN} pip install "pybind11[global]" +${CONDA_RUN} pip install cloudpickle importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" # Determine tensordict installation source based on branch/tag # - release/* branches or release tags: use PyPI stable release @@ -28,5 +29,3 @@ install_tensordict() { } install_tensordict - -export TORCHRL_BUILD_VERSION=0.10.0 diff --git a/.github/scripts/pre-build-script.sh b/.github/scripts/pre-build-script.sh index 4662784fcbb..523099b5538 100644 --- a/.github/scripts/pre-build-script.sh +++ b/.github/scripts/pre-build-script.sh @@ -6,6 +6,7 @@ ${CONDA_RUN} pip install --upgrade setuptools packaging ${CONDA_RUN} pip install "pybind11[global]" +${CONDA_RUN} pip install cloudpickle importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" ${CONDA_RUN} conda install anaconda::cmake -y # Determine tensordict installation source based on branch/tag @@ -16,19 +17,19 @@ install_tensordict() { if [[ "$source" == "stable" ]]; then echo "Installing tensordict from PyPI (stable) - explicit override" - ${CONDA_RUN} pip install tensordict -U + ${CONDA_RUN} pip install tensordict -U --no-deps elif [[ "$source" == "git" ]]; then echo "Installing tensordict from git - explicit override" - ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U + ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U --no-deps elif [[ "$GITHUB_REF_TYPE" == "branch" && "$GITHUB_REF_NAME" == release/* ]]; then echo "Installing tensordict from PyPI (stable) - detected release branch: $GITHUB_REF_NAME" - ${CONDA_RUN} pip install tensordict -U + ${CONDA_RUN} pip install tensordict -U --no-deps elif [[ "$GITHUB_REF_TYPE" == "tag" && "$GITHUB_REF_NAME" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then echo "Installing tensordict from PyPI (stable) - detected release tag: $GITHUB_REF_NAME" - ${CONDA_RUN} pip install tensordict -U + ${CONDA_RUN} pip install tensordict -U --no-deps else echo "Installing tensordict from git - branch: ${GITHUB_REF_NAME:-unknown}, type: ${GITHUB_REF_TYPE:-unknown}" - ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U + ${CONDA_RUN} pip install git+https://github.com/pytorch/tensordict.git -U --no-deps fi } diff --git a/.github/scripts/td_script.sh b/.github/scripts/td_script.sh index 57735501dd7..7a298e23764 100644 --- a/.github/scripts/td_script.sh +++ b/.github/scripts/td_script.sh @@ -1,10 +1,11 @@ #!/bin/bash -export TORCHRL_BUILD_VERSION=0.12.0 +export TORCHRL_BUILD_VERSION="${BUILD_VERSION:-0.12.0}" ${CONDA_RUN} pip install --upgrade setuptools packaging # Always install pybind11 - required for building C++ extensions ${CONDA_RUN} pip install "pybind11[global]" +${CONDA_RUN} pip install cloudpickle importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" # Check if ARCH is set to aarch64 ARCH=${ARCH:-} # This sets ARCH to an empty string if it's not defined diff --git a/.github/scripts/version_script.bat b/.github/scripts/version_script.bat index d9cb0b6ec72..b2b109760a4 100644 --- a/.github/scripts/version_script.bat +++ b/.github/scripts/version_script.bat @@ -1,5 +1,9 @@ @echo off -set TORCHRL_BUILD_VERSION=0.12.0 +if "%BUILD_VERSION%" == "" ( + set TORCHRL_BUILD_VERSION=0.12.0 +) else ( + set TORCHRL_BUILD_VERSION=%BUILD_VERSION% +) echo TORCHRL_BUILD_VERSION is set to %TORCHRL_BUILD_VERSION% @echo on diff --git a/.github/unittest/linux_libs/scripts_gym/run_all.sh b/.github/unittest/linux_libs/scripts_gym/run_all.sh index 3be9c927959..7c8d5784707 100755 --- a/.github/unittest/linux_libs/scripts_gym/run_all.sh +++ b/.github/unittest/linux_libs/scripts_gym/run_all.sh @@ -74,10 +74,10 @@ export pybind11_DIR printf "* Installing tensordict\n" if [[ "$RELEASE" == 0 ]]; then # Install tensordict dependencies (since we use --no-deps) - uv pip install cloudpickle packaging importlib_metadata orjson "pyvers>=0.1.0,<0.2.0" + uv pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" uv pip install --no-build-isolation --no-deps git+https://github.com/pytorch/tensordict.git else - uv pip install cloudpickle packaging importlib_metadata orjson "pyvers>=0.1.0,<0.2.0" + uv pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" uv pip install --no-deps tensordict fi diff --git a/.github/unittest/linux_libs/scripts_habitat/install.sh b/.github/unittest/linux_libs/scripts_habitat/install.sh index 44b79764951..30352cf5e88 100755 --- a/.github/unittest/linux_libs/scripts_habitat/install.sh +++ b/.github/unittest/linux_libs/scripts_habitat/install.sh @@ -21,24 +21,24 @@ git submodule sync && git submodule update --init --recursive printf "Installing PyTorch with %s\n" "${CU_VERSION}" if [[ "$TORCH_VERSION" == "nightly" ]]; then - pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 -U + pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/"${CU_VERSION}" -U elif [[ "$TORCH_VERSION" == "stable" ]]; then - pip3 install torch --index-url https://download.pytorch.org/whl/cu128 + pip3 install torch --index-url https://download.pytorch.org/whl/"${CU_VERSION}" fi # install tensordict -# install tensordict +pip3 install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then - pip3 install git+https://github.com/pytorch/tensordict.git + pip3 install --no-deps git+https://github.com/pytorch/tensordict.git else - pip3 install tensordict + pip3 install --no-deps tensordict fi # smoke test python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_habitat/run_all.sh b/.github/unittest/linux_libs/scripts_habitat/run_all.sh index f5318539a13..61e74f2e445 100755 --- a/.github/unittest/linux_libs/scripts_habitat/run_all.sh +++ b/.github/unittest/linux_libs/scripts_habitat/run_all.sh @@ -4,7 +4,7 @@ set -euxo pipefail set -v -apt-get update && apt-get upgrade -y +apt-get update apt-get install -y vim git wget cmake ninja-build # OpenGL/EGL dependencies for headless rendering @@ -19,8 +19,6 @@ apt-get install -y pkg-config #apt-get upgrade -y libstdc++6 #apt-get install -y libgcc -apt-get dist-upgrade -y - this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" # from cudagl docker image cp $this_dir/10_nvidia.json /usr/share/glvnd/egl_vendor.d/10_nvidia.json diff --git a/.github/unittest/linux_libs/scripts_meltingpot/environment.yml b/.github/unittest/linux_libs/scripts_meltingpot/environment.yml index e003419348c..5d0301d9518 100644 --- a/.github/unittest/linux_libs/scripts_meltingpot/environment.yml +++ b/.github/unittest/linux_libs/scripts_meltingpot/environment.yml @@ -5,7 +5,6 @@ dependencies: - pip - pip: - cloudpickle - - torch - pytest - pytest-cov - pytest-mock diff --git a/.github/unittest/linux_libs/scripts_meltingpot/install.sh b/.github/unittest/linux_libs/scripts_meltingpot/install.sh index c31a17e8fdd..93a5763f901 100755 --- a/.github/unittest/linux_libs/scripts_meltingpot/install.sh +++ b/.github/unittest/linux_libs/scripts_meltingpot/install.sh @@ -44,17 +44,18 @@ else fi # install tensordict +pip3 install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then - pip3 install git+https://github.com/pytorch/tensordict.git + pip3 install --no-deps git+https://github.com/pytorch/tensordict.git else - pip3 install tensordict + pip3 install --no-deps tensordict fi # smoke test python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" conda install conda-forge::jq -y diff --git a/.github/unittest/linux_libs/scripts_minari/install.sh b/.github/unittest/linux_libs/scripts_minari/install.sh index c8f55da7582..4480f9b8435 100755 --- a/.github/unittest/linux_libs/scripts_minari/install.sh +++ b/.github/unittest/linux_libs/scripts_minari/install.sh @@ -43,7 +43,7 @@ else fi # install tensordict and its deps -uv pip install cloudpickle packaging importlib_metadata orjson "pyvers>=0.1.0,<0.2.0" +uv pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then uv pip install --no-deps git+https://github.com/pytorch/tensordict.git else diff --git a/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh b/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh index a8e1523dc22..619d7236255 100755 --- a/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh +++ b/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh @@ -57,17 +57,18 @@ fi conda install -y cmake # install tensordict +pip3 install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then - pip3 install git+https://github.com/pytorch/tensordict.git + pip3 install --no-deps git+https://github.com/pytorch/tensordict.git else - pip3 install tensordict + pip3 install --no-deps tensordict fi # smoke test python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" # Extracted from run_test.sh to run once. diff --git a/.github/unittest/linux_libs/scripts_vmas/environment.yml b/.github/unittest/linux_libs/scripts_vmas/environment.yml index 58c4a00400c..231a8b22bb8 100644 --- a/.github/unittest/linux_libs/scripts_vmas/environment.yml +++ b/.github/unittest/linux_libs/scripts_vmas/environment.yml @@ -10,8 +10,6 @@ dependencies: - numpy - pyglet==1.5.27 - six - - torch - - vmas - zipp - pytest - pytest-cov diff --git a/.github/unittest/linux_libs/scripts_vmas/install.sh b/.github/unittest/linux_libs/scripts_vmas/install.sh index 1b717a16e41..0017a0f19de 100755 --- a/.github/unittest/linux_libs/scripts_vmas/install.sh +++ b/.github/unittest/linux_libs/scripts_vmas/install.sh @@ -44,15 +44,18 @@ else fi # install tensordict +pip3 install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then - pip3 install git+https://github.com/pytorch/tensordict.git + pip3 install --no-deps git+https://github.com/pytorch/tensordict.git else - pip3 install tensordict + pip3 install --no-deps tensordict fi +pip3 install vmas + # smoke test python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" diff --git a/.github/unittest/linux_olddeps/scripts_gym_0_13/install.sh b/.github/unittest/linux_olddeps/scripts_gym_0_13/install.sh index abc9572124e..474eb4977a2 100755 --- a/.github/unittest/linux_olddeps/scripts_gym_0_13/install.sh +++ b/.github/unittest/linux_olddeps/scripts_gym_0_13/install.sh @@ -50,9 +50,11 @@ fi if [[ "$RELEASE" == 0 ]]; then # pybind11 headers are required to build tensordict's C++ extension. python -m pip install "pybind11[global]" - python -m pip install git+https://github.com/pytorch/tensordict.git + python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" + python -m pip install --no-deps git+https://github.com/pytorch/tensordict.git else - python -m pip install tensordict + python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" + python -m pip install --no-deps tensordict fi # smoke test @@ -63,5 +65,5 @@ python -c "import tensordict; print(f'tensordict: {tensordict.__version__}')" python -m pip install "setuptools<82" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" diff --git a/.github/unittest/linux_optdeps/scripts/run_all.sh b/.github/unittest/linux_optdeps/scripts/run_all.sh index 1277681090b..9dce18e7545 100755 --- a/.github/unittest/linux_optdeps/scripts/run_all.sh +++ b/.github/unittest/linux_optdeps/scripts/run_all.sh @@ -135,14 +135,15 @@ python -c "import functorch" #fi # install tensordict +pip3 install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then - pip3 install git+https://github.com/pytorch/tensordict.git + pip3 install --no-deps git+https://github.com/pytorch/tensordict.git else - pip3 install tensordict + pip3 install --no-deps tensordict fi printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/llm/scripts_llm/install.sh b/.github/unittest/llm/scripts_llm/install.sh index 5c821e15c87..6ed03e64f2e 100644 --- a/.github/unittest/llm/scripts_llm/install.sh +++ b/.github/unittest/llm/scripts_llm/install.sh @@ -46,18 +46,19 @@ git submodule sync && git submodule update --init --recursive #fi # install tensordict +pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" if [[ "$RELEASE" == 0 ]]; then pip install "pybind11[global]" ninja - pip install git+https://github.com/pytorch/tensordict.git + pip install --no-deps git+https://github.com/pytorch/tensordict.git else - pip install tensordict + pip install --no-deps tensordict fi # smoke test python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/llm/scripts_llm/run_test.sh b/.github/unittest/llm/scripts_llm/run_test.sh index b1c785c3a69..a6abf752d09 100644 --- a/.github/unittest/llm/scripts_llm/run_test.sh +++ b/.github/unittest/llm/scripts_llm/run_test.sh @@ -6,7 +6,7 @@ eval "$(./conda/bin/conda shell.bash hook)" conda activate ./env apt-get update && apt-get install -y git gcc cmake -ln -s /usr/bin/swig3.0 /usr/bin/swig +ln -s /usr/bin/swig3.0 /usr/bin/swig 2>/dev/null || true export PYTORCH_TEST_WITH_SLOW='1' export LAZY_LEGACY_OP=False diff --git a/.github/unittest/llm/scripts_sglang/install.sh b/.github/unittest/llm/scripts_sglang/install.sh index 9aa49497028..7dee1b5ae87 100755 --- a/.github/unittest/llm/scripts_sglang/install.sh +++ b/.github/unittest/llm/scripts_sglang/install.sh @@ -68,7 +68,7 @@ fi printf "* Installing tensordict\n" # Install tensordict dependencies first (pyvers is required but --no-deps skips it) -uv pip install cloudpickle packaging importlib_metadata orjson "pyvers>=0.1.0,<0.2.0" +uv pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" uv pip install "pybind11[global]" ninja if [[ "$RELEASE" == 0 ]]; then uv pip install --no-build-isolation --no-deps git+https://github.com/pytorch/tensordict.git @@ -120,6 +120,17 @@ export PATH="$HOME/.deno/bin:$PATH" # Install mcp uv pip install mcp langdetect +# SGLang may resolve a backend-specific torch/triton stack. Reinstall +# TensorDict and TorchRL after that resolution so native extensions are built +# against the final torch wheel present in the environment. +printf "* Reinstalling TensorDict and TorchRL against final backend stack\n" +if [[ "$RELEASE" == 0 ]]; then + uv pip install --reinstall --no-build-isolation --no-deps git+https://github.com/pytorch/tensordict.git +else + uv pip install --reinstall --no-deps tensordict +fi +uv pip install --reinstall -e . --no-build-isolation --no-deps + # Verify installations deno --version || echo "Warning: Deno not installed" diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index ff14531b6d5..8fbf828a3d8 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -87,9 +87,10 @@ jobs: python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 -U python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers accelerate python -m pip install "pybind11[global]" - python3.10 -m pip install git+https://github.com/pytorch/tensordict + python3.10 -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" + python3.10 -m pip install --no-deps git+https://github.com/pytorch/tensordict python3.10 -m pip install safetensors tqdm pandas numpy matplotlib ray - python3.10 setup.py develop + python3.10 -m pip install -e . --no-build-isolation --no-deps bash .github/unittest/helpers/assert_torch_version.sh nightly diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml index 83219234b9b..d2c16c858f6 100644 --- a/.github/workflows/benchmarks_pr.yml +++ b/.github/workflows/benchmarks_pr.yml @@ -88,7 +88,8 @@ jobs: python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 -U python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers accelerate ray python3.10 -m pip install "pybind11[global]" - python3.10 -m pip install git+https://github.com/pytorch/tensordict + python3.10 -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" + python3.10 -m pip install --no-deps git+https://github.com/pytorch/tensordict python3.10 -m pip install safetensors tqdm pandas numpy matplotlib bash .github/unittest/helpers/assert_torch_version.sh nightly @@ -104,7 +105,7 @@ jobs: git fetch --no-tags --depth=1 "https://github.com/${PR_BASE_REPO}.git" "${PR_BASE_SHA}" git checkout --detach "${PR_BASE_SHA}" rm -rf build - python3.10 setup.py develop + python3.10 -m pip install -e . --no-build-isolation --no-deps if [ "${{ matrix.device }}" = "GPU" ]; then # test import and fail early if the GPU runner did not expose CUDA @@ -130,7 +131,7 @@ jobs: fi git checkout --detach "${PR_HEAD_SHA}" rm -rf build - python3.10 setup.py develop + python3.10 -m pip install -e . --no-build-isolation --no-deps if [ "${{ matrix.device }}" = "GPU" ]; then # test import and fail early if the GPU runner did not expose CUDA diff --git a/.github/workflows/build-wheels-aarch64-linux.yml b/.github/workflows/build-wheels-aarch64-linux.yml index df71fd5ce5a..5ea3a5482dd 100644 --- a/.github/workflows/build-wheels-aarch64-linux.yml +++ b/.github/workflows/build-wheels-aarch64-linux.yml @@ -64,6 +64,7 @@ jobs: with-cpu: ${{ inputs.with-cpu || (github.event_name == 'workflow_dispatch' && (inputs.build-cpu && 'enable' || 'disable')) || 'enable' }} channel: ${{ inputs.channel || '' }} use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }} + python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]' build: needs: generate-matrix strategy: diff --git a/.github/workflows/build-wheels-linux.yml b/.github/workflows/build-wheels-linux.yml index 615bf184b48..00e3f00e9d4 100644 --- a/.github/workflows/build-wheels-linux.yml +++ b/.github/workflows/build-wheels-linux.yml @@ -81,6 +81,7 @@ jobs: test-infra-ref: ${{ inputs.test-infra-ref || 'main' }} channel: ${{ inputs.channel || '' }} use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }} + python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]' # For workflow_dispatch: convert boolean to enable/disable string # For workflow_call: use the string input directly # Default: build Linux CUDA wheels now that prioritized replay has CUDA kernels. diff --git a/.github/workflows/build-wheels-m1.yml b/.github/workflows/build-wheels-m1.yml index 690e2878967..2d5749a3cc4 100644 --- a/.github/workflows/build-wheels-m1.yml +++ b/.github/workflows/build-wheels-m1.yml @@ -61,6 +61,7 @@ jobs: test-infra-ref: ${{ inputs.test-infra-ref || 'main' }} channel: ${{ inputs.channel || '' }} use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }} + python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]' # macOS only supports CPU builds with-cpu: ${{ inputs.with-cpu || (github.event_name == 'workflow_dispatch' && (inputs.build-cpu && 'enable' || 'disable')) || 'enable' }} build: diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml index ae414c3251c..375476e1e84 100644 --- a/.github/workflows/build-wheels-windows.yml +++ b/.github/workflows/build-wheels-windows.yml @@ -71,6 +71,7 @@ jobs: test-infra-ref: ${{ inputs.test-infra-ref || 'main' }} channel: ${{ inputs.channel || '' }} use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }} + python-versions: '["3.10", "3.11", "3.12", "3.13", "3.14"]' # For workflow_dispatch: convert boolean to enable/disable string # For workflow_call: use the string input directly # Default: disable CUDA (torchrl is CPU-only), enable CPU (Windows has no ROCm) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index cc4ecd3cb93..014e8f3dee0 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -66,12 +66,14 @@ jobs: if [[ ${{ github.event_name }} == push && (${{ github.ref_type }} == tag || (${{ github.ref_type }} == branch && ${{ github.ref_name }} == release/*)) ]]; then # Use stable PyTorch for releases (tags and release/* branches) python -m pip install torch torchvision --quiet --root-user-action=ignore - python -m pip install tensordict --quiet --root-user-action=ignore + python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" --quiet --root-user-action=ignore + python -m pip install tensordict --no-deps --quiet --root-user-action=ignore python -m pip install torchcodec --quiet --root-user-action=ignore else # Use nightly PyTorch for main, nightly, and PRs python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U --quiet --root-user-action=ignore - python -m pip install git+https://github.com/pytorch/tensordict.git --quiet --root-user-action=ignore + python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" --quiet --root-user-action=ignore + python -m pip install git+https://github.com/pytorch/tensordict.git --no-deps --quiet --root-user-action=ignore torchcodec_dir=$(mktemp -d) git clone --depth 1 https://github.com/pytorch/torchcodec.git "$torchcodec_dir" python_base="$(python -c 'import sys; print(sys.base_prefix)')" @@ -86,7 +88,7 @@ jobs: python -m pip install -r docs/requirements.txt --quiet --root-user-action=ignore # 7. Install TorchRL - python -m pip install -e . --no-build-isolation + python -m pip install -e . --no-build-isolation --no-deps # 8. Assert PyTorch version matches expectations if [[ ${{ github.event_name }} == push && (${{ github.ref_type }} == tag || (${{ github.ref_type }} == branch && ${{ github.ref_name }} == release/*)) ]]; then diff --git a/.github/workflows/nightly_build.yml b/.github/workflows/nightly_build.yml index 9e098419a11..018df463564 100644 --- a/.github/workflows/nightly_build.yml +++ b/.github/workflows/nightly_build.yml @@ -68,9 +68,6 @@ jobs: export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" python -mpip install build wheel ./build_nightly.sh - find dist -name '*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: Upload wheel for the test-wheel job uses: actions/upload-artifact@v4 with: @@ -186,6 +183,10 @@ jobs: env: PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }} run: | + if [ -z "${PYPI_TOKEN:-}" ]; then + echo "PYPI_TOKEN is not set; skipping PyPI upload." + exit 0 + fi export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" python -mpip install twine python -m twine upload \ @@ -335,6 +336,10 @@ jobs: env: PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }} run: | + if [ -z "${PYPI_TOKEN:-}" ]; then + echo "PYPI_TOKEN is not set; skipping PyPI upload." + exit 0 + fi python -mpip install twine python -m twine upload \ --username __token__ \ diff --git a/.github/workflows/test-linux-habitat.yml b/.github/workflows/test-linux-habitat.yml index bc686fe5fd3..580849525e9 100644 --- a/.github/workflows/test-linux-habitat.yml +++ b/.github/workflows/test-linux-habitat.yml @@ -32,7 +32,7 @@ jobs: with: runner: linux.g5.12xlarge.nvidia.gpu repository: pytorch/rl - docker-image: "nvidia/cuda:12.1.1-devel-ubuntu22.04" + docker-image: "nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04" gpu-arch-type: cuda gpu-arch-version: ${{ matrix.cuda_arch_version }} timeout: 90 diff --git a/docs/source/reference/objectives_common.rst b/docs/source/reference/objectives_common.rst index 3bdf05f174f..ceb111c6969 100644 --- a/docs/source/reference/objectives_common.rst +++ b/docs/source/reference/objectives_common.rst @@ -28,6 +28,7 @@ Value Estimators TD1Estimator TDLambdaEstimator GAE + VTrace MultiAgentGAE .. currentmodule:: torchrl.objectives diff --git a/packaging/build_wheels.sh b/packaging/build_wheels.sh index 6ad2cc96754..09b7e2a6821 100755 --- a/packaging/build_wheels.sh +++ b/packaging/build_wheels.sh @@ -28,10 +28,7 @@ else fi if [[ "$OSTYPE" == "msys" ]]; then - IS_WHEEL=1 "$script_dir/windows/internal/vc_env_helper.bat" python -m build --wheel + IS_WHEEL=1 "$script_dir/windows/internal/vc_env_helper.bat" python -m build --wheel --no-isolation else - python -m build --wheel - if [[ "$(uname)" != Darwin ]]; then - rename "linux_x86_64" "manylinux1_x86_64" dist/*.whl - fi + python -m build --wheel --no-isolation fi diff --git a/packaging/prepare_nightly_build.py b/packaging/prepare_nightly_build.py index 4206e7c30aa..f7c56992b33 100644 --- a/packaging/prepare_nightly_build.py +++ b/packaging/prepare_nightly_build.py @@ -24,11 +24,11 @@ def prepare_nightly_build(): with open(pyproject_path) as f: content = f.read() - # Replace tensordict dependency with tensordict-nightly using regex - # This pattern matches "tensordict" followed by any version constraints - tensordict_pattern = r"tensordict[^,\]]*" + # Replace only the quoted project dependency, not tool sections such as + # ``[tool.uv.sources]``. + tensordict_pattern = r'"tensordict[^"]*"' if re.search(tensordict_pattern, content): - content = re.sub(tensordict_pattern, "tensordict-nightly", content) + content = re.sub(tensordict_pattern, '"tensordict-nightly"', content) logger.info("Replaced tensordict with tensordict-nightly in pyproject.toml") else: logger.info("tensordict dependency not found in pyproject.toml") diff --git a/packaging/verify_nightly_version.py b/packaging/verify_nightly_version.py index df6807431c2..7b750669a7d 100644 --- a/packaging/verify_nightly_version.py +++ b/packaging/verify_nightly_version.py @@ -48,7 +48,9 @@ print(f"Checking tensordict version: {tensordict_version}") # Check if it's a nightly version (either date format or contains 'd' followed by date) - if re.match(date_pattern, tensordict_version) or "d2025" in tensordict_version: + if re.match(date_pattern, tensordict_version) or re.search( + r"d20\d{6}", tensordict_version + ): print(f"✓ tensordict version {tensordict_version} appears to be nightly") else: # Check if it's a stable version that should not be used in nightly builds diff --git a/setup.py b/setup.py index 24d08bec506..66bffcd8cec 100644 --- a/setup.py +++ b/setup.py @@ -341,7 +341,9 @@ def _version_with_local_sha(base_version: str) -> str: def set_version(): # Prefer explicit build version if provided by build tooling. if "SETUPTOOLS_SCM_PRETEND_VERSION" not in os.environ: - override = os.environ.get("TORCHRL_BUILD_VERSION") + override = os.environ.get("TORCHRL_BUILD_VERSION") or os.environ.get( + "BUILD_VERSION" + ) if override: os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = override.strip() else: diff --git a/test/objectives/test_loss_module.py b/test/objectives/test_loss_module.py index 1b3e1e958b2..01cc2c1f554 100644 --- a/test/objectives/test_loss_module.py +++ b/test/objectives/test_loss_module.py @@ -337,7 +337,17 @@ def _forward_value_estimator_keys(self, **kwargs) -> None: ) ) + [ - ["cuda", torch.float, "cuda:0"], + pytest.param( + "cuda", + torch.float, + "cuda:0", + marks=[ + pytest.mark.gpu, + pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" + ), + ], + ), ["double", torch.double, "cpu"], [torch.double, torch.double, "cpu"], [torch.half, torch.half, "cpu"], diff --git a/test/objectives/test_values.py b/test/objectives/test_values.py index 9bdd2aecbd8..1649c4cdf64 100644 --- a/test/objectives/test_values.py +++ b/test/objectives/test_values.py @@ -146,6 +146,23 @@ def test_value_chunk_options_are_mutually_exclusive(self): num_chunk=3, ) + def test_num_chunks_splits_into_requested_number_of_chunks(self): + value_net = TensorDictModule( + nn.Linear(3, 1), + in_keys=["obs"], + out_keys=["state_value"], + ) + estimator = GAE( + gamma=0.9, + lmbda=0.95, + value_network=value_net, + num_chunks=3, + ) + td = TensorDict({"obs": torch.randn(4, 3)}, [4]) + assert len(estimator._split_value_net_input(td)) == 3 + empty_td = TensorDict({"obs": torch.randn(0, 3)}, [0]) + assert estimator._split_value_net_input(empty_td) == (empty_td,) + @pytest.mark.parametrize("vectorized", [False, True]) @pytest.mark.parametrize("deactivate_vmap", [False, True]) @pytest.mark.parametrize("method", ["forward", "value_estimate"]) diff --git a/test/test_collectors.py b/test/test_collectors.py index 78781824fad..9eaf666a982 100644 --- a/test/test_collectors.py +++ b/test/test_collectors.py @@ -6147,9 +6147,12 @@ def test_profile_config_should_profile_worker(self): False, pytest.param( True, - marks=pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA not available" - ), + marks=[ + pytest.mark.gpu, + pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" + ), + ], ), ], ) diff --git a/torchrl/objectives/value/advantages.py b/torchrl/objectives/value/advantages.py index 99eb567216f..81c67b77334 100644 --- a/torchrl/objectives/value/advantages.py +++ b/torchrl/objectives/value/advantages.py @@ -103,6 +103,21 @@ class ValueEstimatorBase(TensorDictModuleBase): If only the value estimate is needed, the :meth:`ValueFunctionBase.value_estimate` should be used instead. + Keyword Args: + value_chunk_size (int, optional): if set, splits value-network calls + into chunks of this many elements along the leading dimension. + Defaults to ``None``. + num_chunks (int, optional): if set, splits value-network calls into + this many chunks along the leading dimension. Mutually exclusive + with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. + Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. + shifted_budget (int, optional): number of extra value-network time slots + used when ``shifted=True``. ``1`` uses a ``T+1`` + budget, ``2`` can represent one internal reset plus the rollout + boundary without dropping samples, and so on. Defaults to ``1``. + """ @dataclass @@ -310,10 +325,12 @@ def _split_value_net_input( self, data: TensorDictBase ) -> tuple[TensorDictBase, ...]: if self.num_chunks is not None: + if data.shape[0] == 0: + return (data,) num_chunks = min(self.num_chunks, data.shape[0]) if num_chunks == 1: return (data,) - return tuple(data.chunk(num_chunks, dim=0)) + return tuple(data.tensor_split(num_chunks, dim=0)) chunk_size = self.value_chunk_size if chunk_size is None or data.numel() <= chunk_size: return (data,) @@ -940,6 +957,8 @@ class TD0Estimator(ValueEstimatorBase): this many chunks along the leading dimension. Mutually exclusive with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1217,6 +1236,8 @@ class TD1Estimator(ValueEstimatorBase): this many chunks along the leading dimension. Mutually exclusive with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1506,6 +1527,8 @@ class TDLambdaEstimator(ValueEstimatorBase): this many chunks along the leading dimension. Mutually exclusive with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -1835,6 +1858,8 @@ class GAE(ValueEstimatorBase): this many chunks along the leading dimension. Mutually exclusive with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout @@ -2428,6 +2453,8 @@ class VTrace(ValueEstimatorBase): this many chunks along the leading dimension. Mutually exclusive with ``value_chunk_size``. ``num_chunk`` is accepted as an alias. Defaults to ``None``. + num_chunk (int, optional): alias for ``num_chunks``. Cannot be set + together with a different ``num_chunks`` value. Defaults to ``None``. shifted_budget (int, optional): number of extra value-network time slots used when ``shifted=True``. ``1`` uses a ``T+1`` budget, ``2`` can represent one internal reset plus the rollout From d0eb34ac8d96d2d52af1916e01f96073dc500aa4 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 10:07:55 -0700 Subject: [PATCH 05/11] [CI] Pin MuJoCo below 3.9 for dm_control --- .github/unittest/linux/scripts/run_all.sh | 8 ++++---- .github/unittest/tutorials/scripts/run_all.sh | 4 ++-- .github/workflows/benchmarks.yml | 2 +- .github/workflows/benchmarks_pr.yml | 2 +- docs/requirements.txt | 4 ++-- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/unittest/linux/scripts/run_all.sh b/.github/unittest/linux/scripts/run_all.sh index 8d6475f42c0..976ddc3175b 100755 --- a/.github/unittest/linux/scripts/run_all.sh +++ b/.github/unittest/linux/scripts/run_all.sh @@ -140,7 +140,7 @@ uv_pip_install \ # labmaze (dm_control dependency) doesn't have Python 3.13+ wheels if [[ "$PYTHON_VERSION" != "3.13" && "$PYTHON_VERSION" != "3.14" ]]; then echo "installing dm_control" - uv_pip_install dm_control + uv_pip_install "dm_control>=1.0.41" "mujoco>=3.8.1,<3.9.0" fi # Install ray for Python < 3.14 (ray doesn't support Python 3.14 yet) @@ -158,7 +158,7 @@ fi # Install mujoco for Python < 3.14 (mujoco doesn't have Python 3.14 wheels yet) if [[ "$PYTHON_VERSION" != "3.14" ]]; then echo "installing mujoco" - uv_pip_install "mujoco>=3.3.7" + uv_pip_install "mujoco>=3.8.1,<3.9.0" fi # Install gymnasium @@ -168,9 +168,9 @@ if [[ "$PYTHON_VERSION" == "3.14" ]]; then uv_pip_install "gymnasium>=1.1" elif [[ "$PYTHON_VERSION" == "3.12" ]]; then uv_pip_install ale-py sympy - uv_pip_install "gymnasium[mujoco]>=1.1" "mo-gymnasium[mujoco]" + uv_pip_install "gymnasium[mujoco]>=1.1" "mo-gymnasium[mujoco]" "mujoco>=3.8.1,<3.9.0" else - uv_pip_install "gymnasium[atari,mujoco]>=1.1" "mo-gymnasium[mujoco]" + uv_pip_install "gymnasium[atari,mujoco]>=1.1" "mo-gymnasium[mujoco]" "mujoco>=3.8.1,<3.9.0" fi # sanity check diff --git a/.github/unittest/tutorials/scripts/run_all.sh b/.github/unittest/tutorials/scripts/run_all.sh index 4e4b8a267b1..bb77eb35862 100755 --- a/.github/unittest/tutorials/scripts/run_all.sh +++ b/.github/unittest/tutorials/scripts/run_all.sh @@ -144,8 +144,8 @@ uv pip install --no-progress \ psutil \ hydra-core \ "imageio==2.26.0" \ - dm_control \ - mujoco \ + "dm_control>=1.0.41" \ + "mujoco>=3.8.1,<3.9.0" \ av \ coverage \ vmas \ diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 8fbf828a3d8..c2cb503c047 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -85,7 +85,7 @@ jobs: fi python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 -U - python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers accelerate + python3.10 -m pip install ninja pytest pytest-benchmark "mujoco>=3.8.1,<3.9.0" "dm_control>=1.0.41" "gym[accept-rom-license,atari]" transformers accelerate python -m pip install "pybind11[global]" python3.10 -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" python3.10 -m pip install --no-deps git+https://github.com/pytorch/tensordict diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml index d2c16c858f6..8b79e8ec017 100644 --- a/.github/workflows/benchmarks_pr.yml +++ b/.github/workflows/benchmarks_pr.yml @@ -86,7 +86,7 @@ jobs: export PYTHON_INCLUDE_DIR=/usr/include/python3.10 python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 -U - python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers accelerate ray + python3.10 -m pip install ninja pytest pytest-benchmark "mujoco>=3.8.1,<3.9.0" "dm_control>=1.0.41" "gym[accept-rom-license,atari]" transformers accelerate ray python3.10 -m pip install "pybind11[global]" python3.10 -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" python3.10 -m pip install --no-deps git+https://github.com/pytorch/tensordict diff --git a/docs/requirements.txt b/docs/requirements.txt index c566cc44fc4..04defeea58f 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -14,8 +14,8 @@ docutils sphinx_design torchvision -dm_control>=1.0.0 -mujoco>=3.0.0 +dm_control>=1.0.41 +mujoco>=3.8.1,<3.9.0 gymnasium[classic_control,atari,mujoco] pygame tqdm From 2c9ff427739fb165476b3852345b0350888f7681 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 11:05:18 -0700 Subject: [PATCH 06/11] [CI] Stabilize Windows and optional dependency installs --- .github/unittest/linux_libs/scripts_habitat/setup_env.sh | 2 +- .github/unittest/linux_libs/scripts_robohive/setup_env.sh | 2 +- .github/unittest/windows_optdepts/scripts/unittest.sh | 7 ++++--- test/transforms/test_observation_transforms.py | 6 ++---- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh index 9312eab0180..9c00c5c60af 100755 --- a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh +++ b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh @@ -80,7 +80,7 @@ conda env update --file "${this_dir}/environment.yml" --prune # 5. Install habitat-sim from source (conda packages don't support Python 3.10+) # Install build dependencies -pip3 install ninja numpy +pip3 install ninja numpy scikit-build-core # Clone and build habitat-sim from source cd "${root_dir}" diff --git a/.github/unittest/linux_libs/scripts_robohive/setup_env.sh b/.github/unittest/linux_libs/scripts_robohive/setup_env.sh index 684b2f061bb..abf1633e3af 100755 --- a/.github/unittest/linux_libs/scripts_robohive/setup_env.sh +++ b/.github/unittest/linux_libs/scripts_robohive/setup_env.sh @@ -82,7 +82,7 @@ conda env update --file "${this_dir}/environment.yml" --prune conda install conda-forge::ffmpeg -y -pip install robohive +pip install robohive --no-deps python -m robohive_init diff --git a/.github/unittest/windows_optdepts/scripts/unittest.sh b/.github/unittest/windows_optdepts/scripts/unittest.sh index 61b77fe45e6..ee851890791 100755 --- a/.github/unittest/windows_optdepts/scripts/unittest.sh +++ b/.github/unittest/windows_optdepts/scripts/unittest.sh @@ -102,13 +102,14 @@ fi echo "=== Installing build dependencies ===" conda install anaconda::cmake -y python -m pip install "pybind11[global]" +python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.0,<0.3.0" # install tensordict echo "=== Installing tensordict ===" if [[ "$RELEASE" == 0 ]]; then - python -m pip install git+https://github.com/pytorch/tensordict + python -m pip install --no-deps git+https://github.com/pytorch/tensordict else - pip3 install tensordict + python -m pip install --no-deps tensordict fi # smoke test @@ -122,7 +123,7 @@ echo "=== Setting up CUDA environment ===" source "$this_dir/set_cuda_envs.sh" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps whatsinside=$(ls -rtlh ./torchrl) echo $whatsinside diff --git a/test/transforms/test_observation_transforms.py b/test/transforms/test_observation_transforms.py index d0bb6c930fa..9a44fa84b51 100644 --- a/test/transforms/test_observation_transforms.py +++ b/test/transforms/test_observation_transforms.py @@ -2851,11 +2851,9 @@ def test_trans_serial_env_check(self): ) self._smoke_one_step(env) - def test_trans_parallel_env_check(self): + def test_trans_parallel_env_check(self, maybe_fork_ParallelEnv): env = TransformedEnv( - ParallelEnv( - 2, lambda: ContinuousActionVecMockEnv(), mp_start_method="fork" - ), + maybe_fork_ParallelEnv(2, lambda: ContinuousActionVecMockEnv()), NextObservationDelta(in_keys=["observation"]), ) try: From 0afb6449f1a422b39f9b6b8af92cc022e327f250 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 11:13:30 -0700 Subject: [PATCH 07/11] [CI] Assert TensorDict source before tests --- .../assert_torch_tensordict_versions.sh | 87 +++++++++++++++++++ .github/unittest/linux/scripts/run_all.sh | 1 + .../unittest/linux/scripts/run_setup_test.sh | 1 + .../linux_libs/scripts_brax/run_all.sh | 1 + .../linux_libs/scripts_habitat/run_all.sh | 1 + .../linux_libs/scripts_minari/run_all.sh | 1 + .../scripts_robohive/install_and_run_test.sh | 1 + .../scripts_gym_0_13/batch_scripts.sh | 1 + .../unittest/linux_optdeps/scripts/run_all.sh | 1 + .../unittest/linux_sota/scripts/run_all.sh | 1 + .github/unittest/tutorials/scripts/run_all.sh | 2 + .../windows_optdepts/scripts/unittest.sh | 1 + .github/workflows/benchmarks.yml | 1 + .github/workflows/benchmarks_pr.yml | 1 + .github/workflows/docs.yml | 2 + .github/workflows/nightly_build.yml | 2 + .github/workflows/test-linux-libs.yml | 16 ++++ .github/workflows/test-linux-llm.yml | 12 +++ .../workflows/validate-test-partitioning.yml | 1 + 19 files changed, 134 insertions(+) create mode 100755 .github/unittest/helpers/assert_torch_tensordict_versions.sh diff --git a/.github/unittest/helpers/assert_torch_tensordict_versions.sh b/.github/unittest/helpers/assert_torch_tensordict_versions.sh new file mode 100755 index 00000000000..9acf65fbfd2 --- /dev/null +++ b/.github/unittest/helpers/assert_torch_tensordict_versions.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +# Assert that the installed PyTorch and TensorDict builds match the CI mode. +# +# Usage: +# bash assert_torch_tensordict_versions.sh nightly [auto|main|stable] +# bash assert_torch_tensordict_versions.sh stable [auto|main|stable] +# +# By default (auto), nightly PyTorch expects TensorDict from main/git and stable +# PyTorch expects a stable TensorDict release. Some jobs intentionally pin stable +# PyTorch while testing TensorDict main; pass "main" explicitly for those jobs. +set -euo pipefail + +expected_torch="${1:?Usage: assert_torch_tensordict_versions.sh nightly|stable [auto|main|stable]}" +expected_tensordict="${2:-auto}" + +PYTHON="${PYTHON:-$(command -v python3 || command -v python)}" + +"$PYTHON" - "$expected_torch" "$expected_tensordict" <<'PY' +from __future__ import annotations + +import importlib.metadata +import re +import sys + +import tensordict +import torch + +expected_torch = sys.argv[1] +expected_tensordict = sys.argv[2] + +if expected_torch not in {"nightly", "stable"}: + raise RuntimeError( + "Expected PyTorch mode must be 'nightly' or 'stable', " + f"got {expected_torch!r}." + ) + +if expected_tensordict == "auto": + expected_tensordict = "main" if expected_torch == "nightly" else "stable" +elif expected_tensordict not in {"main", "stable"}: + raise RuntimeError( + "Expected TensorDict source must be 'auto', 'main', or 'stable', " + f"got {expected_tensordict!r}." + ) + +torch_version = torch.__version__ +torch_is_nightly = "dev" in torch_version +expected_torch_is_nightly = expected_torch == "nightly" +if torch_is_nightly != expected_torch_is_nightly: + raise RuntimeError( + f"PyTorch version mismatch: expected {expected_torch} but got " + f"torch {torch_version}." + ) + +try: + tensordict_version = tensordict.__version__ +except AttributeError: + tensordict_version = importlib.metadata.version("tensordict") + +try: + tensordict_dist = importlib.metadata.distribution("tensordict") + tensordict_direct_url = tensordict_dist.read_text("direct_url.json") or "" +except importlib.metadata.PackageNotFoundError: + tensordict_direct_url = "" + +tensordict_from_main = bool( + re.search(r"(?:\.dev|\+g|dev)", tensordict_version) + or "github.com/pytorch/tensordict" in tensordict_direct_url +) + +if expected_tensordict == "main" and not tensordict_from_main: + raise RuntimeError( + "TensorDict source mismatch: expected TensorDict from main/git but got " + f"tensordict {tensordict_version}." + ) + +if expected_tensordict == "stable" and tensordict_from_main: + raise RuntimeError( + "TensorDict source mismatch: expected a stable TensorDict release but got " + f"tensordict {tensordict_version}." + ) + +print(f"PyTorch version check OK: torch {torch_version} (expected {expected_torch})") +print( + "TensorDict source check OK: " + f"tensordict {tensordict_version} (expected {expected_tensordict})" +) +PY diff --git a/.github/unittest/linux/scripts/run_all.sh b/.github/unittest/linux/scripts/run_all.sh index 976ddc3175b..dd9d5c06288 100755 --- a/.github/unittest/linux/scripts/run_all.sh +++ b/.github/unittest/linux/scripts/run_all.sh @@ -325,6 +325,7 @@ export PYTORCH_TEST_WITH_SLOW='1' python -m torch.utils.collect_env bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" Xvfb :99 -screen 0 1024x768x24 & diff --git a/.github/unittest/linux/scripts/run_setup_test.sh b/.github/unittest/linux/scripts/run_setup_test.sh index 195f52b0ae9..daacaadc304 100644 --- a/.github/unittest/linux/scripts/run_setup_test.sh +++ b/.github/unittest/linux/scripts/run_setup_test.sh @@ -80,5 +80,6 @@ else fi bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" pytest -q test/test_setup.py -vv diff --git a/.github/unittest/linux_libs/scripts_brax/run_all.sh b/.github/unittest/linux_libs/scripts_brax/run_all.sh index 550c6f0bcc5..dc1b333d193 100755 --- a/.github/unittest/linux_libs/scripts_brax/run_all.sh +++ b/.github/unittest/linux_libs/scripts_brax/run_all.sh @@ -9,5 +9,6 @@ this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" bash ${this_dir}/setup_env.sh bash ${this_dir}/install.sh PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" bash ${this_dir}/run_test.sh bash ${this_dir}/post_process.sh diff --git a/.github/unittest/linux_libs/scripts_habitat/run_all.sh b/.github/unittest/linux_libs/scripts_habitat/run_all.sh index 61e74f2e445..9d9c7f1c0d4 100755 --- a/.github/unittest/linux_libs/scripts_habitat/run_all.sh +++ b/.github/unittest/linux_libs/scripts_habitat/run_all.sh @@ -27,6 +27,7 @@ bash ${this_dir}/setup_env.sh bash ${this_dir}/install.sh PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" bash ${this_dir}/run_test.sh bash ${this_dir}/post_process.sh diff --git a/.github/unittest/linux_libs/scripts_minari/run_all.sh b/.github/unittest/linux_libs/scripts_minari/run_all.sh index 1b26e655f9f..d6ccac46796 100755 --- a/.github/unittest/linux_libs/scripts_minari/run_all.sh +++ b/.github/unittest/linux_libs/scripts_minari/run_all.sh @@ -11,6 +11,7 @@ echo "Running minari tests with uv-based setup..." source "${this_dir}/setup_env.sh" source "${this_dir}/install.sh" PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" source "${this_dir}/run_test.sh" source "${this_dir}/post_process.sh" diff --git a/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh b/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh index 619d7236255..3dddcf3b51b 100755 --- a/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh +++ b/.github/unittest/linux_libs/scripts_robohive/install_and_run_test.sh @@ -81,6 +81,7 @@ git config --global --add safe.directory '*' root_dir="$(git rev-parse --show-toplevel)" bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" env_dir="${root_dir}/env" lib_dir="${env_dir}/lib" diff --git a/.github/unittest/linux_olddeps/scripts_gym_0_13/batch_scripts.sh b/.github/unittest/linux_olddeps/scripts_gym_0_13/batch_scripts.sh index 431ec985593..3e0cd8ddd62 100755 --- a/.github/unittest/linux_olddeps/scripts_gym_0_13/batch_scripts.sh +++ b/.github/unittest/linux_olddeps/scripts_gym_0_13/batch_scripts.sh @@ -7,4 +7,5 @@ DIR="$(cd "$(dirname "$0")" && pwd)" $DIR/install.sh PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" stable +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" stable main $DIR/run_test.sh diff --git a/.github/unittest/linux_optdeps/scripts/run_all.sh b/.github/unittest/linux_optdeps/scripts/run_all.sh index 9dce18e7545..5408abf17a3 100755 --- a/.github/unittest/linux_optdeps/scripts/run_all.sh +++ b/.github/unittest/linux_optdeps/scripts/run_all.sh @@ -160,6 +160,7 @@ export LAZY_LEGACY_OP=False python -m torch.utils.collect_env bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" # Avoid error: "fatal: unsafe repository" git config --global --add safe.directory '*' root_dir="$(git rev-parse --show-toplevel)" diff --git a/.github/unittest/linux_sota/scripts/run_all.sh b/.github/unittest/linux_sota/scripts/run_all.sh index 9fc97fac6ec..d68b4a0f67e 100755 --- a/.github/unittest/linux_sota/scripts/run_all.sh +++ b/.github/unittest/linux_sota/scripts/run_all.sh @@ -148,6 +148,7 @@ uv pip install -e . --no-build-isolation --no-deps # ================================ Run tests ========================================= # bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" bash ${this_dir}/run_test.sh diff --git a/.github/unittest/tutorials/scripts/run_all.sh b/.github/unittest/tutorials/scripts/run_all.sh index bb77eb35862..496a8bc1bfd 100755 --- a/.github/unittest/tutorials/scripts/run_all.sh +++ b/.github/unittest/tutorials/scripts/run_all.sh @@ -171,6 +171,8 @@ uv pip install --no-progress -e . --no-build-isolation --no-deps # ==================================================================================== # # ================================ Run tests ========================================= # +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" + bash ${this_dir}/run_test.sh # ==================================================================================== # diff --git a/.github/unittest/windows_optdepts/scripts/unittest.sh b/.github/unittest/windows_optdepts/scripts/unittest.sh index ee851890791..0b3e091194c 100755 --- a/.github/unittest/windows_optdepts/scripts/unittest.sh +++ b/.github/unittest/windows_optdepts/scripts/unittest.sh @@ -150,6 +150,7 @@ echo "=== Collecting environment info ===" python -m torch.utils.collect_env bash "${root_dir}/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" +bash "${root_dir}/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" echo "=== Starting pytest execution ===" echo "Current working directory: $(pwd)" diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index c2cb503c047..c967f21eb30 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -93,6 +93,7 @@ jobs: python3.10 -m pip install -e . --no-build-isolation --no-deps bash .github/unittest/helpers/assert_torch_version.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly if [ "${{ matrix.device }}" = "GPU" ]; then # test import and fail early if the GPU runner did not expose CUDA diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml index 8b79e8ec017..72e5de509f5 100644 --- a/.github/workflows/benchmarks_pr.yml +++ b/.github/workflows/benchmarks_pr.yml @@ -93,6 +93,7 @@ jobs: python3.10 -m pip install safetensors tqdm pandas numpy matplotlib bash .github/unittest/helpers/assert_torch_version.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly - name: Run baseline benchmarks run: | set -euxo pipefail diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 014e8f3dee0..6bb9939ea86 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -93,8 +93,10 @@ jobs: # 8. Assert PyTorch version matches expectations if [[ ${{ github.event_name }} == push && (${{ github.ref_type }} == tag || (${{ github.ref_type }} == branch && ${{ github.ref_name }} == release/*)) ]]; then bash .github/unittest/helpers/assert_torch_version.sh stable + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh stable else bash .github/unittest/helpers/assert_torch_version.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly fi # 9. Set sanitize version diff --git a/.github/workflows/nightly_build.yml b/.github/workflows/nightly_build.yml index 018df463564..e8ab6521390 100644 --- a/.github/workflows/nightly_build.yml +++ b/.github/workflows/nightly_build.yml @@ -145,6 +145,7 @@ jobs: export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" python -m torch.utils.collect_env bash .github/unittest/helpers/assert_torch_version.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly python -c "import torchrl; print(torchrl.__version__);from torchrl.data import ReplayBuffer" EXIT_STATUS=0 pytest test/smoke_test.py -v --durations 200 @@ -303,6 +304,7 @@ jobs: mkdir test-reports python -m torch.utils.collect_env bash .github/unittest/helpers/assert_torch_version.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly python -c "import torchrl; print(torchrl.__version__);from torchrl.data import ReplayBuffer" EXIT_STATUS=0 pytest test/smoke_test.py -v --durations 200 diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml index 1a246200820..ed64e649d30 100644 --- a/.github/workflows/test-linux-libs.yml +++ b/.github/workflows/test-linux-libs.yml @@ -199,6 +199,7 @@ jobs: bash .github/unittest/linux_libs/scripts_envpool/setup_env.sh bash .github/unittest/linux_libs/scripts_envpool/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_envpool/run_test.sh bash .github/unittest/linux_libs/scripts_envpool/post_process.sh @@ -234,6 +235,7 @@ jobs: bash .github/unittest/linux_libs/scripts_gen-dgrl/setup_env.sh bash .github/unittest/linux_libs/scripts_gen-dgrl/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_gen-dgrl/run_test.sh bash .github/unittest/linux_libs/scripts_gen-dgrl/post_process.sh @@ -325,6 +327,7 @@ jobs: bash .github/unittest/linux_libs/scripts_jumanji/setup_env.sh bash .github/unittest/linux_libs/scripts_jumanji/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_jumanji/run_test.sh bash .github/unittest/linux_libs/scripts_jumanji/post_process.sh @@ -361,6 +364,7 @@ jobs: bash .github/unittest/linux_libs/scripts_meltingpot/setup_env.sh bash .github/unittest/linux_libs/scripts_meltingpot/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_meltingpot/run_test.sh bash .github/unittest/linux_libs/scripts_meltingpot/post_process.sh @@ -400,6 +404,7 @@ jobs: bash .github/unittest/linux_libs/scripts_open_spiel/setup_env.sh bash .github/unittest/linux_libs/scripts_open_spiel/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_open_spiel/run_test.sh bash .github/unittest/linux_libs/scripts_open_spiel/post_process.sh @@ -439,6 +444,7 @@ jobs: bash .github/unittest/linux_libs/scripts_chess/setup_env.sh bash .github/unittest/linux_libs/scripts_chess/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_chess/run_test.sh bash .github/unittest/linux_libs/scripts_chess/post_process.sh @@ -478,6 +484,7 @@ jobs: bash .github/unittest/linux_libs/scripts_unity_mlagents/setup_env.sh bash .github/unittest/linux_libs/scripts_unity_mlagents/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_unity_mlagents/run_test.sh bash .github/unittest/linux_libs/scripts_unity_mlagents/post_process.sh @@ -546,6 +553,7 @@ jobs: bash .github/unittest/linux_libs/scripts_openx/setup_env.sh bash .github/unittest/linux_libs/scripts_openx/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_openx/run_test.sh bash .github/unittest/linux_libs/scripts_openx/post_process.sh @@ -582,6 +590,7 @@ jobs: bash .github/unittest/linux_libs/scripts_pettingzoo/setup_env.sh bash .github/unittest/linux_libs/scripts_pettingzoo/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_pettingzoo/run_test.sh bash .github/unittest/linux_libs/scripts_pettingzoo/post_process.sh @@ -622,6 +631,7 @@ jobs: bash .github/unittest/linux_libs/scripts_procgen/setup_env.sh bash .github/unittest/linux_libs/scripts_procgen/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_procgen/run_test.sh bash .github/unittest/linux_libs/scripts_procgen/post_process.sh @@ -662,6 +672,7 @@ jobs: bash .github/unittest/linux_libs/scripts_safety_gymnasium/setup_env.sh bash .github/unittest/linux_libs/scripts_safety_gymnasium/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_safety_gymnasium/run_test.sh bash .github/unittest/linux_libs/scripts_safety_gymnasium/post_process.sh @@ -732,6 +743,7 @@ jobs: bash .github/unittest/linux_libs/scripts_roboset/setup_env.sh bash .github/unittest/linux_libs/scripts_roboset/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_roboset/run_test.sh bash .github/unittest/linux_libs/scripts_roboset/post_process.sh @@ -769,6 +781,7 @@ jobs: bash .github/unittest/linux_libs/scripts_sklearn/setup_env.sh bash .github/unittest/linux_libs/scripts_sklearn/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_sklearn/run_test.sh bash .github/unittest/linux_libs/scripts_sklearn/post_process.sh @@ -809,6 +822,7 @@ jobs: bash .github/unittest/linux_libs/scripts_smacv2/setup_env.sh bash .github/unittest/linux_libs/scripts_smacv2/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_smacv2/run_test.sh bash .github/unittest/linux_libs/scripts_smacv2/post_process.sh @@ -845,6 +859,7 @@ jobs: bash .github/unittest/linux_libs/scripts_vd4rl/setup_env.sh bash .github/unittest/linux_libs/scripts_vd4rl/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_vd4rl/run_test.sh bash .github/unittest/linux_libs/scripts_vd4rl/post_process.sh @@ -885,6 +900,7 @@ jobs: bash .github/unittest/linux_libs/scripts_vmas/setup_env.sh bash .github/unittest/linux_libs/scripts_vmas/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" bash .github/unittest/linux_libs/scripts_vmas/run_test.sh bash .github/unittest/linux_libs/scripts_vmas/post_process.sh diff --git a/.github/workflows/test-linux-llm.yml b/.github/workflows/test-linux-llm.yml index fdf4e407eee..4ec7b831b4c 100644 --- a/.github/workflows/test-linux-llm.yml +++ b/.github/workflows/test-linux-llm.yml @@ -55,6 +55,12 @@ jobs: bash .github/unittest/llm/scripts_llm/setup_env.sh bash .github/unittest/llm/scripts_llm/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + if [[ "$RELEASE" == 0 ]]; then + TENSORDICT_EXPECTATION=main + else + TENSORDICT_EXPECTATION=stable + fi + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" "$TENSORDICT_EXPECTATION" bash .github/unittest/llm/scripts_llm/run_test.sh bash .github/unittest/llm/scripts_llm/post_process.sh @@ -94,5 +100,11 @@ jobs: bash .github/unittest/llm/scripts_sglang/setup_env.sh bash .github/unittest/llm/scripts_sglang/install.sh PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION" + if [[ "$RELEASE" == 0 ]]; then + TENSORDICT_EXPECTATION=main + else + TENSORDICT_EXPECTATION=stable + fi + PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION" "$TENSORDICT_EXPECTATION" bash .github/unittest/llm/scripts_sglang/run_test.sh bash .github/unittest/llm/scripts_sglang/post_process.sh diff --git a/.github/workflows/validate-test-partitioning.yml b/.github/workflows/validate-test-partitioning.yml index 182767ec2b4..9c7886e09ba 100644 --- a/.github/workflows/validate-test-partitioning.yml +++ b/.github/workflows/validate-test-partitioning.yml @@ -35,6 +35,7 @@ jobs: pip install pytest tensordict pip install torch --index-url https://download.pytorch.org/whl/cpu bash .github/unittest/helpers/assert_torch_version.sh stable + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh stable - name: Validate test partitioning env: From 0d4aaf81dea30accc0d4cf65af2803afcba99dee Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 11:56:23 -0700 Subject: [PATCH 08/11] [CI] Preserve TensorDict main in libs jobs --- .github/unittest/linux_libs/scripts_ataridqn/install.sh | 2 +- .github/unittest/linux_libs/scripts_botorch/install.sh | 2 +- .github/unittest/linux_libs/scripts_brax/install.sh | 2 +- .github/unittest/linux_libs/scripts_chess/install.sh | 2 +- .github/unittest/linux_libs/scripts_d4rl/install.sh | 2 +- .github/unittest/linux_libs/scripts_envpool/install.sh | 2 +- .github/unittest/linux_libs/scripts_gen-dgrl/install.sh | 2 +- .github/unittest/linux_libs/scripts_genesis/install.sh | 2 +- .github/unittest/linux_libs/scripts_habitat/run_all.sh | 1 + .github/unittest/linux_libs/scripts_jumanji/install.sh | 2 +- .github/unittest/linux_libs/scripts_open_spiel/install.sh | 2 +- .github/unittest/linux_libs/scripts_openx/install.sh | 2 +- .github/unittest/linux_libs/scripts_pettingzoo/install.sh | 2 +- .github/unittest/linux_libs/scripts_robohive/setup_env.sh | 1 + .github/unittest/linux_libs/scripts_roboset/install.sh | 2 +- .github/unittest/linux_libs/scripts_sklearn/install.sh | 2 +- .github/unittest/linux_libs/scripts_smacv2/install.sh | 2 +- .github/unittest/linux_libs/scripts_torch_geometric/install.sh | 2 +- .github/unittest/linux_libs/scripts_unity_mlagents/install.sh | 2 +- .github/unittest/linux_libs/scripts_vd4rl/install.sh | 2 +- 20 files changed, 20 insertions(+), 18 deletions(-) diff --git a/.github/unittest/linux_libs/scripts_ataridqn/install.sh b/.github/unittest/linux_libs/scripts_ataridqn/install.sh index d1f138b4ed2..aa37674e7e9 100755 --- a/.github/unittest/linux_libs/scripts_ataridqn/install.sh +++ b/.github/unittest/linux_libs/scripts_ataridqn/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_botorch/install.sh b/.github/unittest/linux_libs/scripts_botorch/install.sh index 395e15ea99e..b52fe75480b 100755 --- a/.github/unittest/linux_libs/scripts_botorch/install.sh +++ b/.github/unittest/linux_libs/scripts_botorch/install.sh @@ -47,7 +47,7 @@ pip install git+https://github.com/pytorch/tensordict.git --progress-bar off python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_brax/install.sh b/.github/unittest/linux_libs/scripts_brax/install.sh index e9fb769f494..3fa858dbe49 100755 --- a/.github/unittest/linux_libs/scripts_brax/install.sh +++ b/.github/unittest/linux_libs/scripts_brax/install.sh @@ -50,7 +50,7 @@ pip install git+https://github.com/pytorch/tensordict.git --progress-bar off python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_chess/install.sh b/.github/unittest/linux_libs/scripts_chess/install.sh index cd3200f3a1e..9b4bc780bfe 100755 --- a/.github/unittest/linux_libs/scripts_chess/install.sh +++ b/.github/unittest/linux_libs/scripts_chess/install.sh @@ -71,7 +71,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_d4rl/install.sh b/.github/unittest/linux_libs/scripts_d4rl/install.sh index 0b26fc87e66..83ef51228c5 100755 --- a/.github/unittest/linux_libs/scripts_d4rl/install.sh +++ b/.github/unittest/linux_libs/scripts_d4rl/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_envpool/install.sh b/.github/unittest/linux_libs/scripts_envpool/install.sh index 8235f3f1b59..3343da63074 100755 --- a/.github/unittest/linux_libs/scripts_envpool/install.sh +++ b/.github/unittest/linux_libs/scripts_envpool/install.sh @@ -52,4 +52,4 @@ else fi printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps diff --git a/.github/unittest/linux_libs/scripts_gen-dgrl/install.sh b/.github/unittest/linux_libs/scripts_gen-dgrl/install.sh index d1f138b4ed2..aa37674e7e9 100755 --- a/.github/unittest/linux_libs/scripts_gen-dgrl/install.sh +++ b/.github/unittest/linux_libs/scripts_gen-dgrl/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_genesis/install.sh b/.github/unittest/linux_libs/scripts_genesis/install.sh index 79e2c0be0b3..244bfb316fa 100644 --- a/.github/unittest/linux_libs/scripts_genesis/install.sh +++ b/.github/unittest/linux_libs/scripts_genesis/install.sh @@ -48,7 +48,7 @@ fi python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_habitat/run_all.sh b/.github/unittest/linux_libs/scripts_habitat/run_all.sh index 9d9c7f1c0d4..e21c3cb7816 100755 --- a/.github/unittest/linux_libs/scripts_habitat/run_all.sh +++ b/.github/unittest/linux_libs/scripts_habitat/run_all.sh @@ -9,6 +9,7 @@ apt-get install -y vim git wget cmake ninja-build # OpenGL/EGL dependencies for headless rendering apt-get install -y libglfw3 libglfw3-dev libgl1-mesa-glx libosmesa6 libosmesa6-dev libglew-dev +apt-get install -y libxinerama-dev libxcursor-dev libxi-dev libxrandr-dev libxxf86vm-dev apt-get install -y libglvnd0 libgl1 libglx0 libegl1 libgles2 apt-get install -y libegl1-mesa-dev libgles2-mesa-dev diff --git a/.github/unittest/linux_libs/scripts_jumanji/install.sh b/.github/unittest/linux_libs/scripts_jumanji/install.sh index 3302f15fba7..48896af5d36 100755 --- a/.github/unittest/linux_libs/scripts_jumanji/install.sh +++ b/.github/unittest/linux_libs/scripts_jumanji/install.sh @@ -54,7 +54,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_open_spiel/install.sh b/.github/unittest/linux_libs/scripts_open_spiel/install.sh index d1c63386b29..b058a3c6096 100755 --- a/.github/unittest/linux_libs/scripts_open_spiel/install.sh +++ b/.github/unittest/linux_libs/scripts_open_spiel/install.sh @@ -54,7 +54,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_openx/install.sh b/.github/unittest/linux_libs/scripts_openx/install.sh index eb1e1c60ba1..bea627bcb36 100755 --- a/.github/unittest/linux_libs/scripts_openx/install.sh +++ b/.github/unittest/linux_libs/scripts_openx/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_pettingzoo/install.sh b/.github/unittest/linux_libs/scripts_pettingzoo/install.sh index 8a71f05120c..34761b74606 100755 --- a/.github/unittest/linux_libs/scripts_pettingzoo/install.sh +++ b/.github/unittest/linux_libs/scripts_pettingzoo/install.sh @@ -54,5 +54,5 @@ fi python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_robohive/setup_env.sh b/.github/unittest/linux_libs/scripts_robohive/setup_env.sh index abf1633e3af..07f8996ff2d 100755 --- a/.github/unittest/linux_libs/scripts_robohive/setup_env.sh +++ b/.github/unittest/linux_libs/scripts_robohive/setup_env.sh @@ -83,6 +83,7 @@ conda env update --file "${this_dir}/environment.yml" --prune conda install conda-forge::ffmpeg -y pip install robohive --no-deps +pip install GitPython python -m robohive_init diff --git a/.github/unittest/linux_libs/scripts_roboset/install.sh b/.github/unittest/linux_libs/scripts_roboset/install.sh index 0b26fc87e66..83ef51228c5 100755 --- a/.github/unittest/linux_libs/scripts_roboset/install.sh +++ b/.github/unittest/linux_libs/scripts_roboset/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_sklearn/install.sh b/.github/unittest/linux_libs/scripts_sklearn/install.sh index 0b26fc87e66..83ef51228c5 100755 --- a/.github/unittest/linux_libs/scripts_sklearn/install.sh +++ b/.github/unittest/linux_libs/scripts_sklearn/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_smacv2/install.sh b/.github/unittest/linux_libs/scripts_smacv2/install.sh index 8a71f05120c..34761b74606 100755 --- a/.github/unittest/linux_libs/scripts_smacv2/install.sh +++ b/.github/unittest/linux_libs/scripts_smacv2/install.sh @@ -54,5 +54,5 @@ fi python -c "import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_torch_geometric/install.sh b/.github/unittest/linux_libs/scripts_torch_geometric/install.sh index 204cf25e3cb..3b0de9d876d 100755 --- a/.github/unittest/linux_libs/scripts_torch_geometric/install.sh +++ b/.github/unittest/linux_libs/scripts_torch_geometric/install.sh @@ -51,7 +51,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_unity_mlagents/install.sh b/.github/unittest/linux_libs/scripts_unity_mlagents/install.sh index d1c63386b29..b058a3c6096 100755 --- a/.github/unittest/linux_libs/scripts_unity_mlagents/install.sh +++ b/.github/unittest/linux_libs/scripts_unity_mlagents/install.sh @@ -54,7 +54,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" diff --git a/.github/unittest/linux_libs/scripts_vd4rl/install.sh b/.github/unittest/linux_libs/scripts_vd4rl/install.sh index eb1e1c60ba1..bea627bcb36 100755 --- a/.github/unittest/linux_libs/scripts_vd4rl/install.sh +++ b/.github/unittest/linux_libs/scripts_vd4rl/install.sh @@ -55,7 +55,7 @@ fi python -c "import functorch;import tensordict" printf "* Installing torchrl\n" -python -m pip install -e . --no-build-isolation +python -m pip install -e . --no-build-isolation --no-deps # smoke test python -c "import torchrl" From 6039be4995e74c24946be3dbede760943a4947a0 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 12:53:36 -0700 Subject: [PATCH 09/11] [CI] Allow nightly TensorDict in wheel smoke tests --- .../assert_torch_tensordict_versions.sh | 24 ++++++++++++++----- .../linux_libs/scripts_habitat/run_all.sh | 12 ++++++++-- .../scripts_robohive/environment.yml | 1 + .github/workflows/nightly_build.yml | 4 ++-- 4 files changed, 31 insertions(+), 10 deletions(-) diff --git a/.github/unittest/helpers/assert_torch_tensordict_versions.sh b/.github/unittest/helpers/assert_torch_tensordict_versions.sh index 9acf65fbfd2..ce47beea03a 100755 --- a/.github/unittest/helpers/assert_torch_tensordict_versions.sh +++ b/.github/unittest/helpers/assert_torch_tensordict_versions.sh @@ -2,15 +2,17 @@ # Assert that the installed PyTorch and TensorDict builds match the CI mode. # # Usage: -# bash assert_torch_tensordict_versions.sh nightly [auto|main|stable] -# bash assert_torch_tensordict_versions.sh stable [auto|main|stable] +# bash assert_torch_tensordict_versions.sh nightly [auto|main|nightly|stable] +# bash assert_torch_tensordict_versions.sh stable [auto|main|nightly|stable] # # By default (auto), nightly PyTorch expects TensorDict from main/git and stable # PyTorch expects a stable TensorDict release. Some jobs intentionally pin stable # PyTorch while testing TensorDict main; pass "main" explicitly for those jobs. +# Nightly wheel jobs install the dated TensorDict nightly package; pass "nightly" +# explicitly for those jobs. set -euo pipefail -expected_torch="${1:?Usage: assert_torch_tensordict_versions.sh nightly|stable [auto|main|stable]}" +expected_torch="${1:?Usage: assert_torch_tensordict_versions.sh nightly|stable [auto|main|nightly|stable]}" expected_tensordict="${2:-auto}" PYTHON="${PYTHON:-$(command -v python3 || command -v python)}" @@ -36,9 +38,10 @@ if expected_torch not in {"nightly", "stable"}: if expected_tensordict == "auto": expected_tensordict = "main" if expected_torch == "nightly" else "stable" -elif expected_tensordict not in {"main", "stable"}: +elif expected_tensordict not in {"main", "nightly", "stable"}: raise RuntimeError( - "Expected TensorDict source must be 'auto', 'main', or 'stable', " + "Expected TensorDict source must be 'auto', 'main', 'nightly', or " + "'stable', " f"got {expected_tensordict!r}." ) @@ -66,6 +69,7 @@ tensordict_from_main = bool( re.search(r"(?:\.dev|\+g|dev)", tensordict_version) or "github.com/pytorch/tensordict" in tensordict_direct_url ) +tensordict_from_nightly = bool(re.match(r"^\d{4}\.\d{2}\.\d{2}", tensordict_version)) if expected_tensordict == "main" and not tensordict_from_main: raise RuntimeError( @@ -73,7 +77,15 @@ if expected_tensordict == "main" and not tensordict_from_main: f"tensordict {tensordict_version}." ) -if expected_tensordict == "stable" and tensordict_from_main: +if expected_tensordict == "nightly" and not tensordict_from_nightly: + raise RuntimeError( + "TensorDict source mismatch: expected a TensorDict nightly release but " + f"got tensordict {tensordict_version}." + ) + +if expected_tensordict == "stable" and ( + tensordict_from_main or tensordict_from_nightly +): raise RuntimeError( "TensorDict source mismatch: expected a stable TensorDict release but got " f"tensordict {tensordict_version}." diff --git a/.github/unittest/linux_libs/scripts_habitat/run_all.sh b/.github/unittest/linux_libs/scripts_habitat/run_all.sh index e21c3cb7816..01df46cba0a 100755 --- a/.github/unittest/linux_libs/scripts_habitat/run_all.sh +++ b/.github/unittest/linux_libs/scripts_habitat/run_all.sh @@ -27,8 +27,16 @@ cp $this_dir/10_nvidia.json /usr/share/glvnd/egl_vendor.d/10_nvidia.json bash ${this_dir}/setup_env.sh bash ${this_dir}/install.sh -PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION" -PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION" +# habitat-lab currently resolves a stable PyTorch build after habitat-sim is +# installed. Keep validating TensorDict source, but accept stable PyTorch in +# this workflow even on non-release branches. +if [[ "$RELEASE" == 0 ]]; then + TENSORDICT_EXPECTATION=main +else + TENSORDICT_EXPECTATION=stable +fi +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" stable +PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" stable "$TENSORDICT_EXPECTATION" bash ${this_dir}/run_test.sh bash ${this_dir}/post_process.sh diff --git a/.github/unittest/linux_libs/scripts_robohive/environment.yml b/.github/unittest/linux_libs/scripts_robohive/environment.yml index a667b841525..753cb29e3f7 100644 --- a/.github/unittest/linux_libs/scripts_robohive/environment.yml +++ b/.github/unittest/linux_libs/scripts_robohive/environment.yml @@ -24,6 +24,7 @@ dependencies: - pybind11[global] - pyyaml - scipy + - flatten_dict - hydra-core - patchelf - mujoco==2.3.3 diff --git a/.github/workflows/nightly_build.yml b/.github/workflows/nightly_build.yml index e8ab6521390..2d6d12e554d 100644 --- a/.github/workflows/nightly_build.yml +++ b/.github/workflows/nightly_build.yml @@ -145,7 +145,7 @@ jobs: export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH" python -m torch.utils.collect_env bash .github/unittest/helpers/assert_torch_version.sh nightly - bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly nightly python -c "import torchrl; print(torchrl.__version__);from torchrl.data import ReplayBuffer" EXIT_STATUS=0 pytest test/smoke_test.py -v --durations 200 @@ -304,7 +304,7 @@ jobs: mkdir test-reports python -m torch.utils.collect_env bash .github/unittest/helpers/assert_torch_version.sh nightly - bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly + bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly nightly python -c "import torchrl; print(torchrl.__version__);from torchrl.data import ReplayBuffer" EXIT_STATUS=0 pytest test/smoke_test.py -v --durations 200 From bdf7ff0f36a24c62968ccb67036d6f2b2e5fbd7f Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 13:16:02 -0700 Subject: [PATCH 10/11] [CI] Tolerate missing nightly version metadata --- packaging/verify_nightly_version.py | 46 +++++++++++++++++++---------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/packaging/verify_nightly_version.py b/packaging/verify_nightly_version.py index 7b750669a7d..24f016bdf45 100644 --- a/packaging/verify_nightly_version.py +++ b/packaging/verify_nightly_version.py @@ -22,22 +22,30 @@ version = torchrl.__version__ print(f"Checking version: {version}") -# Check that version is not the major version (0.9.0) -if re.match(r"^\d+\.\d+\.\d+$", version): - raise ValueError(f"Version should not be the major version: {version}") - -# Check that version matches date format (YYYY.M.D) date_pattern = r"^\d{4}\.\d{1,2}\.\d{1,2}$" -if not re.match(date_pattern, version): - raise ValueError(f"Version should match date format YYYY.M.D, got: {version}") +if not isinstance(version, str) or not version: + print( + "WARNING: torchrl.__version__ is missing; skipping nightly version " + "format validation." + ) +else: + # Check that version is not the major version (0.9.0) + if re.match(r"^\d+\.\d+\.\d+$", version): + raise ValueError(f"Version should not be the major version: {version}") -# Verify it's today's date -today = date.today() -expected_version = f"{today.year}.{today.month}.{today.day}" -if version != expected_version: - raise ValueError(f"Version should be today date {expected_version}, got: {version}") + # Check that version matches date format (YYYY.M.D) + if not re.match(date_pattern, version): + raise ValueError(f"Version should match date format YYYY.M.D, got: {version}") -print(f"✓ Version {version} is correctly formatted as nightly date") + # Verify it's today's date + today = date.today() + expected_version = f"{today.year}.{today.month}.{today.day}" + if version != expected_version: + raise ValueError( + f"Version should be today date {expected_version}, got: {version}" + ) + + print(f"Version {version} is correctly formatted as nightly date") # Check that tensordict-nightly is installed (not stable tensordict) try: @@ -48,10 +56,15 @@ print(f"Checking tensordict version: {tensordict_version}") # Check if it's a nightly version (either date format or contains 'd' followed by date) - if re.match(date_pattern, tensordict_version) or re.search( + if not isinstance(tensordict_version, str) or not tensordict_version: + print( + "WARNING: tensordict.__version__ is missing; skipping nightly " + "version format validation." + ) + elif re.match(date_pattern, tensordict_version) or re.search( r"d20\d{6}", tensordict_version ): - print(f"✓ tensordict version {tensordict_version} appears to be nightly") + print(f"tensordict version {tensordict_version} appears to be nightly") else: # Check if it's a stable version that should not be used in nightly builds if tensordict_version.startswith("0.9."): @@ -59,7 +72,8 @@ f"tensordict should be nightly, not stable version: {tensordict_version}" ) print( - f"⚠ tensordict version {tensordict_version} - please verify this is nightly" + f"WARNING: tensordict version {tensordict_version} - please verify " + "this is nightly" ) except ImportError: From 15bc6b5aee7796381b86e1c1a83be8217014e3e7 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 27 May 2026 13:48:58 -0700 Subject: [PATCH 11/11] [CI] Stabilize RoboHive and RoboSet libs jobs --- .../linux_libs/scripts_robohive/environment.yml | 1 + .../unittest/linux_libs/scripts_roboset/setup_env.sh | 10 +++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/unittest/linux_libs/scripts_robohive/environment.yml b/.github/unittest/linux_libs/scripts_robohive/environment.yml index 753cb29e3f7..a911d4380c7 100644 --- a/.github/unittest/linux_libs/scripts_robohive/environment.yml +++ b/.github/unittest/linux_libs/scripts_robohive/environment.yml @@ -25,6 +25,7 @@ dependencies: - pyyaml - scipy - flatten_dict + - termcolor - hydra-core - patchelf - mujoco==2.3.3 diff --git a/.github/unittest/linux_libs/scripts_roboset/setup_env.sh b/.github/unittest/linux_libs/scripts_roboset/setup_env.sh index a7afd646646..c99398ed394 100755 --- a/.github/unittest/linux_libs/scripts_roboset/setup_env.sh +++ b/.github/unittest/linux_libs/scripts_roboset/setup_env.sh @@ -66,4 +66,12 @@ cat "${this_dir}/environment.yml" pip3 install pip --upgrade -conda env update --file "${this_dir}/environment.yml" --prune +for attempt in 1 2 3; do + if conda env update --file "${this_dir}/environment.yml" --prune; then + break + fi + if [ "${attempt}" = "3" ]; then + exit 1 + fi + sleep $((attempt * 30)) +done