Tree

class torchrl.data.Tree(count: 'int | torch.Tensor' = None, wins: 'int | torch.Tensor' = None, index: 'torch.Tensor | None' = None, hash: 'int | None' = None, node_id: 'int | None' = None, rollout: 'TensorDict | None' = None, node_data: 'TensorDict | None' = None, subtree: 'Tree' = None, _parent: 'weakref.ref | list[weakref.ref] | None' = None, specs: 'Composite | None' = None, *, batch_size, device=None, names=None)

property branching_action: torch.Tensor | TensorDictBase | None

返回分支到此特定節點的動作。

返回:

一個張量、tensordict 或 None，如果節點沒有父節點。

另請參閱

當 rollout 資料包含單個步長時，這將等於 prev_action。

另請參閱

樹中與給定節點（或觀察）關聯的 所有 動作.

cat(dim: int = 0, *, out=None)

將tensordicts沿給定維度連線成一個tensordict。

此呼叫等同於呼叫 torch.cat()，但與 torch.compile 相容。

property device: device

檢索張量類的裝置型別。

dumps(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → Self

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

edges() → list[tuple[int, int]]

檢索樹中的邊列表。

每條邊都表示為一個節點 ID 的元組：父節點 ID 和子節點 ID。樹使用廣度優先搜尋 (BFS) 進行遍歷，以確保所有邊都被訪問。

返回:

一個元組列表，其中每個元組包含一個父節點 ID 和一個子節點 ID。

classmethod fields()

返回一個描述此資料類的欄位的元組。欄位型別為 Field。

接受一個數據類或其例項。元組元素為 Field 型別。

from_any(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)

Recursively converts any object to a TensorDict.

注意

from_any 比常規的 TensorDict 建構函式限制更少。它可以使用自定義啟發式方法將資料結構（如 dataclasses 或 tuples）轉換為 tensordict。此方法可能會產生一些額外的開銷，並在對映策略方面涉及更多主觀選擇。

注意

This method recursively converts the input object to a TensorDict. If the object is already a TensorDict (or any similar tensor collection object), it will be returned as is.

引數:

obj – The object to be converted.

關鍵字引數:

• auto_batch_size (bool, optional) – 如果 True，將自動計算 batch size。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• device (torch.device, optional) – 將建立 TensorDict 的裝置。

• batch_size (torch.Size, optional) – TensorDict 的批次大小。與 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

• 資料類透過 from_dataclass() (資料類將被轉換為 TensorDict 例項，而不是 tensorclass)。

• 命名元組透過 from_namedtuple()。

• 透過 from_dict() 的字典。

• 元組透過 from_tuple()。

• NumPy 的結構化陣列透過 from_struct_array()。

• HDF5 物件透過 from_h5()。

from_dataclass(*, dest_cls: Type | None = None, auto_batch_size: bool = False, batch_dims: int | None = None, as_tensorclass: bool = False, device: torch.device | None = None, batch_size: torch.Size | None = None)

Converts a dataclass into a TensorDict instance.

引數:

dataclass – The dataclass instance to be converted.

關鍵字引數:

• dest_cls (tensorclass, optional) – 用於對映資料的 tensorclass 型別。如果未提供，則建立一個新類。如果 obj 是一個型別或 as_tensorclass 為 False，則無效。

• auto_batch_size (bool, optional) – 如果 True，將自動確定並應用 batch size 到生成的 TensorDict。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• as_tensorclass (bool, optional) – 如果 True，則將轉換委託給自由函式 from_dataclass()，並返回一個 tensorclass（tensorclass()）型別或例項，而不是 TensorDict。預設為 False。

• device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

A TensorDict instance derived from the provided dataclass, unless as_tensorclass is True, in which case a tensor-compatible class or instance is returned.

丟擲:

TypeError – 如果提供的輸入不是資料類例項。

警告

This method is distinct from the free function from_dataclass and serves a different purpose. While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance.

注意

• 此方法建立一個新的 TensorDict 例項，其鍵對應於輸入 dataclass 的欄位。

• 結果 TensorDict 中的每個鍵都使用 `cls.from_any` 方法進行初始化。

• auto_batch_size 選項允許自動確定批次大小並將其應用於結果 TensorDict。

from_h5(*, mode: str = 'r', auto_batch_size: bool = False, batch_dims: int | None = None, batch_size: torch.Size | None = None)

從 h5 檔案建立 PersistentTensorDict。

引數:

filename (str) – h5 檔案的路徑。

關鍵字引數:

• mode (str, optional) – 讀取模式。預設為 "r"。

• auto_batch_size (bool, optional) – 如果 True，將自動計算 batch size。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入 h5 檔案的 PersistentTensorDict 表示。

示例

>>> td = TensorDict.from_h5("path/to/file.h5")
>>> print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

from_modules(*, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)

為 vmap 的 ensemable 學習/特徵期望應用檢索多個模組的引數。

引數:

modules (nn.Module 序列) – 要從中獲取引數的模組。如果模組結構不同，則需要懶惰堆疊（請參閱下面的 `lazy_stack` 引數）。

關鍵字引數:

• as_module (bool, optional) – 如果 True，則返回一個 TensorDictParams 例項，可用於將引數儲存在 torch.nn.Module 中。預設為 False。

• lock (bool, optional) – 如果 True，則鎖定的結果 tensordict。預設為 True。

• use_state_dict (bool, optional) –

  如果為 True，則將使用模組的 state-dict 並將其解壓到具有模型樹結構的 TensorDict 中。預設為 False。

  注意

  這在使用 state-dict hook 時尤其有用。

• lazy_stack (bool, optional) –

  是否密集堆疊或懶惰堆疊引數。預設為 False（密集堆疊）。

  注意

  lazy_stack 和 as_module 是互斥的特性。

  警告

  懶惰輸出和非懶惰輸出之間有一個重要的區別：非懶惰輸出將使用所需的批次大小重新例項化引數，而 `lazy_stack` 將僅將引數表示為懶惰堆疊。這意味著，雖然原始引數可以安全地傳遞給最佳化器（當 `lazy_stack=True` 時），但在設定為 `True` 時需要傳遞新引數。

  警告

  雖然使用 lazy stack 來保留原始引數引用可能很誘人，但請記住，每次呼叫 get() 時，lazy stack 都會執行堆疊操作。這將需要記憶體（引數大小的 N 倍，如果構建了圖，則更多）和計算時間。這也意味著最佳化器將包含更多引數，像 step() 或 zero_grad() 這樣的操作的執行將花費更長的時間。總的來說，lazy_stack 應僅限於極少數用例。

• expand_identical (bool, optional) – 如果 True 且相同的引數（相同標識）被堆疊到自身，則將返回該引數的擴充套件版本。在 lazy_stack=True 時忽略此引數。

示例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> torch.manual_seed(0)
>>> empty_module = nn.Linear(3, 4, device="meta")
>>> n_models = 2
>>> modules = [nn.Linear(3, 4) for _ in range(n_models)]
>>> params = TensorDict.from_modules(*modules)
>>> print(params)
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)
>>> # example of batch execution
>>> def exec_module(params, x):
...     with params.to_module(empty_module):
...         return empty_module(x)
>>> x = torch.randn(3)
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> # since lazy_stack = False, backprop leaves the original params untouched
>>> y.sum().backward()
>>> assert params["weight"].grad.norm() > 0
>>> assert modules[0].weight.grad is None

