RandomCropTensorDict

class torchrl.envs.transforms.RandomCropTensorDict(sub_seq_len: int, sample_dim: int = - 1, mask_key: NestedKey | None = None)

用于 ReplayBuffer 和模块的轨迹子采样器。

在输入 tensordict 的最后一个维度上收集指定长度的子序列。这可以用于从 ReplayBuffer 中采样的轨迹中获取裁剪后的轨迹。

此转换器主要设计用于回放缓冲区和模块。目前，它不能用作环境转换器。如果需要此功能，请随时通过 issue 请求。

参数:

• sub_seq_len (int) – 要采样的子轨迹的长度

• sample_dim (int, optional) – 应该发生裁剪的维度。应优先使用负维度，以使转换器对具有不同批次维度的 tensordict 具有鲁棒性。默认为 -1（TorchRL 中的默认时间维度）。

• mask_key (NestedKey) – 如果提供，这表示在进行采样时要查找的掩码键。如果提供，将只返回有效元素。假设掩码是一个布尔张量，其前半部分为 True 值，后半部分为 False 值，中间没有混合。 RandomCropTensorDict 不会检查这一点，因此不正确的掩码可能导致任何错误而未被察觉。默认值：None（无掩码键）。

forward(tensordict: TensorDictBase) → TensorDictBase

读取输入 tensordict，并对选定的键应用转换。

默认情况下，此方法

• 直接调用 _apply_transform()。

• 不调用 _step() 或 _call()。

此方法不在 env.step 的任何地方被调用。但是，它在 sample() 中被调用。

注意

forward 还使用 dispatch 通过将参数名称转换为键来处理常规关键字参数。

示例

>>> class TransformThatMeasuresBytes(Transform):
...     '''Measures the number of bytes in the tensordict, and writes it under `"bytes"`.'''
...     def __init__(self):
...         super().__init__(in_keys=[], out_keys=["bytes"])
...
...     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
...         bytes_in_td = tensordict.bytes()
...         tensordict["bytes"] = bytes
...         return tensordict
>>> t = TransformThatMeasuresBytes()
>>> env = env.append_transform(t) # works within envs
>>> t(TensorDict(a=0))  # Works offline too.