Shortcuts

Source code for torchrl.data.replay_buffers.samplers

# 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

import json
import textwrap
import warnings
from abc import ABC, abstractmethod
from collections import OrderedDict
from copy import copy, deepcopy
from multiprocessing.context import get_spawning_popen
from pathlib import Path
from typing import Any, Dict, Tuple, Union

import numpy as np
import torch

from tensordict import MemoryMappedTensor, TensorDict
from tensordict.utils import NestedKey

from torchrl._extension import EXTENSION_WARNING

from torchrl._utils import _replace_last, logger
from torchrl.data.replay_buffers.storages import Storage, StorageEnsemble, TensorStorage
from torchrl.data.replay_buffers.utils import _is_int, unravel_index

try:
    from torchrl._torchrl import (
        MinSegmentTreeFp32,
        MinSegmentTreeFp64,
        SumSegmentTreeFp32,
        SumSegmentTreeFp64,
    )
except ImportError:
    warnings.warn(EXTENSION_WARNING)

_EMPTY_STORAGE_ERROR = "Cannot sample from an empty storage."


[docs]class Sampler(ABC): """A generic sampler base class for composable Replay Buffers.""" # Some samplers - mainly those without replacement - # need to keep track of the number of remaining batches _remaining_batches = int(torch.iinfo(torch.int64).max) # The RNG is set by the replay buffer _rng: torch.Generator | None = None @abstractmethod def sample(self, storage: Storage, batch_size: int) -> Tuple[Any, dict]: ... def add(self, index: int) -> None: return def extend(self, index: torch.Tensor) -> None: return def update_priority( self, index: Union[int, torch.Tensor], priority: Union[float, torch.Tensor], *, storage: Storage | None = None, ) -> dict | None: warnings.warn( f"Calling update_priority() on a sampler {type(self).__name__} that is not prioritized. Make sure this is the indented behavior." ) return def mark_update( self, index: Union[int, torch.Tensor], *, storage: Storage | None = None ) -> None: return @property def default_priority(self) -> float: return 1.0 @abstractmethod def state_dict(self) -> Dict[str, Any]: ... @abstractmethod def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ... @property def ran_out(self) -> bool: # by default, samplers never run out return False @abstractmethod def _empty(self): ... @abstractmethod def dumps(self, path): ... @abstractmethod def loads(self, path): ... def __repr__(self): return f"{self.__class__.__name__}()" def __getstate__(self): state = copy(self.__dict__) state["_rng"] = None return state
[docs]class RandomSampler(Sampler): """A uniformly random sampler for composable replay buffers. Args: batch_size (int, optional): if provided, the batch size to be used by the replay buffer when calling :meth:`~.ReplayBuffer.sample`. """ def sample(self, storage: Storage, batch_size: int) -> Tuple[torch.Tensor, dict]: if len(storage) == 0: raise RuntimeError(_EMPTY_STORAGE_ERROR) index = storage._rand_given_ndim(batch_size) return index, {} def _empty(self): pass def dumps(self, path): # no op ... def loads(self, path): # no op ... def state_dict(self) -> Dict[str, Any]: return {} def load_state_dict(self, state_dict: Dict[str, Any]) -> None: return
[docs]class SamplerWithoutReplacement(Sampler): """A data-consuming sampler that ensures that the same sample is not present in consecutive batches. Args: drop_last (bool, optional): if ``True``, the last incomplete sample (if any) will be dropped. If ``False``, this last sample will be kept and (unlike with torch dataloaders) completed with other samples from a fresh indices permutation. Defaults to ``False``. shuffle (bool, optional): if ``False``, the items are not randomly permuted. This enables to iterate over the replay buffer in the order the data was collected. Defaults to ``True``. *Caution*: If the size of the storage changes in between two calls, the samples will be re-shuffled (as we can't generally keep track of which samples have been sampled before and which haven't). Similarly, it is expected that the storage content remains the same in between two calls, but this is not enforced. When the sampler reaches the end of the list of available indices, a new sample order will be generated and the resulting indices will be completed with this new draw, which can lead to duplicated indices, unless the :obj:`drop_last` argument is set to ``True``. """ def __init__(self, drop_last: bool = False, shuffle: bool = True): self._sample_list = None self.len_storage = 0 self.drop_last = drop_last self._ran_out = False self.shuffle = shuffle def dumps(self, path): path = Path(path) path.mkdir(exist_ok=True) TensorDict(self.state_dict()).memmap(path) def loads(self, path): sd = TensorDict.load_memmap(path).to_dict() self.load_state_dict(sd) def _get_sample_list(self, storage: Storage, len_storage: int, batch_size: int): if storage is None: device = self._sample_list.device else: device = storage.device if hasattr(storage, "device") else None if self.shuffle: _sample_list = torch.randperm( len_storage, device=device, generator=self._rng ) else: _sample_list = torch.arange(len_storage, device=device) self._sample_list = _sample_list if self.drop_last: self._remaining_batches = self._sample_list.numel() // batch_size else: self._remaining_batches = -(self._sample_list.numel() // -batch_size) def _single_sample(self, len_storage, batch_size): index = self._sample_list[:batch_size] self._sample_list = self._sample_list[batch_size:] if self.drop_last: self._remaining_batches = self._sample_list.numel() // batch_size else: self._remaining_batches = -(self._sample_list.numel() // -batch_size) # check if we have enough elements for one more batch, assuming same batch size # will be used each time sample is called if self._sample_list.shape[0] == 0 or ( self.drop_last and len(self._sample_list) < batch_size ): self.ran_out = True self._get_sample_list( storage=None, len_storage=len_storage, batch_size=batch_size ) else: self.ran_out = False return index def _storage_len(self, storage): return len(storage) def sample( self, storage: Storage, batch_size: int ) -> Tuple[Any, dict]: # noqa: F811 len_storage = self._storage_len(storage) if len_storage == 0: raise RuntimeError(_EMPTY_STORAGE_ERROR) if not len_storage: raise RuntimeError("An empty storage was passed") if self.len_storage != len_storage or self._sample_list is None: self._get_sample_list(storage, len_storage, batch_size=batch_size) if len_storage < batch_size and self.drop_last: raise ValueError( f"The batch size ({batch_size}) is greater than the storage capacity ({len_storage}). " "This makes it impossible to return a sample without repeating indices. " "Consider changing the sampler class or turn the 'drop_last' argument to False." ) self.len_storage = len_storage index = self._single_sample(len_storage, batch_size) if storage.ndim > 1: index = unravel_index(index, storage.shape) # we 'always' return the indices. The 'drop_last' just instructs the # sampler to turn to `ran_out = True` whenever the next sample # will be too short. This will be read by the replay buffer # as a signal for an early break of the __iter__(). return index, {} @property def ran_out(self): return self._ran_out @ran_out.setter def ran_out(self, value): self._ran_out = value def _empty(self): self._sample_list = None self.len_storage = 0 self._ran_out = False def state_dict(self) -> Dict[str, Any]: return OrderedDict( len_storage=self.len_storage, _sample_list=self._sample_list, drop_last=self.drop_last, _ran_out=self._ran_out, ) def load_state_dict(self, state_dict: Dict[str, Any]) -> None: self.len_storage = state_dict["len_storage"] self._sample_list = state_dict["_sample_list"] self.drop_last = state_dict["drop_last"] self._ran_out = state_dict["_ran_out"] def __repr__(self): if self._sample_list is not None: perc = len(self._sample_list) / self.len_storage * 100 else: perc = 0.0 return f"{self.__class__.__name__}({perc: 4.4f}% sampled)"
[docs]class PrioritizedSampler(Sampler): """Prioritized sampler for replay buffer. Presented in "Schaul, T.; Quan, J.; Antonoglou, I.; and Silver, D. 2015. Prioritized experience replay." (https://arxiv.org/abs/1511.05952) Args: max_capacity (int): maximum capacity of the buffer. alpha (float): exponent α determines how much prioritization is used, with α = 0 corresponding to the uniform case. beta (float): importance sampling negative exponent. eps (float, optional): delta added to the priorities to ensure that the buffer does not contain null priorities. Defaults to 1e-8. reduction (str, optional): the reduction method for multidimensional tensordicts (ie stored trajectory). Can be one of "max", "min", "median" or "mean". max_priority_within_buffer (bool, optional): if ``True``, the max-priority is tracked within the buffer. When ``False``, the max-priority tracks the maximum value since the instantiation of the sampler. Examples: >>> from torchrl.data.replay_buffers import ReplayBuffer, LazyTensorStorage, PrioritizedSampler >>> from tensordict import TensorDict >>> rb = ReplayBuffer(storage=LazyTensorStorage(10), sampler=PrioritizedSampler(max_capacity=10, alpha=1.0, beta=1.0)) >>> priority = torch.tensor([0, 1000]) >>> data_0 = TensorDict({"reward": 0, "obs": [0], "action": [0], "priority": priority[0]}, []) >>> data_1 = TensorDict({"reward": 1, "obs": [1], "action": [2], "priority": priority[1]}, []) >>> rb.add(data_0) >>> rb.add(data_1) >>> rb.update_priority(torch.tensor([0, 1]), priority=priority) >>> sample, info = rb.sample(10, return_info=True) >>> print(sample) TensorDict( fields={ action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.int64, is_shared=False), obs: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.int64, is_shared=False), priority: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False), reward: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([10]), device=cpu, is_shared=False) >>> print(info) {'_weight': array([1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11, 1.e-11], dtype=float32), 'index': array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])} .. note:: Using a :class:`~torchrl.data.replay_buffers.TensorDictReplayBuffer` can smoothen the process of updating the priorities: >>> from torchrl.data.replay_buffers import TensorDictReplayBuffer as TDRB, LazyTensorStorage, PrioritizedSampler >>> from tensordict import TensorDict >>> rb = TDRB( ... storage=LazyTensorStorage(10), ... sampler=PrioritizedSampler(max_capacity=10, alpha=1.0, beta=1.0), ... priority_key="priority", # This kwarg isn't present in regular RBs ... ) >>> priority = torch.tensor([0, 1000]) >>> data_0 = TensorDict({"reward": 0, "obs": [0], "action": [0], "priority": priority[0]}, []) >>> data_1 = TensorDict({"reward": 1, "obs": [1], "action": [2], "priority": priority[1]}, []) >>> data = torch.stack([data_0, data_1]) >>> rb.extend(data) >>> rb.update_priority(data) # Reads the "priority" key as indicated in the constructor >>> sample, info = rb.sample(10, return_info=True) >>> print(sample['index']) # The index is packed with the tensordict tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) """ def __init__( self, max_capacity: int, alpha: float, beta: float, eps: float = 1e-8, dtype: torch.dtype = torch.float, reduction: str = "max", max_priority_within_buffer: bool = False, ) -> None: if alpha < 0: raise ValueError( f"alpha must be greater or equal than 0, got alpha={alpha}" ) if beta < 0: raise ValueError(f"beta must be greater or equal to 0, got beta={beta}") self._max_capacity = max_capacity self._alpha = alpha self._beta = beta self._eps = eps self.reduction = reduction self.dtype = dtype self._max_priority_within_buffer = max_priority_within_buffer self._init() def __repr__(self): return f"{self.__class__.__name__}(alpha={self._alpha}, beta={self._beta}, eps={self._eps}, reduction={self.reduction})" @property 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( f"Samplers of type {type(self)} cannot be shared between processes." ) return super().__getstate__() def _init(self): if self.dtype in (torch.float, torch.FloatType, torch.float32): self._sum_tree = SumSegmentTreeFp32(self._max_capacity) self._min_tree = MinSegmentTreeFp32(self._max_capacity) elif self.dtype in (torch.double, torch.DoubleTensor, torch.float64): self._sum_tree = SumSegmentTreeFp64(self._max_capacity) self._min_tree = MinSegmentTreeFp64(self._max_capacity) else: raise NotImplementedError( f"dtype {self.dtype} not supported by PrioritizedSampler" ) self._max_priority = None def _empty(self): self._init() @property def _max_priority(self): max_priority_index = self.__dict__.get("_max_priority") if max_priority_index is None: return (None, None) return max_priority_index @_max_priority.setter def _max_priority(self, value): self.__dict__["_max_priority"] = value def _maybe_erase_max_priority(self, index): if not self._max_priority_within_buffer: return max_priority_index = self._max_priority[1] if max_priority_index is None: return def check_index(index=index, max_priority_index=max_priority_index): if isinstance(index, torch.Tensor): # index can be 1d or 2d if index.ndim == 1: is_overwritten = (index == max_priority_index).any() else: is_overwritten = (index == max_priority_index).all(-1).any() elif isinstance(index, int): is_overwritten = index == max_priority_index elif isinstance(index, slice): # This won't work if called recursively is_overwritten = max_priority_index in range( index.indices(self._max_capacity) ) elif isinstance(index, tuple): is_overwritten = isinstance(max_priority_index, tuple) if is_overwritten: for idx, mpi in zip(index, max_priority_index): is_overwritten &= check_index(idx, mpi) else: raise TypeError(f"index of type {type(index)} is not recognized.") return is_overwritten is_overwritten = check_index() if is_overwritten: self._max_priority = None @property def default_priority(self) -> float: mp = self._max_priority[0] if mp is None: mp = 1 return (mp + self._eps) ** self._alpha def sample(self, storage: Storage, batch_size: int) -> torch.Tensor: if len(storage) == 0: raise RuntimeError(_EMPTY_STORAGE_ERROR) p_sum = self._sum_tree.query(0, len(storage)) p_min = self._min_tree.query(0, len(storage)) if p_sum <= 0: raise RuntimeError("non-positive p_sum") if p_min <= 0: raise RuntimeError("non-positive p_min") # For some undefined reason, only np.random works here. # All PT attempts fail, even when subsequently transformed into numpy if self._rng is None: mass = np.random.uniform(0.0, p_sum, size=batch_size) else: mass = torch.rand(batch_size, generator=self._rng) * p_sum # mass = torch.zeros(batch_size, dtype=torch.double).uniform_(0.0, p_sum) # mass = torch.rand(batch_size).mul_(p_sum) index = self._sum_tree.scan_lower_bound(mass) index = torch.as_tensor(index) if not index.ndim: index = index.unsqueeze(0) index.clamp_max_(len(storage) - 1) weight = torch.as_tensor(self._sum_tree[index]) # get indices where weight is 0 zero_weight = weight == 0 index = index while zero_weight.any(): index = torch.where(zero_weight, index - 1, index) if (index < 0).any(): raise RuntimeError("Failed to find a suitable index") weight = torch.as_tensor(self._sum_tree[index]) zero_weight = weight == 0 # Importance sampling weight formula: # w_i = (p_i / sum(p) * N) ^ (-beta) # weight_i = w_i / max(w) # weight_i = (p_i / sum(p) * N) ^ (-beta) / # ((min(p) / sum(p) * N) ^ (-beta)) # weight_i = ((p_i / sum(p) * N) / (min(p) / sum(p) * N)) ^ (-beta) # weight_i = (p_i / min(p)) ^ (-beta) # weight = np.power(weight / (p_min + self._eps), -self._beta) weight = torch.pow(weight / p_min, -self._beta) if storage.ndim > 1: index = unravel_index(index, storage.shape) return index, {"_weight": weight} return index, {"_weight": weight} def add(self, index: torch.Tensor | int) -> None: super().add(index) self._maybe_erase_max_priority(index) def extend(self, index: torch.Tensor | tuple) -> None: super().extend(index) self._maybe_erase_max_priority(index)
[docs] @torch.no_grad() def update_priority( self, index: Union[int, torch.Tensor], priority: Union[float, torch.Tensor], *, storage: TensorStorage | None = None, ) -> None: # noqa: D417 """Updates the priority of the data pointed by the index. Args: index (int or torch.Tensor): indexes of the priorities to be updated. priority (Number or torch.Tensor): new priorities of the indexed elements. Keyword Args: storage (Storage, optional): a storage used to map the Nd index size to the 1d size of the sum_tree and min_tree. Only required whenever ``index.ndim > 2``. """ priority = torch.as_tensor(priority, device=torch.device("cpu")).detach() index = torch.as_tensor(index, dtype=torch.long, device=torch.device("cpu")) # we need to reshape priority if it has more than one element or if it has # a different shape than index if priority.numel() > 1 and priority.shape != index.shape: try: priority = priority.reshape(index.shape[:1]) except Exception as err: raise RuntimeError( "priority should be a number or an iterable of the same " f"length as index. Got priority of shape {priority.shape} and index " f"{index.shape}." ) from err elif priority.numel() <= 1: priority = priority.squeeze() # MaxValueWriter will set -1 for items in the data that we don't want # to update. We therefore have to keep only the non-negative indices. if _is_int(index): if index == -1: return else: if index.ndim > 1: if storage is None: raise RuntimeError( "storage should be provided to Sampler.update_priority when the storage has more " "than one dimension." ) try: shape = storage.shape except AttributeError: raise AttributeError( "Could not retrieve the storage shape. If your storage is not a TensorStorage subclass " "or its shape isn't accessible via the shape attribute, submit an issue on GitHub." ) index = torch.as_tensor(np.ravel_multi_index(index.unbind(-1), shape)) valid_index = index >= 0 if not valid_index.any(): return if not valid_index.all(): index = index[valid_index] if priority.ndim: priority = priority[valid_index] max_p, max_p_idx = priority.max(dim=0) cur_max_priority, cur_max_priority_index = self._max_priority if cur_max_priority is None or max_p > cur_max_priority: cur_max_priority, cur_max_priority_index = self._max_priority = ( max_p, index[max_p_idx] if index.ndim else index, ) priority = torch.pow(priority + self._eps, self._alpha) self._sum_tree[index] = priority self._min_tree[index] = priority if ( self._max_priority_within_buffer and cur_max_priority_index is not None and (index == cur_max_priority_index).any() ): maxval, maxidx = torch.tensor( [self._sum_tree[i] for i in range(self._max_capacity)] ).max(0) self._max_priority = (maxval, maxidx)
def mark_update( self, index: Union[int, torch.Tensor], *, storage: Storage | None = None ) -> None: self.update_priority(index, self.default_priority, storage=storage) def state_dict(self) -> Dict[str, Any]: return { "_alpha": self._alpha, "_beta": self._beta, "_eps": self._eps, "_max_priority": self._max_priority, "_sum_tree": deepcopy(self._sum_tree), "_min_tree": deepcopy(self._min_tree), } def load_state_dict(self, state_dict: Dict[str, Any]) -> None: self._alpha = state_dict["_alpha"] self._beta = state_dict["_beta"] self._eps = state_dict["_eps"] self._max_priority = state_dict["_max_priority"] self._sum_tree = state_dict.pop("_sum_tree") self._min_tree = state_dict.pop("_min_tree") def dumps(self, path): path = Path(path).absolute() path.mkdir(exist_ok=True) try: mm_st = MemoryMappedTensor.from_filename( shape=(self._max_capacity,), dtype=torch.float64, filename=path / "sumtree.memmap", ) mm_mt = MemoryMappedTensor.from_filename( shape=(self._max_capacity,), dtype=torch.float64, filename=path / "mintree.memmap", ) except FileNotFoundError: mm_st = MemoryMappedTensor.empty( (self._max_capacity,), dtype=torch.float64, filename=path / "sumtree.memmap", ) mm_mt = MemoryMappedTensor.empty( (self._max_capacity,), dtype=torch.float64, filename=path / "mintree.memmap", ) mm_st.copy_( torch.as_tensor([self._sum_tree[i] for i in range(self._max_capacity)]) ) mm_mt.copy_( torch.as_tensor([self._min_tree[i] for i in range(self._max_capacity)]) ) with open(path / "sampler_metadata.json", "w") as file: json.dump( { "_alpha": self._alpha, "_beta": self._beta, "_eps": self._eps, "_max_priority": self._max_priority, "_max_capacity": self._max_capacity, }, file, ) def loads(self, path): path = Path(path).absolute() with open(path / "sampler_metadata.json", "r") as file: metadata = json.load(file) self._alpha = metadata["_alpha"] self._beta = metadata["_beta"] self._eps = metadata["_eps"] self._max_priority = metadata["_max_priority"] _max_capacity = metadata["_max_capacity"] if _max_capacity != self._max_capacity: raise RuntimeError( f"max capacity of loaded metadata ({_max_capacity}) differs from self._max_capacity ({self._max_capacity})." ) mm_st = MemoryMappedTensor.from_filename( shape=(self._max_capacity,), dtype=torch.float64, filename=path / "sumtree.memmap", ) mm_mt = MemoryMappedTensor.from_filename( shape=(self._max_capacity,), dtype=torch.float64, filename=path / "mintree.memmap", ) for i, elt in enumerate(mm_st.tolist()): self._sum_tree[i] = elt for i, elt in enumerate(mm_mt.tolist()): self._min_tree[i] = elt
[docs]class SliceSampler(Sampler): """Samples slices of data along the first dimension, given start and stop signals. This class samples sub-trajectories with replacement. For a version without replacement, see :class:`~torchrl.data.replay_buffers.samplers.SliceSamplerWithoutReplacement`. Keyword Args: num_slices (int): the number of slices to be sampled. The batch-size must be greater or equal to the ``num_slices`` argument. Exclusive with ``slice_len``. slice_len (int): the length of the slices to be sampled. The batch-size must be greater or equal to the ``slice_len`` argument and divisible by it. Exclusive with ``num_slices``. end_key (NestedKey, optional): the key indicating the end of a trajectory (or episode). Defaults to ``("next", "done")``. traj_key (NestedKey, optional): the key indicating the trajectories. Defaults to ``"episode"`` (commonly used across datasets in TorchRL). ends (torch.Tensor, optional): a 1d boolean tensor containing the end of run signals. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. If provided, it is assumed that the storage is at capacity and that if the last element of the ``ends`` tensor is ``False``, the same trajectory spans across end and beginning. trajectories (torch.Tensor, optional): a 1d integer tensor containing the run ids. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. If provided, it is assumed that the storage is at capacity and that if the last element of the trajectory tensor is identical to the first, the same trajectory spans across end and beginning. cache_values (bool, optional): to be used with static datasets. Will cache the start and end signal of the trajectory. This can be safely used even if the trajectory indices change during calls to :class:`~torchrl.data.ReplayBuffer.extend` as this operation will erase the cache. .. warning:: ``cache_values=True`` will not work if the sampler is used with a storage that is extended by another buffer. For instance: >>> buffer0 = ReplayBuffer(storage=storage, ... sampler=SliceSampler(num_slices=8, cache_values=True), ... writer=ImmutableWriter()) >>> buffer1 = ReplayBuffer(storage=storage, ... sampler=other_sampler) >>> # Wrong! Does not erase the buffer from the sampler of buffer0 >>> buffer1.extend(data) .. warning:: ``cache_values=True`` will not work as expected if the buffer is shared between processes and one process is responsible for writing and one process for sampling, as erasing the cache can only be done locally. truncated_key (NestedKey, optional): If not ``None``, this argument indicates where a truncated signal should be written in the output data. This is used to indicate to value estimators where the provided trajectory breaks. Defaults to ``("next", "truncated")``. This feature only works with :class:`~torchrl.data.replay_buffers.TensorDictReplayBuffer` instances (otherwise the truncated key is returned in the info dictionary returned by the :meth:`~torchrl.data.replay_buffers.ReplayBuffer.sample` method). strict_length (bool, optional): if ``False``, trajectories of length shorter than `slice_len` (or `batch_size // num_slices`) will be allowed to appear in the batch. If ``True``, trajectories shorted than required will be filtered out. Be mindful that this can result in effective `batch_size` shorter than the one asked for! Trajectories can be split using :func:`~torchrl.collectors.split_trajectories`. Defaults to ``True``. compile (bool or dict of kwargs, optional): if ``True``, the bottleneck of the :meth:`~sample` method will be compiled with :func:`~torch.compile`. Keyword arguments can also be passed to torch.compile with this arg. Defaults to ``False``. span (bool, int, Tuple[bool | int, bool | int], optional): if provided, the sampled trajectory will span across the left and/or the right. This means that possibly fewer elements will be provided than what was required. A boolean value means that at least one element will be sampled per trajectory. An integer `i` means that at least `slice_len - i` samples will be gathered for each sampled trajectory. Using tuples allows a fine grained control over the span on the left (beginning of the stored trajectory) and on the right (end of the stored trajectory). .. note:: To recover the trajectory splits in the storage, :class:`~torchrl.data.replay_buffers.samplers.SliceSampler` will first attempt to find the ``traj_key`` entry in the storage. If it cannot be found, the ``end_key`` will be used to reconstruct the episodes. Examples: >>> import torch >>> from tensordict import TensorDict >>> from torchrl.data.replay_buffers import LazyMemmapStorage, TensorDictReplayBuffer >>> from torchrl.data.replay_buffers.samplers import SliceSampler >>> torch.manual_seed(0) >>> rb = TensorDictReplayBuffer( ... storage=LazyMemmapStorage(1_000_000), ... sampler=SliceSampler(cache_values=True, num_slices=10), ... batch_size=320, ... ) >>> episode = torch.zeros(1000, dtype=torch.int) >>> episode[:300] = 1 >>> episode[300:550] = 2 >>> episode[550:700] = 3 >>> episode[700:] = 4 >>> data = TensorDict( ... { ... "episode": episode, ... "obs": torch.randn((3, 4, 5)).expand(1000, 3, 4, 5), ... "act": torch.randn((20,)).expand(1000, 20), ... "other": torch.randn((20, 50)).expand(1000, 20, 50), ... }, [1000] ... ) >>> rb.extend(data) >>> sample = rb.sample() >>> print("sample:", sample) >>> print("episodes", sample.get("episode").unique()) episodes tensor([1, 2, 3, 4], dtype=torch.int32) :class:`~torchrl.data.replay_buffers.SliceSampler` is default-compatible with most of TorchRL's datasets: Examples: >>> import torch >>> >>> from torchrl.data.datasets import RobosetExperienceReplay >>> from torchrl.data import SliceSampler >>> >>> torch.manual_seed(0) >>> num_slices = 10 >>> dataid = list(RobosetExperienceReplay.available_datasets)[0] >>> data = RobosetExperienceReplay(dataid, batch_size=320, sampler=SliceSampler(num_slices=num_slices)) >>> for batch in data: ... batch = batch.reshape(num_slices, -1) ... break >>> print("check that each batch only has one episode:", batch["episode"].unique(dim=1)) check that each batch only has one episode: tensor([[19], [14], [ 8], [10], [13], [ 4], [ 2], [ 3], [22], [ 8]]) """ def __init__( self, *, num_slices: int = None, slice_len: int = None, end_key: NestedKey | None = None, traj_key: NestedKey | None = None, ends: torch.Tensor | None = None, trajectories: torch.Tensor | None = None, cache_values: bool = False, truncated_key: NestedKey | None = ("next", "truncated"), strict_length: bool = True, compile: bool | dict = False, span: bool | int | Tuple[bool | int, bool | int] = False, ): self.num_slices = num_slices self.slice_len = slice_len self.end_key = end_key self.traj_key = traj_key self.truncated_key = truncated_key self.cache_values = cache_values self._fetch_traj = True self.strict_length = strict_length self._cache = {} if isinstance(span, (bool, int)): span = (span, span) self.span = span if trajectories is not None: if traj_key is not None or end_key: raise RuntimeError( "`trajectories` and `end_key` or `traj_key` are exclusive arguments." ) if ends is not None: raise RuntimeError("trajectories and ends are exclusive arguments.") if not cache_values: raise RuntimeError( "To be used, trajectories requires `cache_values` to be set to `True`." ) vals = self._find_start_stop_traj( trajectory=trajectories, at_capacity=True, ) self._cache["stop-and-length"] = vals elif ends is not None: if traj_key is not None or end_key: raise RuntimeError( "`ends` and `end_key` or `traj_key` are exclusive arguments." ) if trajectories is not None: raise RuntimeError("trajectories and ends are exclusive arguments.") if not cache_values: raise RuntimeError( "To be used, ends requires `cache_values` to be set to `True`." ) vals = self._find_start_stop_traj(end=ends, at_capacity=True) self._cache["stop-and-length"] = vals else: if end_key is None: end_key = ("next", "done") if traj_key is None: traj_key = "episode" self.end_key = end_key self.traj_key = traj_key if not ((num_slices is None) ^ (slice_len is None)): raise TypeError( "Either num_slices or slice_len must be not None, and not both. " f"Got num_slices={num_slices} and slice_len={slice_len}." ) self.compile = bool(compile) if self.compile: if isinstance(compile, dict): kwargs = compile else: kwargs = {} self._get_index = torch.compile(self._get_index, **kwargs) def __getstate__(self): if get_spawning_popen() is not None and self.cache_values: logger.warning( f"It seems you are sharing a {type(self).__name__} across processes with" f"cache_values=True. " f"While this isn't forbidden and could perfectly work if your dataset " f"is unaltered on both processes, remember that calling extend/add on" f"one process will NOT erase the cache on another process's sampler, " f"which will cause synchronization issues." ) state = super().__getstate__() state["_cache"] = {} return state def extend(self, index: torch.Tensor) -> None: super().extend(index) if self.cache_values: self._cache.clear() def add(self, index: torch.Tensor) -> None: super().add(index) if self.cache_values: self._cache.clear() def __repr__(self): return ( f"{self.__class__.__name__}(num_slices={self.num_slices}, " f"slice_len={self.slice_len}, " f"end_key={self.end_key}, " f"traj_key={self.traj_key}, " f"truncated_key={self.truncated_key}, " f"strict_length={self.strict_length})" ) @classmethod def _find_start_stop_traj( cls, *, trajectory=None, end=None, at_capacity: bool, cursor=None ): if trajectory is not None: # slower # _, stop_idx = torch.unique_consecutive(trajectory, return_counts=True) # stop_idx = stop_idx.cumsum(0) - 1 # even slower # t = trajectory.unsqueeze(0) # w = torch.tensor([1, -1], dtype=torch.int).view(1, 1, 2) # stop_idx = torch.conv1d(t, w).nonzero() # faster end = trajectory[:-1] != trajectory[1:] if not at_capacity: end = torch.cat([end, torch.ones_like(end[:1])], 0) else: end = torch.cat([end, trajectory[-1:] != trajectory[:1]], 0) length = trajectory.shape[0] else: # We presume that not done at the end means that the traj spans across end and beginning of storage length = end.shape[0] if not at_capacity: end = end.clone() end[length - 1] = True ndim = end.ndim if at_capacity: # we must have at least one end by traj to individuate trajectories # so if no end can be found we set it manually if cursor is not None: if isinstance(cursor, torch.Tensor): cursor = cursor[-1].item() elif isinstance(cursor, range): cursor = cursor[-1] if not _is_int(cursor): raise RuntimeError( "cursor should be an integer or a 1d tensor or a range." ) end = torch.index_fill( end, index=torch.tensor(cursor, device=end.device, dtype=torch.long), dim=0, value=1, ) if not end.any(0).all(): mask = ~end.any(0, True) mask = torch.cat([torch.zeros_like(end[:-1]), mask]) end = torch.masked_fill(mask, end, 1) if ndim == 0: raise RuntimeError( "Expected the end-of-trajectory signal to be at least 1-dimensional." ) return cls._end_to_start_stop(length=length, end=end) @staticmethod def _end_to_start_stop(end, length): # Using transpose ensures the start and stop are sorted the same way stop_idx = end.transpose(0, -1).nonzero() stop_idx[:, [0, -1]] = stop_idx[:, [-1, 0]].clone() # First build the start indices as the stop + 1, we'll shift it later start_idx = stop_idx.clone() start_idx[:, 0] += 1 start_idx[:, 0] %= end.shape[0] # shift start: to do this, we check when the non-first dim indices are identical # and get a mask like [False, True, True, False, True, ...] where False means # that there's a switch from one dim to another (ie, a switch from one element of the batch # to another). We roll this one step along the time dimension and these two # masks provide us with the indices of the permutation matrix we need # to apply to start_idx. if start_idx.shape[0] > 1: start_idx_mask = (start_idx[1:, 1:] == start_idx[:-1, 1:]).all(-1) m1 = torch.cat([torch.zeros_like(start_idx_mask[:1]), start_idx_mask]) m2 = torch.cat([start_idx_mask, torch.zeros_like(start_idx_mask[:1])]) start_idx_replace = torch.empty_like(start_idx) start_idx_replace[m1] = start_idx[m2] start_idx_replace[~m1] = start_idx[~m2] start_idx = start_idx_replace else: # In this case we have only one start and stop has already been set pass lengths = stop_idx[:, 0] - start_idx[:, 0] + 1 lengths[lengths <= 0] = lengths[lengths <= 0] + length return start_idx, stop_idx, lengths def _start_to_end(self, st: torch.Tensor, length: int): arange = torch.arange(length, device=st.device, dtype=st.dtype) ndims = st.shape[-1] - 1 if st.ndim else 0 if ndims: arange = torch.stack([arange] + [torch.zeros_like(arange)] * ndims, -1) else: arange = arange.unsqueeze(-1) if st.shape != arange.shape: # we do this to make sure that we're not broadcasting the start # wrong as a tensor with shape [N] can't be expanded to [N, 1] # without getting an error st = st.expand_as(arange) return arange + st def _tensor_slices_from_startend(self, seq_length, start, storage_length): # start is a 2d tensor resulting from nonzero() # seq_length is a 1d tensor indicating the desired length of each sequence if isinstance(seq_length, int): arange = torch.arange(seq_length, device=start.device, dtype=start.dtype) ndims = start.shape[-1] - 1 if (start.ndim - 1) else 0 if ndims: arange_reshaped = torch.empty( arange.shape + torch.Size([ndims + 1]), device=start.device, dtype=start.dtype, ) arange_reshaped[..., 0] = arange arange_reshaped[..., 1:] = 0 else: arange_reshaped = arange.unsqueeze(-1) arange_expanded = arange_reshaped.expand( torch.Size([start.shape[0]]) + arange_reshaped.shape ) if start.shape != arange_expanded.shape: n_missing_dims = arange_expanded.dim() - start.dim() start_expanded = start[ (slice(None),) + (None,) * n_missing_dims ].expand_as(arange_expanded) result = (start_expanded + arange_expanded).flatten(0, 1) else: # when padding is needed result = torch.cat( [ self._start_to_end(_start, _seq_len) for _start, _seq_len in zip(start, seq_length) ] ) result[:, 0] = result[:, 0] % storage_length return result def _get_stop_and_length(self, storage, fallback=True): if self.cache_values and "stop-and-length" in self._cache: return self._cache.get("stop-and-length") if self._fetch_traj: # We first try with the traj_key try: if isinstance(storage, TensorStorage): trajectory = storage[:][self._used_traj_key] else: try: trajectory = storage[:][self.traj_key] except Exception: raise RuntimeError( "Could not get a tensordict out of the storage, which is required for SliceSampler to compute the trajectories." ) vals = self._find_start_stop_traj( trajectory=trajectory.clone(), at_capacity=storage._is_full, cursor=getattr(storage, "_last_cursor", None), ) if self.cache_values: self._cache["stop-and-length"] = vals return vals except KeyError: if fallback: self._fetch_traj = False return self._get_stop_and_length(storage, fallback=False) raise else: try: try: done = storage[:][self.end_key] except Exception: raise RuntimeError( "Could not get a tensordict out of the storage, which is required for SliceSampler to compute the trajectories." ) vals = self._find_start_stop_traj( end=done.squeeze()[: len(storage)], at_capacity=storage._is_full, cursor=getattr(storage, "_last_cursor", None), ) if self.cache_values: self._cache["stop-and-length"] = vals return vals except KeyError: if fallback: self._fetch_traj = True return self._get_stop_and_length(storage, fallback=False) raise def _adjusted_batch_size(self, batch_size): if self.num_slices is not None: if batch_size % self.num_slices != 0: raise RuntimeError( f"The batch-size must be divisible by the number of slices, got " f"batch_size={batch_size} and num_slices={self.num_slices}." ) seq_length = batch_size // self.num_slices num_slices = self.num_slices else: if batch_size % self.slice_len != 0: raise RuntimeError( f"The batch-size must be divisible by the slice length, got " f"batch_size={batch_size} and slice_len={self.slice_len}." ) seq_length = self.slice_len num_slices = batch_size // self.slice_len return seq_length, num_slices def sample(self, storage: Storage, batch_size: int) -> Tuple[torch.Tensor, dict]: # pick up as many trajs as we need start_idx, stop_idx, lengths = self._get_stop_and_length(storage) # we have to make sure that the number of dims of the storage # is the same as the stop/start signals since we will # use these to sample the storage if start_idx.shape[1] != storage.ndim: raise RuntimeError( f"Expected the end-of-trajectory signal to be " f"{storage.ndim}-dimensional. Got a {start_idx.shape[1]} tensor " "instead." ) seq_length, num_slices = self._adjusted_batch_size(batch_size) storage_length = storage.shape[0] return self._sample_slices( lengths, start_idx, stop_idx, seq_length, num_slices, storage_length=storage_length, storage=storage, ) def _sample_slices( self, lengths: torch.Tensor, start_idx: torch.Tensor, stop_idx: torch.Tensor, seq_length: int, num_slices: int, storage_length: int, traj_idx: torch.Tensor | None = None, *, storage, ) -> Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]]: # start_idx and stop_idx are 2d tensors organized like a non-zero def get_traj_idx(maxval): return torch.randint( maxval, (num_slices,), device=lengths.device, generator=self._rng ) if (lengths < seq_length).any(): if self.strict_length: idx = lengths >= seq_length if not idx.any(): raise RuntimeError( f"Did not find a single trajectory with sufficient length (length range: {lengths.min()} - {lengths.max()} / required={seq_length}))." ) if ( isinstance(seq_length, torch.Tensor) and seq_length.shape == lengths.shape ): seq_length = seq_length[idx] lengths_idx = lengths[idx] start_idx = start_idx[idx] stop_idx = stop_idx[idx] if traj_idx is None: traj_idx = get_traj_idx(lengths_idx.shape[0]) else: # Here we must filter out the indices that correspond to trajectories # we don't want to keep. That could potentially lead to an empty sample. # The difficulty with this adjustment is that traj_idx points to a full # sequences of lengths, but we filter out part of it so we must # convert traj_idx to a boolean mask, index this mask with the # valid indices and then recover the nonzero. idx_mask = torch.zeros_like(idx) idx_mask[traj_idx] = True traj_idx = idx_mask[idx].nonzero().squeeze(-1) if not traj_idx.numel(): raise RuntimeError( "None of the provided indices pointed to a trajectory of " "sufficient length. Consider using strict_length=False for the " "sampler instead." ) num_slices = traj_idx.shape[0] del idx lengths = lengths_idx else: if traj_idx is None: traj_idx = get_traj_idx(lengths.shape[0]) else: num_slices = traj_idx.shape[0] # make seq_length a tensor with values clamped by lengths seq_length = lengths[traj_idx].clamp_max(seq_length) else: if traj_idx is None: traj_idx = get_traj_idx(lengths.shape[0]) else: num_slices = traj_idx.shape[0] return self._get_index( lengths=lengths, start_idx=start_idx, stop_idx=stop_idx, num_slices=num_slices, seq_length=seq_length, storage_length=storage_length, traj_idx=traj_idx, storage=storage, ) def _get_index( self, lengths: torch.Tensor, start_idx: torch.Tensor, stop_idx: torch.Tensor, seq_length: int, num_slices: int, storage_length: int, traj_idx: torch.Tensor | None = None, *, storage, ) -> Tuple[torch.Tensor, dict]: # end_point is the last possible index for start last_indexable_start = lengths[traj_idx] - seq_length + 1 if not self.span[1]: end_point = last_indexable_start elif self.span[1] is True: end_point = lengths[traj_idx] + 1 else: span_left = self.span[1] if span_left >= seq_length: raise ValueError( "The right and left span must be strictly lower than the sequence length" ) end_point = lengths[traj_idx] - span_left if not self.span[0]: start_point = 0 elif self.span[0] is True: start_point = -seq_length + 1 else: span_right = self.span[0] if span_right >= seq_length: raise ValueError( "The right and left span must be strictly lower than the sequence length" ) start_point = -span_right relative_starts = ( torch.rand(num_slices, device=lengths.device, generator=self._rng) * (end_point - start_point) ).floor().to(start_idx.dtype) + start_point if self.span[0]: out_of_traj = relative_starts < 0 if out_of_traj.any(): # a negative start means sampling fewer elements seq_length = torch.where( ~out_of_traj, seq_length, seq_length + relative_starts ) relative_starts = torch.where(~out_of_traj, relative_starts, 0) if self.span[1]: out_of_traj = relative_starts + seq_length > lengths[traj_idx] if out_of_traj.any(): # a negative start means sampling fewer elements seq_length = torch.minimum( seq_length, lengths[traj_idx] - relative_starts ) starts = torch.cat( [ (start_idx[traj_idx, 0] + relative_starts).unsqueeze(1), start_idx[traj_idx, 1:], ], 1, ) index = self._tensor_slices_from_startend(seq_length, starts, storage_length) if self.truncated_key is not None: truncated_key = self.truncated_key done_key = _replace_last(truncated_key, "done") terminated_key = _replace_last(truncated_key, "terminated") truncated = torch.zeros( (index.shape[0], 1), dtype=torch.bool, device=index.device ) if isinstance(seq_length, int): truncated.view(num_slices, -1)[:, -1] = 1 else: truncated[seq_length.cumsum(0) - 1] = 1 index = index.to(torch.long).unbind(-1) st_index = storage[index] try: done = st_index[done_key] | truncated except KeyError: done = truncated.clone() try: terminated = st_index[terminated_key] except KeyError: terminated = torch.zeros_like(truncated) return index, { truncated_key: truncated, done_key: done, terminated_key: terminated, } index = index.to(torch.long).unbind(-1) return index, {} @property def _used_traj_key(self): return self.__dict__.get("__used_traj_key", self.traj_key) @_used_traj_key.setter def _used_traj_key(self, value): self.__dict__["__used_traj_key"] = value @property def _used_end_key(self): return self.__dict__.get("__used_end_key", self.end_key) @_used_end_key.setter def _used_end_key(self, value): self.__dict__["__used_end_key"] = value def _empty(self): pass def dumps(self, path): # no op - cache does not need to be saved ... def loads(self, path): # no op ... def state_dict(self) -> Dict[str, Any]: return {} def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ...
[docs]class SliceSamplerWithoutReplacement(SliceSampler, SamplerWithoutReplacement): """Samples slices of data along the first dimension, given start and stop signals, without replacement. This class is to be used with static replay buffers or in between two replay buffer extensions. Extending the replay buffer will reset the the sampler, and continuous sampling without replacement is currently not allowed. Keyword Args: drop_last (bool, optional): if ``True``, the last incomplete sample (if any) will be dropped. If ``False``, this last sample will be kept. Defaults to ``False``. num_slices (int): the number of slices to be sampled. The batch-size must be greater or equal to the ``num_slices`` argument. Exclusive with ``slice_len``. slice_len (int): the length of the slices to be sampled. The batch-size must be greater or equal to the ``slice_len`` argument and divisible by it. Exclusive with ``num_slices``. end_key (NestedKey, optional): the key indicating the end of a trajectory (or episode). Defaults to ``("next", "done")``. traj_key (NestedKey, optional): the key indicating the trajectories. Defaults to ``"episode"`` (commonly used across datasets in TorchRL). ends (torch.Tensor, optional): a 1d boolean tensor containing the end of run signals. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. trajectories (torch.Tensor, optional): a 1d integer tensor containing the run ids. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. truncated_key (NestedKey, optional): If not ``None``, this argument indicates where a truncated signal should be written in the output data. This is used to indicate to value estimators where the provided trajectory breaks. Defaults to ``("next", "truncated")``. This feature only works with :class:`~torchrl.data.replay_buffers.TensorDictReplayBuffer` instances (otherwise the truncated key is returned in the info dictionary returned by the :meth:`~torchrl.data.replay_buffers.ReplayBuffer.sample` method). strict_length (bool, optional): if ``False``, trajectories of length shorter than `slice_len` (or `batch_size // num_slices`) will be allowed to appear in the batch. If ``True``, trajectories shorted than required will be filtered out. Be mindful that this can result in effective `batch_size` shorter than the one asked for! Trajectories can be split using :func:`~torchrl.collectors.split_trajectories`. Defaults to ``True``. shuffle (bool, optional): if ``False``, the order of the trajectories is not shuffled. Defaults to ``True``. compile (bool or dict of kwargs, optional): if ``True``, the bottleneck of the :meth:`~sample` method will be compiled with :func:`~torch.compile`. Keyword arguments can also be passed to torch.compile with this arg. Defaults to ``False``. .. note:: To recover the trajectory splits in the storage, :class:`~torchrl.data.replay_buffers.samplers.SliceSamplerWithoutReplacement` will first attempt to find the ``traj_key`` entry in the storage. If it cannot be found, the ``end_key`` will be used to reconstruct the episodes. Examples: >>> import torch >>> from tensordict import TensorDict >>> from torchrl.data.replay_buffers import LazyMemmapStorage, TensorDictReplayBuffer >>> from torchrl.data.replay_buffers.samplers import SliceSamplerWithoutReplacement >>> >>> rb = TensorDictReplayBuffer( ... storage=LazyMemmapStorage(1000), ... # asking for 10 slices for a total of 320 elements, ie, 10 trajectories of 32 transitions each ... sampler=SliceSamplerWithoutReplacement(num_slices=10), ... batch_size=320, ... ) >>> episode = torch.zeros(1000, dtype=torch.int) >>> episode[:300] = 1 >>> episode[300:550] = 2 >>> episode[550:700] = 3 >>> episode[700:] = 4 >>> data = TensorDict( ... { ... "episode": episode, ... "obs": torch.randn((3, 4, 5)).expand(1000, 3, 4, 5), ... "act": torch.randn((20,)).expand(1000, 20), ... "other": torch.randn((20, 50)).expand(1000, 20, 50), ... }, [1000] ... ) >>> rb.extend(data) >>> sample = rb.sample() >>> # since we want trajectories of 32 transitions but there are only 4 episodes to >>> # sample from, we only get 4 x 32 = 128 transitions in this batch >>> print("sample:", sample) >>> print("trajectories in sample", sample.get("episode").unique()) :class:`~torchrl.data.replay_buffers.SliceSamplerWithoutReplacement` is default-compatible with most of TorchRL's datasets, and allows users to consume datasets in a dataloader-like fashion: Examples: >>> import torch >>> >>> from torchrl.data.datasets import RobosetExperienceReplay >>> from torchrl.data import SliceSamplerWithoutReplacement >>> >>> torch.manual_seed(0) >>> num_slices = 10 >>> dataid = list(RobosetExperienceReplay.available_datasets)[0] >>> data = RobosetExperienceReplay(dataid, batch_size=320, ... sampler=SliceSamplerWithoutReplacement(num_slices=num_slices)) >>> # the last sample is kept, since drop_last=False by default >>> for i, batch in enumerate(data): ... print(batch.get("episode").unique()) tensor([ 5, 6, 8, 11, 12, 14, 16, 17, 19, 24]) tensor([ 1, 2, 7, 9, 10, 13, 15, 18, 21, 22]) tensor([ 0, 3, 4, 20, 23]) """ def __init__( self, *, num_slices: int | None = None, slice_len: int | None = None, drop_last: bool = False, end_key: NestedKey | None = None, traj_key: NestedKey | None = None, ends: torch.Tensor | None = None, trajectories: torch.Tensor | None = None, truncated_key: NestedKey | None = ("next", "truncated"), strict_length: bool = True, shuffle: bool = True, compile: bool | dict = False, ): SliceSampler.__init__( self, num_slices=num_slices, slice_len=slice_len, end_key=end_key, traj_key=traj_key, cache_values=True, truncated_key=truncated_key, strict_length=strict_length, ends=ends, trajectories=trajectories, compile=compile, ) SamplerWithoutReplacement.__init__(self, drop_last=drop_last, shuffle=shuffle) def __repr__(self): if self._sample_list is not None: perc = len(self._sample_list) / self.len_storage * 100 else: perc = 0 return ( f"{self.__class__.__name__}(" f"num_slices={self.num_slices}, " f"slice_len={self.slice_len}, " f"end_key={self.end_key}, " f"traj_key={self.traj_key}, " f"truncated_key={self.truncated_key}, " f"strict_length={self.strict_length}," f"{perc}% sampled)" ) def _empty(self): self._cache = {} SamplerWithoutReplacement._empty(self) def _storage_len(self, storage): return self._storage_len_buffer def sample( self, storage: Storage, batch_size: int ) -> Tuple[Tuple[torch.Tensor, ...], dict]: start_idx, stop_idx, lengths = self._get_stop_and_length(storage) # we have to make sure that the number of dims of the storage # is the same as the stop/start signals since we will # use these to sample the storage if start_idx.shape[1] != storage.ndim: raise RuntimeError( f"Expected the end-of-trajectory signal to be " f"{storage.ndim}-dimensional. Got a {start_idx.shape[1]} tensor " "instead." ) self._storage_len_buffer = len(start_idx) # first get indices of the trajectories we want to retrieve seq_length, num_slices = self._adjusted_batch_size(batch_size) indices, _ = SamplerWithoutReplacement.sample(self, storage, num_slices) storage_length = storage.shape[0] # traj_idx will either be a single tensor or a tuple that can be reorganized # like a non-zero through stacking. def tuple_to_tensor(traj_idx, lengths=lengths): if isinstance(traj_idx, tuple): traj_idx = torch.arange(len(storage), device=lengths.device).view( storage.shape )[traj_idx] return traj_idx idx, info = self._sample_slices( lengths, start_idx, stop_idx, seq_length, num_slices, storage_length, traj_idx=tuple_to_tensor(indices), storage=storage, ) return idx, info def state_dict(self) -> Dict[str, Any]: return SamplerWithoutReplacement.state_dict(self) def load_state_dict(self, state_dict: Dict[str, Any]) -> None: return SamplerWithoutReplacement.load_state_dict(self, state_dict)
[docs]class PrioritizedSliceSampler(SliceSampler, PrioritizedSampler): """Samples slices of data along the first dimension, given start and stop signals, using prioritized sampling. This class samples sub-trajectories with replacement following a priority weighting presented in "Schaul, T.; Quan, J.; Antonoglou, I.; and Silver, D. 2015. Prioritized experience replay." (https://arxiv.org/abs/1511.05952) For more info see :class:`~torchrl.data.replay_buffers.samplers.SliceSampler` and :class:`~torchrl.data.replay_buffers.samplers.PrioritizedSampler`. .. warning:: PrioritizedSliceSampler will look at the priorities of the individual transitions and sample the start points accordingly. This means that transitions with a low priority may as well appear in the samples if they follow another of higher priority, and transitions with a high priority but closer to the end of a trajectory may never be sampled if they cannot be used as start points. Currently, it is the user responsibility to aggregate priorities across items of a trajectory using :meth:`~.update_priority`. Args: alpha (float): exponent α determines how much prioritization is used, with α = 0 corresponding to the uniform case. beta (float): importance sampling negative exponent. eps (float, optional): delta added to the priorities to ensure that the buffer does not contain null priorities. Defaults to 1e-8. reduction (str, optional): the reduction method for multidimensional tensordicts (i.e., stored trajectory). Can be one of "max", "min", "median" or "mean". Keyword Args: num_slices (int): the number of slices to be sampled. The batch-size must be greater or equal to the ``num_slices`` argument. Exclusive with ``slice_len``. slice_len (int): the length of the slices to be sampled. The batch-size must be greater or equal to the ``slice_len`` argument and divisible by it. Exclusive with ``num_slices``. end_key (NestedKey, optional): the key indicating the end of a trajectory (or episode). Defaults to ``("next", "done")``. traj_key (NestedKey, optional): the key indicating the trajectories. Defaults to ``"episode"`` (commonly used across datasets in TorchRL). ends (torch.Tensor, optional): a 1d boolean tensor containing the end of run signals. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. trajectories (torch.Tensor, optional): a 1d integer tensor containing the run ids. To be used whenever the ``end_key`` or ``traj_key`` is expensive to get, or when this signal is readily available. Must be used with ``cache_values=True`` and cannot be used in conjunction with ``end_key`` or ``traj_key``. cache_values (bool, optional): to be used with static datasets. Will cache the start and end signal of the trajectory. This can be safely used even if the trajectory indices change during calls to :class:`~torchrl.data.ReplayBuffer.extend` as this operation will erase the cache. .. warning:: ``cache_values=True`` will not work if the sampler is used with a storage that is extended by another buffer. For instance: >>> buffer0 = ReplayBuffer(storage=storage, ... sampler=SliceSampler(num_slices=8, cache_values=True), ... writer=ImmutableWriter()) >>> buffer1 = ReplayBuffer(storage=storage, ... sampler=other_sampler) >>> # Wrong! Does not erase the buffer from the sampler of buffer0 >>> buffer1.extend(data) .. warning:: ``cache_values=True`` will not work as expected if the buffer is shared between processes and one process is responsible for writing and one process for sampling, as erasing the cache can only be done locally. truncated_key (NestedKey, optional): If not ``None``, this argument indicates where a truncated signal should be written in the output data. This is used to indicate to value estimators where the provided trajectory breaks. Defaults to ``("next", "truncated")``. This feature only works with :class:`~torchrl.data.replay_buffers.TensorDictReplayBuffer` instances (otherwise the truncated key is returned in the info dictionary returned by the :meth:`~torchrl.data.replay_buffers.ReplayBuffer.sample` method). strict_length (bool, optional): if ``False``, trajectories of length shorter than `slice_len` (or `batch_size // num_slices`) will be allowed to appear in the batch. If ``True``, trajectories shorted than required will be filtered out. Be mindful that this can result in effective `batch_size` shorter than the one asked for! Trajectories can be split using :func:`~torchrl.collectors.split_trajectories`. Defaults to ``True``. compile (bool or dict of kwargs, optional): if ``True``, the bottleneck of the :meth:`~sample` method will be compiled with :func:`~torch.compile`. Keyword arguments can also be passed to torch.compile with this arg. Defaults to ``False``. span (bool, int, Tuple[bool | int, bool | int], optional): if provided, the sampled trajectory will span across the left and/or the right. This means that possibly fewer elements will be provided than what was required. A boolean value means that at least one element will be sampled per trajectory. An integer `i` means that at least `slice_len - i` samples will be gathered for each sampled trajectory. Using tuples allows a fine grained control over the span on the left (beginning of the stored trajectory) and on the right (end of the stored trajectory). max_priority_within_buffer (bool, optional): if ``True``, the max-priority is tracked within the buffer. When ``False``, the max-priority tracks the maximum value since the instantiation of the sampler. Defaults to ``False``. Examples: >>> import torch >>> from torchrl.data.replay_buffers import TensorDictReplayBuffer, LazyMemmapStorage, PrioritizedSliceSampler >>> from tensordict import TensorDict >>> sampler = PrioritizedSliceSampler(max_capacity=9, num_slices=3, alpha=0.7, beta=0.9) >>> rb = TensorDictReplayBuffer(storage=LazyMemmapStorage(9), sampler=sampler, batch_size=6) >>> data = TensorDict( ... { ... "observation": torch.randn(9,16), ... "action": torch.randn(9, 1), ... "episode": torch.tensor([0,0,0,1,1,1,2,2,2], dtype=torch.long), ... "steps": torch.tensor([0,1,2,0,1,2,0,1,2], dtype=torch.long), ... ("next", "observation"): torch.randn(9,16), ... ("next", "reward"): torch.randn(9,1), ... ("next", "done"): torch.tensor([0,0,1,0,0,1,0,0,1], dtype=torch.bool).unsqueeze(1), ... }, ... batch_size=[9], ... ) >>> rb.extend(data) >>> sample, info = rb.sample(return_info=True) >>> print("episode", sample["episode"].tolist()) episode [2, 2, 2, 2, 1, 1] >>> print("steps", sample["steps"].tolist()) steps [1, 2, 0, 1, 1, 2] >>> print("weight", info["_weight"].tolist()) weight [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] >>> priority = torch.tensor([0,3,3,0,0,0,1,1,1]) >>> rb.update_priority(torch.arange(0,9,1), priority=priority) >>> sample, info = rb.sample(return_info=True) >>> print("episode", sample["episode"].tolist()) episode [2, 2, 2, 2, 2, 2] >>> print("steps", sample["steps"].tolist()) steps [1, 2, 0, 1, 0, 1] >>> print("weight", info["_weight"].tolist()) weight [9.120110917137936e-06, 9.120110917137936e-06, 9.120110917137936e-06, 9.120110917137936e-06, 9.120110917137936e-06, 9.120110917137936e-06] """ def __init__( self, max_capacity: int, alpha: float, beta: float, eps: float = 1e-8, dtype: torch.dtype = torch.float, reduction: str = "max", *, num_slices: int = None, slice_len: int = None, end_key: NestedKey | None = None, traj_key: NestedKey | None = None, ends: torch.Tensor | None = None, trajectories: torch.Tensor | None = None, cache_values: bool = False, truncated_key: NestedKey | None = ("next", "truncated"), strict_length: bool = True, compile: bool | dict = False, span: bool | int | Tuple[bool | int, bool | int] = False, max_priority_within_buffer: bool = False, ): SliceSampler.__init__( self, num_slices=num_slices, slice_len=slice_len, end_key=end_key, traj_key=traj_key, cache_values=cache_values, truncated_key=truncated_key, strict_length=strict_length, ends=ends, trajectories=trajectories, compile=compile, span=span, ) PrioritizedSampler.__init__( self, max_capacity=max_capacity, alpha=alpha, beta=beta, eps=eps, dtype=dtype, reduction=reduction, max_priority_within_buffer=max_priority_within_buffer, ) if self.span[0]: # Span left is hard to achieve because we need to sample 'negative' starts, but to sample # the start we rely on PrioritizedSampler which has no idea it's looking at trajectories. # # Another way to go about this would be to stochastically decrease the seq_length to # accommodate this but that would require to over-sample the starts too. # warnings.warn( f"Left spanning is disabled for {type(self).__name__} and will be automatically turned off. " f"If this feature is required, please file an issue on torchrl GitHub repo." ) self.span = (0, self.span[1]) def __repr__(self): return ( f"{self.__class__.__name__}(" f"num_slices={self.num_slices}, " f"slice_len={self.slice_len}, " f"end_key={self.end_key}, " f"traj_key={self.traj_key}, " f"truncated_key={self.truncated_key}, " f"strict_length={self.strict_length}," f"alpha={self._alpha}, " f"beta={self._beta}, " f"eps={self._eps}" ) def __getstate__(self): state = SliceSampler.__getstate__(self) state.update(PrioritizedSampler.__getstate__(self)) return state def mark_update( self, index: Union[int, torch.Tensor], *, storage: Storage | None = None ) -> None: return PrioritizedSampler.mark_update(self, index, storage=storage) def _padded_indices(self, shapes, arange) -> torch.Tensor: # this complex mumbo jumbo creates a left padded tensor with valid indices on the right, e.g. # tensor([[ 0, 1, 2, 3, 4], # [-1, -1, 5, 6, 7], # [-1, 8, 9, 10, 11]]) # where the -1 items on the left are padded values num_groups = shapes.shape[0] max_group_len = shapes.max() pad_lengths = max_group_len - shapes # Get all the start and end indices within arange for each group group_ends = shapes.cumsum(0) group_starts = torch.empty_like(group_ends) group_starts[0] = 0 group_starts[1:] = group_ends[:-1] pad = torch.empty( (num_groups, max_group_len), dtype=arange.dtype, device=arange.device ) for pad_row, group_start, group_end, pad_len in zip( pad, group_starts, group_ends, pad_lengths ): pad_row[:pad_len] = -1 pad_row[pad_len:] = arange[group_start:group_end] return pad def _preceding_stop_idx(self, storage, lengths, seq_length, start_idx): preceding_stop_idx = self._cache.get("preceding_stop_idx") if preceding_stop_idx is not None: return preceding_stop_idx arange = torch.arange(storage.shape.numel()) shapes = lengths.view(-1, 1).cpu() if not shapes.sum() - 1 == arange[-1]: raise RuntimeError("Wrong shapes / arange configuration") if not self.strict_length: # First, remove the starts from the arange # We do this because each traj can be sampled all_but_starts = torch.ones(arange.shape, dtype=torch.bool) starts = lengths.cumsum(0) starts = torch.cat([torch.zeros_like(starts[:1]), starts[:-1]]) all_but_starts[starts] = False arange = arange[all_but_starts] shapes = shapes - 1 pad = self._padded_indices(shapes, arange) _, span_right = self.span[0], self.span[1] if span_right and isinstance(span_right, bool): preceding_stop_idx = pad[:, -1:] else: # Mask the rightmost values of that padded tensor preceding_stop_idx = pad[:, -seq_length + 1 + span_right :] preceding_stop_idx = preceding_stop_idx[preceding_stop_idx >= 0] if storage._is_full: preceding_stop_idx = ( preceding_stop_idx + np.ravel_multi_index( tuple(start_idx[0].tolist()), storage._total_shape ) ) % storage._total_shape.numel() if self.cache_values: self._cache["preceding_stop_idx"] = preceding_stop_idx return preceding_stop_idx def sample(self, storage: Storage, batch_size: int) -> Tuple[torch.Tensor, dict]: # Sample `batch_size` indices representing the start of a slice. # The sampling is based on a weight vector. start_idx, stop_idx, lengths = self._get_stop_and_length(storage) seq_length, num_slices = self._adjusted_batch_size(batch_size) preceding_stop_idx = self._preceding_stop_idx( storage, lengths, seq_length, start_idx ) if storage.ndim > 1: # we need to convert indices of the permuted, flatten storage to indices in a flatten storage (not permuted) # This is because the lengths come as they would for a permuted storage preceding_stop_idx = unravel_index( preceding_stop_idx, (storage.shape[-1], *storage.shape[:-1]) ) preceding_stop_idx = (preceding_stop_idx[-1], *preceding_stop_idx[:-1]) preceding_stop_idx = torch.as_tensor( np.ravel_multi_index(preceding_stop_idx, storage.shape) ) # force to not sample index at the end of a trajectory vals = torch.tensor(self._sum_tree[preceding_stop_idx.cpu().numpy()]) self._sum_tree[preceding_stop_idx.cpu().numpy()] = 0.0 # and no need to update self._min_tree starts, info = PrioritizedSampler.sample( self, storage=storage, batch_size=batch_size // seq_length ) self._sum_tree[preceding_stop_idx.cpu().numpy()] = vals # We must truncate the seq_length if (1) not strict length or (2) span[1] if self.span[1] or not self.strict_length: if not isinstance(starts, torch.Tensor): starts_tensor = torch.stack(list(starts), dim=-1).to(stop_idx.device) else: starts_tensor = starts.unsqueeze(1).to(stop_idx.device) # Find the stop that comes after the start index # say start_tensor has shape [N, X] and stop_idx has shape [M, X] # diff will have shape [M, N, X] stop_idx_corr = stop_idx.clone() stop_idx_corr[:, 0] = torch.where( stop_idx[:, 0] < start_idx[:, 0], stop_idx[:, 0] + storage._len_along_dim0, stop_idx[:, 0], ) diff = stop_idx_corr.unsqueeze(1) - starts_tensor.unsqueeze(0) # filter out all items that don't belong to the same dim in the storage mask = (diff[:, :, 1:] != 0).any(-1) diff = diff[:, :, 0] diff[mask] = diff.max() + 1 diff = diff.reshape(-1, starts_tensor.shape[0]) # We remove all neg values from consideration diff[diff < 0] = diff.max() + 1 # Take the arg min along dim 0 (thereby reducing dim M) idx = diff.argmin(dim=0) stops = stop_idx_corr[idx, 0] # TODO: here things may not work bc we could have spanning trajs, # though I cannot show that it breaks in the tests if starts_tensor.ndim > 1: starts_tensor = starts_tensor[:, 0] seq_length = (stops - starts_tensor + 1).clamp_max(seq_length) if (seq_length <= 0).any(): raise RuntimeError( "failed to compute seq_length, please report this bug" ) if isinstance(starts, tuple): starts = torch.stack(starts, -1) # starts = torch.as_tensor(starts, device=lengths.device) info["_weight"] = torch.as_tensor(info["_weight"], device=lengths.device) # extends starting indices of each slice with sequence_length to get indices of all steps index = self._tensor_slices_from_startend( seq_length, starts, storage_length=storage.shape[0] ) # repeat the weight of each slice to match the number of steps info["_weight"] = torch.repeat_interleave(info["_weight"], seq_length) if self.truncated_key is not None: # following logics borrowed from SliceSampler truncated_key = self.truncated_key done_key = _replace_last(truncated_key, "done") terminated_key = _replace_last(truncated_key, "terminated") truncated = torch.zeros( (index.shape[0], 1), dtype=torch.bool, device=index.device ) if isinstance(seq_length, int): truncated.view(num_slices, -1)[:, -1] = 1 else: truncated[seq_length.cumsum(0) - 1] = 1 index = index.to(torch.long).unbind(-1) st_index = storage[index] try: done = st_index[done_key] | truncated except KeyError: done = truncated.clone() try: terminated = st_index[terminated_key] except KeyError: terminated = torch.zeros_like(truncated) info.update( { truncated_key: truncated, done_key: done, terminated_key: terminated, } ) return index, info return index.to(torch.long).unbind(-1), info def _empty(self): # no op for SliceSampler PrioritizedSampler._empty(self) def dumps(self, path): # no op for SliceSampler PrioritizedSampler.dumps(self, path) def loads(self, path): # no op for SliceSampler return PrioritizedSampler.loads(self, path) def state_dict(self): # no op for SliceSampler return PrioritizedSampler.state_dict(self) def add(self, index: torch.Tensor) -> None: PrioritizedSampler.add(self, index) return SliceSampler.add(self, index) def extend(self, index: torch.Tensor) -> None: PrioritizedSampler.extend(self, index) return SliceSampler.extend(self, index)
[docs]class SamplerEnsemble(Sampler): """An ensemble of samplers. This class is designed to work with :class:`~torchrl.data.replay_buffers.replay_buffers.ReplayBufferEnsemble`. It contains the samplers as well as the sampling strategy hyperparameters. Args: samplers (sequence of Sampler): the samplers to make the composite sampler. Keyword Args: p (list or tensor of probabilities, optional): if provided, indicates the weights of each dataset during sampling. sample_from_all (bool, optional): if ``True``, each dataset will be sampled from. This is not compatible with the ``p`` argument. Defaults to ``False``. num_buffer_sampled (int, optional): the number of buffers to sample. if ``sample_from_all=True``, this has no effect, as it defaults to the number of buffers. If ``sample_from_all=False``, buffers will be sampled according to the probabilities ``p``. .. warning:: The indices provided in the info dictionary are placed in a :class:`~tensordict.TensorDict` with keys ``index`` and ``buffer_ids`` that allow the upper :class:`~torchrl.data.ReplayBufferEnsemble` and :class:`~torchrl.data.StorageEnsemble` objects to retrieve the data. This format is different from with other samplers which usually return indices as regular tensors. """ def __init__( self, *samplers, p=None, sample_from_all=False, num_buffer_sampled=None ): self._rng_private = None self._samplers = samplers self.sample_from_all = sample_from_all if sample_from_all and p is not None: raise RuntimeError( "Cannot pass both `p` argument and `sample_from_all=True`." ) self.p = p self.num_buffer_sampled = num_buffer_sampled @property def _rng(self): return self._rng_private @_rng.setter def _rng(self, value): self._rng_private = value for sampler in self._samplers: sampler._rng = value @property def p(self): return self._p @p.setter def p(self, value): if not isinstance(value, torch.Tensor) and value is not None: value = torch.tensor(value) if value is not None: value = value / value.sum().clamp_min(1e-6) self._p = value @property def num_buffer_sampled(self): value = self.__dict__.get("_num_buffer_sampled", None) if value is None: value = self.__dict__["_num_buffer_sampled"] = len(self._samplers) return value @num_buffer_sampled.setter def num_buffer_sampled(self, value): self.__dict__["_num_buffer_sampled"] = value def sample(self, storage, batch_size): if batch_size % self.num_buffer_sampled > 0: raise ValueError if not isinstance(storage, StorageEnsemble): raise TypeError sub_batch_size = batch_size // self.num_buffer_sampled if self.sample_from_all: samples, infos = zip( *[ sampler.sample(storage, sub_batch_size) for storage, sampler in zip(storage._storages, self._samplers) ] ) buffer_ids = torch.arange(len(samples)) else: if self.p is None: buffer_ids = torch.randint( len(self._samplers), (self.num_buffer_sampled,), generator=self._rng, device=getattr(storage, "device", None), ) else: buffer_ids = torch.multinomial(self.p, self.num_buffer_sampled, True) samples, infos = zip( *[ self._samplers[i].sample(storage._storages[i], sub_batch_size) for i in buffer_ids.tolist() ] ) samples = [ sample if isinstance(sample, torch.Tensor) else torch.stack(sample, -1) for sample in samples ] if all(samples[0].shape == sample.shape for sample in samples[1:]): samples_stack = torch.stack(samples) else: samples_stack = torch.nested.nested_tensor(list(samples)) samples = TensorDict( { "index": samples_stack, "buffer_ids": buffer_ids, }, batch_size=[self.num_buffer_sampled], ) infos = torch.stack( [ TensorDict.from_dict(info, batch_dims=samples.ndim - 1) if info else TensorDict({}, []) for info in infos ] ) return samples, infos def dumps(self, path: Path): path = Path(path).absolute() for i, sampler in enumerate(self._samplers): sampler.dumps(path / str(i)) def loads(self, path: Path): path = Path(path).absolute() for i, sampler in enumerate(self._samplers): sampler.loads(path / str(i)) def state_dict(self) -> Dict[str, Any]: state_dict = OrderedDict() for i, sampler in enumerate(self._samplers): state_dict[str(i)] = sampler.state_dict() return state_dict def load_state_dict(self, state_dict: Dict[str, Any]) -> None: for i, sampler in enumerate(self._samplers): sampler.load_state_dict(state_dict[str(i)]) def _empty(self): raise NotImplementedError _INDEX_ERROR = "Expected an index of type torch.Tensor, range, np.ndarray, int, slice or ellipsis, got {} instead." def __getitem__(self, index): if isinstance(index, tuple): if index[0] is Ellipsis: index = (slice(None), index[1:]) result = self[index[0]] if len(index) > 1: raise IndexError( f"Tuple of length greater than 1 are not accepted to index samplers of type {type(self)}." ) return result if isinstance(index, slice) and index == slice(None): return self if isinstance(index, (list, range, np.ndarray)): index = torch.as_tensor(index) if isinstance(index, torch.Tensor): if index.ndim > 1: raise RuntimeError( f"Cannot index a {type(self)} with tensor indices that have more than one dimension." ) if index.is_floating_point(): raise TypeError( "A floating point index was recieved when an integer dtype was expected." ) if isinstance(index, int) or (not isinstance(index, slice) and len(index) == 0): try: index = int(index) except Exception: raise IndexError(self._INDEX_ERROR.format(type(index))) try: return self._samplers[index] except IndexError: raise IndexError(self._INDEX_ERROR.format(type(index))) if isinstance(index, torch.Tensor): index = index.tolist() samplers = [self._samplers[i] for i in index] else: # slice samplers = self._samplers[index] p = self._p[index] return SamplerEnsemble( *samplers, p=p, sample_from_all=self.sample_from_all, num_buffer_sampled=self.num_buffer_sampled, ) def __len__(self): return len(self._samplers) def __repr__(self): samplers = textwrap.indent(f"samplers={self._samplers}", " " * 4) return f"{self.__class__.__name__}(\n{samplers})"

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources