[BugFix] CompositeSpec nested key deletion (pytorch#1059)

Co-authored-by: vmoens <vincentmoens@gmail.com>
ai4co · Apr 13, 2023 · 7e21e6f · 7e21e6f
1 parent d318220
commit 7e21e6f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/test/test_specs.py b/test/test_specs.py
@@ -581,6 +581,13 @@ def test_nested_composite_spec_setitem(self, is_complete, device, dtype):
         )
         assert ts["nested_cp"]["nested_cp"]["obs"] is None
 
+    def test_nested_composite_spec_delitem(self, is_complete, device, dtype):
+        ts = self._composite_spec(is_complete, device, dtype)
+        ts["nested_cp"] = self._composite_spec(is_complete, device, dtype)
+        ts["nested_cp"]["nested_cp"] = self._composite_spec(is_complete, device, dtype)
+        del ts["nested_cp", "nested_cp", "obs"]
+        assert ("nested_cp", "nested_cp", "obs") not in ts.keys(True, True)
+
     def test_nested_composite_spec_update(self, is_complete, device, dtype):
         ts = self._composite_spec(is_complete, device, dtype)
         ts["nested_cp"] = self._composite_spec(is_complete, device, dtype)

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -2257,6 +2257,19 @@ def __iter__(self):
             yield k
 
     def __delitem__(self, key: str) -> None:
+        if isinstance(key, tuple) and len(key) > 1:
+            del self._specs[key[0]][key[1:]]
+            return
+        elif isinstance(key, tuple):
+            del self._specs[key[0]]
+            return
+        elif not isinstance(key, str):
+            raise TypeError(
+                f"Got key of type {type(key)} when a string or a tuple of strings was expected."
+            )
+
+        if key in {"shape", "device", "dtype", "space"}:
+            raise AttributeError(f"CompositeSpec has no key {key}")
         del self._specs[key]
 
     def encode(self, vals: Dict[str, Any]) -> Dict[str, torch.Tensor]:

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -2102,7 +2102,7 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
             )
 
         if isinstance(observation_spec, CompositeSpec) and len(
-            [key for key in self.in_keys if key not in observation_spec]
+            [key for key in self.in_keys if key not in observation_spec.keys(True)]
         ):
             raise ValueError(
                 "CatTensor got a list of keys that does not match the keys in observation_spec. "
@@ -2136,7 +2136,8 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
         )
         if self._del_keys:
             for key in self.keys_to_exclude:
-                del observation_spec[key]
+                if key in observation_spec.keys(True):
+                    del observation_spec[key]
         return observation_spec
 
     def __repr__(self) -> str: