[BugFix] ReplayBuffer's storage now signal back when changes happen

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,53 @@ 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 rb not in storage.attached_entities
+    assert prb in storage.attached_entities
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

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,17 @@ def sample(self, batch_size: int) -> Tuple[Any, dict]:
 
         return ret
 
+    def mark_update(self, index) -> None:
+        """Marks a given storage index as having changed.
+
+        Derived classes can deal with this however appropriate,
+        forwarding this call to whichever parts are needed.
+
+        Args:
+            index: The modified index from storage.
+        """
+        return self.update_priority(index, self._sampler.default_priority)
+
 
 class TensorDictReplayBuffer(ReplayBuffer):
     """TensorDict-specific wrapper around the ReplayBuffer class.

torchrl/data/replay_buffers/samplers.py

@@ -28,7 +28,7 @@ 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
 

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,28 @@ 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 buffer to this storage.
+
+        Replay Buffers that read from this storage must call this
+        method to attach themselves. This guarantees that when data
+        in the storage changes, all relevant pieces of the buffer are
+        aware of it even if the storage is shared with another buffer
+        (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 i in self.attached_entities:
+            i.mark_update(index)
+        return ret
 
     def __iter__(self):
         for i in range(len(self)):