ReplayDatasetAdapter#

class torchrl.data.datasets.ReplayDatasetAdapter(*args, **kwargs)[source]#

Import and export an external dataset format for use with replay buffers.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data import TensorStorage
>>> from torchrl.data.datasets import ReplayDatasetAdapter
>>> class ColumnAdapter:
...     def from_external(self, columns):
...         size = next(iter(columns.values())).shape[0]
...         return TensorStorage(TensorDict(columns, [size])).as_dataset()
...     def to_external(self, dataset):
...         keys = dataset[0].keys()
...         return {
...             key: torch.stack([dataset[index][key] for index in range(len(dataset))])
...             for key in keys
...         }
>>> adapter: ReplayDatasetAdapter[dict[str, torch.Tensor]] = ColumnAdapter()
>>> external = {"observation": torch.arange(2)}
>>> round_trip = adapter.to_external(adapter.from_external(external))
>>> torch.equal(round_trip["observation"], external["observation"])
True

ReplayDatasetAdapter#

Docs

Tutorials

Resources