BatchSubSampler

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

用於線上 RL 最先進實作的資料子取樣器。

此類別會從剛從環境收集的一整批資料中取出部分樣本。

參數:

• batch_size (int) – 要收集的子批次大小。提供的批次大小必須等於輸出 tensordict 中的項目總數，該 tensordict 的大小將為 [batch_size // sub_traj_len, sub_traj_len]。

• sub_traj_len (int, optional) – 子樣本必須在線上設定中具有的軌跡長度。預設值為 -1（即採用軌跡的完整長度）

• min_sub_traj_len (int, optional) – 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')

在預設位置將 hook 註冊到訓練器中。

參數:

• trainer (Trainer) – 訓練器 (trainer)，hook必須在此註冊。

• name (str) – hook的名稱。

注意

若要將hook註冊到非預設位置，請使用 register_op()。