SliceSamplerWithoutReplacement

class torchrl.data.replay_buffers.SliceSamplerWithoutReplacement(*, num_slices: int | None = None, slice_len: int | None = None, drop_last: bool = False, end_key: NestedKey | None = None, traj_key: NestedKey | None = None, ends: torch.Tensor | None = None, trajectories: torch.Tensor | None = None, truncated_key: NestedKey | None = ('next', 'truncated'), strict_length: bool = True, shuffle: bool = True, compile: bool | dict = False)

在第一个维度上采样数据片段，给定开始和停止信号，并且不进行替换。

此类用于静态回放缓冲区或两个回放缓冲区扩展之间。扩展回放缓冲区将重置采样器，并且当前不允许连续不进行替换的采样。

关键字参数：

• drop_last (bool, 可选) – 如果True，则将丢弃最后一个不完整的样本（如果有）。如果False，则保留此最后一个样本。默认为False。

• 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 结合使用。

• trajectories (torch.Tensor, 可选) – 一个 1d 整数张量，包含运行 ID。在end_key 或traj_key 获取成本很高或此信号已准备就绪时使用。必须与cache_values=True 一起使用，并且不能与end_key 或traj_key 结合使用。

• truncated_key (NestedKey, 可选) – 如果不为None，则此参数指示在输出数据中应写入截断信号的位置。这用于指示值估计器轨迹在哪里断开。默认为("next", "truncated")。此功能仅适用于TensorDictReplayBuffer 实例（否则，截断键将在sample() 方法返回的信息字典中返回）。

• strict_length (布尔值, 可选) – 如果为 False，则允许长度短于 slice_len（或 batch_size // num_slices）的轨迹出现在批次中。如果为 True，则会过滤掉短于所需长度的轨迹。请注意，这可能导致有效的 batch_size 短于请求的！可以使用 split_trajectories() 分割轨迹。默认为 True。

• shuffle (布尔值, 可选) – 如果为 False，则轨迹的顺序不会被打乱。默认为 True。

• compile (布尔值 或 关键字参数字典, 可选) – 如果为 True，则 sample() 方法的瓶颈部分将使用 compile() 进行编译。还可以通过此参数将关键字参数传递给 torch.compile。默认为 False。

注意

要恢复存储中的轨迹分割，SliceSamplerWithoutReplacement 将首先尝试在存储中查找 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 SliceSamplerWithoutReplacement
>>>
>>> rb = TensorDictReplayBuffer(
...     storage=LazyMemmapStorage(1000),
...     # asking for 10 slices for a total of 320 elements, ie, 10 trajectories of 32 transitions each
...     sampler=SliceSamplerWithoutReplacement(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()
>>> # since we want trajectories of 32 transitions but there are only 4 episodes to
>>> # sample from, we only get 4 x 32 = 128 transitions in this batch
>>> print("sample:", sample)
>>> print("trajectories in sample", sample.get("episode").unique())

SliceSamplerWithoutReplacement 与大多数 TorchRL 的数据集默认兼容，并允许用户以类似数据加载器的方式使用数据集。

示例

>>> import torch
>>>
>>> from torchrl.data.datasets import RobosetExperienceReplay
>>> from torchrl.data import SliceSamplerWithoutReplacement
>>>
>>> torch.manual_seed(0)
>>> num_slices = 10
>>> dataid = list(RobosetExperienceReplay.available_datasets)[0]
>>> data = RobosetExperienceReplay(dataid, batch_size=320,
...     sampler=SliceSamplerWithoutReplacement(num_slices=num_slices))
>>> # the last sample is kept, since drop_last=False by default
>>> for i, batch in enumerate(data):
...     print(batch.get("episode").unique())
tensor([ 5,  6,  8, 11, 12, 14, 16, 17, 19, 24])
tensor([ 1,  2,  7,  9, 10, 13, 15, 18, 21, 22])
tensor([ 0,  3,  4, 20, 23])