[Feature] Add scheduler for alpha/beta parameters of PrioritizedSampl…

…er (#2452)

Co-authored-by: Vincent Moens <[email protected]>

test/test_rb.py

@@ -59,6 +59,11 @@
     SliceSampler,
     SliceSamplerWithoutReplacement,
 )
+from torchrl.data.replay_buffers.scheduler import (
+    LinearScheduler,
+    SchedulerList,
+    StepScheduler,
+)
 
 from torchrl.data.replay_buffers.storages import (
     LazyMemmapStorage,
@@ -100,6 +105,7 @@
     VecNorm,
 )
 
+
 OLD_TORCH = parse(torch.__version__) < parse("2.0.0")
 _has_tv = importlib.util.find_spec("torchvision") is not None
 _has_gym = importlib.util.find_spec("gym") is not None
@@ -3041,6 +3047,77 @@ def test_prioritized_slice_sampler_episodes(device):
     ), "after priority update, only episode 1 and 3 are expected to be sampled"
 
 
+@pytest.mark.parametrize("alpha", [0.6, torch.tensor(1.0)])
+@pytest.mark.parametrize("beta", [0.7, torch.tensor(0.1)])
+@pytest.mark.parametrize("gamma", [0.1])
+@pytest.mark.parametrize("total_steps", [200])
+@pytest.mark.parametrize("n_annealing_steps", [100])
+@pytest.mark.parametrize("anneal_every_n", [10, 159])
+@pytest.mark.parametrize("alpha_min", [0, 0.2])
+@pytest.mark.parametrize("beta_max", [1, 1.4])
+def test_prioritized_parameter_scheduler(
+    alpha,
+    beta,
+    gamma,
+    total_steps,
+    n_annealing_steps,
+    anneal_every_n,
+    alpha_min,
+    beta_max,
+):
+    rb = TensorDictPrioritizedReplayBuffer(
+        alpha=alpha, beta=beta, storage=ListStorage(max_size=1000)
+    )
+    data = TensorDict({"data": torch.randn(1000, 5)}, batch_size=1000)
+    rb.extend(data)
+    alpha_scheduler = LinearScheduler(
+        rb, param_name="alpha", final_value=alpha_min, num_steps=n_annealing_steps
+    )
+    beta_scheduler = StepScheduler(
+        rb,
+        param_name="beta",
+        gamma=gamma,
+        n_steps=anneal_every_n,
+        max_value=beta_max,
+        mode="additive",
+    )
+
+    scheduler = SchedulerList(schedulers=(alpha_scheduler, beta_scheduler))
+
+    alpha = alpha if torch.is_tensor(alpha) else torch.tensor(alpha)
+    alpha_min = torch.tensor(alpha_min)
+    expected_alpha_vals = torch.linspace(alpha, alpha_min, n_annealing_steps + 1)
+    expected_alpha_vals = torch.nn.functional.pad(
+        expected_alpha_vals, (0, total_steps - n_annealing_steps), value=alpha_min
+    )
+
+    expected_beta_vals = [beta]
+    annealing_steps = total_steps // anneal_every_n
+    gammas = torch.arange(0, annealing_steps + 1, dtype=torch.float32) * gamma
+    expected_beta_vals = (
+        (beta + gammas).repeat_interleave(anneal_every_n).clip(None, beta_max)
+    )
+    for i in range(total_steps):
+        curr_alpha = rb.sampler.alpha
+        torch.testing.assert_close(
+            curr_alpha
+            if torch.is_tensor(curr_alpha)
+            else torch.tensor(curr_alpha).float(),
+            expected_alpha_vals[i],
+            msg=f"expected {expected_alpha_vals[i]}, got {curr_alpha}",
+        )
+        curr_beta = rb.sampler.beta
+        torch.testing.assert_close(
+            curr_beta
+            if torch.is_tensor(curr_beta)
+            else torch.tensor(curr_beta).float(),
+            expected_beta_vals[i],
+            msg=f"expected {expected_beta_vals[i]}, got {curr_beta}",
+        )
+        rb.sample(20)
+        scheduler.step()
+
+
 class TestEnsemble:
     def _make_data(self, data_type):
         if data_type is torch.Tensor:

torchrl/data/replay_buffers/samplers.py

@@ -395,6 +395,22 @@ def __repr__(self):
     def max_size(self):
         return self._max_capacity
 
