Repository: pytorch/rl

Issue:
[BUG] Using LazyTensorStorage messes up the indexing of the replay-buffer
## Describe the bug
When using `LazyTensorStorage` and the replay buffer isn't fully filled, the indexing of the buffer goes wrong.

## To Reproduce
The following code snippet demonstrates the problem:

```python
import torch
from torchrl.data import ReplayBuffer, LazyTensorStorage

buffer = ReplayBuffer(storage=LazyTensorStorage(1_000))
for i in range(500):
    elem = torch.empty(10)
    buffer.add(elem)

print(len(buffer)) # 500, correct
print(len(buffer[10:])) # 990, wrong!!! should be 490
```

## Expected behavior
Changing the storage of the replay buffer shouldn't change its behaviour. The default storage works as expected, but the LazyTensorStorage makes a mistake when the buffer isn't fully filled.

## System info
Describe the characteristic of your environment:
 * Python 3.11.6
 * Apple M1 MacBook Pro
 * 16GB memory
 * MacOs 14.1.1

```python
import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)

0.2.1 1.26.2 3.11.6 (v3.11.6:8b6ee5ba3b, Oct  2 2023, 11:18:21) [Clang 13.0.0 (clang-1300.0.29.30)] darwin
```

## Checklist

- [ x] I have checked that there is no similar issue in the repo (**required**)
- [ x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [ x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Thanks for reporting this.
I agree that this behaviour should be fixed.
For context, printing `len(buffer)` and `len(buffer[idx])` are actually two different things:
- `len(buffer)` tells you what is the total size of the replay buffer, seen as a full set of tools (a storage, a writer, a sampler and its transforms).
- As soon as you index the buffer, it instantiates a sample and you lose the buffer attributes (storage, writer, ...). If you print the `buffer` or `buffer[idx]` you will get drastically different objects.

I can understand that this is a bit counterintuitive. The reason we do this is that as soon as you index the buffer, we suppose that you want to access its data. Also, it's hard to write a `__getitem__` that would return another buffer that is a view of the original one with consistent storage, writer etc. (for instance, a view of a prioritized sampler where we rebuild the segment trees would be quite challenging to write!). 

I will provide a fix asap

Resolving patch:
diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
--- a/torchrl/data/replay_buffers/storages.py
+++ b/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_()


Test patch:
diff --git a/test/test_rb.py b/test/test_rb.py
--- a/test/test_rb.py
+++ b/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."""