[BugFix] Initialising the classes LazyTensorStorage with a nested TensorDict raises error

test/test_rb.py

@@ -12,12 +12,8 @@
 import pytest
 import torch
 from _utils_internal import get_available_devices
-from tensordict.tensordict import assert_allclose_td, TensorDictBase, TensorDict
-from torchrl.data import (
-    PrioritizedReplayBuffer,
-    ReplayBuffer,
-    TensorDictReplayBuffer,
-)
+from tensordict.tensordict import assert_allclose_td, TensorDict, TensorDictBase
+from torchrl.data import PrioritizedReplayBuffer, ReplayBuffer, TensorDictReplayBuffer
 from torchrl.data.replay_buffers import (
     rb_prototype,
     samplers,
@@ -32,25 +28,25 @@
 )
 from torchrl.data.replay_buffers.writers import RoundRobinWriter
 from torchrl.envs.transforms.transforms import (
-    CatTensors,
-    FlattenObservation,
-    SqueezeTransform,
-    ToTensorImage,
-    RewardClipping,
     BinarizeReward,
-    Resize,
-    CenterCrop,
-    UnsqueezeTransform,
-    GrayScale,
-    ObservationNorm,
     CatFrames,
-    RewardScaling,
-    DoubleToFloat,
-    VecNorm,
+    CatTensors,
+    CenterCrop,
     DiscreteActionProjection,
+    DoubleToFloat,
     FiniteTensorDictCheck,
+    FlattenObservation,
+    GrayScale,
     gSDENoise,
+    ObservationNorm,
     PinMemoryTransform,
+    Resize,
+    RewardClipping,
+    RewardScaling,
+    SqueezeTransform,
+    ToTensorImage,
+    UnsqueezeTransform,
+    VecNorm,
 )
 
 _has_tv = importlib.util.find_spec("torchvision") is not None
@@ -198,6 +194,34 @@ def test_index(self, rb_type, sampler, writer, storage, size):
         assert b
 
 
+@pytest.mark.parametrize("max_size", [1000])
+@pytest.mark.parametrize("shape", [[3, 4]])
+@pytest.mark.parametrize("storage", [LazyTensorStorage, LazyMemmapStorage])
+class TestStorages:
+    def _get_nested_td(self, shape):
+        nested_td = TensorDict(
+            {
+                "key1": torch.ones(*shape),
+                "key2": torch.ones(*shape),
+                "next": TensorDict(
+                    {
+                        "key1": torch.ones(*shape),
+                        "key2": torch.ones(*shape),
+                    },
+                    shape,
+                ),
+            },
+            shape,
+        )
+        return nested_td
+
+    def test_init(self, max_size, shape, storage):
+        td = self._get_nested_td(shape)
+        mystorage = storage(max_size=max_size)
+        mystorage._init(td)
+        assert mystorage._storage.shape == (max_size, *shape)
+
+
 @pytest.mark.parametrize("priority_key", ["pk", "td_error"])
 @pytest.mark.parametrize("contiguous", [True, False])
 @pytest.mark.parametrize("device", get_available_devices())

torchrl/data/replay_buffers/storages.py

@@ -7,11 +7,11 @@
 import os
 from collections import OrderedDict
 from copy import copy
-from typing import Any, Sequence, Union, Dict
+from typing import Any, Dict, Sequence, Union
 
 import torch
 from tensordict.memmap import MemmapTensor
-from tensordict.tensordict import TensorDictBase, TensorDict
+from tensordict.tensordict import TensorDict, TensorDictBase
 
 from torchrl._utils import _CKPT_BACKEND
 from torchrl.data.replay_buffers.utils import INT_CLASSES
@@ -236,20 +236,13 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
                 dtype=data.dtype,
             )
         else:
-            out = TensorDict({}, [self.max_size, *data.shape])
-            print("The storage is being created: ")
-            for key, tensor in data.items():
-                if isinstance(tensor, TensorDictBase):
-                    out[key] = (
-                        tensor.expand(self.max_size).clone().to(self.device).zero_()
-                    )
-                else:
-                    out[key] = torch.empty(
-                        self.max_size,
-                        *tensor.shape,
-                        device=self.device,
-                        dtype=tensor.dtype,
-                    )
+            out = (
+                data.expand(self.max_size, *data.shape)
+                .to_tensordict()
+                .zero_()
+                .clone()
+                .to(self.device)
+            )
 
         self._storage = out
         self.initialized = True