# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
from collections.abc import Mapping
from itertools import islice
from typing import Any, TYPE_CHECKING
import torch
from tensordict import is_tensor_collection, LazyStackedTensorDict
from torch.utils.data import Dataset, get_worker_info, IterableDataset
from torch.utils.data._utils.collate import default_collate
if TYPE_CHECKING:
from torch.utils.data._utils.worker import WorkerInfo
from torchrl.data.replay_buffers.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.storages import Storage
__all__ = ["ReplayBufferDataset", "StorageDataset", "tensordict_collate"]
[docs]
def tensordict_collate(batch: Any) -> Any:
"""Collate function for a :class:`torch.utils.data.DataLoader` reading TorchRL storages or buffers.
A batch that is already a tensor, a tensor collection or a tuple of them,
as fetched by :class:`StorageDataset` or yielded by
:class:`ReplayBufferDataset`, is returned unchanged. A list of samples, as
produced by per-item storages or by composing datasets with
:class:`torch.utils.data.ConcatDataset`, is stacked: tensor collections
lazily when their shapes differ, tensors densely, and mappings or tuples
element-wise. The default torch collation iterates a tensordict over its
batch dimension and cannot be used.
Args:
batch (Tensor, TensorDictBase or list): the fetched batch.
Examples:
>>> import torch
>>> from tensordict import TensorDict
>>> from torch.utils.data import DataLoader
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer, tensordict_collate
>>> rb = ReplayBuffer(storage=LazyTensorStorage(100))
>>> _ = rb.extend(TensorDict({"obs": torch.arange(100)}, [100]))
>>> loader = DataLoader(
... rb.storage.as_dataset(),
... batch_size=4,
... shuffle=True,
... collate_fn=tensordict_collate,
... )
>>> next(iter(loader))["obs"].shape
torch.Size([4])
"""
if not isinstance(batch, list) or not batch:
return batch
first = batch[0]
if is_tensor_collection(first):
return LazyStackedTensorDict.maybe_dense_stack(batch)
if isinstance(first, torch.Tensor):
return torch.stack(batch)
if isinstance(first, Mapping):
return {key: tensordict_collate([item[key] for item in batch]) for key in first}
if isinstance(first, tuple):
return tuple(tensordict_collate(list(items)) for items in zip(*batch))
return default_collate(batch)
def _check_fork_safe(storage: Storage) -> None:
if not storage._fork_safe():
raise RuntimeError(
f"This DataLoader worker was forked with a private copy of a "
f"{type(storage).__name__} allocated after the dataset was created, "
"and would read stale rows. Create the dataset once the storage "
"holds data, or use a LazyMemmapStorage or the spawn start method."
)
[docs]
class StorageDataset(Dataset):
"""A map-style :class:`torch.utils.data.Dataset` reading a TorchRL storage.
The dataset has one item per storage entry and a
:class:`torch.utils.data.DataLoader` reads it with any torch sampler and
worker processes. ``__getitems__`` fetches an index batch with a single
:meth:`~torchrl.data.replay_buffers.Storage.get` call, so the loader
receives one batch rather than a list of items: pass
:func:`tensordict_collate` as ``collate_fn``, which returns such batches
unchanged and stacks lists of items. Storages with more than one dimension
are read through :meth:`~torchrl.data.replay_buffers.Storage.flatten`.
Reading the storage directly with a DataLoader, without this adapter,
fetches items one by one and passes the list to the collate function.
DataLoader workers read the storage content live: by default, creating the
dataset moves CPU tensor storage to shared memory and memory-mapped storages are
read through their files, so rows written after the workers start are
visible to them under every start method. Reads are not synchronized with
writes, and a row written while a worker reads it can come back partially
updated. A :class:`~torchrl.data.replay_buffers.ListStorage` cannot be
sent to spawned workers and forked workers read a snapshot of it. Only the
storage is sent to the workers, not the buffers attached to it.
Args:
storage (Storage): the storage to read. Must be one-dimensional.
Keyword Args:
share_memory (bool, optional): whether to move CPU tensor storage to
shared memory when creating the dataset, so DataLoader workers
observe later writes. Defaults to ``True``. Set to ``False`` for
reads in the calling process or when storage sharing is already
arranged.
Examples:
>>> import torch
>>> from tensordict import TensorDict
>>> from torch.utils.data import DataLoader
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer, tensordict_collate
>>> rb = ReplayBuffer(storage=LazyTensorStorage(100))
>>> _ = rb.extend(TensorDict({"obs": torch.arange(100)}, [100]))
>>> dataset = rb.storage.as_dataset()
>>> len(dataset)
100
>>> loader = DataLoader(
... dataset, batch_size=4, shuffle=True, collate_fn=tensordict_collate
... )
>>> next(iter(loader))["obs"].shape
torch.Size([4])
"""
def __init__(self, storage: Storage, *, share_memory: bool = True) -> None:
if storage.ndim > 1:
raise RuntimeError(
f"A {type(storage).__name__} with ndim={storage.ndim} is not a "
"flat dataset. Read storage.flatten().as_dataset() instead."
)
if share_memory:
storage._share_memory_()
self.storage = storage
self._worker_checked = False
def __getstate__(self) -> dict[str, Any]:
storage = self.storage
storage_state = storage.__getstate__()
storage_state["_attached_entities_list"] = []
return {**self.__dict__, "storage": (type(storage), storage_state)}
def __setstate__(self, state: dict[str, Any]) -> None:
storage_cls, storage_state = state.pop("storage")
storage = storage_cls.__new__(storage_cls)
storage.__setstate__(storage_state)
self.__dict__.update(state, storage=storage)
def _check_worker(self) -> None:
if not self._worker_checked and get_worker_info() is not None:
_check_fork_safe(self.storage)
self._worker_checked = True
def __len__(self) -> int:
return len(self.storage)
def __getitem__(self, index: int | list[int]) -> Any:
self._check_worker()
return self.storage[index]
def __getitems__(self, indices: list[int]) -> Any:
self._check_worker()
return self.storage.get(indices)
def __repr__(self) -> str:
return f"{type(self).__name__}(storage={self.storage!r})"
[docs]
class ReplayBufferDataset(IterableDataset):
"""A :class:`torch.utils.data.IterableDataset` streaming batches from a replay buffer.
Iterating the dataset iterates the buffer: every item is one batch of
``replay_buffer.batch_size`` elements with the buffer transforms applied.
Pass the dataset to a :class:`torch.utils.data.DataLoader` with
``batch_size=None`` and :func:`tensordict_collate` to sample in worker
processes. Each worker holds its own copy of the buffer, so the sampler and
the transforms run in the worker and ``num_batches`` is split between
workers. The storage content is shared rather than copied and workers
observe later writes as described in :class:`StorageDataset`.
Buffer prefetching is disabled in workers and prefetched batches are never
serialized to them, the DataLoader prefetches instead. A buffer built with
a :class:`torch.Generator` is reseeded once per worker from the worker
seed, so sampling in workers is reproducible when the DataLoader is seeded
(``torch.manual_seed`` or ``DataLoader(generator=...)``).
Samplers whose
:attr:`~torchrl.data.replay_buffers.Sampler.requires_shared_state` is
``True``, which is every sampler except those that declare their draws
stateless such as :class:`~torchrl.data.replay_buffers.RandomSampler` and
:class:`~torchrl.data.replay_buffers.SliceSampler`, are rejected when
workers are used. So is a :class:`~torchrl.data.RateLimitedReplayBuffer`
that has not been shared with :meth:`~torchrl.data.ReplayBuffer.share`,
since each worker would otherwise spend its own copy of the sample budget.
Args:
replay_buffer (ReplayBuffer): the buffer to sample from. Its
``batch_size`` must be set.
Keyword Args:
num_batches (int or None, optional): number of batches yielded by one
iterator, shared between DataLoader workers. ``None`` streams
batches until the sampler runs out, which samplers with replacement
never do. Defaults to ``None``.
Examples:
>>> import torch
>>> from tensordict import TensorDict
>>> from torch.utils.data import DataLoader
>>> from torchrl.data import (
... LazyMemmapStorage,
... SliceSampler,
... TensorDictReplayBuffer,
... tensordict_collate,
... )
>>> rb = TensorDictReplayBuffer(
... storage=LazyMemmapStorage(1000),
... sampler=SliceSampler(num_slices=4, traj_key="episode", cache_values=True),
... batch_size=32,
... )
>>> _ = rb.extend(
... TensorDict(
... {"obs": torch.randn(1000, 3), "episode": torch.arange(1000) // 50},
... [1000],
... )
... )
>>> loader = DataLoader(
... rb.as_dataset(num_batches=8),
... batch_size=None,
... num_workers=4,
... persistent_workers=True,
... collate_fn=tensordict_collate,
... )
>>> batches = list(loader)
>>> len(batches), batches[0]["obs"].shape
(8, torch.Size([32, 3]))
"""
def __init__(
self, replay_buffer: ReplayBuffer, *, num_batches: int | None = None
) -> None:
replay_buffer.storage._share_memory_()
self.replay_buffer = replay_buffer
self.num_batches = num_batches
self._worker_seed = None
def __getstate__(self) -> dict[str, Any]:
replay_buffer = self.replay_buffer
self._check_shared_state(replay_buffer)
return {
**self.__dict__,
"replay_buffer": (
type(replay_buffer),
replay_buffer._state_without_prefetch(),
),
}
def __setstate__(self, state: dict[str, Any]) -> None:
buffer_cls, buffer_state = state.pop("replay_buffer")
replay_buffer = buffer_cls.__new__(buffer_cls)
replay_buffer.__setstate__(buffer_state)
self.__dict__.update(state, replay_buffer=replay_buffer)
def _check_shared_state(self, replay_buffer: ReplayBuffer) -> None:
sampler = replay_buffer.sampler
if sampler.requires_shared_state:
raise RuntimeError(
f"{type(sampler).__name__} keeps sampling state that DataLoader "
"workers cannot share. Use num_workers=0 or a sampler without "
"cross-process state."
)
if replay_buffer._requires_shared_state:
raise RuntimeError(
f"{type(replay_buffer).__name__} keeps sampling state that "
"DataLoader workers cannot share unless the buffer is shared. "
"Call share() on the buffer or use num_workers=0."
)
def _setup_worker(self) -> WorkerInfo | None:
worker = get_worker_info()
if worker is None:
return None
replay_buffer = self.replay_buffer
self._check_shared_state(replay_buffer)
_check_fork_safe(replay_buffer.storage)
replay_buffer._reset_worker_state()
rng = replay_buffer._rng
if rng is not None and self._worker_seed != worker.seed:
replay_buffer.set_rng(
torch.Generator(device=rng.device).manual_seed(worker.seed)
)
self._worker_seed = worker.seed
return worker
def __iter__(self):
worker = self._setup_worker()
num_batches = self.num_batches
if worker is not None and num_batches is not None:
num_batches = len(range(worker.id, num_batches, worker.num_workers))
yield from islice(self.replay_buffer, num_batches)
def __repr__(self) -> str:
return (
f"{type(self).__name__}(replay_buffer={self.replay_buffer!r}, "
f"num_batches={self.num_batches})"
)
