SliceSamplerWithoutReplacement¶
- class torchrl.data.replay_buffers.SliceSamplerWithoutReplacement(*, num_slices: Optional[int] = None, slice_len: Optional[int] = None, drop_last: bool = False, end_key: Optional[NestedKey] = None, traj_key: Optional[NestedKey] = None, ends: Optional[Tensor] = None, trajectories: Optional[Tensor] = None, truncated_key: tensordict._nestedkey.NestedKey | None = ('next', 'truncated'), strict_length: bool = True, shuffle: bool = True, compile: bool | dict = False)[source]¶
沿第一个维度采样数据切片,给定起始和停止信号,无需替换。
此类与静态回放缓冲区一起使用,或在两个回放缓冲区扩展之间使用。扩展回放缓冲区将重置采样器,目前不允许连续无替换采样。
- 关键字参数:
drop_last (bool, optional) – 如果
True
,则将删除最后一个不完整的样本(如果有)。如果False
,则保留最后一个样本。默认为False
。num_slices (int) – 要采样的切片数量。批大小必须大于或等于
num_slices
参数。与slice_len
互斥。slice_len (int) – 要采样的切片长度。批大小必须大于或等于
slice_len
参数,并且可以被其整除。与num_slices
互斥。end_key (NestedKey, optional) – 指示轨迹(或 episode)结束的键。默认为
("next", "done")
。traj_key (NestedKey, optional) – 指示轨迹的键。默认为
"episode"
(TorchRL 中数据集常用的键)。ends (torch.Tensor, optional) – 一个 1d 布尔张量,包含运行结束信号。当
end_key
或traj_key
获取成本高昂,或者当此信号已准备就绪时使用。必须与cache_values=True
一起使用,并且不能与end_key
或traj_key
结合使用。trajectories (torch.Tensor, optional) – 一个 1d 整数张量,包含运行 ID。当
end_key
或traj_key
获取成本高昂,或者当此信号已准备就绪时使用。必须与cache_values=True
一起使用,并且不能与end_key
或traj_key
结合使用。truncated_key (NestedKey, optional) – 如果不为
None
,则此参数指示截断信号应写入输出数据的位置。这用于向值估计器指示提供的轨迹在哪里中断。默认为("next", "truncated")
。此功能仅适用于TensorDictReplayBuffer
实例(否则截断键将返回到sample()
方法返回的 info 字典中)。strict_length (bool, optional) – 如果
False
,则允许长度短于 slice_len(或 batch_size // num_slices)的轨迹出现在批次中。如果True
,则会过滤掉短于要求的轨迹。请注意,这可能会导致有效的 batch_size 小于要求的批大小!可以使用split_trajectories()
拆分轨迹。默认为True
。shuffle (bool, optional) – 如果
False
,则轨迹的顺序不会被打乱。默认为True
。compile (bool or dict of kwargs, optional) – 如果
True
,则sample()
方法的瓶颈将使用compile()
进行编译。关键字参数也可以通过此参数传递给 torch.compile。默认为False
。
注意
为了恢复存储中的轨迹分割,
SliceSamplerWithoutReplacement
将首先尝试在存储中查找traj_key
条目。如果找不到,将使用end_key
重构 episode。示例
>>> 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 的大多数数据集兼容,并允许用户以类似 dataloader 的方式使用数据集示例
>>> 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])