[BugFix] Fix composite setitem

ghstack-source-id: f33b49b
Pull Request resolved: #2778

(cherry picked from commit c2a149d)

Author: Vincent Moens

torchrl/data/tensor_specs.py

@@ -4818,7 +4818,7 @@ def get(self, item, default=NO_DEFAULT):
     def __setitem__(self, key, value):
         dest = self
         if isinstance(key, tuple) and len(key) > 1:
-            while key[0] not in self.keys():
+            while key[0] not in dest.keys():
                 dest[key[0]] = dest = Composite(shape=self.shape, device=self.device)
                 if len(key) > 2:
                     key = key[1:]

torchrl/envs/transforms/transforms.py

@@ -8015,9 +8015,11 @@ def transform_input_spec(self, input_spec: Composite) -> Composite:
                 input_spec["full_action_spec"][out_key] = input_spec[
                     "full_action_spec"
                 ][action_key].clone()
-                if not self.create_copy:
+        if not self.create_copy:
+            for action_key in self.parent.action_keys:
+                if action_key in self.in_keys_inv:
                     del input_spec["full_action_spec"][action_key]
-        for state_key in self.parent.full_state_spec.keys(True):
+        for state_key in self.parent.full_state_spec.keys(True, True):
             if state_key in self.in_keys_inv:
                 for i, out_key in enumerate(self.out_keys_inv):  # noqa: B007
                     if self.in_keys_inv[i] == state_key:
@@ -8028,7 +8030,9 @@ def transform_input_spec(self, input_spec: Composite) -> Composite:
                 input_spec["full_state_spec"][out_key] = input_spec["full_state_spec"][
                     state_key
                 ].clone()
-                if not self.create_copy:
+        if not self.create_copy:
+            for state_key in self.parent.full_state_spec.keys(True, True):
+                if state_key in self.in_keys_inv:
                     del input_spec["full_state_spec"][state_key]
         return input_spec