torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/data/replay_buffers/replay_buffers.py

@@ -0,0 +1,2376 @@
+# 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 collections
+import contextlib
+import json
+import multiprocessing
+import textwrap
+import threading
+import warnings
+from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+try:
+    from torch.compiler import is_compiling
+except ImportError:
+    from torch._dynamo import is_compiling
+
+from functools import partial, wraps
+from typing import TYPE_CHECKING, TypeVar
+
+from tensordict import (
+    is_tensor_collection,
+    is_tensorclass,
+    LazyStackedTensorDict,
+    NestedKey,
+    TensorDict,
+    TensorDictBase,
+    unravel_key,
+)
+from tensordict.nn.utils import _set_dispatch_td_nn_modules
+from tensordict.utils import expand_as_right, expand_right
+from torch import Tensor
+from torch.utils._pytree import tree_map
+
+from torchrl._utils import accept_remote_rref_udf_invocation, rl_warnings
+from torchrl.data.replay_buffers.samplers import (
+    PrioritizedSampler,
+    RandomSampler,
+    Sampler,
+    SamplerEnsemble,
+)
+from torchrl.data.replay_buffers.storages import (
+    _get_default_collate,
+    _stack_anything,
+    ListStorage,
+    Storage,
+    StorageEnsemble,
+)
+from torchrl.data.replay_buffers.utils import (
+    _is_int,
+    _reduce,
+    _to_numpy,
+    _to_torch,
+    INT_CLASSES,
+    pin_memory_output,
+)
+from torchrl.data.replay_buffers.writers import (
+    RoundRobinWriter,
+    TensorDictRoundRobinWriter,
+    Writer,
+    WriterEnsemble,
+)
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.envs.transforms.transforms import _InvertTransform, Transform
+
+T = TypeVar("T")
+if TYPE_CHECKING:
+    from typing import Self
+else:
+    Self = T
+
+
+def _maybe_delay_init(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if self._delayed_init and not self.initialized:
+            self._init()
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
+class ReplayBuffer:
+    """A generic, composable replay buffer class.
+
+    Keyword Args:
+        storage (Storage, Callable[[], Storage], optional): the storage to be used.
+            If a callable is passed, it is used as constructor for the storage.
+            If none is provided a default :class:`~torchrl.data.replay_buffers.ListStorage` with
+            ``max_size`` of ``1_000`` will be created.
+        sampler (Sampler, Callable[[], Sampler], optional): the sampler to be used.
+            If a callable is passed, it is used as constructor for the sampler.
+            If none is provided, a default :class:`~torchrl.data.replay_buffers.RandomSampler`
+            will be used.
+        writer (Writer, Callable[[], Writer], optional): the writer to be used.
+            If a callable is passed, it is used as constructor for the writer.
+            If none is provided a default :class:`~torchrl.data.replay_buffers.RoundRobinWriter`
+            will be used.
+        collate_fn (callable, optional): merges a list of samples to form a
+            mini-batch of Tensor(s)/outputs.  Used when using batched
+            loading from a map-style dataset. The default value will be decided
+            based on the storage type.
+        pin_memory (bool): whether pin_memory() should be called on the rb
+            samples.
+        prefetch (int, optional): number of next batches to be prefetched
+            using multithreading. Defaults to None (no prefetching).
+        transform (Transform or Callable[[Any], Any], optional): Transform to be executed when
+            :meth:`sample` is called.
+            To chain transforms use the :class:`~torchrl.envs.Compose` class.
+            Transforms should be used with :class:`tensordict.TensorDict`
+            content. A generic callable can also be passed if the replay buffer
+            is used with PyTree structures (see example below).
+            Unlike storages, writers and samplers, transform constructors must
+            be passed as separate keyword argument :attr:`transform_factory`,
+            as it is impossible to distinguish a constructor from a transform.
+        transform_factory (Callable[[], Callable], optional): a factory for the
+            transform. Exclusive with :attr:`transform`.
+        batch_size (int, optional): the batch size to be used when sample() is
+            called.
+
+            .. note::
+              The batch-size can be specified at construction time via the
+              ``batch_size`` argument, or at sampling time. The former should
+              be preferred whenever the batch-size is consistent across the
+              experiment. If the batch-size is likely to change, it can be
+              passed to the :meth:`sample` method. This option is
+              incompatible with prefetching (since this requires to know the
+              batch-size in advance) as well as with samplers that have a
+              ``drop_last`` argument.
+
+        dim_extend (int, optional): indicates the dim to consider for
+            extension when calling :meth:`extend`. Defaults to ``storage.ndim-1``.
+            When using ``dim_extend > 0``, we recommend using the ``ndim``
+            argument in the storage instantiation if that argument is
+            available, to let storages know that the data is
+            multi-dimensional and keep consistent notions of storage-capacity
+            and batch-size during sampling.
+
+            .. note:: This argument has no effect on :meth:`add` and
+                therefore should be used with caution when both :meth:`add`
+                and :meth:`extend` are used in a codebase. For example:
+
+                    >>> data = torch.zeros(3, 4)
+                    >>> rb = ReplayBuffer(
+                    ...     storage=LazyTensorStorage(10, ndim=2),
+                    ...     dim_extend=1)
+                    >>> # these two approaches are equivalent:
+                    >>> for d in data.unbind(1):
+                    ...     rb.add(d)
+                    >>> rb.extend(data)
+
+        generator (torch.Generator, optional): a generator to use for sampling.
+            Using a dedicated generator for the replay buffer can allow a fine-grained control
+            over seeding, for instance keeping the global seed different but the RB seed identical
+            for distributed jobs.
+            Defaults to ``None`` (global default generator).
+
+            .. warning:: As of now, the generator has no effect on the transforms.
+        shared (bool, optional): whether the buffer will be shared using multiprocessing or not.
+            Defaults to ``False``.
+        compilable (bool, optional): whether the writer is compilable.
+            If ``True``, the writer cannot be shared between multiple processes.
+            Defaults to ``False``.
+        delayed_init (bool, optional): whether to initialize storage, writer, sampler and transform
+            the first time the buffer is used rather than during construction.
+            This is useful when the replay buffer needs to be pickled and sent to remote workers,
+            particularly when using transforms with modules that require gradients.
+            If not specified, defaults to ``True`` when ``transform_factory`` is provided,
+            and ``False`` otherwise.
+
+    Examples:
+        >>> import torch
+        >>>
+        >>> from torchrl.data import ReplayBuffer, ListStorage
+        >>>
+        >>> torch.manual_seed(0)
+        >>> rb = ReplayBuffer(
+        ...     storage=ListStorage(max_size=1000),
+        ...     batch_size=5,
+        ... )
+        >>> # populate the replay buffer and get the item indices
+        >>> data = range(10)
+        >>> indices = rb.extend(data)
+        >>> # sample will return as many elements as specified in the constructor
+        >>> sample = rb.sample()
+        >>> print(sample)
+        tensor([4, 9, 3, 0, 3])
+        >>> # Passing the batch-size to the sample method overrides the one in the constructor
+        >>> sample = rb.sample(batch_size=3)
+        >>> print(sample)
+        tensor([9, 7, 3])
+        >>> # one cans sample using the ``sample`` method or iterate over the buffer
+        >>> for i, batch in enumerate(rb):
+        ...     print(i, batch)
+        ...     if i == 3:
+        ...         break
+        0 tensor([7, 3, 1, 6, 6])
+        1 tensor([9, 8, 6, 6, 8])
+        2 tensor([4, 3, 6, 9, 1])
+        3 tensor([4, 4, 1, 9, 9])
+
+    Replay buffers accept *any* kind of data. Not all storage types
+    will work, as some expect numerical data only, but the default
+    :class:`~torchrl.data.ListStorage` will:
+
+    Examples:
+        >>> torch.manual_seed(0)
+        >>> buffer = ReplayBuffer(storage=ListStorage(100), collate_fn=lambda x: x)
+        >>> indices = buffer.extend(["a", 1, None])
+        >>> buffer.sample(3)
+        [None, 'a', None]
+
+    The :class:`~torchrl.data.replay_buffers.TensorStorage`, :class:`~torchrl.data.replay_buffers.LazyMemmapStorage`
+    and :class:`~torchrl.data.replay_buffers.LazyTensorStorage` also work
+    with any PyTree structure (a PyTree is a nested structure of arbitrary depth made of dicts,
+    lists or tuples where the leaves are tensors) provided that it only contains
+    tensor data.
+
+    Examples:
+        >>> from torch.utils._pytree import tree_map
+        >>> def transform(x):
+        ...     # Zeros all the data in the pytree
+        ...     return tree_map(lambda y: y * 0, x)
+        >>> rb = ReplayBuffer(storage=LazyMemmapStorage(100), transform=transform)
+        >>> data = {
+        ...     "a": torch.randn(3),
+        ...     "b": {"c": (torch.zeros(2), [torch.ones(1)])},
+        ...     30: -torch.ones(()),
+        ... }
+        >>> rb.add(data)
+        >>> # The sample has a similar structure to the data (with a leading dimension of 10 for each tensor)
+        >>> s = rb.sample(10)
+        >>> # let's check that our transform did its job:
+        >>> def assert0(x):
+        >>>     assert (x == 0).all()
+        >>> tree_map(assert0, s)
+
+    """
+
+    def __init__(
+        self,
+        *,
+        storage: Storage | Callable[[], Storage] | None = None,
+        sampler: Sampler | Callable[[], Sampler] | None = None,
+        writer: Writer | Callable[[], Writer] | None = None,
+        collate_fn: Callable | None = None,
+        pin_memory: bool = False,
+        prefetch: int | None = None,
+        transform: Transform | Callable | None = None,  # noqa-F821
+        transform_factory: Callable[[], Transform | Callable]
+        | None = None,  # noqa-F821
+        batch_size: int | None = None,
+        dim_extend: int | None = None,
+        checkpointer: StorageCheckpointerBase  # noqa: F821
+        | Callable[[], StorageCheckpointerBase]  # noqa: F821
+        | None = None,  # noqa: F821
+        generator: torch.Generator | None = None,
+        shared: bool = False,
+        compilable: bool | None = None,
+        delayed_init: bool | None = None,
+    ) -> None:
+        self._delayed_init = delayed_init
+        self._initialized = False
+
+        # Store init parameters for potential delayed initialization
+        self._init_storage = storage
+        self._init_sampler = sampler
+        self._init_writer = writer
+        self._init_collate_fn = collate_fn
+        self._init_transform = transform
+        self._init_transform_factory = transform_factory
+        self._init_checkpointer = checkpointer
+        self._init_generator = generator
+        self._init_compilable = compilable
+
+        if transform is not None and transform_factory is not None:
+            raise TypeError(
+                f"transform and transform_factory are mutually exclusive. "
+                f"Got transform={transform} and transform_factory={transform_factory}."
+            )
+
+        # Auto-detect delayed_init when transform_factory is provided
+        if transform_factory is not None and delayed_init is None:
+            delayed_init = True
+        elif delayed_init is None:
+            delayed_init = False
+
+        # Update _delayed_init after auto-detection
+        self._delayed_init = delayed_init
+
+        self._pin_memory = pin_memory
+        self._prefetch = bool(prefetch)
+        self._prefetch_cap = prefetch or 0
+        self._prefetch_queue = collections.deque()
+        self._batch_size = batch_size
+
+        if batch_size is None and prefetch:
+            raise ValueError(
+                "Dynamic batch-size specification is incompatible "
+                "with multithreaded sampling. "
+                "When using prefetch, the batch-size must be specified in "
+                "advance. "
+            )
+
+        if dim_extend is not None and dim_extend < 0:
+            raise ValueError("dim_extend must be a positive value.")
+        self._dim_extend = dim_extend
+
+        if self._prefetch_cap:
+            self._prefetch_executor = ThreadPoolExecutor(max_workers=self._prefetch_cap)
+
+        if shared and prefetch:
+            raise ValueError("Cannot share prefetched replay buffers.")
+        self.shared = shared
+        self.share(self.shared)
+
+        self._replay_lock = threading.RLock()
+        self._futures_lock = threading.RLock()
+
+        # If not delayed, initialize immediately
+        if not self._delayed_init:
+            self._init()
+
+    def _init(self) -> None:
+        """Initialize the replay buffer components.
+
+        This method is called either immediately during __init__ (if delayed_init=False)
+        or on first use of the buffer (if delayed_init=True).
+        """
+        if self._initialized:
+            return
+
+        self._initialized = True
+        try:
+            # Initialize storage
+            self._storage = self._maybe_make_storage(
+                self._init_storage, compilable=self._init_compilable
+            )
+            self._storage.attach(self)
+
+            # Initialize sampler
+            self._sampler = self._maybe_make_sampler(self._init_sampler)
+
+            # Initialize writer
+            self._writer = self._maybe_make_writer(self._init_writer)
+            self._writer.register_storage(self._storage)
+
+            # Initialize collate function
+            self._get_collate_fn(self._init_collate_fn)
+
+            # Initialize transform
+            self._transform = self._maybe_make_transform(
+                self._init_transform, self._init_transform_factory
+            )
+
+            # Check batch_size compatibility with sampler
+            if (
+                self._batch_size is None
+                and hasattr(self._sampler, "drop_last")
+                and self._sampler.drop_last
+            ):
+                raise ValueError(
+                    "Samplers with drop_last=True must work with a predictable batch-size. "
+                    "Please pass the batch-size to the ReplayBuffer constructor."
+                )
+
+            # Set dim_extend properly now that storage is initialized
+            if self._dim_extend is None:
+                if self._storage is not None:
+                    ndim = self._storage.ndim
+                    self._dim_extend = ndim - 1
+                else:
+                    self._dim_extend = 1
+
+            # Set checkpointer and generator
+            self._storage.checkpointer = self._init_checkpointer
+            self.set_rng(generator=self._init_generator)
+
+            # Initialize prioritized sampler if needed
+            self._initialize_prioritized_sampler()
+
+            # Remove init parameters
+            self._init_storage = None
+            self._init_sampler = None
+            self._init_writer = None
+            self._init_collate_fn = None
+            self._init_transform = None
+            self._init_transform_factory = None
+            self._init_checkpointer = None
+            self._init_generator = None
+            self._init_compilable = None
+        except Exception as e:
+            self._initialized = False
+            raise e
+
+    @property
+    def initialized(self) -> bool:
+        """Whether the replay buffer has been initialized."""
+        return self._initialized
+
+    def _initialize_prioritized_sampler(self) -> None:
+        """Initialize priority trees for existing data when using PrioritizedSampler.
+
+        This method ensures that when a PrioritizedSampler is used with storage that
+        already contains data, the priority trees are properly populated with default
+        priorities for all existing entries.
+        """
+        from .samplers import PrioritizedSampler
+
+        if isinstance(self._sampler, PrioritizedSampler) and len(self._storage) > 0:
+            # Set default priorities for all existing data
+            indices = torch.arange(len(self._storage), dtype=torch.long)
+            default_priorities = torch.full(
+                (len(self._storage),), self._sampler.default_priority, dtype=torch.float
+            )
+            self._sampler.update_priority(indices, default_priorities)
+
+    def _maybe_make_storage(
+        self, storage: Storage | Callable[[], Storage] | None, compilable
+    ) -> Storage:
+        if storage is None:
+            return ListStorage(max_size=1_000, compilable=compilable)
+        elif isinstance(storage, Storage):
+            return storage
+        elif callable(storage):
+            storage = storage()
+        if not isinstance(storage, Storage):
+            raise TypeError(
+                "storage must be either a Storage or a callable returning a storage instance."
+            )
+        return storage
+
+    def _maybe_make_sampler(
+        self, sampler: Sampler | Callable[[], Sampler] | None
+    ) -> Sampler:
+        if sampler is None:
+            return RandomSampler()
+        elif isinstance(sampler, Sampler):
+            return sampler
+        elif callable(sampler):
+            sampler = sampler()
+        if not isinstance(sampler, Sampler):
+            raise TypeError(
+                "sampler must be either a Sampler or a callable returning a sampler instance."
+            )
+        return sampler
+
+    def _maybe_make_writer(
+        self, writer: Writer | Callable[[], Writer] | None
+    ) -> Writer:
+        if writer is None:
+            return RoundRobinWriter()
+        elif isinstance(writer, Writer):
+            return writer
+        elif callable(writer):
+            writer = writer()
+        if not isinstance(writer, Writer):
+            raise TypeError(
+                "writer must be either a Writer or a callable returning a writer instance."
+            )
+        return writer
+
+    def _maybe_make_transform(
+        self,
+        transform: Transform | Callable[[], Transform] | None,
+        transform_factory: Callable | None,
+    ) -> Transform:
+        from torchrl.envs.transforms.transforms import (
+            _CallableTransform,
+            Compose,
+            Transform,
+        )
+
+        if transform_factory is not None:
+            if transform is not None:
+                raise TypeError(
+                    "transform and transform_factory cannot be used simultaneously"
+                )
+            transform = transform_factory()
+        if transform is None:
+            transform = Compose()
+        elif not isinstance(transform, Compose):
+            if not isinstance(transform, Transform) and callable(transform):
+                transform = _CallableTransform(transform)
+            elif not isinstance(transform, Transform):
+                raise RuntimeError(
+                    "transform must be either a Transform instance or a callable."
+                )
+            transform = Compose(transform)
+        transform.eval()
+        return transform
+
+    def share(self, shared: bool = True) -> Self:
+        self.shared = shared
+        if self.shared:
+            self._write_lock = multiprocessing.Lock()
+        else:
+            self._write_lock = contextlib.nullcontext()
+        return self
+
+    @_maybe_delay_init
+    def set_rng(self, generator) -> None:
+        self._rng = generator
+        self._storage._rng = generator
+        self._sampler._rng = generator
+        self._writer._rng = generator
+
+    @property
+    def dim_extend(self):
+        return self._dim_extend
+
+    @property
+    def batch_size(self):
+        """The batch size of the replay buffer.
+
+        The batch size can be overriden by setting the `batch_size` parameter in the :meth:`sample` method.
+
+        It defines both the number of samples returned by :meth:`sample` and the number of samples that are
+        yielded by the :class:`ReplayBuffer` iterator.
+        """
+        return self._batch_size
+
+    @dim_extend.setter
+    def dim_extend(self, value):
+        if (
+            hasattr(self, "_dim_extend")
+            and self._dim_extend is not None
+            and self._dim_extend != value
+        ):
+            raise RuntimeError(
+                "dim_extend cannot be reset. Please create a new replay buffer."
+            )
+
+        if value is None:
+            if self._initialized and self._storage is not None:
+                ndim = self._storage.ndim
+                value = ndim - 1
+            else:
+                value = 1
+
+        self._dim_extend = value
+
+    def _transpose(self, data):
+        if is_tensor_collection(data):
+            return data.transpose(self.dim_extend, 0)
+        return tree_map(lambda x: x.transpose(self.dim_extend, 0), data)
+
+    def _get_collate_fn(self, collate_fn):
+        self._collate_fn = (
+            collate_fn
+            if collate_fn is not None
+            else _get_default_collate(
+                self._storage, _is_tensordict=isinstance(self, TensorDictReplayBuffer)
+            )
+        )
+
+    @_maybe_delay_init
+    def set_storage(self, storage: Storage, collate_fn: Callable | None = None):
+        """Sets a new storage in the replay buffer and returns the previous storage.
+
+        Args:
+            storage (Storage): the new storage for the buffer.
+            collate_fn (callable, optional): if provided, the collate_fn is set to this
+                value. Otherwise it is reset to a default value.
+
+        """
+        prev_storage = self._storage
+        self._storage = storage
+        self._get_collate_fn(collate_fn)
+
+        return prev_storage
+
+    @_maybe_delay_init
+    def set_writer(self, writer: Writer):
+        """Sets a new writer in the replay buffer and returns the previous writer."""
+        prev_writer = self._writer
+        self._writer = writer
+        self._writer.register_storage(self._storage)
+        return prev_writer
+
+    @_maybe_delay_init
+    def set_sampler(self, sampler: Sampler):
+        """Sets a new sampler in the replay buffer and returns the previous sampler."""
+        prev_sampler = self._sampler
+        self._sampler = sampler
+        return prev_sampler
+
+    @_maybe_delay_init
+    def __len__(self) -> int:
+        with self._replay_lock:
+            return len(self._storage)
+
+    def _getattr(self, attr):
+        # To access properties in remote settings, see RayReplayBuffer.write_count for instance
+        return getattr(self, attr)
+
+    def _setattr(self, attr, value):
+        # To set properties in remote settings
+        setattr(self, attr, value)
+        return None  # explicit return for remote calls
+
+    @property
+    @_maybe_delay_init
+    def write_count(self) -> int:
+        """The total number of items written so far in the buffer through add and extend."""
+        return self._writer._write_count
+
+    def __repr__(self) -> str:
+        from torchrl.envs.transforms import Compose
+
+        storage = textwrap.indent(f"storage={getattr(self, '_storage', None)}", " " * 4)
+        writer = textwrap.indent(f"writer={getattr(self, '_writer', None)}", " " * 4)
+        sampler = textwrap.indent(f"sampler={getattr(self, '_sampler', None)}", " " * 4)
+        if getattr(self, "_transform", None) is not None and not (
+            isinstance(self._transform, Compose)
+            and not len(getattr(self, "_transform", None))
+        ):
+            transform = textwrap.indent(
+                f"transform={getattr(self, '_transform', None)}", " " * 4
+            )
+            transform = f"\n{self._transform}, "
+        else:
+            transform = ""
+        batch_size = textwrap.indent(
+            f"batch_size={getattr(self, '_batch_size', None)}", " " * 4
+        )
+        collate_fn = textwrap.indent(
+            f"collate_fn={getattr(self, '_collate_fn', None)}", " " * 4
+        )
+        return f"{self.__class__.__name__}(\n{storage}, \n{sampler}, \n{writer}, {transform}\n{batch_size}, \n{collate_fn})"
+
+    @_maybe_delay_init
+    @pin_memory_output
+    def __getitem__(self, index: int | torch.Tensor | NestedKey) -> Any:
+        if isinstance(index, str) or (isinstance(index, tuple) and unravel_key(index)):
+            return self[:][index]
+        if isinstance(index, tuple):
+            if len(index) == 1:
+                return self[index[0]]
+            else:
+                return self[:][index]
+        index = _to_numpy(index)
+
+        if self.dim_extend > 0:
+            index = (slice(None),) * self.dim_extend + (index,)
+            with self._replay_lock:
+                data = self._storage[index]
+            data = self._transpose(data)
+        else:
+            with self._replay_lock:
+                data = self._storage[index]
+
+        if not isinstance(index, INT_CLASSES):
+            data = self._collate_fn(data)
+
+        if self._transform is not None and len(self._transform):
+            with data.unlock_() if is_tensor_collection(
+                data
+            ) else contextlib.nullcontext():
+                data = self._transform(data)
+
+        return data
+
+    @_maybe_delay_init
+    def __setitem__(self, index, value) -> None:
+        if isinstance(index, str) or (isinstance(index, tuple) and unravel_key(index)):
+            self[:][index] = value
+            return
+        if isinstance(index, tuple):
+            if len(index) == 1:
+                self[index[0]] = value
+            else:
+                self[:][index] = value
+            return
+        index = _to_numpy(index)
+
+        if self._transform is not None and len(self._transform):
+            value = self._transform.inv(value)
+
+        if self.dim_extend > 0:
+            index = (slice(None),) * self.dim_extend + (index,)
+            with self._replay_lock:
+                self._storage[index] = self._transpose(value)
+        else:
+            with self._replay_lock:
+                self._storage[index] = value
+        return
+
+    @_maybe_delay_init
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "_storage": self._storage.state_dict(),
+            "_sampler": self._sampler.state_dict(),
+            "_writer": self._writer.state_dict(),
+            "_transforms": self._transform.state_dict(),
+            "_batch_size": self._batch_size,
+            "_rng": (self._rng.get_state().clone(), str(self._rng.device))
+            if self._rng is not None
+            else None,
+        }
+
+    @_maybe_delay_init
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._storage.load_state_dict(state_dict["_storage"])
+        self._sampler.load_state_dict(state_dict["_sampler"])
+        self._writer.load_state_dict(state_dict["_writer"])
+        self._transform.load_state_dict(state_dict["_transforms"])
+        self._batch_size = state_dict["_batch_size"]
+        rng = state_dict.get("_rng")
+        if rng is not None:
+            state, device = rng
+            rng = torch.Generator(device=device)
+            rng.set_state(state)
+            self.set_rng(generator=rng)
+
+    @_maybe_delay_init
+    def dumps(self, path):
+        """Saves the replay buffer on disk at the specified path.
+
+        Args:
+            path (Path or str): path where to save the replay buffer.
+
+        Examples:
+            >>> import tempfile
+            >>> import tqdm
+            >>> from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
+            >>> from torchrl.data.replay_buffers.samplers import PrioritizedSampler, RandomSampler
+            >>> import torch
+            >>> from tensordict import TensorDict
+            >>> # Build and populate the replay buffer
+            >>> S = 1_000_000
+            >>> sampler = PrioritizedSampler(S, 1.1, 1.0)
+            >>> # sampler = RandomSampler()
+            >>> storage = LazyMemmapStorage(S)
+            >>> rb = TensorDictReplayBuffer(storage=storage, sampler=sampler)
+            >>>
+            >>> for _ in tqdm.tqdm(range(100)):
+            ...     td = TensorDict({"obs": torch.randn(100, 3, 4), "next": {"obs": torch.randn(100, 3, 4)}, "td_error": torch.rand(100)}, [100])
+            ...     rb.extend(td)
+            ...     sample = rb.sample(32)
+            ...     rb.update_tensordict_priority(sample)
+            >>> # save and load the buffer
+            >>> with tempfile.TemporaryDirectory() as tmpdir:
+            ...     rb.dumps(tmpdir)
+            ...
+            ...     sampler = PrioritizedSampler(S, 1.1, 1.0)
+            ...     # sampler = RandomSampler()
+            ...     storage = LazyMemmapStorage(S)
+            ...     rb_load = TensorDictReplayBuffer(storage=storage, sampler=sampler)
+            ...     rb_load.loads(tmpdir)
+            ...     assert len(rb) == len(rb_load)
+
+        """
+        path = Path(path).absolute()
+        path.mkdir(exist_ok=True)
+        self._storage.dumps(path / "storage")
+        self._sampler.dumps(path / "sampler")
+        self._writer.dumps(path / "writer")
+        if self._rng is not None:
+            rng_state = TensorDict(
+                rng_state=self._rng.get_state().clone(),
+                device=self._rng.device,
+            )
+            rng_state.memmap(path / "rng_state")
+
+        # fall back on state_dict for transforms
+        transform_sd = self._transform.state_dict()
+        if transform_sd:
+            torch.save(transform_sd, path / "transform.t")
+        with open(path / "buffer_metadata.json", "w") as file:
+            json.dump({"batch_size": self._batch_size}, file)
+
+    @_maybe_delay_init
+    def loads(self, path):
+        """Loads a replay buffer state at the given path.
+
+        The buffer should have matching components and be saved using :meth:`dumps`.
+
+        Args:
+            path (Path or str): path where the replay buffer was saved.
+
+        See :meth:`dumps` for more info.
+
+        """
+        path = Path(path).absolute()
+        self._storage.loads(path / "storage")
+        self._sampler.loads(path / "sampler")
+        self._writer.loads(path / "writer")
+        if (path / "rng_state").exists():
+            rng_state = TensorDict.load_memmap(path / "rng_state")
+            rng = torch.Generator(device=rng_state.device)
+            rng.set_state(rng_state["rng_state"])
+            self.set_rng(rng)
+        # fall back on state_dict for transforms
+        if (path / "transform.t").exists():
+            self._transform.load_state_dict(torch.load(path / "transform.t"))
+        with open(path / "buffer_metadata.json") as file:
+            metadata = json.load(file)
+        self._batch_size = metadata["batch_size"]
+
+    @_maybe_delay_init
+    def save(self, *args, **kwargs):
+        """Alias for :meth:`dumps`."""
+        return self.dumps(*args, **kwargs)
+
+    @_maybe_delay_init
+    def dump(self, *args, **kwargs):
+        """Alias for :meth:`dumps`."""
+        return self.dumps(*args, **kwargs)
+
+    @_maybe_delay_init
+    def load(self, *args, **kwargs):
+        """Alias for :meth:`loads`."""
+        return self.loads(*args, **kwargs)
+
+    @_maybe_delay_init
+    def register_save_hook(self, hook: Callable[[Any], Any]):
+        """Registers a save hook for the storage.
+
+        .. note:: Hooks are currently not serialized when saving a replay buffer: they must
+            be manually re-initialized every time the buffer is created.
+
+        """
+        self._storage.register_save_hook(hook)
+
+    @_maybe_delay_init
+    def register_load_hook(self, hook: Callable[[Any], Any]):
+        """Registers a load hook for the storage.
+
+        .. note:: Hooks are currently not serialized when saving a replay buffer: they must
+            be manually re-initialized every time the buffer is created.
+
+        """
+        self._storage.register_load_hook(hook)
+
+    @_maybe_delay_init
+    def add(self, data: Any) -> int:
+        """Add a single element to the replay buffer.
+
+        Args:
+            data (Any): data to be added to the replay buffer
+
+        Returns:
+            index where the data lives in the replay buffer.
+        """
+        if self._transform is not None and len(self._transform):
+            with _set_dispatch_td_nn_modules(is_tensor_collection(data)):
+                make_none = False
+                # Transforms usually expect a time batch dimension when called within a RB, so we unsqueeze the data temporarily
+                is_tc = is_tensor_collection(data)
+                cm = data.unsqueeze(-1) if is_tc else contextlib.nullcontext(data)
+                new_data = None
+                with cm as data_unsq:
+                    data_unsq_r = self._transform.inv(data_unsq)
+                    if is_tc and data_unsq_r is not None:
+                        # this is a no-op whenever the result matches the input
+                        new_data = data_unsq_r.squeeze(-1)
+                    else:
+                        make_none = data_unsq_r is None
+                data = new_data if new_data is not None else data
+                if make_none:
+                    data = None
+        if data is None:
+            return torch.zeros((0, self._storage.ndim), dtype=torch.long)
+        if rl_warnings() and is_tensor_collection(data) and data.ndim:
+            warnings.warn(
+                f"Using `add()` with a TensorDict that has batch_size={data.batch_size}. "
+                f"Use `extend()` to add multiple elements, or `add()` with a single element (batch_size=torch.Size([])). "
+                "You can silence this warning by setting the `RL_WARNINGS` environment variable to `'0'`."
+            )
+
+        return self._add(data)
+
+    def _add(self, data):
+        with self._replay_lock, self._write_lock:
+            index = self._writer.add(data)
+            self._sampler.add(index)
+        return index
+
+    def _extend(self, data: Sequence, *, update_priority: bool = True) -> torch.Tensor:
+        is_comp = is_compiling()
+        nc = contextlib.nullcontext()
+        with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
+            if self.dim_extend > 0:
+                data = self._transpose(data)
+            index = self._writer.extend(data)
+            self._sampler.extend(index)
+        return index
+
+    @_maybe_delay_init
+    def extend(
+        self, data: Sequence, *, update_priority: bool | None = None
+    ) -> torch.Tensor:
+        """Extends the replay buffer with one or more elements contained in an iterable.
+
+        If present, the inverse transforms will be called.`
+
+        Args:
+            data (iterable): collection of data to be added to the replay
+                buffer.
+
+        Keyword Args:
+            update_priority (bool, optional): Whether to update the priority of the data. Defaults to True.
+                Without effect in this class. See :meth:`~torchrl.data.TensorDictReplayBuffer.extend` for more details.
+
+        Returns:
+            Indices of the data added to the replay buffer.
+
+        .. warning:: :meth:`~torchrl.data.replay_buffers.ReplayBuffer.extend` can have an
+          ambiguous signature when dealing with lists of values, which should be interpreted
+          either as PyTree (in which case all elements in the list will be put in a slice
+          in the stored PyTree in the storage) or a list of values to add one at a time.
+          To solve this, TorchRL makes the clear-cut distinction between list and tuple:
+          a tuple will be viewed as a PyTree, a list (at the root level) will be interpreted
+          as a stack of values to add one at a time to the buffer.
+          For :class:`~torchrl.data.replay_buffers.ListStorage` instances, only
+          unbound elements can be provided (no PyTrees).
+
+        """
+        if update_priority is not None:
+            raise NotImplementedError(
+                "update_priority is not supported in this class. See :meth:`~torchrl.data.TensorDictReplayBuffer.extend` for more details."
+            )
+        if self._transform is not None and len(self._transform):
+            with _set_dispatch_td_nn_modules(is_tensor_collection(data)):
+                data = self._transform.inv(data)
+        if data is None:
+            return torch.zeros((0, self._storage.ndim), dtype=torch.long)
+        return self._extend(data, update_priority=update_priority)
+
+    @_maybe_delay_init
+    def update_priority(
+        self,
+        index: int | torch.Tensor | tuple[torch.Tensor],
+        priority: int | torch.Tensor,
+    ) -> None:
+        if isinstance(index, tuple):
+            index = torch.stack(index, -1)
+        priority = torch.as_tensor(priority)
+        if self.dim_extend > 0 and priority.ndim > 1:
+            priority = self._transpose(priority).flatten()
+            # priority = priority.flatten()
+        with self._replay_lock, self._write_lock:
+            self._sampler.update_priority(index, priority, storage=self.storage)
+
+    @pin_memory_output
+    def _sample(self, batch_size: int) -> tuple[Any, dict]:
+        is_comp = is_compiling()
+        nc = contextlib.nullcontext()
+        with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
+            index, info = self._sampler.sample(self._storage, batch_size)
+            info["index"] = index
+            data = self._storage.get(index)
+        if not isinstance(index, INT_CLASSES):
+            data = self._collate_fn(data)
+        if self._transform is not None and len(self._transform):
+            is_td = is_tensor_collection(data)
+            with data.unlock_() if is_td else contextlib.nullcontext(), _set_dispatch_td_nn_modules(
+                is_td
+            ):
+                data = self._transform(data)
+
+        return data, info
+
+    @_maybe_delay_init
+    def empty(self, empty_write_count: bool = True):
+        """Empties the replay buffer and reset cursor to 0.
+
+        Args:
+            empty_write_count (bool, optional): Whether to empty the write_count attribute. Defaults to `True`.
+        """
+        self._writer._empty(empty_write_count=empty_write_count)
+        self._sampler._empty()
+        self._storage._empty()
+
+    @_maybe_delay_init
+    def sample(self, batch_size: int | None = None, return_info: bool = False) -> Any:
+        """Samples a batch of data from the replay buffer.
+
+        Uses Sampler to sample indices, and retrieves them from Storage.
+
+        Args:
+            batch_size (int, optional): size of data to be collected. If none
+                is provided, this method will sample a batch-size as indicated
+                by the sampler.
+            return_info (bool): whether to return info. If True, the result
+                is a tuple (data, info). If False, the result is the data.
+
+        Returns:
+            A batch of data selected in the replay buffer.
+            A tuple containing this batch and info if return_info flag is set to True.
+        """
+        if (
+            batch_size is not None
+            and self._batch_size is not None
+            and batch_size != self._batch_size
+        ):
+            warnings.warn(
+                f"Got conflicting batch_sizes in constructor ({self._batch_size}) "
+                f"and `sample` ({batch_size}). Refer to the ReplayBuffer documentation "
+                "for a proper usage of the batch-size arguments. "
+                "The batch-size provided to the sample method "
+                "will prevail."
+            )
+        elif batch_size is None and self._batch_size is not None:
+            batch_size = self._batch_size
+        elif batch_size is None:
+            raise RuntimeError(
+                "batch_size not specified. You can specify the batch_size when "
+                "constructing the replay buffer, or pass it to the sample method. "
+                "Refer to the ReplayBuffer documentation "
+                "for a proper usage of the batch-size arguments."
+            )
+        if not self._prefetch:
+            result = self._sample(batch_size)
+        else:
+            with self._futures_lock:
+                while (
+                    len(self._prefetch_queue)
+                    < min(self._sampler._remaining_batches, self._prefetch_cap)
+                    and not self._sampler.ran_out
+                ) or not len(self._prefetch_queue):
+                    fut = self._prefetch_executor.submit(self._sample, batch_size)
+                    self._prefetch_queue.append(fut)
+                result = self._prefetch_queue.popleft().result()
+
+        if return_info:
+            out, info = result
+            if getattr(self.storage, "device", None) is not None:
+                device = self.storage.device
+                info = tree_map(lambda x: x.to(device) if hasattr(x, "to") else x, info)
+            return out, info
+        return result[0]
+
+    @_maybe_delay_init
+    def mark_update(self, index: int | torch.Tensor) -> None:
+        self._sampler.mark_update(index, storage=self._storage)
+
+    @_maybe_delay_init
+    def append_transform(
+        self, transform: Transform, *, invert: bool = False  # noqa-F821
+    ) -> ReplayBuffer:  # noqa: D417
+        """Appends transform at the end.
+
+        Transforms are applied in order when `sample` is called.
+
+        Args:
+            transform (Transform): The transform to be appended
+
+        Keyword Args:
+            invert (bool, optional): if ``True``, the transform will be inverted (forward calls will be called
+                during writing and inverse calls during reading). Defaults to ``False``.
+
+        Example:
+            >>> rb = ReplayBuffer(storage=LazyMemmapStorage(10), batch_size=4)
+            >>> data = TensorDict({"a": torch.zeros(10)}, [10])
+            >>> def t(data):
+            ...     data += 1
+            ...     return data
+            >>> rb.append_transform(t, invert=True)
+            >>> rb.extend(data)
+            >>> assert (data == 1).all()
+
+        """
+        from torchrl.envs.transforms.transforms import _CallableTransform, Transform
+
+        if not isinstance(transform, Transform) and callable(transform):
+            transform = _CallableTransform(transform)
+        if invert:
+            transform = _InvertTransform(transform)
+        transform.eval()
+        self._transform.append(transform)
+        return self
+
+    @_maybe_delay_init
+    def insert_transform(
+        self,
+        index: int,
+        transform: Transform,  # noqa-F821
+        *,
+        invert: bool = False,
+    ) -> ReplayBuffer:  # noqa: D417
+        """Inserts transform.
+
+        Transforms are executed in order when `sample` is called.
+
+        Args:
+            index (int): Position to insert the transform.
+            transform (Transform): The transform to be appended
+
+        Keyword Args:
+            invert (bool, optional): if ``True``, the transform will be inverted (forward calls will be called
+                during writing and inverse calls during reading). Defaults to ``False``.
+
+        """
+        transform.eval()
+        if invert:
+            transform = _InvertTransform(transform)
+        self._transform.insert(index, transform)
+        return self
+
+    _iterator = None
+
+    @_maybe_delay_init
+    def next(self):
+        """Returns the next item in the replay buffer.
+
+        This method is used to iterate over the replay buffer in contexts where __iter__ is not available,
+        such as :class:`~torchrl.data.replay_buffers.RayReplayBuffer`.
+        """
+        try:
+            if self._iterator is None:
+                self._iterator = iter(self)
+            out = next(self._iterator)
+            # if any, we don't want the device ref to be passed in distributed settings
+            if out is not None and (out.device != "cpu"):
+                out = out.copy().clear_device_()
+            return out
+        except StopIteration:
+            self._iterator = None
+            return None
+
+    @_maybe_delay_init
+    def __iter__(self):
+        if self._sampler.ran_out:
+            self._sampler.ran_out = False
+        if self._batch_size is None:
+            raise RuntimeError(
+                "Cannot iterate over the replay buffer. "
+                "Batch_size was not specified during construction of the replay buffer."
+            )
+        while not self._sampler.ran_out or (
+            self._prefetch and len(self._prefetch_queue)
+        ):
+            yield self.sample()
+
+    @_maybe_delay_init
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        if getattr(self, "_rng", None) is not None:
+            rng_state = TensorDict(
+                rng_state=self._rng.get_state().clone(),
+                device=self._rng.device,
+            )
+            state["_rng"] = rng_state
+        _replay_lock = state.pop("_replay_lock", None)
+        _futures_lock = state.pop("_futures_lock", None)
+        if _replay_lock is not None:
+            state["_replay_lock_placeholder"] = None
+        if _futures_lock is not None:
+            state["_futures_lock_placeholder"] = None
+        # Remove non-picklable prefetch objects - they will be recreated on unpickle
+        _prefetch_queue = state.pop("_prefetch_queue", None)
+        _prefetch_executor = state.pop("_prefetch_executor", None)
+        if _prefetch_queue is not None:
+            state["_prefetch_queue_placeholder"] = None
+        if _prefetch_executor is not None:
+            state["_prefetch_executor_placeholder"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]):
+        rngstate = None
+        if "_rng" in state:
+            rngstate = state["_rng"]
+            if rngstate is not None:
+                rng = torch.Generator(device=rngstate.device)
+                rng.set_state(rngstate["rng_state"])
+
+        if "_replay_lock_placeholder" in state:
+            state.pop("_replay_lock_placeholder")
+            _replay_lock = threading.RLock()
+            state["_replay_lock"] = _replay_lock
+        if "_futures_lock_placeholder" in state:
+            state.pop("_futures_lock_placeholder")
+            _futures_lock = threading.RLock()
+            state["_futures_lock"] = _futures_lock
+        # Recreate prefetch objects after unpickling if they were present
+        if "_prefetch_queue_placeholder" in state:
+            state.pop("_prefetch_queue_placeholder")
+            state["_prefetch_queue"] = collections.deque()
+        if "_prefetch_executor_placeholder" in state:
+            state.pop("_prefetch_executor_placeholder")
+            state["_prefetch_executor"] = ThreadPoolExecutor(
+                max_workers=state["_prefetch_cap"]
+            )
+        self.__dict__.update(state)
+        if rngstate is not None:
+            self.set_rng(rng)
+
+    @property
+    @_maybe_delay_init
+    def sampler(self) -> Sampler:
+        """The sampler of the replay buffer.
+
+        The sampler must be an instance of :class:`~torchrl.data.replay_buffers.Sampler`.
+
+        """
+        return self._sampler
+
+    @property
+    @_maybe_delay_init
+    def writer(self) -> Writer:
+        """The writer of the replay buffer.
+
+        The writer must be an instance of :class:`~torchrl.data.replay_buffers.Writer`.
+
+        """
+        return self._writer
+
+    @property
+    @_maybe_delay_init
+    def storage(self) -> Storage:
+        """The storage of the replay buffer.
+
+        The storage must be an instance of :class:`~torchrl.data.replay_buffers.Storage`.
+
+        """
+        return self._storage
+
+    @property
+    @_maybe_delay_init
+    def transform(self) -> Transform:
+        """The transform of the replay buffer.
+
+        The transform must be an instance of :class:`~torchrl.envs.transforms.Transform`.
+        """
+        return self._transform
+
+
+class PrioritizedReplayBuffer(ReplayBuffer):
+    """Prioritized replay buffer.
+
+    All arguments are keyword-only arguments.
+
+    Presented in "Schaul, T.; Quan, J.; Antonoglou, I.; and Silver, D. 2015.
+    Prioritized experience replay." (https://arxiv.org/abs/1511.05952)
+
+    Args:
+        alpha (:obj:`float`): exponent α determines how much prioritization is used,
+            with α = 0 corresponding to the uniform case.
+        beta (:obj:`float`): importance sampling negative exponent.
+        eps (:obj:`float`): delta added to the priorities to ensure that the buffer
+            does not contain null priorities.
+        storage (Storage, optional): the storage to be used. If none is provided
+            a default :class:`~torchrl.data.replay_buffers.ListStorage` with
+            ``max_size`` of ``1_000`` will be created.
+        sampler (Sampler, optional): the sampler to be used. If none is provided,
+            a default :class:`~torchrl.data.replay_buffers.PrioritizedSampler` with
+            ``alpha``, ``beta``, and ``eps`` will be created.
+        collate_fn (callable, optional): merges a list of samples to form a
+            mini-batch of Tensor(s)/outputs.  Used when using batched
+            loading from a map-style dataset. The default value will be decided
+            based on the storage type.
+        pin_memory (bool): whether pin_memory() should be called on the rb
+            samples.
+        prefetch (int, optional): number of next batches to be prefetched
+            using multithreading. Defaults to None (no prefetching).
+        transform (Transform, optional): Transform to be executed when
+            sample() is called.
+            To chain transforms use the :class:`~torchrl.envs.Compose` class.
+            Transforms should be used with :class:`tensordict.TensorDict`
+            content. If used with other structures, the transforms should be
+            encoded with a ``"data"`` leading key that will be used to
+            construct a tensordict from the non-tensordict content.
+        batch_size (int, optional): the batch size to be used when sample() is
+            called.
+
+            .. note:: The batch-size can be specified at construction time via the
+              ``batch_size`` argument, or at sampling time. The former should
+              be preferred whenever the batch-size is consistent across the
+              experiment. If the batch-size is likely to change, it can be
+              passed to the :meth:`sample` method. This option is
+              incompatible with prefetching (since this requires to know the
+              batch-size in advance) as well as with samplers that have a
+              ``drop_last`` argument.
+
+        dim_extend (int, optional): indicates the dim to consider for
+            extension when calling :meth:`extend`. Defaults to ``storage.ndim-1``.
+            When using ``dim_extend > 0``, we recommend using the ``ndim``
+            argument in the storage instantiation if that argument is
+            available, to let storages know that the data is
+            multi-dimensional and keep consistent notions of storage-capacity
+            and batch-size during sampling.
+
+            .. note:: This argument has no effect on :meth:`add` and
+                therefore should be used with caution when both :meth:`add`
+                and :meth:`extend` are used in a codebase. For example:
+
+                    >>> data = torch.zeros(3, 4)
+                    >>> rb = ReplayBuffer(
+                    ...     storage=LazyTensorStorage(10, ndim=2),
+                    ...     dim_extend=1)
+                    >>> # these two approaches are equivalent:
+                    >>> for d in data.unbind(1):
+                    ...     rb.add(d)
+                    >>> rb.extend(data)
+
+        delayed_init (bool, optional): whether to initialize storage, writer, sampler and transform
+            the first time the buffer is used rather than during construction.
+            This is useful when the replay buffer needs to be pickled and sent to remote workers,
+            particularly when using transforms with modules that require gradients.
+            If not specified, defaults to ``True`` when ``transform_factory`` is provided,
+            and ``False`` otherwise.
+
+    .. note::
+        Generic prioritized replay buffers (ie. non-tensordict backed) require
+        calling :meth:`~.sample` with the ``return_info`` argument set to
+        ``True`` to have access to the indices, and hence update the priority.
+        Using :class:`tensordict.TensorDict` and the related
+        :class:`~torchrl.data.TensorDictPrioritizedReplayBuffer` simplifies this
+        process.
+
+    Examples:
+        >>> import torch
+        >>>
+        >>> from torchrl.data import ListStorage, PrioritizedReplayBuffer
+        >>>
+        >>> torch.manual_seed(0)
+        >>>
+        >>> rb = PrioritizedReplayBuffer(alpha=0.7, beta=0.9, storage=ListStorage(10))
+        >>> data = range(10)
+        >>> rb.extend(data)
+        >>> sample = rb.sample(3)
+        >>> print(sample)
+        tensor([1, 0, 1])
+        >>> # get the info to find what the indices are
+        >>> sample, info = rb.sample(5, return_info=True)
+        >>> print(sample, info)
+        tensor([2, 7, 4, 3, 5]) {'priority_weight': array([1., 1., 1., 1., 1.], dtype=float32), 'index': array([2, 7, 4, 3, 5])}
+        >>> # update priority
+        >>> priority = torch.ones(5) * 5
+        >>> rb.update_priority(info["index"], priority)
+        >>> # and now a new sample, the weights should be updated
+        >>> sample, info = rb.sample(5, return_info=True)
+        >>> print(sample, info)
+        tensor([2, 5, 2, 2, 5]) {'priority_weight': array([0.36278465, 0.36278465, 0.36278465, 0.36278465, 0.36278465],
+              dtype=float32), 'index': array([2, 5, 2, 2, 5])}
+
+    """
+
+    def __init__(
+        self,
+        *,
+        alpha: float,
+        beta: float,
+        eps: float = 1e-8,
+        dtype: torch.dtype = torch.float,
+        storage: Storage | None = None,
+        sampler: Sampler | None = None,
+        collate_fn: Callable | None = None,
+        pin_memory: bool = False,
+        prefetch: int | None = None,
+        transform: Transform | None = None,  # noqa-F821
+        batch_size: int | None = None,
+        dim_extend: int | None = None,
+        delayed_init: bool = False,
+    ) -> None:
+        if storage is None:
+            storage = ListStorage(max_size=1_000)
+        if sampler is None:
+            sampler = PrioritizedSampler(storage.max_size, alpha, beta, eps, dtype)
+        super().__init__(
+            storage=storage,
+            sampler=sampler,
+            collate_fn=collate_fn,
+            pin_memory=pin_memory,
+            prefetch=prefetch,
+            transform=transform,
+            batch_size=batch_size,
+            dim_extend=dim_extend,
+            delayed_init=delayed_init,
+        )
+
+
+class TensorDictReplayBuffer(ReplayBuffer):
+    """TensorDict-specific wrapper around the :class:`~torchrl.data.ReplayBuffer` class.
+
+    Keyword Args:
+        storage (Storage, Callable[[], Storage], optional): the storage to be used.
+            If a callable is passed, it is used as constructor for the storage.
+            If none is provided a default :class:`~torchrl.data.replay_buffers.ListStorage` with
+            ``max_size`` of ``1_000`` will be created.
+        sampler (Sampler, Callable[[], Sampler], optional): the sampler to be used.
+            If a callable is passed, it is used as constructor for the sampler.
+            If none is provided, a default :class:`~torchrl.data.replay_buffers.RandomSampler`
+            will be used.
+        writer (Writer, Callable[[], Writer], optional): the writer to be used.
+            If a callable is passed, it is used as constructor for the writer.
+            If none is provided a default :class:`~torchrl.data.replay_buffers.TensorDictRoundRobinWriter`
+            will be used.
+        collate_fn (callable, optional): merges a list of samples to form a
+            mini-batch of Tensor(s)/outputs.  Used when using batched
+            loading from a map-style dataset. The default value will be decided
+            based on the storage type.
+        pin_memory (bool): whether pin_memory() should be called on the rb
+            samples.
+        prefetch (int, optional): number of next batches to be prefetched
+            using multithreading. Defaults to None (no prefetching).
+        transform (Transform or Callable[[Any], Any], optional): Transform to be executed when
+            :meth:`sample` is called.
+            To chain transforms use the :class:`~torchrl.envs.Compose` class.
+            Transforms should be used with :class:`tensordict.TensorDict`
+            content. A generic callable can also be passed if the replay buffer
+            is used with PyTree structures (see example below).
+            Unlike storages, writers and samplers, transform constructors must
+            be passed as separate keyword argument :attr:`transform_factory`,
+            as it is impossible to distinguish a constructor from a transform.
+        transform_factory (Callable[[], Callable], optional): a factory for the
+            transform. Exclusive with :attr:`transform`.
+        batch_size (int, optional): the batch size to be used when sample() is
+            called.
+
+            .. note::
+              The batch-size can be specified at construction time via the
+              ``batch_size`` argument, or at sampling time. The former should
+              be preferred whenever the batch-size is consistent across the
+              experiment. If the batch-size is likely to change, it can be
+              passed to the :meth:`~.sample` method. This option is
+              incompatible with prefetching (since this requires to know the
+              batch-size in advance) as well as with samplers that have a
+              ``drop_last`` argument.
+
+        priority_key (str, optional): the key at which priority is assumed to
+            be stored within TensorDicts added to this ReplayBuffer.
+            This is to be used when the sampler is of type
+            :class:`~torchrl.data.PrioritizedSampler`.
+            Defaults to ``"td_error"``.
+        dim_extend (int, optional): indicates the dim to consider for
+            extension when calling :meth:`~.extend`. Defaults to ``storage.ndim-1``.
+            When using ``dim_extend > 0``, we recommend using the ``ndim``
+            argument in the storage instantiation if that argument is
+            available, to let storages know that the data is
+            multi-dimensional and keep consistent notions of storage-capacity
+            and batch-size during sampling.
+
+            .. note:: This argument has no effect on :meth:`~.add` and
+                therefore should be used with caution when both :meth:`~.add`
+                and :meth:`~.extend` are used in a codebase. For example:
+
+                    >>> data = torch.zeros(3, 4)
+                    >>> rb = ReplayBuffer(
+                    ...     storage=LazyTensorStorage(10, ndim=2),
+                    ...     dim_extend=1)
+                    >>> # these two approaches are equivalent:
+                    >>> for d in data.unbind(1):
+                    ...     rb.add(d)
+                    >>> rb.extend(data)
+
+        generator (torch.Generator, optional): a generator to use for sampling.
+            Using a dedicated generator for the replay buffer can allow a fine-grained control
+            over seeding, for instance keeping the global seed different but the RB seed identical
+            for distributed jobs.
+            Defaults to ``None`` (global default generator).
+
+            .. warning:: As of now, the generator has no effect on the transforms.
+        shared (bool, optional): whether the buffer will be shared using multiprocessing or not.
+            Defaults to ``False``.
+        compilable (bool, optional): whether the writer is compilable.
+            If ``True``, the writer cannot be shared between multiple processes.
+            Defaults to ``False``.
+        delayed_init (bool, optional): whether to initialize storage, writer, sampler and transform
+            the first time the buffer is used rather than during construction.
+            This is useful when the replay buffer needs to be pickled and sent to remote workers,
+            particularly when using transforms with modules that require gradients.
+            If not specified, defaults to ``True`` when ``transform_factory`` is provided,
+            and ``False`` otherwise.
+
+    Examples:
+        >>> import torch
+        >>>
+        >>> from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
+        >>> from tensordict import TensorDict
+        >>>
+        >>> torch.manual_seed(0)
+        >>>
+        >>> rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=5)
+        >>> data = TensorDict({"a": torch.ones(10, 3), ("b", "c"): torch.zeros(10, 1, 1)}, [10])
+        >>> rb.extend(data)
+        >>> sample = rb.sample(3)
+        >>> # samples keep track of the index
+        >>> print(sample)
+        TensorDict(
+            fields={
+                a: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                b: TensorDict(
+                    fields={
+                        c: Tensor(shape=torch.Size([3, 1, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                    batch_size=torch.Size([3]),
+                    device=cpu,
+                    is_shared=False),
+                index: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([3]),
+            device=cpu,
+            is_shared=False)
+        >>> # we can iterate over the buffer
+        >>> for i, data in enumerate(rb):
+        ...     print(i, data)
+        ...     if i == 2:
+        ...         break
+        0 TensorDict(
+            fields={
+                a: Tensor(shape=torch.Size([5, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                b: TensorDict(
+                    fields={
+                        c: Tensor(shape=torch.Size([5, 1, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                    batch_size=torch.Size([5]),
+                    device=cpu,
+                    is_shared=False),
+                index: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([5]),
+            device=cpu,
+            is_shared=False)
+        1 TensorDict(
+            fields={
+                a: Tensor(shape=torch.Size([5, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                b: TensorDict(
+                    fields={
+                        c: Tensor(shape=torch.Size([5, 1, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                    batch_size=torch.Size([5]),
+                    device=cpu,
+                    is_shared=False),
+                index: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([5]),
+            device=cpu,
+            is_shared=False)
+
+    """
+
+    def __init__(self, *, priority_key: str = "td_error", **kwargs) -> None:
+        writer = kwargs.get("writer", None)
+        if writer is None:
+            kwargs["writer"] = partial(
+                TensorDictRoundRobinWriter, compilable=kwargs.get("compilable")
+            )
+        super().__init__(**kwargs)
+        self.priority_key = priority_key
+
+    def _get_priority_item(self, tensordict: TensorDictBase) -> float:
+        priority = tensordict.get(self.priority_key, None)
+        if self._storage.ndim > 1:
+            # We have to flatten the priority otherwise we'll be aggregating
+            # the priority across batches
+            priority = priority.flatten(0, self._storage.ndim - 1)
+        if priority is None:
+            return self._sampler.default_priority
+        try:
+            if priority.numel() > 1:
+                priority = _reduce(priority, self._sampler.reduction)
+            else:
+                priority = priority.item()
+        except ValueError:
+            raise ValueError(
+                f"Found a priority key of size"
+                f" {tensordict.get(self.priority_key).shape} but expected "
+                f"scalar value"
+            )
+
+        if self._storage.ndim > 1:
+            priority = priority.unflatten(0, tensordict.shape[: self._storage.ndim])
+
+        return priority
+
+    def _get_priority_vector(self, tensordict: TensorDictBase) -> torch.Tensor:
+        priority = tensordict.get(self.priority_key, None)
+        if priority is None:
+            return torch.tensor(
+                self._sampler.default_priority,
+                dtype=torch.float,
+                device=tensordict.device,
+            ).expand(tensordict.shape[0])
+        if self._storage.ndim > 1 and priority.ndim >= self._storage.ndim:
+            # We have to flatten the priority otherwise we'll be aggregating
+            # the priority across batches
+            priority = priority.flatten(0, self._storage.ndim - 1)
+
+        priority = priority.reshape(priority.shape[0], -1)
+        priority = _reduce(priority, self._sampler.reduction, dim=1)
+
+        if self._storage.ndim > 1:
+            priority = priority.unflatten(0, tensordict.shape[: self._storage.ndim])
+
+        return priority
+
+    @_maybe_delay_init
+    def add(self, data: TensorDictBase) -> int:
+        if self._transform is not None:
+            with _set_dispatch_td_nn_modules(is_tensor_collection(data)):
+                data = self._transform.inv(data)
+        if data is None:
+            return torch.zeros((0, self._storage.ndim), dtype=torch.long)
+
+        index = super()._add(data)
+        if index is not None:
+            if is_tensor_collection(data):
+                self._set_index_in_td(data, index)
+
+            self.update_tensordict_priority(data)
+        return index
+
+    @_maybe_delay_init
+    def extend(
+        self, tensordicts: TensorDictBase, *, update_priority: bool | None = None
+    ) -> torch.Tensor:
+        """Extends the replay buffer with a batch of data.
+
+        Args:
+            tensordicts (TensorDictBase): The data to extend the replay buffer with.
+
+        Keyword Args:
+            update_priority (bool, optional): Whether to update the priority of the data. Defaults to True.
+
+        Returns:
+            The indices of the data that were added to the replay buffer.
+        """
+        if not isinstance(tensordicts, TensorDictBase):
+            raise ValueError(
+                f"{self.__class__.__name__} only accepts TensorDictBase subclasses. tensorclasses "
+                f"and other types are not compatible with that class. "
+                "Please use a regular `ReplayBuffer` instead."
+            )
+        if self._transform is not None:
+            tensordicts = self._transform.inv(tensordicts)
+        if tensordicts is None:
+            return torch.zeros((0, self._storage.ndim), dtype=torch.long)
+
+        index = super()._extend(tensordicts)
+
+        # TODO: to be usable directly, the indices should be flipped but the issue
+        #  is that just doing this results in indices that are not sorted like the original data
+        #  so the actually indices will have to be used on the _storage directly (not on the buffer)
+        self._set_index_in_td(tensordicts, index)
+        if update_priority is None:
+            update_priority = True
+        if update_priority:
+            try:
+                vector = tensordicts.get(self.priority_key)
+                if vector is not None:
+                    self.update_priority(index, vector)
+            except Exception as e:
+                raise RuntimeError(
+                    "Failed to update priority of extended data. You can try to set update_priority=False in the extend method and update the priority manually."
+                ) from e
+        return index
+
+    def _set_index_in_td(self, tensordict, index):
+        if index is None:
+            return
+        if _is_int(index):
+            index = torch.as_tensor(index, device=tensordict.device)
+        elif index.ndim == 2 and index.shape[:1] != tensordict.shape[:1]:
+            for dim in range(2, tensordict.ndim + 1):
+                if index.shape[:1].numel() == tensordict.shape[:dim].numel():
+                    # if index has 2 dims and is in a non-zero format
+                    index = index.unflatten(0, tensordict.shape[:dim])
+                    break
+            else:
+                raise RuntimeError(
+                    f"could not find how to reshape index with shape {index.shape} to fit in tensordict with shape {tensordict.shape}"
+                )
+            tensordict.set("index", index)
+            return
+        tensordict.set("index", expand_as_right(index, tensordict))
+
+    @_maybe_delay_init
+    def update_tensordict_priority(self, data: TensorDictBase) -> None:
+        if not isinstance(self._sampler, PrioritizedSampler):
+            return
+        if data.ndim:
+            priority = self._get_priority_vector(data)
+        else:
+            priority = torch.as_tensor(self._get_priority_item(data))
+        index = data.get("index")
+        if self._storage.ndim > 1 and index.ndim == 2:
+            index = index.unbind(-1)
+        else:
+            while index.shape != priority.shape:
+                # reduce index
+                index = index[..., 0]
+        return self.update_priority(index, priority)
+
+    def sample(
+        self,
+        batch_size: int | None = None,
+        return_info: bool = False,
+        include_info: bool | None = None,
+    ) -> TensorDictBase:
+        """Samples a batch of data from the replay buffer.
+
+        Uses Sampler to sample indices, and retrieves them from Storage.
+
+        Args:
+            batch_size (int, optional): size of data to be collected. If none
+                is provided, this method will sample a batch-size as indicated
+                by the sampler.
+            return_info (bool): whether to return info. If True, the result
+                is a tuple (data, info). If False, the result is the data.
+
+        Returns:
+            A tensordict containing a batch of data selected in the replay buffer.
+            A tuple containing this tensordict and info if return_info flag is set to True.
+        """
+        if include_info is not None:
+            warnings.warn(
+                "include_info is going to be deprecated soon."
+                "The default behavior has changed to `include_info=True` "
+                "to avoid bugs linked to wrongly preassigned values in the "
+                "output tensordict."
+            )
+
+        data, info = super().sample(batch_size, return_info=True)
+        is_tc = is_tensor_collection(data)
+        if is_tc and not is_tensorclass(data) and include_info in (True, None):
+            is_locked = data.is_locked
+            if is_locked:
+                data.unlock_()
+            for key, val in info.items():
+                if key == "index" and isinstance(val, tuple):
+                    val = torch.stack(val, -1)
+                try:
+                    val = _to_torch(val, data.device)
+                    if val.ndim < data.ndim:
+                        val = expand_as_right(val, data)
+                    data.set(key, val)
+                except RuntimeError:
+                    raise RuntimeError(
+                        "Failed to set the metadata (e.g., indices or weights) in the sampled tensordict within TensorDictReplayBuffer.sample. "
+                        "This is probably caused by a shape mismatch (one of the transforms has probably modified "
+                        "the shape of the output tensordict). "
+                        "You can always recover these items from the `sample` method from a regular ReplayBuffer "
+                        "instance with the 'return_info' flag set to True."
+                    )
+            if is_locked:
+                data.lock_()
+        elif not is_tc and include_info in (True, None):
+            raise RuntimeError("Cannot include info in non-tensordict data")
+        if return_info:
+            return data, info
+        return data
+
+    @pin_memory_output
+    def _sample(self, batch_size: int) -> tuple[Any, dict]:
+        is_comp = is_compiling()
+        nc = contextlib.nullcontext()
+        with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
+            index, info = self._sampler.sample(self._storage, batch_size)
+            info["index"] = index
+            data = self._storage.get(index)
+        if not isinstance(index, INT_CLASSES):
+            data = self._collate_fn(data)
+        if self._transform is not None and len(self._transform):
+            with data.unlock_(), _set_dispatch_td_nn_modules(True):
+                data = self._transform(data)
+        return data, info
+
+
+class TensorDictPrioritizedReplayBuffer(TensorDictReplayBuffer):
+    """TensorDict-specific wrapper around the :class:`~torchrl.data.PrioritizedReplayBuffer` class.
+
+    This class returns tensordicts with a new key ``"index"`` that represents
+    the index of each element in the replay buffer. It also provides the
+    :meth:`~.update_tensordict_priority` method that only requires for the
+    tensordict to be passed to it with its new priority value.
+
+    Keyword Args:
+        alpha (:obj:`float`): exponent α determines how much prioritization is used,
+            with α = 0 corresponding to the uniform case.
+        beta (:obj:`float`): importance sampling negative exponent.
+        eps (:obj:`float`): delta added to the priorities to ensure that the buffer
+            does not contain null priorities.
+        storage (Storage, Callable[[], Storage], optional): the storage to be used.
+            If a callable is passed, it is used as constructor for the storage.
+            If none is provided a default :class:`~torchrl.data.replay_buffers.ListStorage` with
+            ``max_size`` of ``1_000`` will be created.
+        collate_fn (callable, optional): merges a list of samples to form a
+            mini-batch of Tensor(s)/outputs.  Used when using batched
+            loading from a map-style dataset. The default value will be decided
+            based on the storage type.
+        pin_memory (bool): whether pin_memory() should be called on the rb
+            samples.
+        prefetch (int, optional): number of next batches to be prefetched
+            using multithreading. Defaults to None (no prefetching).
+        transform (Transform or Callable[[Any], Any], optional): Transform to be executed when
+            :meth:`sample` is called.
+            To chain transforms use the :class:`~torchrl.envs.Compose` class.
+            Transforms should be used with :class:`tensordict.TensorDict`
+            content. A generic callable can also be passed if the replay buffer
+            is used with PyTree structures (see example below).
+            Unlike storages, writers and samplers, transform constructors must
+            be passed as separate keyword argument :attr:`transform_factory`,
+            as it is impossible to distinguish a constructor from a transform.
+        transform_factory (Callable[[], Callable], optional): a factory for the
+            transform. Exclusive with :attr:`transform`.
+        batch_size (int, optional): the batch size to be used when sample() is
+            called.
+
+            .. note::
+              The batch-size can be specified at construction time via the
+              ``batch_size`` argument, or at sampling time. The former should
+              be preferred whenever the batch-size is consistent across the
+              experiment. If the batch-size is likely to change, it can be
+              passed to the :meth:`~.sample` method. This option is
+              incompatible with prefetching (since this requires to know the
+              batch-size in advance) as well as with samplers that have a
+              ``drop_last`` argument.
+
+        priority_key (str, optional): the key at which priority is assumed to
+            be stored within TensorDicts added to this ReplayBuffer.
+            This is to be used when the sampler is of type
+            :class:`~torchrl.data.PrioritizedSampler`.
+            Defaults to ``"td_error"``.
+        reduction (str, optional): the reduction method for multidimensional
+            tensordicts (ie stored trajectories). Can be one of "max", "min",
+            "median" or "mean".
+        dim_extend (int, optional): indicates the dim to consider for
+            extension when calling :meth:`~.extend`. Defaults to ``storage.ndim-1``.
+            When using ``dim_extend > 0``, we recommend using the ``ndim``
+            argument in the storage instantiation if that argument is
+            available, to let storages know that the data is
+            multi-dimensional and keep consistent notions of storage-capacity
+            and batch-size during sampling.
+
+            .. note:: This argument has no effect on :meth:`~.add` and
+                therefore should be used with caution when both :meth:`~.add`
+                and :meth:`~.extend` are used in a codebase. For example:
+
+                    >>> data = torch.zeros(3, 4)
+                    >>> rb = ReplayBuffer(
+                    ...     storage=LazyTensorStorage(10, ndim=2),
+                    ...     dim_extend=1)
+                    >>> # these two approaches are equivalent:
+                    >>> for d in data.unbind(1):
+                    ...     rb.add(d)
+                    >>> rb.extend(data)
+
+        generator (torch.Generator, optional): a generator to use for sampling.
+            Using a dedicated generator for the replay buffer can allow a fine-grained control
+            over seeding, for instance keeping the global seed different but the RB seed identical
+            for distributed jobs.
+            Defaults to ``None`` (global default generator).
+
+            .. warning:: As of now, the generator has no effect on the transforms.
+        shared (bool, optional): whether the buffer will be shared using multiprocessing or not.
+            Defaults to ``False``.
+        compilable (bool, optional): whether the writer is compilable.
+            If ``True``, the writer cannot be shared between multiple processes.
+            Defaults to ``False``.
+        delayed_init (bool, optional): whether to initialize storage, writer, sampler and transform
+            the first time the buffer is used rather than during construction.
+            This is useful when the replay buffer needs to be pickled and sent to remote workers,
+            particularly when using transforms with modules that require gradients.
+            If not specified, defaults to ``True`` when ``transform_factory`` is provided,
+            and ``False`` otherwise.
+
+    Examples:
+        >>> import torch
+        >>>
+        >>> from torchrl.data import LazyTensorStorage, TensorDictPrioritizedReplayBuffer
+        >>> from tensordict import TensorDict
+        >>>
+        >>> torch.manual_seed(0)
+        >>>
+        >>> rb = TensorDictPrioritizedReplayBuffer(alpha=0.7, beta=1.1, storage=LazyTensorStorage(10), batch_size=5)
+        >>> data = TensorDict({"a": torch.ones(10, 3), ("b", "c"): torch.zeros(10, 3, 1)}, [10])
+        >>> rb.extend(data)
+        >>> print("len of rb", len(rb))
+        len of rb 10
+        >>> sample = rb.sample(5)
+        >>> print(sample)
+        TensorDict(
+            fields={
+                priority_weight: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.float32, is_shared=False),
+                a: Tensor(shape=torch.Size([5, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                b: TensorDict(
+                    fields={
+                        c: Tensor(shape=torch.Size([5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                    batch_size=torch.Size([5]),
+                    device=cpu,
+                    is_shared=False),
+                index: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.int64, is_shared=False)},
+            batch_size=torch.Size([5]),
+            device=cpu,
+            is_shared=False)
+        >>> print("index", sample["index"])
+        index tensor([9, 5, 2, 2, 7])
+        >>> # give a high priority to these samples...
+        >>> sample.set("td_error", 100*torch.ones(sample.shape))
+        >>> # and update priority
+        >>> rb.update_tensordict_priority(sample)
+        >>> # the new sample should have a high overlap with the previous one
+        >>> sample = rb.sample(5)
+        >>> print(sample)
+        TensorDict(
+            fields={
+                priority_weight: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.float32, is_shared=False),
+                a: Tensor(shape=torch.Size([5, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                b: TensorDict(
+                    fields={
+                        c: Tensor(shape=torch.Size([5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
+                    batch_size=torch.Size([5]),
+                    device=cpu,
+                    is_shared=False),
+                index: Tensor(shape=torch.Size([5]), device=cpu, dtype=torch.int64, is_shared=False)},
+            batch_size=torch.Size([5]),
+            device=cpu,
+            is_shared=False)
+        >>> print("index", sample["index"])
+        index tensor([2, 5, 5, 9, 7])
+
+    """
+
+    def __init__(
+        self,
+        *,
+        alpha: float,
+        beta: float,
+        priority_key: str = "td_error",
+        eps: float = 1e-8,
+        storage: Storage | None = None,
+        collate_fn: Callable | None = None,
+        pin_memory: bool = False,
+        prefetch: int | None = None,
+        transform: Transform | None = None,  # noqa-F821
+        reduction: str = "max",
+        batch_size: int | None = None,
+        dim_extend: int | None = None,
+        generator: torch.Generator | None = None,
+        shared: bool = False,
+        compilable: bool = False,
+    ) -> None:
+        storage = self._maybe_make_storage(storage, compilable=compilable)
+        sampler = PrioritizedSampler(
+            storage.max_size, alpha, beta, eps, reduction=reduction
+        )
+        super().__init__(
+            priority_key=priority_key,
+            storage=storage,
+            sampler=sampler,
+            collate_fn=collate_fn,
+            pin_memory=pin_memory,
+            prefetch=prefetch,
+            transform=transform,
+            batch_size=batch_size,
+            dim_extend=dim_extend,
+            generator=generator,
+            shared=shared,
+            compilable=compilable,
+        )
+
+
+@accept_remote_rref_udf_invocation
+class RemoteTensorDictReplayBuffer(TensorDictReplayBuffer):
+    """A remote invocation friendly ReplayBuffer class. Public methods can be invoked by remote agents using `torch.rpc` or called locally as normal."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def sample(
+        self,
+        batch_size: int | None = None,
+        include_info: bool | None = None,
+        return_info: bool = False,
+    ) -> TensorDictBase:
+        return super().sample(
+            batch_size=batch_size, include_info=include_info, return_info=return_info
+        )
+
+    def add(self, data: TensorDictBase) -> int:
+        return super().add(data)
+
+    def extend(
+        self, tensordicts: list | TensorDictBase, *, update_priority: bool | None = None
+    ) -> torch.Tensor:
+        return super().extend(tensordicts, update_priority=update_priority)
+
+    def update_priority(
+        self, index: int | torch.Tensor, priority: int | torch.Tensor
+    ) -> None:
+        return super().update_priority(index, priority)
+
+    def update_tensordict_priority(self, data: TensorDictBase) -> None:
+        return super().update_tensordict_priority(data)
+
+
+class InPlaceSampler:
+    """[Deprecated] A sampler to write tennsordicts in-place."""
+
+    def __init__(self, device: DEVICE_TYPING | None = None):
+        raise RuntimeError(
+            "This class has been removed without replacement. In-place sampling should be avoided."
+        )
+
+
+def stack_tensors(list_of_tensor_iterators: list) -> tuple[torch.Tensor]:
+    """Zips a list of iterables containing tensor-like objects and stacks the resulting lists of tensors together.
+
+    Args:
+        list_of_tensor_iterators (list): Sequence containing similar iterators,
+            where each element of the nested iterator is a tensor whose
+            shape match the tensor of other iterators that have the same index.
+
+    Returns:
+         Tuple of stacked tensors.
+
+    Examples:
+         >>> list_of_tensor_iterators = [[torch.ones(3), torch.zeros(1,2)]
+         ...     for _ in range(4)]
+         >>> stack_tensors(list_of_tensor_iterators)
+         (tensor([[1., 1., 1.],
+                 [1., 1., 1.],
+                 [1., 1., 1.],
+                 [1., 1., 1.]]), tensor([[[0., 0.]],
+         <BLANKLINE>
+                 [[0., 0.]],
+         <BLANKLINE>
+                 [[0., 0.]],
+         <BLANKLINE>
+                 [[0., 0.]]]))
+
+    """
+    return tuple(torch.stack(tensors, 0) for tensors in zip(*list_of_tensor_iterators))
+
+
+class ReplayBufferEnsemble(ReplayBuffer):
+    """An ensemble of replay buffers.
+
+    This class allows to read and sample from multiple replay buffers at once.
+    It automatically composes ensemble of storages (:class:`~torchrl.data.replay_buffers.storages.StorageEnsemble`),
+    writers (:class:`~torchrl.data.replay_buffers.writers.WriterEnsemble`) and
+    samplers (:class:`~torchrl.data.replay_buffers.samplers.SamplerEnsemble`).
+
+    .. note::
+      Writing directly to this class is forbidden, but it can be indexed to retrieve
+      the nested nested-buffer and extending it.
+
+    There are two distinct ways of constructing a :class:`~torchrl.data.ReplayBufferEnsemble`:
+    one can either pass a list of replay buffers, or directly pass the components
+    (storage, writers and samplers) like it is done for other replay buffer subclasses.
+
+    Args:
+        rbs (sequence of ReplayBuffer instances, optional): the replay buffers to ensemble.
+        storages (StorageEnsemble, optional): the ensemble of storages, if the replay
+            buffers are not passed.
+        samplers (SamplerEnsemble, optional): the ensemble of samplers, if the replay
+            buffers are not passed.
+        writers (WriterEnsemble, optional): the ensemble of writers, if the replay
+            buffers are not passed.
+        transform (Transform, optional): if passed, this will be the transform
+            of the ensemble of replay buffers. Individual transforms for each
+            replay buffer is retrieved from its parent replay buffer, or directly
+            written in the :class:`~torchrl.data.replay_buffers.storages.StorageEnsemble`
+            object.
+        batch_size (int, optional): the batch-size to use during sampling.
+        collate_fn (callable, optional): the function to use to collate the
+            data after each individual collate_fn has been called and the data
+            is placed in a list (along with the buffer id).
+        collate_fns (list of callables, optional): collate_fn of each nested
+            replay buffer. Retrieved from the :class:`~ReplayBuffer` instances
+            if not provided.
+        p (list of float or Tensor, optional): a list of floating numbers
+            indicating the relative weight of each replay buffer. Can also
+            be passed to torchrl.data.replay_buffers.samplers.SamplerEnsemble`
+            if the buffer is built explicitly.
+        sample_from_all (bool, optional): if ``True``, each dataset will be sampled
+            from. This is not compatible with the ``p`` argument. Defaults to ``False``.
+            Can also be passed to torchrl.data.replay_buffers.samplers.SamplerEnsemble`
+            if the buffer is built explicitly.
+        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``. Can also
+            be passed to torchrl.data.replay_buffers.samplers.SamplerEnsemble`
+            if the buffer is built explicitly.
+        generator (torch.Generator, optional): a generator to use for sampling.
+            Using a dedicated generator for the replay buffer can allow a fine-grained control
+            over seeding, for instance keeping the global seed different but the RB seed identical
+            for distributed jobs.
+            Defaults to ``None`` (global default generator).
+
+            .. warning:: As of now, the generator has no effect on the transforms.
+
+        shared (bool, optional): whether the buffer will be shared using multiprocessing or not.
+            Defaults to ``False``.
+        delayed_init (bool, optional): whether to initialize storage, writer, sampler and transform
+            the first time the buffer is used rather than during construction.
+            This is useful when the replay buffer needs to be pickled and sent to remote workers,
+            particularly when using transforms with modules that require gradients.
+            If not specified, defaults to ``True`` when ``transform_factory`` is provided,
+            and ``False`` otherwise.
+
+    Examples:
+        >>> from torchrl.envs import Compose, ToTensorImage, Resize, RenameTransform
+        >>> from torchrl.data import TensorDictReplayBuffer, ReplayBufferEnsemble, LazyMemmapStorage
+        >>> from tensordict import TensorDict
+        >>> import torch
+        >>> rb0 = TensorDictReplayBuffer(
+        ...     storage=LazyMemmapStorage(10),
+        ...     transform=Compose(
+        ...         ToTensorImage(in_keys=["pixels", ("next", "pixels")]),
+        ...         Resize(32, in_keys=["pixels", ("next", "pixels")]),
+        ...         RenameTransform([("some", "key")], ["renamed"]),
+        ...     ),
+        ... )
+        >>> rb1 = TensorDictReplayBuffer(
+        ...     storage=LazyMemmapStorage(10),
+        ...     transform=Compose(
+        ...         ToTensorImage(in_keys=["pixels", ("next", "pixels")]),
+        ...         Resize(32, in_keys=["pixels", ("next", "pixels")]),
+        ...         RenameTransform(["another_key"], ["renamed"]),
+        ...     ),
+        ... )
+        >>> rb = ReplayBufferEnsemble(
+        ...     rb0,
+        ...     rb1,
+        ...     p=[0.5, 0.5],
+        ...     transform=Resize(33, in_keys=["pixels"], out_keys=["pixels33"]),
+        ... )
+        >>> print(rb)
+        ReplayBufferEnsemble(
+            storages=StorageEnsemble(
+                storages=(<torchrl.data.replay_buffers.storages.LazyMemmapStorage object at 0x13a2ef430>, <torchrl.data.replay_buffers.storages.LazyMemmapStorage object at 0x13a2f9310>),
+                transforms=[Compose(
+                        ToTensorImage(keys=['pixels', ('next', 'pixels')]),
+                        Resize(w=32, h=32, interpolation=InterpolationMode.BILINEAR, keys=['pixels', ('next', 'pixels')]),
+                        RenameTransform(keys=[('some', 'key')])), Compose(
+                        ToTensorImage(keys=['pixels', ('next', 'pixels')]),
+                        Resize(w=32, h=32, interpolation=InterpolationMode.BILINEAR, keys=['pixels', ('next', 'pixels')]),
+                        RenameTransform(keys=['another_key']))]),
+            samplers=SamplerEnsemble(
+                samplers=(<torchrl.data.replay_buffers.samplers.RandomSampler object at 0x13a2f9220>, <torchrl.data.replay_buffers.samplers.RandomSampler object at 0x13a2f9f70>)),
+            writers=WriterEnsemble(
+                writers=(<torchrl.data.replay_buffers.writers.TensorDictRoundRobinWriter object at 0x13a2d9b50>, <torchrl.data.replay_buffers.writers.TensorDictRoundRobinWriter object at 0x13a2f95b0>)),
+        batch_size=None,
+        transform=Compose(
+                Resize(w=33, h=33, interpolation=InterpolationMode.BILINEAR, keys=['pixels'])),
+        collate_fn=<built-in method stack of type object at 0x128648260>)
+        >>> data0 = TensorDict(
+        ...     {
+        ...         "pixels": torch.randint(255, (10, 244, 244, 3)),
+        ...         ("next", "pixels"): torch.randint(255, (10, 244, 244, 3)),
+        ...         ("some", "key"): torch.randn(10),
+        ...     },
+        ...     batch_size=[10],
+        ... )
+        >>> data1 = TensorDict(
+        ...     {
+        ...         "pixels": torch.randint(255, (10, 64, 64, 3)),
+        ...         ("next", "pixels"): torch.randint(255, (10, 64, 64, 3)),
+        ...         "another_key": torch.randn(10),
+        ...     },
+        ...     batch_size=[10],
+        ... )
+        >>> rb[0].extend(data0)
+        >>> rb[1].extend(data1)
+        >>> for _ in range(2):
+        ...     sample = rb.sample(10)
+        ...     assert sample["next", "pixels"].shape == torch.Size([2, 5, 3, 32, 32])
+        ...     assert sample["pixels"].shape == torch.Size([2, 5, 3, 32, 32])
+        ...     assert sample["pixels33"].shape == torch.Size([2, 5, 3, 33, 33])
+        ...     assert sample["renamed"].shape == torch.Size([2, 5])
+
+    """
+
+    _collate_fn_val = None
+
+    def __init__(
+        self,
+        *rbs,
+        storages: StorageEnsemble | None = None,
+        samplers: SamplerEnsemble | None = None,
+        writers: WriterEnsemble | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        batch_size: int | None = None,
+        collate_fn: Callable | None = None,
+        collate_fns: list[Callable] | None = None,
+        p: Tensor = None,
+        sample_from_all: bool = False,
+        num_buffer_sampled: int | None = None,
+        generator: torch.Generator | None = None,
+        shared: bool = False,
+        **kwargs,
+    ):
+
+        if collate_fn is None:
+            collate_fn = _stack_anything
+
+        if rbs:
+            if storages is not None or samplers is not None or writers is not None:
+                raise RuntimeError
+            # Ensure all replay buffers are initialized before creating ensemble
+            for rb in rbs:
+                if (
+                    hasattr(rb, "_delayed_init")
+                    and rb._delayed_init
+                    and not rb.initialized
+                ):
+                    rb._init()
+            storages = StorageEnsemble(
+                *[rb._storage for rb in rbs], transforms=[rb._transform for rb in rbs]
+            )
+            samplers = SamplerEnsemble(
+                *[rb._sampler for rb in rbs],
+                p=p,
+                sample_from_all=sample_from_all,
+                num_buffer_sampled=num_buffer_sampled,
+            )
+            writers = WriterEnsemble(*[rb._writer for rb in rbs])
+            if collate_fns is None:
+                collate_fns = [rb._collate_fn for rb in rbs]
+        else:
+            rbs = None
+            if collate_fns is None:
+                collate_fns = [
+                    _get_default_collate(storage) for storage in storages._storages
+                ]
+        self._rbs = rbs
+        self._collate_fns = collate_fns
+        super().__init__(
+            storage=storages,
+            sampler=samplers,
+            writer=writers,
+            transform=transform,
+            batch_size=batch_size,
+            collate_fn=collate_fn,
+            generator=generator,
+            shared=shared,
+            **kwargs,
+        )
+
+    def _sample(self, *args, **kwargs):
+        sample, info = super()._sample(*args, **kwargs)
+        if isinstance(sample, TensorDictBase):
+            buffer_ids = info.get(("index", "buffer_ids"))
+            info.set(
+                ("index", "buffer_ids"), expand_right(buffer_ids, sample.batch_size)
+            )
+            if isinstance(info, LazyStackedTensorDict):
+                for _info, _sample in zip(
+                    info.unbind(info.stack_dim), sample.unbind(info.stack_dim)
+                ):
+                    _info.batch_size = _sample.batch_size
+                info = torch.stack(info.tensordicts, info.stack_dim)
+            else:
+                info.batch_size = sample.batch_size
+            sample.update(info)
+
+        return sample, info
+
+    @property
+    def _collate_fn(self):
+        def new_collate(samples):
+            samples = [self._collate_fns[i](sample) for (i, sample) in samples]
+            return self._collate_fn_val(samples)
+
+        return new_collate
+
+    @_collate_fn.setter
+    def _collate_fn(self, value):
+        self._collate_fn_val = value
+
+    _INDEX_ERROR = "Expected an index of type torch.Tensor, range, np.ndarray, int, slice or ellipsis, got {} instead."
+
+    def __getitem__(
+        self, index: int | torch.Tensor | tuple | np.ndarray | list | slice | Ellipsis
+    ) -> Any:
+        # accepts inputs:
+        # (int | 1d tensor | 1d list | 1d array | slice | ellipsis | range, int | tensor | list | array | slice | ellipsis | range)
+        # tensor
+        if isinstance(index, tuple):
+            if index[0] is Ellipsis:
+                index = (slice(None), index[1:])
+            rb = self[index[0]]
+            if len(index) > 1:
+                if rb is self:
+                    # then index[0] is an ellipsis/slice(None)
+                    sample = [
+                        (i, storage[index[1:]])
+                        for i, storage in enumerate(self._storage._storages)
+                    ]
+                    return self._collate_fn(sample)
+                if isinstance(rb, ReplayBufferEnsemble):
+                    new_index = (slice(None), *index[1:])
+                    return rb[new_index]
+                return rb[index[1:]]
+            return rb
+        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 received when an integer dtype was expected."
+                )
+        if self._rbs is not None and (
+            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._rbs[index]
+            except IndexError:
+                raise IndexError(self._INDEX_ERROR.format(type(index)))
+
+        if self._rbs is not None:
+            if isinstance(index, torch.Tensor):
+                index = index.tolist()
+                rbs = [self._rbs[i] for i in index]
+                _collate_fns = [self._collate_fns[i] for i in index]
+            else:
+                try:
+                    # slice
+                    rbs = self._rbs[index]
+                    _collate_fns = self._collate_fns[index]
+                except IndexError:
+                    raise IndexError(self._INDEX_ERROR.format(type(index)))
+            p = self._sampler._p[index] if self._sampler._p is not None else None
+            return ReplayBufferEnsemble(
+                *rbs,
+                transform=self._transform,
+                batch_size=self._batch_size,
+                collate_fn=self._collate_fn_val,
+                collate_fns=_collate_fns,
+                sample_from_all=self._sampler.sample_from_all,
+                num_buffer_sampled=self._sampler.num_buffer_sampled,
+                p=p,
+            )
+
+        try:
+            samplers = self._sampler[index]
+            writers = self._writer[index]
+            storages = self._storage[index]
+            if isinstance(index, torch.Tensor):
+                _collate_fns = [self._collate_fns[i] for i in index.tolist()]
+            else:
+                _collate_fns = self._collate_fns[index]
+            p = self._sampler._p[index] if self._sampler._p is not None else None
+
+        except IndexError:
+            raise IndexError(self._INDEX_ERROR.format(type(index)))
+
+        return ReplayBufferEnsemble(
+            samplers=samplers,
+            writers=writers,
+            storages=storages,
+            transform=self._transform,
+            batch_size=self._batch_size,
+            collate_fn=self._collate_fn_val,
+            collate_fns=_collate_fns,
+            sample_from_all=self._sampler.sample_from_all,
+            num_buffer_sampled=self._sampler.num_buffer_sampled,
+            p=p,
+        )
+
+    def __len__(self):
+        return len(self._storage)
+
+    def __repr__(self):
+        storages = textwrap.indent(f"storages={self._storage}", " " * 4)
+        writers = textwrap.indent(f"writers={self._writer}", " " * 4)
+        samplers = textwrap.indent(f"samplers={self._sampler}", " " * 4)
+        return f"ReplayBufferEnsemble(\n{storages}, \n{samplers}, \n{writers}, \nbatch_size={self._batch_size}, \ntransform={self._transform}, \ncollate_fn={self._collate_fn_val})"