Repository: pytorch/rl

Issue:
[BUG] Composable replay buffers with shared storage break when a prioritized exist but is not extended
## Context

We have a prototype of composable replay buffer.
If has 4 component: 
- A storage: either a list or some other form of storage that keeps a register of tensors
- A writer that writes stuff in the storage
- A sampler, which you use to get random items from the storage.

Examples: Writer can be a FIFO or something else, the storage can be on disk or in mem, the sampler can be uniform or prioritized (ie weighted)
The final element is replay buffer itself, which handles the 3 previous elements.

One interesting use case is that we can build multiple RB that share a component: for instance you could have 2 replay buffers with same storage but different sampling strategy.

## Describe the bug


Now say I have 2 replay buffers, one with a uniform sampling and another one with a prioritized sampling strategy. We assume both to share the same storage and writer.

If I update the second (prioritized) the storage should be informed that new elements have come in and it should assign them some default weight value.

But if I update the first (uniform) it has no idea that somewhere there is a prioritized sampler attached to the storage.

When I query my prioritized sampler, it will tell me that there are no weights assigned to some elements (or raise a strange issue).

## To Reproduce

If two replay buffers share the same storage, one is prioritized and the other not, extending the one with random sampling breaks the other:
```python
# let's try another thing: extend rb1

n = 100


storage = LazyMemmapStorage(n)
writer = RoundRobinWriter()
sampler0 = RandomSampler()
sampler1 = PrioritizedSampler(max_capacity=n, alpha=0.7, beta=1.1)

rb0 = ReplayBuffer(storage=storage, writer=writer, sampler=sampler0, collate_fn=lambda x: x)
rb1 = ReplayBuffer(storage=storage, writer=writer, sampler=sampler1, collate_fn=lambda x: x)

data = TensorDict({"a": torch.arange(50)}, [50])
rb0.extend(data)

assert len(rb0) == 50
assert len(storage) == 50
assert len(rb1) == 50

rb0.sample(10) # works
rb1.sample(10) # breaks
```

If in the code snippet above, `rb0.extend(data)` is replaced by `rb1.extend(data)` everything works fine.

## Expected behavior

Extending one or the other replay buffer should work.

## Additional context

Ideally the bug fix should consider the following questions:
(1) Similar bugs might rise in the future, and we don't want a solution that is too tailored to Prioritized replay buffers.
Currently, `ReplayBuffer.add` calls `ReplayBuffer._sampler.add()`, but what that should probably be changed.
Perhaps it should be for the storage to call `add` on all the siblings samplers it has: that would involve that when we register a storage in a buffer, the storage should be informed of its new parent buffer, and look for siblings there.
Something like
```
class Storage:
   def __init__(self, parent: ReplayBuffer):
       self.list_of_parents.append(parent)

   def extend(self, data):
       # do stuff
       ...
       # update all samplers
       for parent in self.list_of_parents:
          parent.sampler.extend(data)
```

(2) another thing to keep in mind is whether it would make sense to have different writers for a single storage. I guess someday someone might want this, but it might also be dangerous.
This can be left as an open issue -- it's not directly linked to this one.



Discussion:

Resolving patch:
diff --git a/torchrl/data/replay_buffers/rb_prototype.py b/torchrl/data/replay_buffers/rb_prototype.py
--- a/torchrl/data/replay_buffers/rb_prototype.py
+++ b/torchrl/data/replay_buffers/rb_prototype.py
@@ -42,6 +42,7 @@ def __init__(
         prefetch: Optional[int] = None,
     ) -> None:
         self._storage = storage if storage is not None else ListStorage(max_size=1_000)
+        self._storage.attach(self)
         self._sampler = sampler if sampler is not None else RandomSampler()
         self._writer = writer if writer is not None else RoundRobinWriter()
         self._writer.register_storage(self._storage)
@@ -155,6 +156,9 @@ def sample(self, batch_size: int) -> Tuple[Any, dict]:
 
         return ret
 
