[BugFix] Fix shape setting in CompositeSpec (#1620)

Co-authored-by: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>

Author: Vincent Moens

test/test_specs.py

@@ -651,6 +651,23 @@ def test_nested_composite_spec_update(self, shape, is_complete, device, dtype):
         }
         assert ts["nested_cp"]["act"] is not None
 
+    def test_change_batch_size(self, shape, is_complete, device, dtype):
+        ts = self._composite_spec(shape, is_complete, device, dtype)
+        ts["nested"] = CompositeSpec(
+            leaf=UnboundedContinuousTensorSpec(shape), shape=shape
+        )
+        ts = ts.expand(3, *shape)
+        assert ts["nested"].shape == (3, *shape)
+        assert ts["nested", "leaf"].shape == (3, *shape)
+        ts.shape = ()
+        # this does not change
+        assert ts["nested"].shape == (3, *shape)
+        assert ts.shape == ()
+        ts["nested"].shape = ()
+        ts.shape = (3,)
+        assert ts.shape == (3,)
+        assert ts["nested"].shape == (3,)
+
 
 @pytest.mark.parametrize("shape", [(), (2, 3)])
 @pytest.mark.parametrize("device", get_default_devices())

torchrl/data/tensor_specs.py

@@ -3135,10 +3135,10 @@ def shape(self, value: torch.Size):
             raise RuntimeError("Cannot modify shape of locked composite spec.")
         for key, spec in self.items():
             if isinstance(spec, CompositeSpec):
-                if spec.shape[: self.ndim] != self.shape:
+                if spec.shape[: len(value)] != value:
                     spec.shape = value
             elif spec is not None:
-                if spec.shape[: self.ndim] != self.shape:
+                if spec.shape[: len(value)] != value:
                     raise ValueError(
                         f"The shape of the spec and the CompositeSpec mismatch during shape resetting: the "
                         f"{self.ndim} first dimensions should match but got self['{key}'].shape={spec.shape} and "

torchrl/envs/libs/gym.py

@@ -2,11 +2,14 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
 import importlib
 import warnings
 from copy import copy
 from types import ModuleType
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Tuple
 from warnings import warn
 
 import numpy as np
@@ -310,7 +313,8 @@ def _gym_to_torchrl_spec_transform(
                 categorical_action_encoding=categorical_action_encoding,
                 remap_state_to_observation=remap_state_to_observation,
             )
-        return CompositeSpec(**spec_out)
+        # the batch-size must be set later
+        return CompositeSpec(spec_out)
     elif isinstance(spec, gym_spaces.dict.Dict):
         return _gym_to_torchrl_spec_transform(
             spec.spaces,
@@ -910,7 +914,7 @@ def info_dict_reader(self, value: callable):
         self._info_dict_reader = value
 
     def _reset(
-        self, tensordict: Optional[TensorDictBase] = None, **kwargs
+        self, tensordict: TensorDictBase | None = None, **kwargs
     ) -> TensorDictBase:
         if self._is_batched:
             # batched (aka 'vectorized') env reset is a bit special: envs are