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

用於軌跡分割的實用函式。

接收一個 tensordict，其中包含一個名為 traj_ids 的鍵，該鍵指示每個軌跡的 id。

由此，構建一個 B x T x … 的零填充 tensordict，其中 B 為批次，T 為最大持續時間。

引數:

rollout_tensordict (TensorDictBase) – 一個在最後一個維度上具有相鄰軌跡的 rollout。

關鍵字引數:

• prefix (NestedKey, optional) – 用於讀取和寫入元資料的 prefix，例如 "traj_ids"（每個軌跡的可選整數 id）以及指示哪些資料有效和無效的 "mask" 條目。如果輸入具有 "collector" 條目，則預設為 "collector"；否則預設為 ()（無 prefix）。prefix 是一個遺留功能，最終將被棄用。如果可能，請優先使用 trajectory_key 或 done_key。

• trajectory_key (NestedKey, optional) – 指向軌跡 id 的鍵。優先於 done_key 和 prefix。如果未提供，則預設為 (prefix, "traj_ids")。

• done_key (NestedKey, optional) – 指向 "done"" 訊號的鍵，如果軌跡無法直接恢復。預設為 "done"。

• as_nested (bool or torch.layout, optional) –

  是否將結果作為巢狀張量返回。預設為 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 佈局一起使用。雖然前者在撰寫本文時具有稍多的功能，但後者由於與 compile() 的相容性更好，將是 PyTorch 團隊未來的主要關注點。

返回:

一個新的 tensordict，其前導維度對應於軌跡。還添加了一個共享 trajectory_key 字首和 tensordict 形狀的 "mask" 布林條目。它指示 tensordict 的有效元素，以及一個 "traj_ids" 條目（如果找不到 trajectory_key）。

示例

>>> 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)