當 lazy_stack=True 時，情況略有不同

>>> params = TensorDict.from_modules(*modules, lazy_stack=True)
>>> print(params)
LazyStackedTensorDict(
    fields={
        bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False,
    stack_dim=0)
>>> # example of batch execution
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> y.sum().backward()
>>> assert modules[0].weight.grad is not None

from_namedtuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)

遞迴地將命名元組轉換為 TensorDict。

引數:

named_tuple – 要轉換的命名元組例項。

關鍵字引數:

• auto_batch_size (bool, optional) – 如果 True，將自動計算 batch size。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入命名元組的 TensorDict 表示。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> nt = data.to_namedtuple()
>>> print(nt)
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
>>> TensorDict.from_namedtuple(nt, auto_batch_size=True)
TensorDict(
    fields={
        a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None),
                a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_pytree(*, batch_size: torch.Size | None = None, auto_batch_size: bool = False, batch_dims: int | None = None)

將 pytree 轉換為 TensorDict 例項。

此方法旨在儘可能保留 pytree 的巢狀結構。

其他非張量鍵將被新增，以跟蹤每個級別的標識，從而提供內建的 pytree 到 tensordict 的雙射轉換 API。

當前接受的類包括列表、元組、命名元組和字典。

注意

對於字典，非 NestedKey 鍵會作為 NonTensorData 例項單獨註冊。

注意

可轉換為張量型別（如 int、float 或 np.ndarray）將被轉換為 torch.Tensor 例項。請注意，此轉換是滿射的：將 tensordict 轉換回 pytree 將無法恢復原始型別。

示例

>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key
>>> class WeirdLookingClass:
...     pass
...
>>> weird_key = WeirdLookingClass()
>>> # Make a pytree with tuple, lists, dict and namedtuple
>>> pytree = (
...     [torch.randint(10, (3,)), torch.zeros(2)],
...     {
...         "tensor": torch.randn(
...             2,
...         ),
...         "td": TensorDict({"one": 1}),
...         weird_key: torch.randint(10, (2,)),
...         "list": [1, 2, 3],
...     },
...     {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()},
... )
>>> # Build a TensorDict from that pytree
>>> td = TensorDict.from_pytree(pytree)
>>> # Recover the pytree
>>> pytree_recon = td.to_pytree()
>>> # Check that the leaves match
>>> def check(v1, v2):
>>>     assert (v1 == v2).all()
>>>
>>> torch.utils._pytree.tree_map(check, pytree, pytree_recon)
>>> assert weird_key in pytree_recon[1]

from_remote_init(group: 'ProcessGroup' | None = None, device: torch.device | None = None) → Self

從遠端傳送的元資料建立新的 tensordict 例項。

此類方法接收由 init_remote 傳送的元資料，建立具有匹配形狀和 dtype 的新 tensordict，然後非同步接收實際的 tensordict 內容。

引數:

• src (int) – 傳送元資料的源程序的 rank。

• group ("ProcessGroup", optional) – 要使用的程序組。預設為 None。

• device (torch.device, 可選) – 用於張量運算的裝置。預設為 None。

返回:

使用接收到的元資料和內容初始化的新 tensordict 例項。

返回型別:

TensorDict

另請參閱

傳送程序應已呼叫 ~.init_remote 來發送元資料和內容。

from_struct_array(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None) → Self

將結構化 numpy 陣列轉換為 TensorDict。

生成的 TensorDict 將與 numpy 陣列共享相同的記憶體內容（這是一次零複製操作）。原地更改結構化 numpy 陣列的值會影響 TensorDict 的內容。

注意

此方法執行零複製操作，這意味著生成的 TensorDict 將與輸入的 numpy 陣列共享相同的記憶體內容。因此，原地更改 numpy 陣列的值會影響 TensorDict 的內容。

引數:

struct_array (np.ndarray) – 要轉換的結構化 numpy 陣列。

關鍵字引數:

• auto_batch_size (bool, optional) – 如果 True，將自動計算 batch size。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• device (torch.device, 可選) –

  將建立 TensorDict 的裝置。預設為 None。

  注意

  更改裝置（即，指定任何非 `None` 或 `"cpu"` 的裝置）將傳輸資料，從而導致返回資料的記憶體位置發生更改。

• batch_size (torch.Size, 可選) – TensorDict 的批次大小。預設為 None。

返回:

輸入的結構化 numpy 陣列的 TensorDict 表示。

示例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()

classmethod from_tensordict(tensordict: TensorDictBase, non_tensordict: dict | None = None, safe: bool = True) → Self

用於例項化新張量類物件的張量類包裝器。

引數:

• tensordict (TensorDictBase) – 張量型別字典

• non_tensordict (dict) – 包含非張量和巢狀張量類物件的字典

• safe (bool) – 如果 tensordict 不是 TensorDictBase 例項，是否引發錯誤

from_tuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)

