SliceSampler¶

根据开始和停止信号，沿第一个维度对数据切片进行采样。

此类对子轨迹进行有放回采样。有关无放回采样的版本，请参阅 SliceSamplerWithoutReplacement。

关键字参数：

num_slices (int) – 要采样的切片数量。批大小必须大于或等于 num_slices 参数。与 slice_len 互斥。
slice_len (int) – 要采样的切片的长度。批大小必须大于或等于 slice_len 参数，并且可以被其整除。与 num_slices 互斥。
end_key (NestedKey, 可选) – 指示轨迹（或情节）结束的键。默认为 ("next", "done")。
traj_key (NestedKey, 可选) – 指示轨迹的键。默认为 "episode"（在 TorchRL 中的数据集中通常使用）。
ends (torch.Tensor, 可选) – 一个 1d 布尔张量，包含运行结束信号。在 end_key 或 traj_key 获取成本很高或此信号已准备就绪时使用。必须与 cache_values=True 一起使用，并且不能与 end_key 或 traj_key 结合使用。如果提供，则假设存储已满，并且如果 ends 张量的最后一个元素为 False，则相同的轨迹跨越结束和开始。
trajectories (torch.Tensor, 可选) – 一个 1d 整数张量，包含运行 ID。在 end_key 或 traj_key 获取成本很高或此信号已准备就绪时使用。必须与 cache_values=True 一起使用，并且不能与 end_key 或 traj_key 结合使用。如果提供，则假设存储已满，并且如果轨迹张量的最后一个元素与第一个相同，则相同的轨迹跨越结束和开始。
cache_values (bool, 可选) –
用于静态数据集。将缓存轨迹的开始和结束信号。即使在调用 extend 期间轨迹索引发生更改时，也可以安全地使用此操作，因为此操作将清除缓存。
警告

如果采样器与由另一个缓冲区扩展的存储一起使用，则cache_values=True 将不起作用。例如
```
>>> buffer0 = ReplayBuffer(storage=storage,
...     sampler=SliceSampler(num_slices=8, cache_values=True),
...     writer=ImmutableWriter())
>>> buffer1 = ReplayBuffer(storage=storage,
...     sampler=other_sampler)
>>> # Wrong! Does not erase the buffer from the sampler of buffer0
>>> buffer1.extend(data)
```
警告

如果缓冲区在进程之间共享，并且一个进程负责写入，另一个进程负责采样，则cache_values=True 不会按预期工作，因为清除缓存只能在本地完成。
truncated_key (NestedKey, 可选) – 如果不是 None，则此参数指示在输出数据中应将截断信号写入何处。这用于指示值估计器提供轨迹中断的位置。默认为 ("next", "truncated")。此功能仅适用于 TensorDictReplayBuffer 实例（否则，截断键将在 sample() 方法返回的信息字典中返回）。
strict_length (bool, 可选) – 如果为 False，则允许长度小于 slice_len（或 batch_size // num_slices）的轨迹出现在批次中。如果为 True，则将过滤掉比所需长度短的轨迹。请注意，这可能导致有效 batch_size 短于请求的 batch_size！可以使用 split_trajectories() 拆分轨迹。默认为 True。
compile (bool 或 kwargs 字典, 可选) – 如果为 True，则 sample() 方法的瓶颈将使用 compile() 编译。也可以使用此参数将关键字参数传递给 torch.compile。默认为 False。
span (bool, int, Tuple[bool | int, bool | int], 可选) – 如果提供，则采样的轨迹将跨越左侧和/或右侧。这意味着可能提供的元素少于所需元素。布尔值表示每个轨迹至少采样一个元素。整数 i 表示每个采样轨迹至少收集 slice_len - i 个样本。使用元组可以对左侧（存储轨迹的开始）和右侧（存储轨迹的结束）的跨度进行细粒度控制。

注意

为了恢复存储中的轨迹分割，SliceSampler 将首先尝试在存储中查找 traj_key 条目。如果找不到，则将使用 end_key 重建片段。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data.replay_buffers import LazyMemmapStorage, TensorDictReplayBuffer
>>> from torchrl.data.replay_buffers.samplers import SliceSampler
>>> torch.manual_seed(0)
>>> rb = TensorDictReplayBuffer(
...     storage=LazyMemmapStorage(1_000_000),
...     sampler=SliceSampler(cache_values=True, num_slices=10),
...     batch_size=320,
... )
>>> episode = torch.zeros(1000, dtype=torch.int)
>>> episode[:300] = 1
>>> episode[300:550] = 2
>>> episode[550:700] = 3
>>> episode[700:] = 4
>>> data = TensorDict(
...     {
...         "episode": episode,
...         "obs": torch.randn((3, 4, 5)).expand(1000, 3, 4, 5),
...         "act": torch.randn((20,)).expand(1000, 20),
...         "other": torch.randn((20, 50)).expand(1000, 20, 50),
...     }, [1000]
... )
>>> rb.extend(data)
>>> sample = rb.sample()
>>> print("sample:", sample)
>>> print("episodes", sample.get("episode").unique())
episodes tensor([1, 2, 3, 4], dtype=torch.int32)

SliceSampler 默认兼容大多数 TorchRL 的数据集

示例

>>> import torch
>>>
>>> from torchrl.data.datasets import RobosetExperienceReplay
>>> from torchrl.data import SliceSampler
>>>
>>> torch.manual_seed(0)
>>> num_slices = 10
>>> dataid = list(RobosetExperienceReplay.available_datasets)[0]
>>> data = RobosetExperienceReplay(dataid, batch_size=320, sampler=SliceSampler(num_slices=num_slices))
>>> for batch in data:
...     batch = batch.reshape(num_slices, -1)
...     break
>>> print("check that each batch only has one episode:", batch["episode"].unique(dim=1))
check that each batch only has one episode: tensor([[19],
        [14],
        [ 8],
        [10],
        [13],
        [ 4],
        [ 2],
        [ 3],
        [22],
        [ 8]])

SliceSampler¶

文档

教程

资源