[BugFix] NonTensor should not convert anything to numpy

ghstack-source-id: 7644f6c
Pull Request resolved: #2771

Author: vmoens

torchrl/data/tensor_specs.py

@@ -2619,6 +2619,19 @@ def unbind(self, dim: int = 0):
             for i in range(self.shape[dim])
         )
 
+    def to_numpy(
+        self, val: torch.Tensor | TensorDictBase, safe: bool = None
+    ) -> np.ndarray | dict:
+        return val
+
+    def encode(
+        self,
+        val: np.ndarray | torch.Tensor | TensorDictBase,
+        *,
+        ignore_device: bool = False,
+    ) -> torch.Tensor | TensorDictBase:
+        return val
+
 
 class _UnboundedMeta(abc.ABCMeta):
     def __call__(cls, *args, **kwargs):
@@ -4918,7 +4931,7 @@ def rand(self, shape: torch.Size = None) -> TensorDictBase:
                 _dict[key] = item.rand(shape)
         # No need to run checks since we know Composite is compliant with
         # TensorDict requirements
-        return TensorDict._new_unsafe(
+        return TensorDict(
             _dict,
             batch_size=_size([*shape, *self.shape]),
             device=self._device,

torchrl/envs/gym_like.py

@@ -342,7 +342,7 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
                     for key, val in TensorDict(obs_dict, []).items(True, True)
                 )
         else:
-            tensordict_out = TensorDict._new_unsafe(
+            tensordict_out = TensorDict(
                 obs_dict,
                 batch_size=tensordict.batch_size,
             )
@@ -376,7 +376,8 @@ def _reset(
 
         source = self.read_obs(obs)
 
-        tensordict_out = TensorDict._new_unsafe(
+        # _new_unsafe cannot be used because it won't wrap non-tensor correctly
+        tensordict_out = TensorDict(
             source=source,
             batch_size=self.batch_size,
         )