BatchSubSampler

class torchrl.trainers.BatchSubSampler(batch_size: int, sub_traj_len: int = 0, min_sub_traj_len: int = 0)

用于在线 RL SOTA 实施的数据子采样器。

此类对从环境中收集的整个数据批次的一部分进行子采样。

参数：

• batch_size (int) – 要收集的子批次大小。提供的批次大小必须等于输出 tensordict 中项目的总数，其大小将为 [batch_size // sub_traj_len, sub_traj_len]。

• sub_traj_len (int, 可选) – 子采样必须在在线设置中具有的轨迹长度。默认值为 -1（即采用轨迹的完整长度）。

• min_sub_traj_len (int, 可选) – sub_traj_len 的最小值，以防批次中的某些元素包含少量步骤。默认值为 -1（即没有最小值）。

示例

>>> td = TensorDict(
...     {
...         key1: torch.stack([torch.arange(0, 10), torch.arange(10, 20)], 0),
...         key2: torch.stack([torch.arange(0, 10), torch.arange(10, 20)], 0),
...     },
...     [2, 10],
... )
>>> trainer.register_op(
...     "process_optim_batch",
...     BatchSubSampler(batch_size=batch_size, sub_traj_len=sub_traj_len),
... )
>>> td_out = trainer._process_optim_batch_hook(td)
>>> assert td_out.shape == torch.Size([batch_size // sub_traj_len, sub_traj_len])

register(trainer: Trainer, name: str = 'batch_subsampler')

在训练器中以默认位置注册钩子。

参数：

• trainer (Trainer) – 必须注册钩子的训练器。

• name (str) – 钩子的名称。

注意

要在默认位置以外的位置注册钩子，请使用 register_op()。