reference/generated/torchrl.data.datasets.RobosetExperienceReplay

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

RobosetExperienceReplay

class torchrl.data.datasets.RobosetExperienceReplay(*args, use_ray_service=False, service_backend=None, service_backend_options=None, **kwargs)

Roboset experience replay dataset.

This class downloads the H5 data from roboset and processes it in a mmap format, which makes indexing (and therefore sampling) faster.

Learn more about roboset here: https://sites.google.com/view/robohive/roboset

The data format follows the TED convention.

Parameters:

• dataset_id (str) – the dataset to be downloaded. Must be part of RobosetExperienceReplay.available_datasets.

• batch_size (int) – Batch-size used during sampling. Can be overridden by data.sample(batch_size) if necessary.

Keyword Arguments:

• root (Path or str, optional) – The Roboset dataset root directory. The actual dataset memory-mapped files will be saved under <root>/<dataset_id>. If none is provided, it defaults to ~/.cache/torchrl/atari.roboset`.

• download (bool or str, optional) – Whether the dataset should be downloaded if not found. Defaults to True. Download can also be passed as "force", in which case the downloaded data will be overwritten.

• sampler (Sampler, optional) – the sampler to be used. If none is provided a default RandomSampler() will be used.

• writer (Writer, optional) – the writer to be used. If none is provided a default ImmutableDatasetWriter 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.

• 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.

• transform (Transform, optional) – Transform to be executed when sample() is called. To chain transforms use the Compose class.

• split_trajs (bool, optional) – if True, the trajectories will be split along the first dimension and padded to have a matching shape. To split the trajectories, the "done" signal will be used, which is recovered via done = truncated | terminated. In other words, it is assumed that any truncated or terminated signal is equivalent to the end of a trajectory. Defaults to False.

Variables:

available_datasets – a list of accepted entries to be downloaded.

Examples

>>> import torch
>>> torch.manual_seed(0)
>>> from torchrl.envs.transforms import ExcludeTransform
>>> from torchrl.data.datasets import RobosetExperienceReplay
>>> d = RobosetExperienceReplay("FK1-v4(expert)/FK1_MicroOpenRandom_v2d-v4", batch_size=32,
...     transform=ExcludeTransform("info", ("next", "info")))  # excluding info dict for conciseness
>>> for batch in d:
...     break
>>> # data is organised by seed and episode, but stored contiguously
>>> print(f"{batch['seed']}, {batch['episode']}")
tensor([2, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 2, 2, 2, 1, 1, 2, 0, 2, 0, 2, 2, 1,
        0, 2, 0, 0, 1, 1, 2, 1]) tensor([17, 20, 18,  9,  6,  1, 12,  6,  2,  6,  8, 15,  8, 21, 17,  3,  9, 20,
        23, 12,  3, 16, 19, 16, 16,  4,  4, 12,  1,  2, 15, 24])
>>> print(batch)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([32, 9]), device=cpu, dtype=torch.float64, is_shared=False),
        done: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        episode: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        index: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([32, 75]), device=cpu, dtype=torch.float64, is_shared=False),
                reward: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float64, is_shared=False),
                terminated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([32]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([32, 75]), device=cpu, dtype=torch.float64, is_shared=False),
        seed: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        time: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.float64, is_shared=False)},
        truncated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([32]),
    device=cpu,
    is_shared=False)

add(data: TensorDictBase) → int

Add a single element to the replay buffer.

Parameters:

data (Any) – data to be added to the replay buffer

Returns:

index where the data lives in the replay buffer.

append_transform(transform: Transform, *, invert: bool = False) → ReplayBuffer

Appends transform at the end.

Transforms are applied in order when sample is called.

Parameters:

transform (Transform) – The transform to be appended

Keyword Arguments:

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()

classmethod as_remote(remote_config=None)

Creates an instance of a remote ray class.

Parameters:

• cls (Python Class) – class to be remotely instantiated.

• remote_config (dict) – the quantity of CPU cores to reserve for this class. Defaults to torchrl.collectors.distributed.ray.DEFAULT_REMOTE_CLASS_CONFIG.

Returns:

A function that creates ray remote class instances.

property batch_size

The batch size of the replay buffer.

The batch size can be overridden by setting the batch_size parameter in the sample() method.

It defines both the number of samples returned by sample() and the number of samples that are yielded by the ReplayBuffer iterator.

client() → T

Return self for the zero-overhead direct backend.

property data_path

Path to the dataset, including split.

property data_path_root

Path to the dataset root.

delete()

Deletes a dataset storage from disk.

dump(*args, **kwargs)

Alias for dumps().

dumps(path)

Saves the replay buffer on disk at the specified path.

Parameters:

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)

empty(empty_write_count: bool = True)

Empties the replay buffer and reset cursor to 0.

Parameters:

empty_write_count (bool, optional) – Whether to empty the write_count attribute. Defaults to True.

extend(tensordicts: TensorDictBase, *, update_priority: bool | None = None) → Tensor

Extends the replay buffer with a batch of data.

Parameters:

tensordicts (TensorDictBase) – The data to extend the replay buffer with.

Keyword Arguments:

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.

property initialized: bool

Whether the replay buffer has been initialized.

insert_transform(index: int, transform: Transform, *, invert: bool = False) → ReplayBuffer

Inserts transform.

Transforms are executed in order when sample is called.

Parameters:

• index (int) – Position to insert the transform.

• transform (Transform) – The transform to be appended

Keyword Arguments:

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.

property is_alive: bool

Whether this direct replay buffer remains available.

load(*args, **kwargs)

Alias for loads().

loads(path)

Loads a replay buffer state at the given path.

The buffer should have matching components and be saved using dumps().

Parameters:

path (Path or str) – path where the replay buffer was saved.

See dumps() for more info.

next()

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 RayReplayBuffer.

preprocess(fn: Callable[[TensorDictBase], TensorDictBase], dim: int = 0, num_workers: int | None = None, *, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = False, pbar: bool = False, mp_start_method: str | None = None, num_frames: int | None = None, dest: str | Path) → TensorStorage

Preprocesses a dataset and returns a new storage with the formatted data.

The data transform must be unitary (work on a single sample of the dataset).

The dataset can subsequently be deleted using delete().

Parameters:

• fn (Callable[[TensorDictBase], TensorDictBase]) – transform to apply to each sample.

• dim (int, optional) – dimension along which the dataset is mapped. Defaults to 0.

• num_workers (int, optional) – number of worker processes to use. Defaults to None.

Keyword Arguments:

• chunksize (int, optional) – chunk size forwarded to map().

• num_chunks (int, optional) – number of chunks forwarded to map().

• pool (multiprocessing.Pool, optional) – worker pool forwarded to map().

• generator (torch.Generator, optional) – random generator forwarded to map().

• max_tasks_per_child (int, optional) – maximum number of tasks per child process forwarded to map().

• worker_threads (int, optional) – number of threads per worker forwarded to map(). Defaults to 1.

• index_with_generator (bool, optional) – whether to index with the generator when mapping. Defaults to False.

• pbar (bool, optional) – whether to display a progress bar. Defaults to False.

• mp_start_method (str, optional) – multiprocessing start method forwarded to map().

• dest (path or equivalent) – a path to the location of the new dataset.

• num_frames (int, optional) – if provided, only the first num_frames will be transformed. This is useful to debug the transform at first.

Returns: A new storage to be used within a ReplayBuffer instance.

Examples

>>> from torchrl.data.datasets import MinariExperienceReplay
>>>
>>> data = MinariExperienceReplay(
...     list(MinariExperienceReplay.available_datasets)[0],
...     batch_size=32
...     )
>>> print(data)
MinariExperienceReplay(
    storages=TensorStorage(TensorDict(
        fields={
            action: MemoryMappedTensor(shape=torch.Size([1000000, 8]), device=cpu, dtype=torch.float32, is_shared=True),
            episode: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.int64, is_shared=True),
            info: TensorDict(
                fields={
                    distance_from_origin: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    forward_reward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                    qpos: MemoryMappedTensor(shape=torch.Size([1000000, 15]), device=cpu, dtype=torch.float64, is_shared=True),
                    qvel: MemoryMappedTensor(shape=torch.Size([1000000, 14]), device=cpu, dtype=torch.float64, is_shared=True),
                    reward_ctrl: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    reward_forward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    reward_survive: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    success: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.bool, is_shared=True),
                    x_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    x_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    y_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                    y_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True)},
                batch_size=torch.Size([1000000]),
                device=cpu,
                is_shared=False),
            next: TensorDict(
                fields={
                    done: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                    info: TensorDict(
                        fields={
                            distance_from_origin: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            forward_reward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            qpos: MemoryMappedTensor(shape=torch.Size([1000000, 15]), device=cpu, dtype=torch.float64, is_shared=True),
                            qvel: MemoryMappedTensor(shape=torch.Size([1000000, 14]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_ctrl: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_forward: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            reward_survive: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            success: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.bool, is_shared=True),
                            x_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            x_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            y_position: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True),
                            y_velocity: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float64, is_shared=True)},
                        batch_size=torch.Size([1000000]),
                        device=cpu,
                        is_shared=False),
                    observation: TensorDict(
                        fields={
                            achieved_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            desired_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                            observation: MemoryMappedTensor(shape=torch.Size([1000000, 27]), device=cpu, dtype=torch.float64, is_shared=True)},
                        batch_size=torch.Size([1000000]),
                        device=cpu,
                        is_shared=False),
                    reward: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.float64, is_shared=True),
                    terminated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                    truncated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True)},
                batch_size=torch.Size([1000000]),
                device=cpu,
                is_shared=False),
            observation: TensorDict(
                fields={
                    achieved_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                    desired_goal: MemoryMappedTensor(shape=torch.Size([1000000, 2]), device=cpu, dtype=torch.float64, is_shared=True),
                    observation: MemoryMappedTensor(shape=torch.Size([1000000, 27]), device=cpu, dtype=torch.float64, is_shared=True)},
                batch_size=torch.Size([1000000]),
                device=cpu,
                is_shared=False)},
        batch_size=torch.Size([1000000]),
        device=cpu,
        is_shared=False)),
    samplers=RandomSampler,
    writers=ImmutableDatasetWriter(),
batch_size=32,
transform=Compose(
),
collate_fn=<function _collate_id at 0x120e21dc0>)
>>> from torchrl.envs import CatTensors, Compose
>>> from tempfile import TemporaryDirectory
>>>
>>> cat_tensors = CatTensors(
...     in_keys=[("observation", "observation"), ("observation", "achieved_goal"),
...              ("observation", "desired_goal")],
...     out_key="obs"
...     )
>>> cat_next_tensors = CatTensors(
...     in_keys=[("next", "observation", "observation"),
...              ("next", "observation", "achieved_goal"),
...              ("next", "observation", "desired_goal")],
...     out_key=("next", "obs")
...     )
>>> t = Compose(cat_tensors, cat_next_tensors)
>>>
>>> def func(td):
...     td = td.select(
...         "action",
...         "episode",
...         ("next", "done"),
...         ("next", "observation"),
...         ("next", "reward"),
...         ("next", "terminated"),
...         ("next", "truncated"),
...         "observation"
...         )
...     td = t(td)
...     return td
>>> with TemporaryDirectory() as tmpdir:
...     new_storage = data.preprocess(func, num_workers=4, pbar=True, mp_start_method="fork", dest=tmpdir)
...     rb = ReplayBuffer(storage=new_storage)
...     print(rb)
ReplayBuffer(
    storage=TensorStorage(
        data=TensorDict(
            fields={
                action: MemoryMappedTensor(shape=torch.Size([1000000, 8]), device=cpu, dtype=torch.float32, is_shared=True),
                episode: MemoryMappedTensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.int64, is_shared=True),
                next: TensorDict(
                    fields={
                        done: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                        obs: MemoryMappedTensor(shape=torch.Size([1000000, 31]), device=cpu, dtype=torch.float64, is_shared=True),
                        observation: TensorDict(
                            fields={
                            },
                            batch_size=torch.Size([1000000]),
                            device=cpu,
                            is_shared=False),
                        reward: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.float64, is_shared=True),
                        terminated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True),
                        truncated: MemoryMappedTensor(shape=torch.Size([1000000, 1]), device=cpu, dtype=torch.bool, is_shared=True)},
                    batch_size=torch.Size([1000000]),
                    device=cpu,
                    is_shared=False),
                obs: MemoryMappedTensor(shape=torch.Size([1000000, 31]), device=cpu, dtype=torch.float64, is_shared=True),
                observation: TensorDict(
                    fields={
                    },
                    batch_size=torch.Size([1000000]),
                    device=cpu,
                    is_shared=False)},
            batch_size=torch.Size([1000000]),
            device=cpu,
            is_shared=False),
        shape=torch.Size([1000000]),
        len=1000000,
        max_size=1000000),
    sampler=RandomSampler(),
    writer=RoundRobinWriter(cursor=0, full_storage=True),
    batch_size=None,
    collate_fn=<function _collate_id at 0x168406fc0>)

query(predicate: Callable[[Trajectory], bool] | None = None, *, trajectory_key: NestedKey | None = None) → list[Trajectory]

Filters the stored trajectories with a query predicate.

Splits the buffer content into trajectories (see iter_trajectories()) and returns those matching the predicate as Trajectory views.

Parameters:

predicate (Callable[[Trajectory], bool], optional) – a TrajectoryPredicate built from traj, or any callable mapping a trajectory to a boolean. Defaults to None (return all trajectories).

Keyword Arguments:

trajectory_key (NestedKey, optional) – entry holding per-transition trajectory ids. Defaults to None (auto-detection from ("collector", "traj_ids"), "traj_ids", "episode" or the done/terminated/truncated flags).

Returns:

A list of matching trajectory views, ordered chronologically (oldest trajectory first; for multi-dimensional storages, grouped by batch coordinate).

The trajectory boundaries are computed from the stored (untransformed) data with the same machinery SliceSampler uses, so samplers and queries always agree on where trajectories start and stop. This includes storages with ndim > 1 (e.g. LazyTensorStorage(..., ndim=2) holding [B, T] batches), whose trajectories are recovered per batch coordinate.

Predicates built from traj report the keys they read via required_keys(); evaluation then only fetches those entries from the storage and only runs the transforms that can affect them. Matching trajectories are extracted in full with the complete transform chain applied, so predicates and results see the same values a sampler would produce. Opaque callables are evaluated against the fully transformed content.

Note

Once the buffer has wrapped around (it is at capacity and older entries have been overwritten), the oldest trajectory may have lost its first transitions to overwriting and will appear truncated at the front. A trajectory written across the wrap point is followed through it and returned whole, in time order.

Examples

>>> from torchrl.data import traj
>>> good_trajs = rb.query((traj.reward.sum() > 100) & (traj.length >= 50))
>>> observations = good_trajs[0].observation

read_all_in_order(end: int | None = None) → Any

Read storage contents in physical order.

This is equivalent to rb[:] when end is None.

Parameters:

end (int, optional) – Number of leading storage entries to read. Defaults to the entire storage slice.

Returns:

A storage slice containing entries [:end].

register_load_hook(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.

register_save_hook(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.

sample(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.

Parameters:

• 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.

• include_info (bool, optional) – deprecated alias for return_info.

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.

property sampler: Sampler

The sampler of the replay buffer.

The sampler must be an instance of Sampler.

save(*args, **kwargs)

Alias for dumps().

property service_backend: str

The canonical deployment backend for this replay buffer.

set_(key, value)

Sets the value of a key across the entire replay buffer in-place.

Parameters:

• key (NestedKey) – the key to set.

• value (torch.Tensor) – the value to write.

Returns:

self

set_at_(key, value, index)

Sets the value of a key at specified indices in the replay buffer.

Parameters:

• key (NestedKey) – the key to set.

• value (torch.Tensor) – the value to write.

• index – the indices where to write the value.

Returns:

self

set_sampler(sampler: Sampler)

Sets a new sampler in the replay buffer and returns the previous sampler.

set_storage(storage: Storage, collate_fn: Callable | None = None)

Sets a new storage in the replay buffer and returns the previous storage.

Parameters:

• 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.

set_writer(writer: Writer)

Sets a new writer in the replay buffer and returns the previous writer.

shutdown(timeout: float | None = None) → None

Mark this direct replay-buffer owner as shut down.

start() → T

Return this already-started direct replay buffer.

stats() → dict[str, int | float | bool]

Returns a cheap, serializable snapshot of the buffer’s operational state.

The snapshot only contains scalar counters and gauges. It never includes the storage content, does not modify the buffer state and is safe to call concurrently with writes and samples. Cumulative counters such as write_count are meant to be converted into rates by an external monitor such as LoggerMonitor.

Calling this method on an uninitialized buffer does not trigger its initialization; an empty snapshot with initialized=False is returned instead (capacity is still reported when the storage already advertises it).

Returns:

• "size": current number of elements in the buffer (mirrors len(buffer));

• "write_count": total number of items written through add and extend (0 for writers that do not track writes, such as ImmutableDatasetWriter);

• "prefetch_queue_size": number of pending prefetched batches;

• "initialized": whether the buffer components are initialized;

• "capacity": maximum number of elements the storage can hold (only present when the storage advertises a max_size);

• "utilization": size / capacity (only present alongside capacity).

Remote clients backed by the distributed transport report a subset of these entries (size and write_count).

Return type:

A dictionary with the following entries

Examples

>>> import torch
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer
>>> rb = ReplayBuffer(storage=LazyTensorStorage(10))
>>> rb.extend(torch.arange(5))
>>> snapshot = rb.stats()
>>> print(snapshot["size"], snapshot["write_count"], snapshot["capacity"])
5 5 10

property storage: Storage

The storage of the replay buffer.

The storage must be an instance of Storage.

property transform: Transform

The transform of the replay buffer.

The transform must be an instance of Transform.

update_(input_dict_or_td, clone=False, *, keys_to_update=None)

Updates the replay buffer in-place with the given dict or TensorDict.

Parameters:

• input_dict_or_td (dict or TensorDictBase) – the data to update with.

• clone (bool, optional) – whether to clone the values before writing. Defaults to False.

• keys_to_update (sequence of NestedKey, optional) – if provided, only these keys will be updated.

Returns:

self

update_if_present(*, index: Tensor, generation: Tensor, patch: Mapping[NestedKey, Tensor] | TensorDictBase, version_key: NestedKey | None = None, version: int | Tensor | None = None, require_newer: bool = False) → ConditionalUpdateResult

Conditionally updates stored records that are still live.

Replay slots are recycled by round-robin writers, so a physical index captured at sampling time can point to a different record by the time an asynchronous computation writes back. This method applies patch only to records whose (index, generation) pair still matches the writer’s current slot generation, skipping records whose slot was reused or emptied since the handle was captured. Skipped records are never modified.

The whole patch is validated (key existence, shape and dtype) before any write happens; a validation failure leaves the storage untouched. Updating a record refreshes its content, not its identity: the same handle keeps working until the slot is rewritten by add, extend or empty.

Generation tracking is opt-in: the buffer must be constructed with a writer that tracks slot generations, e.g. RoundRobinWriter(track_generations=True) (see ref_buffers_generations). Calling this method on a buffer whose writer does not track generations raises a RuntimeError.

Keyword Arguments:

• index (torch.Tensor) – storage indices, as returned by extend() or found in the sample under "index".

• generation (torch.Tensor) – slot generations captured with the indices, as found in the sample under "index_generation".

• patch (mapping of NestedKey to torch.Tensor, or TensorDictBase) – the fields to overwrite for live records. Leading dimension must match the number of records addressed by index.

• version_key (NestedKey, optional) – a stored per-record scalar field holding each record’s current version. When passed (together with version), a generation-live record is only patched if the incoming version compares favorably against the stored one, and the accepted version is written into version_key atomically with the patch. version_key may not appear in patch. Nested keys must be passed in tuple form (("nested", "version")); dotted strings are rejected. Defaults to None (no version comparison).

• version (int or torch.Tensor, optional) – the incoming version, either a scalar (broadcast to every record) or a tensor with one entry per record. Must be passed together with version_key.

• require_newer (bool, optional) – if True, a record is only patched when version > stored; if False, ties are accepted (version >= stored). When the same slot is addressed several times in one call, only the row carrying the highest incoming version is applied (the last such row on ties); the losing rows are reported in version_rejected. Defaults to False.

Returns:

A ConditionalUpdateResult whose updated mask is aligned with the input index order, with updated_count and stale_count conveniences. When version_key is passed, its version_rejected mask marks generation-live records that were rejected by the version comparison (None otherwise).

Raises:

• RuntimeError – if the storage does not support conditional updates (for example ListStorage) or the writer does not track slot generations.

• KeyError – if a patch key (or version_key) does not exist in the storage.

• ValueError – if a patch entry has an incompatible shape or dtype, if only one of version_key / version is passed, if version_key appears in patch or names a non-scalar field, or if it is a dotted string.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data import (
...     LazyTensorStorage,
...     TensorDictReplayBuffer,
...     TensorDictRoundRobinWriter,
... )
>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(10),
...     writer=TensorDictRoundRobinWriter(track_generations=True),
...     batch_size=4,
... )
>>> rb.extend(TensorDict({"obs": torch.zeros(10, 3)}, batch_size=[10]))
>>> sample = rb.sample()
>>> result = rb.update_if_present(
...     index=sample["index"],
...     generation=sample["index_generation"],
...     patch={"obs": torch.ones(4, 3)},
... )
>>> print(result.updated_count, result.stale_count)
4 0

With a version comparison, outdated asynchronous writers lose deterministically:

>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(10),
...     writer=TensorDictRoundRobinWriter(track_generations=True),
...     batch_size=4,
... )
>>> rb.extend(
...     TensorDict(
...         {
...             "obs": torch.zeros(10, 3),
...             "v": torch.full((10,), 5, dtype=torch.int64),
...         },
...         batch_size=[10],
...     )
... )
>>> sample = rb.sample()
>>> result = rb.update_if_present(
...     index=sample["index"],
...     generation=sample["index_generation"],
...     patch={"obs": torch.ones(4, 3)},
...     version_key="v",
...     version=4,
...     require_newer=True,
... )
>>> print(result.updated_count, result.version_rejected_count)
0 4

write_all(data: Any, end: int | None = None) → None

Write data back to storage in physical order.

This is equivalent to rb[:end] = data. If end is None, end defaults to data.shape[0] for tensor collections and len(data) otherwise. If data spans the full storage, this is equivalent to rb[:] = data.

Parameters:

• data – Data to write to storage.

• end (int, optional) – Number of leading storage entries to update. Defaults to data.shape[0] for tensor collections and len(data) otherwise.

property write_count: int

The total number of items written so far in the buffer through add and extend.

property writer: Writer

The writer of the replay buffer.

The writer must be an instance of Writer.