Repository: continualai/avalanche
Issue: GDumb strategy performance issues with Rotated/Permuted MNIST benchmarks

When running GDumb for MNIST variants where the number of classes remains the same (from 0 to 9) in each experience,  accuracy on all old tasks drops significantly with the newer task. I wanted to confirm this is a bug in the Gdumb strategy in avalanche or I am missing something on my part.

Below is the script to reproduce the results for Rotated MNIST:
```
train_transform = Compose([
        ToTensor(),
        Normalize((0.1307,), (0.3081,))
    ])
eval_transform = Compose([
        ToTensor(),
        Normalize((0.1307,), (0.3081,))
    ])
scenario = RotatedMNIST(
    n_experiences=2, 
    train_transform=train_transform, 
    eval_transform=eval_transform,
    seed=1234)

interactive_logger = InteractiveLogger()

eval_plugin = EvaluationPlugin(
    accuracy_metrics(experience=True, stream=True),
    loss_metrics(experience=True, stream=True),
    loggers=[interactive_logger])

model = SimpleMLP(hidden_size=100, hidden_layers=2, num_classes=10, drop_rate=0.)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

strategy = Naive(
            model, optimizer, criterion,
            train_mb_size=16, device=device,
            eval_mb_size=16, train_epochs=15,
            plugins=[GDumbPlugin(mem_size=500)]
        )

print('Starting experiment...')
for experience in scenario.train_stream:
  print("Start training on experience ", experience.current_experience)

  strategy.train(experience)
  print("End training on experience ", experience.current_experience)
  print('Computing accuracy on the test set')
  strategy.eval(scenario.test_stream[:])
```


Below are the results:
After training on Exp0:
Top1_Acc_Stream/eval_phase/test_stream/Task000 = 0.8488
Top1_Acc_Stream/eval_phase/test_stream/Task001 = 0.1460

After training on Exp1:
Top1_Acc_Stream/eval_phase/test_stream/Task000 = 0.1373
Top1_Acc_Stream/eval_phase/test_stream/Task001 = 0.7879

Similar trend is observed for PermutedMNIST as well.

Thank you!

Changed files:
avalanche/training/plugins/replay.py, tests/training/test_plugins.py, tests/training/test_replay.py, avalanche/training/plugins/gdumb.py, avalanche/training/storage_policy.py, tests/test_dataloaders.py, avalanche/benchmarks/utils/data_loader.py, avalanche/training/plugins/cope.py

diff --git a/avalanche/benchmarks/utils/data_loader.py b/avalanche/benchmarks/utils/data_loader.py
index bd641e15..db18912b 100644
--- a/avalanche/benchmarks/utils/data_loader.py
+++ b/avalanche/benchmarks/utils/data_loader.py
@@ -195,7 +195,7 @@ class GroupBalancedInfiniteDataLoader:
                 sampler=infinite_sampler,
                 **kwargs)
             self.dataloaders.append(dl)
-        self.max_len = max([len(d) for d in self.dataloaders])
+        self.max_len = 10 ** 10
 
     def __iter__(self):
         iter_dataloaders = []
diff --git a/avalanche/training/plugins/cope.py b/avalanche/training/plugins/cope.py
index 3fedddb4..d75bca91 100644
--- a/avalanche/training/plugins/cope.py
+++ b/avalanche/training/plugins/cope.py
@@ -8,7 +8,7 @@ from torch.nn.modules import Module
 from avalanche.training.utils import get_last_fc_layer, swap_last_fc_layer
 from avalanche.benchmarks.utils import AvalancheConcatDataset
 from avalanche.training.plugins.strategy_plugin import StrategyPlugin
-from avalanche.training.storage_policy import ClassBalancedStoragePolicy
+from avalanche.training.storage_policy import ClassBalancedBuffer
 from avalanche.benchmarks.utils.data_loader import \\
     ReplayDataLoader
 
@@ -48,9 +48,8 @@ class CoPEPlugin(StrategyPlugin):
         # Operational memory: replay memory
         self.replay_mem = {}
         self.mem_size = mem_size  # replay memory size