+    def mark_update(self, index: Union[int, torch.Tensor]) -> None:
+        self._sampler.mark_update(index)
+
 
 class TensorDictReplayBuffer(ReplayBuffer):
     """TensorDict-specific wrapper around the ReplayBuffer class.
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
--- a/torchrl/data/replay_buffers/samplers.py
+++ b/torchrl/data/replay_buffers/samplers.py
@@ -28,10 +28,13 @@ def extend(self, index: torch.Tensor) -> None:
         pass
 
     def update_priority(
-        self, index: Union[int, torch.Tensor], priority: Union[int, torch.Tensor]
+        self, index: Union[int, torch.Tensor], priority: Union[float, torch.Tensor]
     ) -> dict:
         pass
 
+    def mark_update(self, index: Union[int, torch.Tensor]) -> None:
+        pass
+
     @property
     def default_priority(self) -> float:
         return 1.0
@@ -182,3 +185,6 @@ def update_priority(
         priority = np.power(priority + self._eps, self._alpha)
         self._sum_tree[index] = priority
         self._min_tree[index] = priority
+
+    def mark_update(self, index: Union[int, torch.Tensor]) -> None:
+        self.update_priority(index, self.default_priority)
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
@@ -29,6 +29,9 @@ class Storage:
 
     def __init__(self, max_size: int) -> None:
         self.max_size = int(max_size)
+        # Prototype feature. RBs that use a given instance of Storage should add
+        # themselves to this set.
+        self._attached_entities = set()
 
     @abc.abstractmethod
     def set(self, cursor: int, data: Any):
@@ -38,11 +41,27 @@ def set(self, cursor: int, data: Any):
     def get(self, index: int) -> Any:
         raise NotImplementedError
 
+    def attach(self, buffer: Any) -> None:
+        """This function attaches a sampler to this storage.
+
+        Buffers that read from this storage must be included as an attached
+        entity by calling this method. This guarantees that when data
+        in the storage changes, components are made aware of changes even if the storage
+        is shared with other buffers (eg. Priority Samplers).
+
+        Args:
+            buffer: the object that reads from this storage.
+        """
+        self._attached_entities.add(buffer)
+
     def __getitem__(self, item):
         return self.get(item)
 
     def __setitem__(self, index, value):
-        return self.set(index, value)
+        ret = self.set(index, value)
+        for ent in self._attached_entities:
+            ent.mark_update(index)
+        return ret
 
     def __iter__(self):
         for i in range(len(self)):


Test patch:
diff --git a/test/test_rb.py b/test/test_rb.py
--- a/test/test_rb.py
+++ b/test/test_rb.py
@@ -21,11 +21,13 @@
     TensorDictPrioritizedReplayBuffer,
     writers,
 )
+from torchrl.data.replay_buffers.samplers import PrioritizedSampler, RandomSampler
 from torchrl.data.replay_buffers.storages import (
     LazyMemmapStorage,
     LazyTensorStorage,
     ListStorage,
 )
+from torchrl.data.replay_buffers.writers import RoundRobinWriter
 from torchrl.data.tensordict.tensordict import assert_allclose_td, TensorDictBase
 
 
@@ -542,6 +544,54 @@ def test_rb_trajectories(stack):
     sampled_td_filtered.batch_size = [3, 4]
 
 
+def test_shared_storage_prioritized_sampler():
+
+    n = 100
+
+    storage = LazyMemmapStorage(n)
+    writer = RoundRobinWriter()
+    sampler0 = RandomSampler()
+    sampler1 = PrioritizedSampler(max_capacity=n, alpha=0.7, beta=1.1)
+
+    rb0 = rb_prototype.ReplayBuffer(
+        storage=storage, writer=writer, sampler=sampler0, collate_fn=lambda x: x
+    )
+    rb1 = rb_prototype.ReplayBuffer(
+        storage=storage, writer=writer, sampler=sampler1, collate_fn=lambda x: x
+    )
+
+    data = TensorDict({"a": torch.arange(50)}, [50])
+
+    # Extend rb0. rb1 should be aware of changes to storage.
+    rb0.extend(data)
+
+    assert len(rb0) == 50
+    assert len(storage) == 50
+    assert len(rb1) == 50
+
+    rb0.sample(10)
+    rb1.sample(10)
+
+    assert rb1._sampler._sum_tree.query(0, 10) == 10
+    assert rb1._sampler._sum_tree.query(0, 50) == 50
+    assert rb1._sampler._sum_tree.query(0, 70) == 50
+
+
+def test_legacy_rb_does_not_attach():
+    n = 10
+    storage = LazyMemmapStorage(n)
+    writer = RoundRobinWriter()
+    sampler = RandomSampler()
+    rb = ReplayBuffer(storage=storage, size=n, prefetch=0, collate_fn=lambda x: x)
+    prb = rb_prototype.ReplayBuffer(
+        storage=storage, writer=writer, sampler=sampler, collate_fn=lambda x: x
+    )
+
+    assert len(storage._attached_entities) == 1
+    assert prb in storage._attached_entities
+    assert rb not in storage._attached_entities
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)