[BugFix] Fix non-full TensorStorage indexing (#1730)

test/test_rb.py

@@ -797,6 +797,14 @@ def test_index(self, rbtype, storage, size, prefetch):
             b = b.all()
         assert b
 
+    def test_index_nonfull(self, rbtype, storage, size, prefetch):
+        # checks that indexing the buffer before it's full gives the accurate view of the data
+        rb = self._get_rb(rbtype, storage=storage, size=size, prefetch=prefetch)
+        data = self._get_data(rbtype, size=size - 1)
+        rb.extend(data)
+        assert len(rb[: size - 1]) == size - 1
+        assert len(rb[size - 2 :]) == 1
+
 
 def test_multi_loops():
     """Tests that one can iterate multiple times over a buffer without rep."""

torchrl/data/replay_buffers/storages.py

@@ -448,11 +448,15 @@ def set(  # noqa: F811
         self._storage[cursor] = data
 
     def get(self, index: Union[int, Sequence[int], slice]) -> Any:
+        if self._len < self.max_size:
+            storage = self._storage[: self._len]
+        else:
+            storage = self._storage
         if not self.initialized:
             raise RuntimeError(
                 "Cannot get an item from an unitialized LazyMemmapStorage"
             )
-        out = self._storage[index]
+        out = storage[index]
         if is_tensor_collection(out):
             out = _reset_batch_size(out)
             return out.unlock_()