StorageDataset#

class torchrl.data.replay_buffers.StorageDataset(storage: Storage, *, share_memory: bool = True)[source]#

A map-style torch.utils.data.Dataset reading a TorchRL storage.

The dataset has one item per storage entry and a torch.utils.data.DataLoader reads it with any torch sampler and worker processes. __getitems__ fetches an index batch with a single get() call, so the loader receives one batch rather than a list of items: pass tensordict_collate() as collate_fn, which returns such batches unchanged and stacks lists of items. Storages with more than one dimension are read through 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 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.

Parameters:: storage (Storage) – the storage to read. Must be one-dimensional.
Keyword Arguments:: 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])

StorageDataset#

Docs

Tutorials

Resources