"[BugFix] Refactor _skip_tensordict to avoid update calls (#2802)"

ghstack-source-id: 0f31b87
Pull Request resolved: #2802

Author: Vincent Moens

torchrl/envs/common.py

@@ -1947,31 +1947,35 @@ def state_spec_unbatched(self, spec: Composite):
         spec = spec.expand(self.batch_size + spec.shape)
         self.state_spec = spec
 
-    def _skip_tensordict(self, tensordict):
+    def _skip_tensordict(self, tensordict: TensorDictBase) -> TensorDictBase:
         # Creates a "skip" tensordict, ie a placeholder for when a step is skipped
         next_tensordict = self.full_done_spec.zero()
         next_tensordict.update(self.full_observation_spec.zero())
         next_tensordict.update(self.full_reward_spec.zero())
 
         # Copy the data from tensordict in `next`
-        def select_and_clone(x, y):
+        keys = set()
+
+        def select_and_clone(name, x, y):
+            keys.add(name)
             if y is not None:
                 if y.device == x.device:
                     return x.clone()
                 return x.to(y.device)
 
-        next_tensordict.update(
-            tensordict._fast_apply(
-                select_and_clone,
-                next_tensordict,
-                device=next_tensordict.device,
-                batch_size=next_tensordict.batch_size,
-                default=None,
-                filter_empty=True,
-                is_leaf=_is_leaf_nontensor,
-            )
+        result = tensordict._fast_apply(
+            select_and_clone,
+            next_tensordict,
+            device=next_tensordict.device,
+            batch_size=next_tensordict.batch_size,
+            default=None,
+            filter_empty=True,
+            is_leaf=_is_leaf_nontensor,
+            named=True,
+            nested_keys=True,
         )
-        return next_tensordict
+        result.update(next_tensordict.exclude(*keys).filter_empty_())
+        return result
 
     def step(self, tensordict: TensorDictBase) -> TensorDictBase:
         """Makes a step in the environment.