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

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

(cherry picked from commit e0d3eee)

Author: Vincent Moens

torchrl/envs/common.py

@@ -1933,6 +1933,36 @@ def state_spec_unbatched(self, spec: Composite):
         spec = spec.expand(self.batch_size + spec.shape)
         self.state_spec = spec
 
+    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`
+        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)
+
+        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,
+        )
+        result.update(next_tensordict.exclude(*keys).filter_empty_())
+        return result
+
     def step(self, tensordict: TensorDictBase) -> TensorDictBase:
         """Makes a step in the environment.