-        self.storage_policy = ClassBalancedStoragePolicy(
-            ext_mem=self.replay_mem,
-            mem_size=self.mem_size,
+        self.storage_policy = ClassBalancedBuffer(
+            max_size=self.mem_size,
             adaptive_size=True)
 
         # Operational memory: Prototypical memory
@@ -167,7 +166,7 @@ class CoPEPlugin(StrategyPlugin):
         store observed samples for replay.
         """
         self._update_prototypes()  # Update prototypes
-        self.storage_policy(strategy)  # Update memory
+        self.storage_policy.update(strategy)  # Update memory
 
     @torch.no_grad()
     def _update_prototypes(self):
diff --git a/avalanche/training/plugins/gdumb.py b/avalanche/training/plugins/gdumb.py
index f0ac4d87..07f51946 100644
--- a/avalanche/training/plugins/gdumb.py
+++ b/avalanche/training/plugins/gdumb.py
@@ -18,7 +18,7 @@ from torch.utils.data import TensorDataset
 
 from avalanche.benchmarks.utils import AvalancheConcatDataset, AvalancheDataset
 from avalanche.training.plugins.strategy_plugin import StrategyPlugin
-from avalanche.training.storage_policy import ClassBalancedStoragePolicy
+from avalanche.training.storage_policy import ClassBalancedBuffer
 
 if TYPE_CHECKING:
     from avalanche.training.strategies import BaseStrategy
@@ -43,9 +43,8 @@ class GDumbPlugin(StrategyPlugin):
 
         # model initialization
         self.buffer = {}
-        self.storage_policy = ClassBalancedStoragePolicy(
-            ext_mem=self.buffer,
-            mem_size=self.mem_size,
+        self.storage_policy = ClassBalancedBuffer(
+            max_size=self.mem_size,
             adaptive_size=True
         )
         self.init_model = None
@@ -60,6 +59,5 @@ class GDumbPlugin(StrategyPlugin):
 
     def after_train_dataset_adaptation(self, strategy: "BaseStrategy",
                                        **kwargs):
-        self.storage_policy(strategy, **kwargs)
-        cat_data = AvalancheConcatDataset(self.buffer.values())
-        strategy.adapted_dataset = cat_data
+        self.storage_policy.update(strategy, **kwargs)
+        strategy.adapted_dataset = self.storage_policy.buffer
diff --git a/avalanche/training/plugins/replay.py b/avalanche/training/plugins/replay.py
index 880e96be..54e07069 100644
--- a/avalanche/training/plugins/replay.py
+++ b/avalanche/training/plugins/replay.py
@@ -4,8 +4,8 @@ from avalanche.benchmarks.utils import AvalancheConcatDataset
 from avalanche.benchmarks.utils.data_loader import \\
     ReplayDataLoader
 from avalanche.training.plugins.strategy_plugin import StrategyPlugin
-from avalanche.training.storage_policy import StoragePolicy, \\
-    ExperienceBalancedStoragePolicy
+from avalanche.training.storage_policy import ExemplarsBuffer, \\
+    ExperienceBalancedBuffer
 
 if TYPE_CHECKING:
     from avalanche.training.strategies import BaseStrategy
@@ -33,22 +33,22 @@ class ReplayPlugin(StrategyPlugin):
     """
 
     def __init__(self, mem_size: int = 200,
-                 storage_policy: Optional["StoragePolicy"] = None):
+                 storage_policy: Optional["ExemplarsBuffer"] = None):
         super().__init__()
         self.mem_size = mem_size
 
         if storage_policy is not None:  # Use other storage policy
             self.storage_policy = storage_policy
-            self.ext_mem = storage_policy.ext_mem  # Keep ref
-            assert storage_policy.mem_size == self.mem_size
-
+            assert storage_policy.max_size == self.mem_size
         else:  # Default