將元組轉換為 TensorDict。

引數:

obj – 要轉換的元組例項。

關鍵字引數:

• auto_batch_size (bool, optional) – 如果 True，將自動計算 batch size。預設為 False。

• batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每個級別的完整 batch size）。

• device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入的元組的 TensorDict 表示。

示例

>>> my_tuple = (1, 2, 3)
>>> td = TensorDict.from_tuple(my_tuple)
>>> print(td)
TensorDict(
    fields={
        0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

fromkeys(value: Any = 0)

從鍵列表和單個值建立 tensordict。

引數:

• keys (list of NestedKey) – 指定新字典鍵的可迭代物件。

• value (compatible type, optional) – 所有鍵的值。預設為 0。

property full_action_spec

樹的動作規範。

這是 Tree.specs[‘input_spec’, ‘full_action_spec’] 的別名。

property full_done_spec

樹的完成規範。

這是 Tree.specs[‘output_spec’, ‘full_done_spec’] 的別名。

property full_observation_spec

樹的觀察規範。

這是 Tree.specs[‘output_spec’, ‘full_observation_spec’] 的別名。

property full_reward_spec

樹的獎勵規範。

這是 Tree.specs[‘output_spec’, ‘full_reward_spec’] 的別名。

property full_state_spec

樹的狀態規範。

這是 Tree.specs[‘input_spec’, ‘full_state_spec’] 的別名。

fully_expanded(env: EnvBase) → bool

如果子節點數量等於環境基數，則返回 True。

get(key: NestedKey, *args, **kwargs)

獲取輸入鍵對應的儲存值。

引數:

• key (str, tuple of str) – 要查詢的鍵。如果是字串元組，則等同於鏈式呼叫 getattr。

• default – 如果在張量類中找不到鍵，則返回預設值。

返回:

儲存在輸入鍵下的值

get_vertex_by_hash(hash: int) → Tree

遍歷樹並返回與給定雜湊對應的節點。

get_vertex_by_id(id: int) → Tree

遍歷樹並返回與給定 ID 對應的節點。

property is_terminal: bool | torch.Tensor

如果沒有子節點，則返回 True。

lazy_stack(dim: int = 0, *, out=None, **kwargs)

建立 TensorDicts 的懶惰堆疊。

有關詳細資訊，請參閱 lazy_stack()。

load(*args, **kwargs) → Self

從磁碟載入 tensordict。

此類方法是 load_memmap() 的代理。

load_(prefix: str | Path, *args, **kwargs)

在當前 tensordict 中從磁碟載入 tensordict。

此類方法是 load_memmap_() 的代理。

load_memmap(device: torch.device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None) → Self

從磁碟載入記憶體對映的 tensordict。

引數:

• prefix (str 或 Path to folder) – 應獲取已儲存 tensordict 的資料夾路徑。

• device (torch.device 或 等效項, 可選) – 如果提供，資料將非同步轉換為該裝置。支援 `"meta"` 裝置，在這種情況下，資料不會被載入，而是建立一組空的 "meta" 張量。這對於在不實際開啟任何檔案的情況下了解模型大小和結構很有用。

• non_blocking (bool, optional) – 如果為 True，則在將張量載入到裝置後不會呼叫同步。預設為 False。

• out (TensorDictBase, optional) – 資料應寫入的可選 tensordict。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

此方法還允許載入巢狀的 tensordicts。

示例

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

tensordict 也可以在“meta”裝置上載入，或者作為假張量載入。

示例

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

load_state_dict(state_dict: dict[str, Any], strict=True, assign=False, from_flatten=False)

嘗試將 state_dict 載入到目標張量類中（原地）。

classmethod make_node(data: TensorDictBase, *, device: torch.device | None = None, batch_size: torch.Size | None = None, specs: Composite | None = None) → Tree

給定一些資料建立一個新節點。

max_length()

返回樹中所有有效路徑的最大長度。

路徑的長度定義為路徑中的節點數。如果樹為空，則返回 0。

返回:

樹中所有有效路徑的最大長度。

返回型別:

int

maybe_dense_stack(dim: int = 0, *, out=None, **kwargs)

嘗試使 TensorDicts 密集堆疊，並在需要時回退到懶惰堆疊。

有關詳細資訊，請參閱 maybe_dense_stack()。

memmap(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → Self

將所有張量寫入記憶體對映的 Tensor 中，並放入新的 tensordict。

引數:

• prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則如果 tensordict 中的某個條目已經是儲存在磁碟上的張量並具有關聯檔案，但未按 prefix 指定的位置儲存，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果 True 且 num_threads>0，則該方法將返回 tensordict 的未來物件。

• share_non_tensor (bool, optional) – 如果為 True，則非張量資料將在程序之間共享，並且單個節點內任何工作程序上的寫入操作（例如原地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子節點數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:

返回一個新的 tensordict，其中張量儲存在磁碟上（如果 return_early=False），否則返回一個 TensorDictFuture 例項。

注意

以這種方式序列化對於深度巢狀的 tensordicts 來說可能很慢，因此不建議在訓練迴圈中呼叫此方法。

memmap_(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → Self

將所有張量原地寫入相應的記憶體對映張量。

引數:

• prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則如果 tensordict 中的某個條目已經是儲存在磁碟上的張量並具有關聯檔案，但未按 prefix 指定的位置儲存，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果 True 且 num_threads>0，則該方法將返回 tensordict 的未來物件。可以使用 future.result() 查詢結果 tensordict。

• share_non_tensor (bool, optional) – 如果為 True，則非張量資料將在程序之間共享，並且單個節點內任何工作程序上的寫入操作（例如原地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子節點數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:

如果 return_early=False，則返回 self，否則返回 TensorDictFuture 例項。

注意

以這種方式序列化對於深度巢狀的 tensordicts 來說可能很慢，因此不建議在訓練迴圈中呼叫此方法。

memmap_like(prefix: str | None = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → Self

建立一個無內容的記憶體對映 tensordict，其形狀與原始 tensordict 相同。

引數:

• prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則如果 tensordict 中的某個條目已經是儲存在磁碟上的張量並具有關聯檔案，但未按 prefix 指定的位置儲存，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果 True 且 num_threads>0，則該方法將返回 tensordict 的未來物件。

• share_non_tensor (bool, optional) – 如果為 True，則非張量資料將在程序之間共享，並且單個節點內任何工作程序上的寫入操作（例如原地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子節點數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:

如果 return_early=False，則建立一個新的 TensorDict 例項，其中資料儲存為記憶體對映張量；否則，建立一個 TensorDictFuture 例項。

注意

這是將一組大型緩衝區寫入磁碟的推薦方法，因為 memmap_() 將複製資訊，這對於大型內容來說可能很慢。

示例

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_()

如果記憶體對映的 tensordict 具有 saved_path，則重新整理其內容。

如果沒有任何路徑與之關聯，此方法將引發異常。

property node_observation: torch.Tensor | TensorDictBase

返回與此特定節點關聯的觀察。

這是在出現分支之前定義節點的觀察（或觀察集）。如果節點包含 rollout() 屬性，則節點觀察通常與最後一次操作產生的觀察相同，即 node.rollout[..., -1]["next", "observation"]。

如果樹規範關聯了多個觀察鍵，則會返回 TensorDict 例項。

為了獲得更一致的表示，請參見 node_observations。

property node_observations: torch.Tensor | TensorDictBase

以 TensorDict 格式返回與此特定節點關聯的觀察。

這是在出現分支之前定義節點的觀察（或觀察集）。如果節點包含 rollout() 屬性，則節點觀察通常與最後一次操作產生的觀察相同，即 node.rollout[..., -1]["next", "observation"]。

如果樹規範關聯了多個觀察鍵，則會返回 TensorDict 例項。

為了獲得更一致的表示，請參見 node_observations。

property num_children: int

此節點的子節點數量。

等於 self.subtree 堆疊中元素的數量。

num_vertices(*, count_repeat: bool = False) → int

返回 Tree 中唯一頂點的數量。

關鍵字引數:

count_repeat (bool, optional) –

確定是否計算重複的頂點。

• 如果為 False，則每個唯一頂點只計算一次。

• 如果為 True，則如果頂點出現在不同的路徑中，則多次計算它們。

預設為False。

返回:

Tree 中唯一頂點的數量。

返回型別:

int

property parent: Tree | None

節點的父節點。

如果節點有父節點且該物件仍在 python 工作區中，則此屬性將返回該父節點。

對於重新分支的樹，此屬性可能會返回一個樹堆疊，其中堆疊的每個索引對應於一個不同的父節點。

注意

parent 屬性在內容上匹配，但在身份上不匹配：tensorclass 物件使用相同的張量（即指向相同記憶體位置的張量）進行重建。

返回:

包含父節點資料的 Tree，或者如果父節點資料超出範圍或節點是根節點，則為 None。

plot(backend: str = 'plotly', figure: str = 'tree', info: list[str] = None, make_labels: Callable[[Any, ...], Any] | None = None)

使用指定的後端和圖形型別繪製樹的視覺化。

引數:

• backend – 要使用的繪圖後端。目前僅支援 ‘plotly’。

• figure – 要繪製的圖形型別。可以是 ‘tree’ 或 ‘box’。

• info – 要包含在圖形中的其他資訊列表（目前未使用）。

• make_labels – 用於為圖形生成自定義標籤的可選函式。

丟擲:

NotImplementedError – 如果指定了不受支援的後端或圖形型別。

property prev_action: torch.Tensor | TensorDictBase | None

在此節點觀察生成之前的動作。

返回:

一個張量、tensordict 或 None，如果節點沒有父節點。

另請參閱

當 rollout 資料包含單個步長時，這將等於 branching_action。

另請參閱

樹中與給定節點（或觀察）關聯的 所有 動作.

rollout_from_path(path: tuple[int]) → TensorDictBase | None

沿給定路徑檢索樹中的 rollout 資料。

對於路徑中的每個節點，rollout 資料沿最後一個維度 (dim=-1) 連線。如果沿路徑未找到 rollout 資料，則返回 None。

引數:

path – 表示樹中路徑的整數元組。

返回:

沿路徑連線的 rollout 資料，或在未找到資料時返回 None。

save(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → Self

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

property selected_actions: torch.Tensor | TensorDictBase | None

返回從該節點分支出的所有選定動作的張量。

set(key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False)

設定一個新的鍵值對。

引數:

• key (str, tuple of str) – 要設定的鍵的名稱。如果是字串元組，則等同於鏈式呼叫 getattr，然後進行最終的 setattr。

• value (Any) – 要儲存在張量類中的值

• inplace (bool, optional) – 如果為 True，則 set 將嘗試原地更新值。如果為 False 或鍵不存在，則值將簡單地寫入其目的地。

返回:

self

stack(dim: int = 0, *, out=None)

沿給定維度將 tensordicts 堆疊成一個單一的 tensordict。

此呼叫等效於呼叫 torch.stack()，但與 torch.compile 相容。

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → dict[str, Any]

返回一個 state_dict 字典，可用於儲存和載入張量類的資料。

to_string(node_format_fn=<function Tree.<lambda>>)

生成樹的字串表示。

此函式可以提取樹中每個節點的資訊，因此對除錯很有用。節點逐行列出。每行包含到節點的路徑，然後是使用 :arg:`node_format_fn` 生成的該節點的字串表示。每行根據到達相應節點所需的路徑步數進行縮排。

引數:

node_format_fn (Callable, optional) – 使用者定義的函式，用於為樹的每個節點生成字串。簽名必須為 (Tree) -> Any，並且輸出必須可轉換為字串。如果未提供此引數，則生成的字串是節點 Tree.node_data 屬性轉換為字典。

示例

>>> from torchrl.data import MCTSForest
>>> from tensordict import TensorDict
>>> forest = MCTSForest()
>>> td_root = TensorDict({"observation": 0,})
>>> rollouts_data = [
...     # [(action, obs), ...]
...     [(3, 123), (1, 456)],
...     [(2, 359), (2, 3094)],
...     [(3, 123), (9, 392), (6, 989), (20, 809), (21, 847)],
...     [(1, 75)],
...     [(3, 123), (0, 948)],
...     [(2, 359), (2, 3094), (10, 68)],
...     [(2, 359), (2, 3094), (11, 9045)],
... ]
>>> for rollout_data in rollouts_data:
...     td = td_root.clone().unsqueeze(0)
...     for action, obs in rollout_data:
...         td = td.update(TensorDict({
...             "action": [action],
...             "next": TensorDict({"observation": [obs]}, [1]),
...         }, [1]))
...         forest.extend(td)
...         td = td["next"].clone()
...
>>> tree = forest.get_tree(td_root)
>>> print(tree.to_string())
(0,) {'observation': tensor(123)}
(0, 0) {'observation': tensor(456)}
(0, 1) {'observation': tensor(847)}
(0, 2) {'observation': tensor(948)}
(1,) {'observation': tensor(3094)}
(1, 0) {'observation': tensor(68)}
(1, 1) {'observation': tensor(9045)}
(2,) {'observation': tensor(75)}

to_tensordict(*, retain_none: bool | None = None) → TensorDict

將張量類轉換為常規 TensorDict。

複製所有條目。記憶體對映和共享記憶體張量將被轉換為常規張量。

引數:

retain_none (bool) – 如果為 True，則 None 值將被寫入 tensordict。否則它們將被丟棄。預設值：True。

返回:

包含與張量類相同值的新的 TensorDict 物件。

unbind(dim: int)

返回沿指定維度解綁的索引張量類例項的元組。

結果張量類例項將共享初始張量類例項的儲存。

valid_paths()

生成樹中的所有有效路徑。

有效路徑是從根節點開始並以葉節點結束的子索引序列。每條路徑表示為一個整數元組，其中每個整數對應於一個子節點的索引。

產生:

tuple – 樹中的有效路徑。

vertices(*, key_type: Literal['id', 'hash', 'path'] = 'hash') → dict[int | tuple[int], Tree]

返回包含 Tree 頂點的對映。

關鍵字引數:

key_type (Literal["id", "hash", "path"], optional) –

指定用於頂點的鍵的型別。

• ”id”: 使用頂點 ID 作為鍵。

• ”hash”: 使用頂點的雜湊作為鍵。

• ”path”: 使用頂點的路徑作為鍵。這可能導致字典的長度比

  當使用 "id" 或 "hash" 時，因為相同的節點可能屬於多個軌跡。預設為 "hash"。

預設為空字串，這可能意味著預設行為。

返回:

將鍵對映到 Tree 頂點的字典。

返回型別:

Dict[int | Tuple[int], Tree]

property visits: int | torch.Tensor

返回與此特定節點關聯的訪問次數。

這是 count 屬性的別名。