+    @property
+    def alpha(self):
+        return self._alpha
+
+    @alpha.setter
+    def alpha(self, value):
+        self._alpha = value
+
+    @property
+    def beta(self):
+        return self._beta
+
+    @beta.setter
+    def beta(self, value):
+        self._beta = value
+
     def __getstate__(self):
         if get_spawning_popen() is not None:
             raise RuntimeError(

torchrl/data/replay_buffers/scheduler.py

@@ -0,0 +1,267 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from typing import Any, Callable, Dict
+
+import numpy as np
+
+import torch
+
+from torchrl.data.replay_buffers.replay_buffers import ReplayBuffer
+from torchrl.data.replay_buffers.samplers import Sampler
+
+
+class ParameterScheduler(ABC):
+    """Scheduler to adjust the value of a given parameter of a replay buffer's sampler.
+
+    Scheduler can for example be used to alter the alpha and beta values in the PrioritizedSampler.
+
+    Args:
+        obj (ReplayBuffer or Sampler): the replay buffer or sampler whose sampler to adjust
+        param_name (str): the name of the attribute to adjust, e.g. `beta` to adjust the beta parameter
+        min_value (Union[int, float], optional): a lower bound for the parameter to be adjusted
+            Defaults to `None`.
+        max_value (Union[int, float], optional): an upper bound for the parameter to be adjusted
+            Defaults to `None`.
+
+    """
+
+    def __init__(
+        self,
+        obj: ReplayBuffer | Sampler,
+        param_name: str,
+        min_value: int | float | None = None,
+        max_value: int | float | None = None,
+    ):
+        if not isinstance(obj, (ReplayBuffer, Sampler)):
+            raise TypeError(
+                f"ParameterScheduler only supports Sampler class. Pass either `ReplayBuffer` or `Sampler` object. Got {type(obj)} instead."
+            )
+        self.sampler = obj.sampler if isinstance(obj, ReplayBuffer) else obj
+        self.param_name = param_name
+        self._min_val = min_value or float("-inf")
+        self._max_val = max_value or float("inf")
+        if not hasattr(self.sampler, self.param_name):
+            raise ValueError(
+                f"Provided class {type(obj).__name__} does not have an attribute {param_name}"
+            )
+        initial_val = getattr(self.sampler, self.param_name)
+        if isinstance(initial_val, torch.Tensor):
+            initial_val = initial_val.clone()
+            self.backend = torch
+        else:
+            self.backend = np
+        self.initial_val = initial_val
+        self._step_cnt = 0
+
+    def state_dict(self):
+        """Returns the state of the scheduler as a :class:`dict`.
+
+        It contains an entry for every variable in ``self.__dict__`` which
+        is not the sampler.
+        """
+        sd = dict(self.__dict__)
+        del sd["sampler"]
+        return sd
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        """Load the scheduler's state.
+
+        Args:
+            state_dict (dict): scheduler state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        self.__dict__.update(state_dict)
+
+    def step(self):
+        self._step_cnt += 1
+        # Apply the step function
+        new_value = self._step()
+        # clip value to specified range
+        new_value_clipped = self.backend.clip(new_value, self._min_val, self._max_val)
+        # Set the new value of the parameter dynamically
+        setattr(self.sampler, self.param_name, new_value_clipped)
+
+    @abstractmethod
+    def _step(self):
+        ...
+
+
+class LambdaScheduler(ParameterScheduler):
+    """Sets a parameter to its initial value times a given function.
+
+    Similar to :class:`~torch.optim.LambdaLR`.
+
+    Args:
+        obj (ReplayBuffer or Sampler): the replay buffer whose sampler to adjust (or the sampler itself).
+        param_name (str): the name of the attribute to adjust, e.g. `beta` to adjust the
+            beta parameter.
+        lambda_fn (Callable[[int], float]): A function which computes a multiplicative factor given an integer
+            parameter ``step_count``.
+        min_value (Union[int, float], optional): a lower bound for the parameter to be adjusted
+            Defaults to `None`.
+        max_value (Union[int, float], optional): an upper bound for the parameter to be adjusted
+            Defaults to `None`.
+
+    """
+
+    def __init__(
+        self,
+        obj: ReplayBuffer | Sampler,
+        param_name: str,
+        lambda_fn: Callable[[int], float],
+        min_value: int | float | None = None,
+        max_value: int | float | None = None,
+    ):
+        super().__init__(obj, param_name, min_value, max_value)
+        self.lambda_fn = lambda_fn
+
+    def _step(self):
+        return self.initial_val * self.lambda_fn(self._step_cnt)
+
+
+class LinearScheduler(ParameterScheduler):
+    """A linear scheduler for gradually altering a parameter in an object over a given number of steps.
+
+    This scheduler linearly interpolates between the initial value of the parameter and a final target value.
+
+    Args:
+        obj (ReplayBuffer or Sampler): the replay buffer whose sampler to adjust (or the sampler itself).
+        param_name (str): the name of the attribute to adjust, e.g. `beta` to adjust the
+            beta parameter.
+        final_value (number): The final value that the parameter will reach after the
+            specified number of steps.
+        num_steps (number, optional): The total number of steps over which the parameter
+            will be linearly altered.
+
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> # Assuming sampler uses initial beta = 0.6
+        >>> # beta = 0.7   if step  == 1
+        >>> # beta = 0.8   if step  == 2
+        >>> # beta = 0.9   if step  == 3
+        >>> # beta = 1.0   if step  >= 4
+        >>> scheduler = LinearScheduler(sampler, param_name='beta', final_value=1.0, num_steps=4)
+        >>> for epoch in range(100):
+        >>>     train(...)
+        >>>     validate(...)
+        >>>     scheduler.step()
+    """
+
+    def __init__(
+        self,
+        obj: ReplayBuffer | Sampler,
+        param_name: str,
+        final_value: int | float,
+        num_steps: int,
+    ):
+        super().__init__(obj, param_name)
+        if isinstance(self.initial_val, torch.Tensor):
+            # cast to same type as initial value
+            final_value = torch.tensor(final_value).to(self.initial_val)
+        self.final_val = final_value
+        self.num_steps = num_steps
+        self._delta = (self.final_val - self.initial_val) / self.num_steps
+
+    def _step(self):
+        # Nit: we should use torch.where instead than if/else here to make the scheduler compatible with compile
+        #  without graph breaks
+        if self._step_cnt < self.num_steps:
+            return self.initial_val + (self._delta * self._step_cnt)
+        else:
+            return self.final_val
+
+
+class StepScheduler(ParameterScheduler):
+    """A step scheduler that alters a parameter after every n steps using either multiplicative or additive changes.
+
+    The scheduler can apply:
+    1. Multiplicative changes: `new_val = curr_val * gamma`
+    2. Additive changes: `new_val = curr_val + gamma`
+
+    Args:
+        obj (ReplayBuffer or Sampler): the replay buffer whose sampler to adjust (or the sampler itself).
+        param_name (str): the name of the attribute to adjust, e.g. `beta` to adjust the
+            beta parameter.
+        gamma (int or float, optional): The value by which to adjust the parameter,
+            either in a multiplicative or additive way.
+        n_steps (int, optional): The number of steps after which the parameter should be altered.
+            Defaults to 1.
+        mode (str, optional): The mode of scheduling. Can be either `'multiplicative'` or `'additive'`.
+            Defaults to `'multiplicative'`.
+        min_value (int or float, optional): a lower bound for the parameter to be adjusted.
+            Defaults to `None`.
+        max_value (int or float, optional): an upper bound for the parameter to be adjusted.
+            Defaults to `None`.
+
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> # Assuming sampler uses initial beta = 0.6
+        >>> # beta = 0.6   if  0 <= step < 10
+        >>> # beta = 0.7   if 10 <= step < 20
+        >>> # beta = 0.8   if 20 <= step < 30
+        >>> # beta = 0.9   if 30 <= step < 40
+        >>> # beta = 1.0   if 40 <= step
+        >>> scheduler = StepScheduler(sampler, param_name='beta', gamma=0.1, mode='additive', max_value=1.0)
+        >>> for epoch in range(100):
+        >>>     train(...)
+        >>>     validate(...)
+        >>>     scheduler.step()
+    """
+
+    def __init__(
+        self,
+        obj: ReplayBuffer | Sampler,
+        param_name: str,
+        gamma: int | float = 0.9,
+        n_steps: int = 1,
+        mode: str = "multiplicative",
+        min_value: int | float | None = None,
+        max_value: int | float | None = None,
+    ):
+
+        super().__init__(obj, param_name, min_value, max_value)
+        self.gamma = gamma
+        self.n_steps = n_steps
+        self.mode = mode
+        if mode == "additive":
+            operator = self.backend.add
+        elif mode == "multiplicative":
+            operator = self.backend.multiply
+        else:
+            raise ValueError(
+                f"Invalid mode: {mode}. Choose 'multiplicative' or 'additive'."
+            )
+        self.operator = operator
+
+    def _step(self):
+        """Applies the scheduling logic to alter the parameter value every `n_steps`."""
+        # Check if the current step count is a multiple of n_steps
+        current_val = getattr(self.sampler, self.param_name)
+        # Nit: we should use torch.where instead than if/else here to make the scheduler compatible with compile
+        #  without graph breaks
+        if self._step_cnt % self.n_steps == 0:
+            return self.operator(current_val, self.gamma)
+        else:
+            return current_val
+
+
+class SchedulerList:
+    """Simple container abstracting a list of schedulers."""
+
+    def __init__(self, schedulers: list[ParameterScheduler]) -> None:
+        if isinstance(schedulers, ParameterScheduler):
+            schedulers = [schedulers]
+        self.schedulers = schedulers
+
+    def append(self, scheduler: ParameterScheduler):
+        self.schedulers.append(scheduler)
+
+    def step(self):
+        for scheduler in self.schedulers:
+            scheduler.step()