-            self.ext_mem = {}  # a Dict<task_id, Dataset>
-            self.storage_policy = ExperienceBalancedStoragePolicy(
-                ext_mem=self.ext_mem,
-                mem_size=self.mem_size,
+            self.storage_policy = ExperienceBalancedBuffer(
+                max_size=self.mem_size,
                 adaptive_size=True)
 
+    @property
+    def ext_mem(self):
+        return self.storage_policy.buffer_groups  # a Dict<task_id, Dataset>
+
     def before_training_exp(self, strategy: "BaseStrategy",
                             num_workers: int = 0, shuffle: bool = True,
                             **kwargs):
@@ -56,15 +56,17 @@ class ReplayPlugin(StrategyPlugin):
         Dataloader to build batches containing examples from both memories and
         the training dataset
         """
-        if len(self.ext_mem) == 0:
+        if len(self.storage_policy.buffer) == 0:
+            # first experience. We don't use the buffer, no need to change
+            # the dataloader.
             return
         strategy.dataloader = ReplayDataLoader(
             strategy.adapted_dataset,
-            AvalancheConcatDataset(self.ext_mem.values()),
+            self.storage_policy.buffer,
             oversample_small_tasks=True,
             num_workers=num_workers,
             batch_size=strategy.train_mb_size,
             shuffle=shuffle)
 
     def after_training_exp(self, strategy: "BaseStrategy", **kwargs):
-        self.storage_policy(strategy, **kwargs)
+        self.storage_policy.update(strategy, **kwargs)
diff --git a/avalanche/training/storage_policy.py b/avalanche/training/storage_policy.py
index dcfee353..eed091fe 100644
--- a/avalanche/training/storage_policy.py
+++ b/avalanche/training/storage_policy.py
@@ -6,7 +6,7 @@ import torch
 from numpy import inf
 from torch import cat, Tensor
 from torch.nn import Module
-from torch.utils.data import random_split, DataLoader
+from torch.utils.data import DataLoader
 
 from avalanche.benchmarks.utils import AvalancheDataset, AvalancheSubset, \\
     AvalancheConcatDataset
@@ -16,123 +16,228 @@ if TYPE_CHECKING:
     from .strategies import BaseStrategy
 
 
-class StoragePolicy(ABC):
-    """
-    A policy to store exemplars in a replay memory.
+class ExemplarsBuffer(ABC):
+    """ A buffer that stores exemplars for rehearsal.
+
+    `self.buffer` is an AvalancheDataset of samples collected from the previous
+    experiences. The buffer can be updated by calling `self.update(strategy)`.
 
-    :param ext_mem: The replay memory dictionary to store samples.
-    :param mem_size: max number of total input samples in the replay memory.
+    :param max_size: max number of input samples in the replay memory.
     """
-    def __init__(self, ext_mem: Dict[int, AvalancheDataset], mem_size: int):
-        self.ext_mem = ext_mem
-        self.mem_size = mem_size
+
+    def __init__(self, max_size: int):
+        self.max_size = max_size
+        """ Maximum size of the buffer. """
+        self._buffer = AvalancheConcatDataset([])
+
+    @property
+    def buffer(self) -> AvalancheDataset:
+        """ Buffer of samples. """
+        return self._buffer
+
+    @buffer.setter
+    def buffer(self, new_buffer: AvalancheDataset):
+        self._buffer = new_buffer
 
     @abstractmethod
-    def __call__(self, data_source: AvalancheDataset, **kwargs):
-        """Store exemplars in the replay memory"""
+    def update(self, strategy: 'BaseStrategy', **kwargs):
+        """ Update `self.buffer` using the `strategy` state.
+
+        :param strategy:
+        :param kwargs:
+        :return:
+        """
+        ...
+
+    @abstractmethod
+    def resize(self, strategy: 'BaseStrategy', new_size: int):
+        """ Update the maximum size of the buffer.
+
+        :param strategy:
+        :param new_size:
+        :return:
+        """
+        ...
+
+
+class ReservoirSamplingBuffer(ExemplarsBuffer):
+    def __init__(self, max_size: int):
+        """ Buffer updated with reservoir sampling. """
+        # The algorithm follows
+        # https://en.wikipedia.org/wiki/Reservoir_sampling
+        # We sample a random uniform value in [0, 1] for each sample and
+        # choose the `size` samples with higher values.
+        # This is equivalent to a random selection of `size_samples`
+        # from the entire stream.
+        super().__init__(max_size)
+        # INVARIANT: _buffer_weights is always sorted.
+        self._buffer_weights = torch.zeros(0)
+
+    def update(self, strategy: 'BaseStrategy', **kwargs):
+        """ Update buffer. """
+        self.update_from_dataset(strategy.experience.dataset)
+
+    def update_from_dataset(self, new_data: AvalancheDataset):
+        new_weights = torch.rand(len(new_data))
+
+        cat_weights = torch.cat([new_weights, self._buffer_weights])
+        cat_data = AvalancheConcatDataset([new_data, self.buffer])
+        sorted_weights, sorted_idxs = cat_weights.sort(descending=True)
+
+        buffer_idxs = sorted_idxs[:self.max_size]
+        self.buffer = AvalancheSubset(cat_data, buffer_idxs)
+        self._buffer_weights = sorted_weights[:self.max_size]
+
+    def resize(self, strategy, new_size):
+        """ Update the maximum size of the buffer. """
+        self.max_size = new_size
+        if len(self.buffer) <= self.max_size:
+            return
+        self.buffer = AvalancheSubset(self.buffer, torch.arange(self.max_size))
+        self._buffer_weights = self._buffer_weights[:self.max_size]
+
+
+class BalancedExemplarsBuffer(ExemplarsBuffer):
+    def __init__(self, max_size: int, adaptive_size: bool = True,
+                 total_num_groups=None):
+        """ A buffer that stores exemplars for rehearsal in separate groups.
+
+        The grouping allows to balance the data (by task, experience,
+        classes..). In combination with balanced data loaders, it can be used
+        to sample balanced mini-batches during training.
+
+        `self.buffer_groups` is a dictionary that stores each group as a
+        separate buffer. The buffers are updated by calling
+        `self.update(strategy)`.
+
+        :param max_size: max number of input samples in the replay memory.
+        :param adaptive_size: True if max_size is divided equally over all
+                              observed experiences (keys in replay_mem).
+        :param total_num_groups: If adaptive size is False, the fixed number
+                                of groups to divide capacity over.
+        """
+        super().__init__(max_size)
+        self.adaptive_size = adaptive_size
+        self.total_num_groups = total_num_groups
+        if not self.adaptive_size:
+            assert self.total_num_groups > 0, \\
+                "You need to specify `total_num_groups` if " \\
+                "`adaptive_size=True`."
+        else:
+            assert self.total_num_groups is None, \\
+                "`total_num_groups` is not compatible with " \\
+                "`adaptive_size=False`."
+
+        self.buffer_groups: Dict[int, ExemplarsBuffer] = {}
+        """ Dictionary of buffers. """
+
+    @property
+    def buffer_datasets(self):
+        """ Return group buffers as a list of `AvalancheDataset`s. """
+        return [g.buffer for g in self.buffer_groups.values()]
+
+    def get_group_lengths(self, num_groups):
+        """ Compute groups lengths given the number of groups `num_groups`. """
+        if self.adaptive_size:
+            lengths = [self.max_size // num_groups for _ in range(num_groups)]
+            # distribute remaining size among experiences.
+            rem = self.max_size - sum(lengths)
+            for i in range(rem):
+                lengths[i] += 1
+        else:
+            lengths = [self.max_size // self.total_num_groups for _ in
+                       range(num_groups)]
+        return lengths
+
+    @property
+    def buffer(self):
+        return AvalancheConcatDataset(
+            [g.buffer for g in self.buffer_groups.values()])
+
+    @buffer.setter
+    def buffer(self, new_buffer):
+        assert NotImplementedError(
+            "Cannot set `self.buffer` for this class. "
+            "You should modify `self.buffer_groups instead.")
 
+    @abstractmethod
+    def update(self, strategy: 'BaseStrategy', **kwargs):
+        """ Update `self.buffer_groups` using the `strategy` state.
 
-class ExperienceBalancedStoragePolicy(StoragePolicy):
-    def __init__(self, ext_mem: Dict, mem_size: int, adaptive_size: bool = True,
-                 num_experiences=-1):
+        :param strategy:
+        :param kwargs:
+        :return:
         """
-        Stores samples for replay, equally divided over experiences.
-        Because it is conditioned on the experience, it should be called in
-        the 'after_training_exp' phase.
+        ...
+
+    def resize(self, strategy, new_size):
+        """ Update the maximum size of the buffers. """
+        self.max_size = new_size
+        lens = self.get_group_lengths(len(self.buffer_groups))
+        for ll, buffer in zip(lens, self.buffer_groups.values()):
+            buffer.resize(strategy, ll)
+
+
+class ExperienceBalancedBuffer(BalancedExemplarsBuffer):
+    def __init__(self, max_size: int, adaptive_size: bool = True,
+                 num_experiences=None):
+        """ Rehearsal buffer with samples balanced over experiences.
 
         The number of experiences can be fixed up front or adaptive, based on
         the 'adaptive_size' attribute. When adaptive, the memory is equally
         divided over all the unique observed experiences so far.
 
-        :param ext_mem: The replay memory dictionary to store samples.
-        :param mem_size: max number of total input samples in the replay memory.
+        :param max_size: max number of total input samples in the replay
+            memory.
         :param adaptive_size: True if mem_size is divided equally over all
                               observed experiences (keys in replay_mem).
         :param num_experiences: If adaptive size is False, the fixed number
                                 of experiences to divide capacity over.
         """
-        super().__init__(ext_mem, mem_size)
-        self.adaptive_size = adaptive_size
-        self.num_experiences = num_experiences
+        super().__init__(max_size, adaptive_size, num_experiences)
 
-        if not self.adaptive_size:
-            assert self.num_experiences > 0, \\
-                """When fixed exp mem size, num_experiences should be > 0."""
-
-    def subsample_single(self, data: AvalancheDataset, new_size: int):
-        """ Subsample `data` to match length `new_size`. """
-        removed_els = len(data) - new_size
-        if removed_els > 0:
-            data, _ = random_split(data, [new_size, removed_els])
-        return data
-
-    def subsample_all_groups(self, new_size: int):
-        """ Subsample all groups equally to match total buffer size
-        `new_size`. """
-        groups = list(self.ext_mem.keys())
-        if len(groups) == 0:
-            return  # buffer is empty.
-
-        num_groups = len(groups) if self.adaptive_size else self.num_experiences
-        group_size = new_size // num_groups
-        last_group_size = group_size + (new_size % num_groups)
-
-        for g in groups[:-1]:
-            self.ext_mem[g] = self.subsample_single(self.ext_mem[g], group_size)
-        # last group may be bigger
-        last = self.ext_mem[groups[-1]]
-        self.ext_mem[groups[-1]] = self.subsample_single(last, last_group_size)
-
-    def __call__(self, strategy: "BaseStrategy", **kwargs):
+    def update(self, strategy: "BaseStrategy", **kwargs):
+        new_data = strategy.experience.dataset
         num_exps = strategy.clock.train_exp_counter + 1
-        num_exps = num_exps if self.adaptive_size else self.num_experiences
-        curr_data = strategy.experience.dataset
+        lens = self.get_group_lengths(num_exps)
 
-        # new group may be bigger because of the remainder.
-        group_size = self.mem_size // num_exps
-        new_group_size = group_size + (self.mem_size % num_exps)
+        new_buffer = ReservoirSamplingBuffer(lens[-1])
+        new_buffer.update_from_dataset(new_data)
+        self.buffer_groups[num_exps - 1] = new_buffer
 
-        self.subsample_all_groups(group_size * (num_exps - 1))
-        curr_data = self.subsample_single(curr_data, new_group_size)
-        self.ext_mem[strategy.clock.train_exp_counter + 1] = curr_data
+        for ll, b in zip(lens, self.buffer_groups.values()):
+            b.resize(strategy, ll)
 
-        # buffer size should always equal self.mem_size
-        len_tot = sum(len(el) for el in self.ext_mem.values())
-        assert len_tot == self.mem_size
 
+class ClassBalancedBuffer(BalancedExemplarsBuffer):
+    def __init__(self, max_size: int, adaptive_size: bool = True,
+                 total_num_classes: int = None):
+        """ Stores samples for replay, equally divided over classes.
 
-class ClassBalancedStoragePolicy(StoragePolicy):
-    def __init__(self, ext_mem: Dict, mem_size: int, adaptive_size: bool = True,
-                 total_num_classes: int = -1, selection_strategy:
-                 Optional["ClassExemplarsSelectionStrategy"] = None):
-        """
-        Stores samples for replay, equally divided over classes.
+        There is a separate buffer updated by reservoir sampling for each
+            class.
         It should be called in the 'after_training_exp' phase (see
         ExperienceBalancedStoragePolicy).
         The number of classes can be fixed up front or adaptive, based on
         the 'adaptive_size' attribute. When adaptive, the memory is equally
         divided over all the unique observed classes so far.
-        :param ext_mem: The replay memory dictionary to store samples in.
-        :param mem_size: The max capacity of the replay memory.
+        
+        :param max_size: The max capacity of the replay memory.
         :param adaptive_size: True if mem_size is divided equally over all
                             observed experiences (keys in replay_mem).
         :param total_num_classes: If adaptive size is False, the fixed number
                                   of classes to divide capacity over.
-        :param selection_strategy: The strategy used to select exemplars to 
-                                   keep in memory when cutting it off
         """
-        super().__init__(ext_mem, mem_size)
-        self.selection_strategy = selection_strategy or \\
-            RandomExemplarsSelectionStrategy()
+        if not adaptive_size:
+            assert total_num_classes > 0, \\
+                """When fixed exp mem size, total_num_classes should be > 0."""
+
+        super().__init__(max_size, adaptive_size, total_num_classes)
         self.adaptive_size = adaptive_size
         self.total_num_classes = total_num_classes
         self.seen_classes = set()
 
-        if not self.adaptive_size:
-            assert self.total_num_classes > 0, \\
-                """When fixed exp mem size, total_num_classes should be > 0."""
-
-    def __call__(self, strategy: "BaseStrategy", **kwargs):
+    def update(self, strategy: "BaseStrategy", **kwargs):
         new_data = strategy.experience.dataset
 
         # Get sample idxs per class
@@ -150,67 +255,160 @@ class ClassBalancedStoragePolicy(StoragePolicy):
         # Update seen classes
         self.seen_classes.update(cl_datasets.keys())
 
-        # how many experiences to divide the memory over
-        div_cnt = len(self.seen_classes) if self.adaptive_size \\
-            else self.total_num_classes
-        class_mem_size = self.mem_size // div_cnt
-
-        # Add current classes data to memory
-        for c, c_mem in cl_datasets.items():
-            if c in self.ext_mem:  # Merge data with previous seen data
-                c_mem = AvalancheConcatDataset((c_mem, self.ext_mem[c]))
-            sorted_indices = self.selection_strategy.make_sorted_indices(
-                strategy, c_mem)
-            self.ext_mem[c] = AvalancheSubset(c_mem, sorted_indices)
-
-        # Distribute remaining samples using counts
-        cutoff_per_exp = self.divide_remaining_samples(class_mem_size, div_cnt)
-
-        # Use counts to remove samples from memory
-        self.cutoff_memory(cutoff_per_exp)
-
-    def divide_remaining_samples(self, exp_mem_size: int, div_cnt: int) -> \\
-            Dict[int, int]:
-        # Find number of remaining samples
-        samples_per_exp = {exp: len(mem) for exp, mem in
-                           self.ext_mem.items()}
-        rem_from_exps = {exp: exp_mem_size - memsize for exp, memsize in
-                         samples_per_exp.items() if
-                         exp_mem_size - memsize > 0}
-        rem_from_div = self.mem_size % div_cnt
-        free_mem = sum(rem_from_exps.values()) + rem_from_div
-
-        # Divide the remaining samples randomly over the experiences
-        cutoff_per_exp = {exp: min(exp_mem_size, len(m))
-                          for exp, m in self.ext_mem.items()}
-
-        # Find remaining data samples to divide
-        rem_samples_exp = {exp: memsize - exp_mem_size for exp, memsize in
-                           samples_per_exp.items()
-                           if memsize - exp_mem_size > 0}
-
-        while len(rem_samples_exp) > 0 and free_mem > 0:
-            exp = random.choice(list(rem_samples_exp.keys()))
-            cutoff_per_exp[exp] += 1
-            free_mem -= 1
-            rem_samples_exp[exp] -= 1
-            if rem_samples_exp[exp] <= 0:
-                del rem_samples_exp[exp]
-
-        return cutoff_per_exp
-
-    def cutoff_memory(self, cutoff_per_exp: Dict[int, int]):
-        # No need to reselect at this point, we expect the first selection to
-        # have sorted the exemplars
-        for exp, cutoff in cutoff_per_exp.items():
-            self.ext_mem[exp] = AvalancheSubset(self.ext_mem[exp],
-                                                list(range(cutoff)))
-
-
-class ClassExemplarsSelectionStrategy(ABC):
+        # associate lengths to classes
+        lens = self.get_group_lengths(len(self.seen_classes))
+        class_to_len = {}
+        for class_id, ll in zip(self.seen_classes, lens):
+            class_to_len[class_id] = ll
+
+        # update buffers with new data
+        for class_id, new_data_c in cl_datasets.items():
+            ll = class_to_len[class_id]
+            if class_id in self.buffer_groups:
+                old_buffer_c = self.buffer_groups[class_id]
+                old_buffer_c.update_from_dataset(new_data_c)
+                old_buffer_c.resize(strategy, ll)
+            else:
+                new_buffer = ReservoirSamplingBuffer(ll)
+                new_buffer.update_from_dataset(new_data_c)
+                self.buffer_groups[class_id] = new_buffer
+
+        # resize buffers
+        for class_id, class_buf in self.buffer_groups.items():
+            self.buffer_groups[class_id].resize(strategy,
+                                                class_to_len[class_id])
+
+
+class ParametricBuffer(BalancedExemplarsBuffer):
+    def __init__(self, max_size: int,
+                 groupby=None,
+                 selection_strategy: Optional[
+                     "ExemplarsSelectionStrategy"] = None):
+        """ Stores samples for replay using a custom selection strategy and
+        grouping.
+
+        :param max_size: The max capacity of the replay memory.
+        :param groupby: Grouping mechanism. One of {None, 'class', 'task',
+        'experience'}.
+        :param selection_strategy: The strategy used to select exemplars to
+                                   keep in memory when cutting it off.
+        """
+        super().__init__(max_size)
+        assert groupby in {None, 'task', 'class', 'experience'}, \\
+            "Unknown grouping scheme. Must be one of {None, 'task', " \\
+            "'class', 'experience'}"
+        self.groupby = groupby
+        ss = selection_strategy or RandomExemplarsSelectionStrategy()
+        self.selection_strategy = ss
+        self.seen_groups = set()
+        self._curr_strategy = None
+
+    def update(self, strategy: "BaseStrategy", **kwargs):
+        new_data = strategy.experience.dataset
+        new_groups = self.make_groups(strategy, new_data)
+        self.seen_groups.update(new_groups.keys())
+
+        # associate lengths to classes
+        lens = self.get_group_lengths(len(self.seen_groups))
+        group_to_len = {}
+        for group_id, ll in zip(self.seen_groups, lens):
+            group_to_len[group_id] = ll
+
+        # update buffers with new data
+        for group_id, new_data_g in new_groups.items():
+            ll = group_to_len[group_id]
+            if group_id in self.buffer_groups:
+                old_buffer_g = self.buffer_groups[group_id]
+                old_buffer_g.update_from_dataset(strategy, new_data_g)
+                old_buffer_g.resize(strategy, ll)
+            else:
+                new_buffer = _ParametricSingleBuffer(ll,
+                                                     self.selection_strategy)
+                new_buffer.update_from_dataset(strategy, new_data_g)
+                self.buffer_groups[group_id] = new_buffer
+
+        # resize buffers
+        for group_id, class_buf in self.buffer_groups.items():
+            self.buffer_groups[group_id].resize(strategy,
+                                                group_to_len[group_id])
+
+    def make_groups(self, strategy, data):
+        if self.groupby is None:
+            return {0: data}
+        elif self.groupby == 'task':
+            return self._split_by_task(data)
+        elif self.groupby == 'experience':
+            return self._split_by_experience(strategy, data)
+        elif self.groupby == 'class':
+            return self._split_by_class(data)
+        else:
+            assert False, "Invalid groupby key. Should never get here."
+
+    def _split_by_class(self, data):
+        # Get sample idxs per class
+        class_idxs = {}
+        for idx, target in enumerate(data.targets):
+            if target not in class_idxs:
+                class_idxs[target] = []
+            class_idxs[target].append(idx)
+
+        # Make AvalancheSubset per class
+        new_groups = {}
+        for c, c_idxs in class_idxs.items():
+            new_groups[c] = AvalancheSubset(data, indices=c_idxs)
+        return new_groups
+
+    def _split_by_experience(self, strategy, data):
+        exp_id = strategy.clock.train_exp_counter + 1
+        return {exp_id: data}
+
+    def _split_by_task(self, data):
+        new_groups = {}
+        for task_id in data.task_set:
+            new_groups[task_id] = data.task_set[task_id]
+        return new_groups
+
+
+class _ParametricSingleBuffer(ExemplarsBuffer):
+    def __init__(self, max_size: int,
+                 selection_strategy: Optional[
+                     "ExemplarsSelectionStrategy"] = None):
+        """ A buffer that stores samples for replay using a custom selection
+        strategy.
+
+        This is a private class. Use `ParametricBalancedBuffer` with
+        `groupby=None` to get the same behavior.
+
+        :param max_size: The max capacity of the replay memory.
+        :param selection_strategy: The strategy used to select exemplars to
+                                   keep in memory when cutting it off.
+        """
+        super().__init__(max_size)
+        ss = selection_strategy or RandomExemplarsSelectionStrategy()
+        self.selection_strategy = ss
+        self._curr_strategy = None
+
+    def update(self, strategy: "BaseStrategy", **kwargs):
+        new_data = strategy.experience.dataset
+        self.update_from_dataset(strategy, new_data)
+
+    def update_from_dataset(self, strategy, new_data):
+        self.buffer = AvalancheConcatDataset([self.buffer, new_data])
+        self.resize(strategy, self.max_size)
+
+    def resize(self, strategy, new_size: int):
+        self.max_size = new_size
+        idxs = self.selection_strategy.make_sorted_indices(
+            strategy=strategy,
+            data=self.buffer)
+        self.buffer = AvalancheSubset(self.buffer, idxs[:self.max_size])
+
+
+class ExemplarsSelectionStrategy(ABC):
     """
-    Base class to define how to select class exemplars from a dataset
+    Base class to define how to select a subset of exemplars from a dataset.
     """
+
     @abstractmethod
     def make_sorted_indices(self, strategy: "BaseStrategy",
                             data: AvalancheDataset) -> List[int]:
@@ -219,9 +417,10 @@ class ClassExemplarsSelectionStrategy(ABC):
 
         The last indices will be the first to be removed when cutoff memory.
         """
+        ...
 
 
-class RandomExemplarsSelectionStrategy(ClassExemplarsSelectionStrategy):
+class RandomExemplarsSelectionStrategy(ExemplarsSelectionStrategy):
     """Select the exemplars at random in the dataset"""
 
     def make_sorted_indices(self, strategy: "BaseStrategy",
@@ -231,9 +430,10 @@ class RandomExemplarsSelectionStrategy(ClassExemplarsSelectionStrategy):
         return indices
 
 
-class FeatureBasedExemplarsSelectionStrategy(ClassExemplarsSelectionStrategy,
+class FeatureBasedExemplarsSelectionStrategy(ExemplarsSelectionStrategy,
                                              ABC):
     """Base class to select exemplars from their features"""
+
     def __init__(self, model: Module, layer_name: str):
         self.feature_extractor = FeatureExtractorBackbone(model, layer_name)
 
@@ -299,3 +499,18 @@ class ClosestToCenterSelectionStrategy(FeatureBasedExemplarsSelectionStrategy):
         center = features.mean(dim=0)
         distances = pow(features - center, 2).sum(dim=1)
         return distances.argsort()
+
+
+__all__ = [
+    'ExemplarsBuffer',
+    'ReservoirSamplingBuffer',
+    'BalancedExemplarsBuffer',
+    'ExperienceBalancedBuffer',
+    'ClassBalancedBuffer',
+    'ParametricBuffer',
+    'ExemplarsSelectionStrategy',
+    'RandomExemplarsSelectionStrategy',
+    'FeatureBasedExemplarsSelectionStrategy',
+    'HerdingSelectionStrategy',
+    'ClosestToCenterSelectionStrategy'
+]
diff --git a/tests/test_dataloaders.py b/tests/test_dataloaders.py
index daa17468..66ce1083 100644
--- a/tests/test_dataloaders.py
+++ b/tests/test_dataloaders.py
@@ -107,7 +107,7 @@ class DataLoaderTests(unittest.TestCase):
             adapted_dataset = step.dataset
             dataloader = ReplayDataLoader(
                     adapted_dataset,
-                    AvalancheConcatDataset(replayPlugin.ext_mem.values()),
+                    replayPlugin.storage_policy.buffer,
                     oversample_small_tasks=True,
                     num_workers=0,
                     batch_size=batch_size,
diff --git a/tests/training/test_plugins.py b/tests/training/test_plugins.py
index 02135b80..9877ec00 100644
--- a/tests/training/test_plugins.py
+++ b/tests/training/test_plugins.py
@@ -13,10 +13,7 @@ from torch.utils.data import TensorDataset
 from avalanche.benchmarks import nc_benchmark
 from avalanche.logging import TextLogger
 from avalanche.models import SimpleMLP
-from avalanche.training.plugins import EvaluationPlugin
-from avalanche.training.plugins import StrategyPlugin, ReplayPlugin
-from avalanche.training.storage_policy import \\
-    ExperienceBalancedStoragePolicy, ClassBalancedStoragePolicy
+from avalanche.training.plugins import StrategyPlugin
 from avalanche.training.plugins.lr_scheduling import LRSchedulerPlugin
 from avalanche.training.strategies import Naive
 
diff --git a/tests/training/test_replay.py b/tests/training/test_replay.py
index 1e749872..d2397eeb 100644
--- a/tests/training/test_replay.py
+++ b/tests/training/test_replay.py
@@ -12,9 +12,9 @@ from avalanche.benchmarks.utils import AvalancheDataset, AvalancheDatasetType, \\
 from avalanche.models import SimpleMLP
 from avalanche.training.plugins import ReplayPlugin
 from avalanche.training.storage_policy import \\
-    ExperienceBalancedStoragePolicy, ClassBalancedStoragePolicy, \\
-    ClassExemplarsSelectionStrategy, \\
-    HerdingSelectionStrategy, ClosestToCenterSelectionStrategy
+    ExperienceBalancedBuffer, ClassBalancedBuffer, \\
+    ExemplarsSelectionStrategy, \\
+    HerdingSelectionStrategy, ClosestToCenterSelectionStrategy, ParametricBuffer
 from avalanche.training.strategies import Naive, BaseStrategy
 from tests.unit_tests_utils import get_fast_benchmark
 
@@ -22,9 +22,11 @@ from tests.unit_tests_utils import get_fast_benchmark
 class ReplayTest(unittest.TestCase):
     def test_replay_balanced_memory(self):
         mem_size = 25
-        policies = [None,
-                    ExperienceBalancedStoragePolicy({}, mem_size=mem_size),
-                    ClassBalancedStoragePolicy({}, mem_size=mem_size)]
+        policies = [
+            None,
+            ExperienceBalancedBuffer(max_size=mem_size),
+            ClassBalancedBuffer(max_size=mem_size)
+        ]
         for policy in policies:
             self._test_replay_balanced_memory(policy, mem_size)
 
@@ -45,21 +47,19 @@ class ReplayTest(unittest.TestCase):
             n_seen_data += len(step.dataset)
             mem_fill = min(mem_size, n_seen_data)
             cl_strategy.train(step)
-            ext_mem = replayPlugin.ext_mem
             lengths = []
-            for task_id in ext_mem.keys():
-                lengths.append(len(ext_mem[task_id]))
+            for d in replayPlugin.storage_policy.buffer_datasets:
+                lengths.append(len(d))
             self.assertEqual(sum(lengths), mem_fill)  # Always fully filled
 
     def test_balancing(self):
-        p1 = ExperienceBalancedStoragePolicy({}, 100, adaptive_size=True)
-        p2 = ClassBalancedStoragePolicy({}, 100, adaptive_size=True)
+        p1 = ExperienceBalancedBuffer(100, adaptive_size=True)
+        p2 = ClassBalancedBuffer(100, adaptive_size=True)
 
         for policy in [p1, p2]:
             self.assert_balancing(policy)
 
     def assert_balancing(self, policy):
-        ext_mem = policy.ext_mem
         benchmark = get_fast_benchmark(use_task_labels=True)
         replay = ReplayPlugin(mem_size=100, storage_policy=policy)
         model = SimpleMLP(num_classes=benchmark.n_classes)
@@ -73,70 +73,37 @@ class ReplayTest(unittest.TestCase):
 
         for exp in benchmark.train_stream:
             cl_strategy.train(exp)
-            print(list(ext_mem.keys()), [len(el) for el in ext_mem.values()])
+
+            ext_mem = policy.buffer_groups
+            ext_mem_data = policy.buffer_datasets
+            print(list(ext_mem.keys()), [len(el) for el in ext_mem_data])
 
             # buffer size should equal self.mem_size if data is large enough
-            len_tot = sum([len(el) for el in ext_mem.values()])
-            assert len_tot == policy.mem_size
+            len_tot = sum([len(el) for el in ext_mem_data])
+            assert len_tot == policy.max_size
 
 
-class ClassBalancePolicyTest(unittest.TestCase):
+class ParametricBufferTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.memory = {}
-        self.policy = ClassBalancedStoragePolicy(self.memory, mem_size=4)
-
-    def test_store_alone_with_enough_memory(self):
-        order = [2, 0, 1, 3]
-
-        self.observe_exemplars({0: list(range(4))}, selection_order=order)
-
-        self.assert_memory_equal({0: order})
-
-    def test_store_alone_without_enough_memory(self):
-        order = [6, 4, 2, 0, 1, 3, 5]
-
-        self.observe_exemplars({0: list(range(7))}, selection_order=order)
-
-        self.assert_memory_equal({0: order[:4]})
-
-    def test_store_multiple_with_enough_memory(self):
-        self.observe_exemplars({0: [0, 1], 1: [10, 11]}, selection_order=[1, 0])
-
-        self.assert_memory_equal({0: [1, 0], 1: [11, 10]})
-
-    def test_store_multiple_without_enough_memory(self):
-        self.observe_exemplars({0: [0, 1, 2], 1: [10, 11, 12]},
-                               selection_order=[2, 0, 1])
-
-        self.assert_memory_equal({0: [2, 0], 1: [12, 10]})
-
-    def test_sequence(self):
-        # 1st observation
-        self.observe_exemplars({0: [0, 1], 1: [10, 11]}, selection_order=[1, 0])
-        self.assert_memory_equal({0: [1, 0], 1: [11, 10]})
-
-        # 2nd observation
-        self.observe_exemplars({2: [20, 21, 22], 3: [30, 31, 32]},
-                               selection_order=[2, 1, 0])
-        self.assert_memory_equal({0: [1], 1: [11], 2: [22], 3: [32]})
-
-    def observe_exemplars(self, class2exemplars: Dict[int, List[int]],
-                          selection_order: List[int]):
-        self.policy.selection_strategy = FixedSelectionStrategy(selection_order)
-        x = tensor(
-            [i for exemplars in class2exemplars.values() for i in exemplars])
-        y = tensor(
-            [class_id for class_id, exemplars in class2exemplars.items() for _
-             in exemplars]).long()
-        dataset = AvalancheTensorDataset(
-            x, y, dataset_type=AvalancheDatasetType.CLASSIFICATION)
-
-        self.policy(MagicMock(experience=MagicMock(dataset=dataset)))
-
-    def assert_memory_equal(self, class2exemplars: Dict[int, List[int]]):
-        self.assertEqual(class2exemplars,
-                         {class_id: [x.tolist() for x, *_ in memory] for
-                          class_id, memory in self.memory.items()})
+        self.benchmark = get_fast_benchmark(use_task_labels=True)
+
+    def test_groupings(self):
+        policies = [
+            ParametricBuffer(max_size=10, groupby=None),
+            ParametricBuffer(max_size=10, groupby='class'),
+            ParametricBuffer(max_size=10, groupby='task'),
+            ParametricBuffer(max_size=10, groupby='experience'),
+        ]
+        for p in policies:
+            self._test_policy(p)
+
+    def _test_policy(self, policy):
+        for exp in self.benchmark.train_stream:
+            policy.update(MagicMock(experience=exp))
+            groups_lens = [(g_id, len(g_data.buffer)) for g_id, g_data in
+                           policy.buffer_groups.items()]
+            print(groups_lens)
+        print("DONE.")
 
 
 class SelectionStrategyTest(unittest.TestCase):
@@ -195,7 +162,7 @@ class AbsLayer(Module):
         return torch.abs(x).reshape((-1, 1))
 
 
-class FixedSelectionStrategy(ClassExemplarsSelectionStrategy):
+class FixedSelectionStrategy(ExemplarsSelectionStrategy):
     """This is a fake strategy used for testing the policy behavior"""
 
     def __init__(self, indices: List[int]):