split_trajectories¶

torchrl.collectors.utils.split_trajectories(rollout_tensordict: TensorDictBase, *, prefix=None, trajectory_key: NestedKey | None = None, done_key: NestedKey | None = None, as_nested: bool = False) → TensorDictBase[源代码]¶

一个用于轨迹分离的实用函数。

采用具有 traj_ids 键的 tensordict，该键指示每个轨迹的 ID。

从那里，构建一个 B x T x … 零填充 tensordict，其中 B 个批次在最大持续时间 T 上

参数：

rollout_tensordict (TensorDictBase) – 沿最后一个维度具有相邻轨迹的 rollout。

关键字参数：

prefix (NestedKey, 可选) – 用于读取和写入元数据的前缀，例如 "traj_ids" （每个轨迹的可选整型 ID）和 "mask" 条目，指示哪些数据有效，哪些无效。如果输入具有 "collector" 条目，则默认为 "collector" ，否则默认为 () （无前缀）。prefix 作为一项遗留功能保留，最终将被弃用。尽可能优先使用 trajectory_key 或 done_key 。
trajectory_key (NestedKey, 可选) – 指向轨迹 ID 的键。优先于 done_key 和 prefix 。如果未提供，则默认为 (prefix, "traj_ids") 。
done_key (NestedKey, 可选) – 指向 "done"" 信号的键，如果无法直接恢复轨迹。默认为 "done" 。
as_nested (bool 或 torch.layout, 可选) –
是否将结果作为嵌套张量返回。默认为 False 。如果提供了 torch.layout ，它将用于构建嵌套张量，否则将使用默认布局。

注意

使用 split_trajectories(tensordict, as_nested=True).to_padded_tensor(mask=mask_key) 应与 as_nested=False 产生完全相同的结果。由于这是一个实验性功能，并且依赖于 nested_tensors，而该 API 在将来可能会发生更改，因此我们将其作为可选功能。使用 as_nested=True 运行时应该更快。

注意

提供布局可以让用户控制是否要使用 torch.strided 或 torch.jagged 布局来使用嵌套张量。虽然前者在撰写本文时具有更多功能，但后者将在将来成为 PyTorch 团队的重点，因为它与 compile() 具有更好的兼容性。

返回：

一个新的 tensordict，其前导维度对应于轨迹。共享 trajectory_key 前缀的 "mask" 布尔条目和 tensordict 形状也会被添加。它指示 tensordict 的有效元素，以及 trajectory_key 无法找到的 "traj_ids" 条目。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> from torchrl.collectors.utils import split_trajectories
>>> obs = torch.cat([torch.arange(10), torch.arange(5)])
>>> obs_ = torch.cat([torch.arange(1, 11), torch.arange(1, 6)])
>>> done = torch.zeros(15, dtype=torch.bool)
>>> done[9] = True
>>> trajectory_id = torch.cat([torch.zeros(10, dtype=torch.int32),
...     torch.ones(5, dtype=torch.int32)])
>>> data = TensorDict({"obs": obs, ("next", "obs"): obs_, ("next", "done"): done, "trajectory": trajectory_id}, batch_size=[15])
>>> data_split = split_trajectories(data, done_key="done")
>>> print(data_split)
TensorDict(
    fields={
        mask: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
                obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([2, 10]),
            device=None,
            is_shared=False),
        obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
        traj_ids: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
        trajectory: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int32, is_shared=False)},
    batch_size=torch.Size([2, 10]),
    device=None,
    is_shared=False)
>>> # check that split_trajectories got the trajectories right with the done signal
>>> assert (data_split["traj_ids"] == data_split["trajectory"]).all()
>>> print(data_split["mask"])
tensor([[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True, False, False, False, False, False]])
>>> data_split = split_trajectories(data, trajectory_key="trajectory")
>>> print(data_split)
TensorDict(
    fields={
        mask: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
                obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([2, 10]),
            device=None,
            is_shared=False),
        obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
        trajectory: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int32, is_shared=False)},
    batch_size=torch.Size([2, 10]),
    device=None,
    is_shared=False)

split_trajectories¶

文档

教程

资源