[Refactor] Make all leaves in tensorclass part of _tensordict, exce…

…pt for NonTensorData (#841)

tensordict/_td.py

@@ -794,15 +794,16 @@ def __setitem__(
                         ) from err
 
             keys = set(self.keys())
-            if any(key not in keys for key in value.keys()):
-                subtd = self._get_sub_tensordict(index)
-            for key, item in value.items():
-                if key in keys:
+            subtd = None
+            for value_key, item in value.items():
+                if value_key in keys:
                     self._set_at_str(
-                        key, item, index, validated=False, non_blocking=False
+                        value_key, item, index, validated=False, non_blocking=False
                     )
                 else:
-                    subtd.set(key, item, inplace=True, non_blocking=False)
+                    if subtd is None:
+                        subtd = self._get_sub_tensordict(index)
+                    subtd.set(value_key, item, inplace=True, non_blocking=False)
         else:
             for key in self.keys():
                 self.set_at_(key, value, index)
@@ -2083,12 +2084,35 @@ def _set_tuple(
         )
         return self
 
+    _SHARED_INPLACE_ERROR = (
+        "You're attempting to update a leaf in-place with a shared "
+        "tensordict, but the new value does not match the previous. "
+        "If you're using NonTensorData, see the class documentation "
+        "to see how to properly pre-allocate memory in shared contexts."
+    )
+
     def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool):
         if not validated:
             value = self._validate_value(value, check_shape=False)
             validated = True
         tensor_in = self._get_str(key, NO_DEFAULT)
 
+        if is_non_tensor(value) and not (self._is_shared or self._is_memmap):
+            dest = self._get_str(key, NO_DEFAULT)
+            is_diff = dest[idx].tolist() != value.tolist()
+            if is_diff:
+                dest_val = dest.maybe_to_stack()
+                dest_val[idx] = value
+                if dest_val is not dest:
+                    self._set_str(
+                        key,
+                        dest_val,
+                        validated=True,
+                        inplace=False,
+                        ignore_lock=True,
+                    )
+            return
+
         if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple):
             warn(
                 "Multiple indexing can lead to unexpected behaviours when "
@@ -2103,12 +2127,7 @@ def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool):
             )
             if tensor_in is not tensor_out:
                 if self._is_shared or self._is_memmap:
-                    raise RuntimeError(
-                        "You're attempting to update a leaf in-place with a shared "
-                        "tensordict, but the new value does not match the previous. "
-                        "If you're using NonTensorData, see the class documentation "
-                        "to see how to properly pre-allocate memory in shared contexts."
-                    )
+                    raise RuntimeError(self._SHARED_INPLACE_ERROR)
                 # this happens only when a NonTensorData becomes a NonTensorStack
                 # so it is legitimate (there is no in-place modification of a tensor
                 # that was expected to happen but didn't).
@@ -2301,16 +2320,7 @@ def _memmap_(
             raise RuntimeError(
                 "memmap and shared memory are mutually exclusive features."
             )
-        dest = (
-            self
-            if inplace
-            else TensorDict(
-                {},
-                batch_size=self.batch_size,
-                names=self.names if self._has_names() else None,
-                device=torch.device("cpu"),
-            )
-        )
+        dest = self if inplace else self.empty(device=torch.device("cpu"))
 
         # We must set these attributes before memmapping because we need the metadata
         # to match the tensordict content.