StackedComposite

class torchrl.data.StackedComposite(*args, **kwargs)

A lazy representation of a stack of composite specs.

Stacks composite specs together along one dimension. When random samples are drawn, a LazyStackedTensorDict is returned.

Indexing is allowed but only along the stack dimension.

This class is aimed to be used in multi-task and multi-agent settings, where heterogeneous specs may occur (same semantic but different shape).

assert_is_in(value: Tensor) → None

断言一个张量是否属于该区域（box），否则抛出异常。

参数:

value (torch.Tensor) – 要检查的值。

cardinality(*args, **kwargs) → Any

规格的基数。

这指的是规格中可能出现的结果的数量。假设复合规格的基数是所有可能结果的笛卡尔积。

clear_device_()

清除 Composite 的设备。

clone() → T

克隆 Composite spec。

锁定的 spec 不会产生锁定的克隆。

contains(item: torch.Tensor | TensorDictBase) → bool

If the value val could have been generated by the TensorSpec, returns True, otherwise False.

See is_in() for more information.

cpu()

将 TensorSpec 转换为“cpu”设备。

cuda(device=None)

将 TensorSpec 转换为“cuda”设备。

property device: Union[device, str, int]

规格的设备。

Only Composite specs can have a None device. All leaves must have a non-null device.

empty()

创建一个像 self 一样的 spec，但没有条目。

encode(val: np.ndarray | list | torch.Tensor | TensorDictBase, *, ignore_device: bool = False) → torch.Tensor | TensorDictBase

使用指定的规格对值进行编码，并返回相应的张量。

此方法用于返回易于映射到 TorchRL 所需域的值（例如 numpy 数组）的环境。如果值已经是张量，则规格不会更改其值，而是按原样返回。

参数:

val (np.ndarray 或 torch.Tensor) – 要编码为张量的值。

关键字参数:

ignore_device (bool, optional) – if True, the spec device will be ignored. This is used to group tensor casting within a call to TensorDict(..., device="cuda") which is faster.

返回:

符合所需张量规格的 torch.Tensor。

enumerate(use_mask: bool = False) → TensorDictBase

返回可以从 TensorSpec 获得的所有样本。

样本将沿第一个维度堆叠。

此方法仅为离散规格实现。

参数:

use_mask (bool, optional) – If True and the spec has a mask, samples that are masked are excluded. Default is False.

erase_memoize_cache() → None

清除用于缓存 encode 执行的缓存。

另请参阅

memoize_encode().

expand(*shape)

返回一个具有扩展形状的新 Spec。

参数:

*shape (tuple or iterable of int) – the new shape of the Spec. Must be broadcastable with the current shape: its length must be at least as long as the current shape length, and its last values must be compliant too; ie they can only differ from it if the current dimension is a singleton.

flatten(start_dim: int, end_dim: int) → T

展平一个 TensorSpec。

Check flatten() for more information on this method.

get(item, default=_NoDefault.ZERO)

从 Composite 中获取一个条目。

如果条目不存在，可以传入一个默认值。

classmethod implements_for_spec(torch_function: Callable) → Callable

为 TensorSpec 注册一个 torch 函数覆盖。

index(index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase) → torch.Tensor | TensorDictBase

索引输入张量。

This method is to be used with specs that encode one or more categorical variables (e.g., OneHot or Categorical), such that indexing of a tensor with a sample can be done without caring about the actual representation of the index.

参数:

• index (int, torch.Tensor, slice or list) – index of the tensor

• tensor_to_index – 要索引的张量

返回:

被索引的张量

示例

>>> from torchrl.data import OneHot
>>> import torch
>>>
>>> one_hot = OneHot(n=100)
>>> categ = one_hot.to_categorical_spec()
>>> idx_one_hot = torch.zeros((100,), dtype=torch.bool)
>>> idx_one_hot[50] = 1
>>> print(one_hot.index(idx_one_hot, torch.arange(100)))
tensor(50)
>>> idx_categ = one_hot.to_categorical(idx_one_hot)
>>> print(categ.index(idx_categ, torch.arange(100)))
tensor(50)

is_empty(recurse: bool = False)

Composite spec 是否包含 spec。

参数:

recurse (bool) – whether to recursively assess if the spec is empty. If True, will return True if there are no leaves. If False (default) will return whether there is any spec defined at the root level.

is_in(value) → bool

If the value val could have been generated by the TensorSpec, returns True, otherwise False.

More precisely, the is_in methods checks that the value val is within the limits defined by the space attribute (the box), and that the dtype, device, shape potentially other metadata match those of the spec. If any of these checks fails, the is_in method will return False.

参数:

val (torch.Tensor) – 要检查的值。

返回:

布尔值，指示值是否属于 TensorSpec 区域。

items(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Callable[[type], bool] | None = None, step_mdp_static_only: bool = False) → _CompositeSpecItemsView

Composite 的条目。

参数:

• include_nested (bool, optional) – if False, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next"]. Default is ``False`, i.e. nested keys will not be returned.

• leaves_only (bool, optional) – if False, the values returned will contain every level of nesting, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next", ("next", "obs")]. Default is False.

关键字参数:

• is_leaf (callable, optional) – 读取一个类型并返回一个布尔值，指示该类型是否应被视为叶子。默认情况下，所有非 Composite 节点都被视为叶子。

• step_mdp_static_only (bool, optional) – if True, only keys that are static under step_mdp will be returned. Default is False.

keys(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Callable[[type], bool] | None = None, step_mdp_static_only: bool = False) → _CompositeSpecKeysView

Composite 的键。

keys 参数反映了 tensordict.TensorDict 的参数。

参数:

• include_nested (bool, optional) – if False, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next"]. Default is ``False`, i.e. nested keys will not be returned.

• leaves_only (bool, optional) – if False, the values returned will contain every level of nesting, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next", ("next", "obs")]. Default is False.

关键字参数:

• is_leaf (callable, optional) – 读取一个类型并返回一个布尔值，指示该类型是否应被视为叶子。默认情况下，所有非 Composite 节点都被视为叶子。

• step_mdp_static_only (bool, optional) – if True, only keys that are static under step_mdp will be returned. Default is False.

lock_(recurse: bool | None = None) → None

锁定 Composite 并防止修改其内容。

The recurse argument control whether the lock will be propagated to sub-specs. The current default is False but it will be turned to True for consistency with the TensorDict API in v0.8.

示例

>>> shape = [3, 4, 5]
>>> spec = Composite(
...         a=Composite(
...         b=Composite(shape=shape[:3], device="cpu"), shape=shape[:2]
...     ),
...     shape=shape[:1],
... )
>>> spec["a"] = spec["a"].clone()
>>> recurse = False
>>> spec.lock_(recurse=recurse)
>>> try:
...     spec["a"] = spec["a"].clone()
... except RuntimeError:
...     print("failed!")
failed!
>>> try:
...     spec["a", "b"] = spec["a", "b"].clone()
...     print("succeeded!")
... except RuntimeError:
...     print("failed!")
succeeded!
>>> recurse = True
>>> spec.lock_(recurse=recurse)
>>> try:
...     spec["a", "b"] = spec["a", "b"].clone()
...     print("succeeded!")
... except RuntimeError:
...     print("failed!")
failed!

make_neg_dim(dim: int)

将特定维度转换为 -1。

memoize_encode(mode: bool = True) → None

创建 encode 方法的缓存可调用序列，以加快其执行速度。

这应该只在输入类型、形状等在给定规格的调用之间预期一致时使用。

参数:

mode (bool, optional) – 是否使用缓存。默认为 True。

另请参阅

the cache can be erased via erase_memoize_cache().

property ndim

规格形状的维数。

相当于 len(spec.shape)。

ndimension()

规格形状的维数。

相当于 len(spec.shape)。

one(shape: Optional[Size] = None) → TensorDictBase

返回盒中的一个填充一的张量。

注意

Even though there is no guarantee that 1 belongs to the spec domain, this method will not raise an exception when this condition is violated. The primary use case of one is to generate empty data buffers, not meaningful data.

参数:

shape (torch.Size) – 一维张量的形状

返回:

在 TensorSpec 区域中采样的填充一的张量。

ones(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Proxy to one().

pop(key: NestedKey, default: Any = _NoDefault.ZERO) → Any

从 composite spec 中删除并返回与指定键关联的值。

此方法在 composite spec 中搜索指定键，将其删除，然后返回其关联的值。如果找不到键，则返回提供的默认值（如果指定），否则引发 KeyError。

参数:

• key (NestedKey) – 要从 composite spec 中删除的键。它可以是单个键或嵌套键。

• default (Any, optional) – 如果在 composite spec 中找不到指定的键，则返回的值。如果未提供且找不到键，则会引发 KeyError。

返回:

从 composite spec 中删除的与指定键关联的值。

返回类型:

任何

抛出:

KeyError – 如果在 composite spec 中找不到指定的键且未提供默认值。

project(val: TensorDictBase) → TensorDictBase

如果输入张量不在 TensorSpec 区域内，则根据定义的启发式方法将其映射回该区域。

参数:

val (torch.Tensor) – 要映射到区域的张量。

返回:

属于 TensorSpec 区域的 torch.Tensor。

rand(shape: Optional[Size] = None) → TensorDictBase

返回规格定义的区域中的随机张量。

采样将在区域内均匀进行，除非区域无界，在这种情况下将绘制正态值。

参数:

shape (torch.Size) – 随机张量的形状

返回:

在 TensorSpec 区域中采样的随机张量。

reshape(*shape) → T

重塑一个 TensorSpec。

Check reshape() for more information on this method.

sample(shape: torch.Size = None) → torch.Tensor | TensorDictBase

返回规格定义的区域中的随机张量。

See rand() for details.

separates(*keys: NestedKey, default: Optional[Any] = None) → Composite

通过将指定的键及其关联值提取到新的复合规格中来拆分复合规格。

此方法将迭代提供的键，将它们从当前复合规格中删除，然后将它们添加到新的复合规格中。如果找不到键，将使用指定的默认值。将返回新的复合规格。

参数:

• *keys (NestedKey) – 要从复合规格中提取的一个或多个键。每个键可以是单个键，也可以是嵌套键。

• default (Any, optional) – 如果在复合规格中找不到指定的键，则使用的值。默认为 None。

返回:

包含提取的键及其关联值的新复合规格。

返回类型:

Composite

注意

如果找不到任何指定的键，则该方法返回 None。

set(name: str, spec: TensorSpec) → StackedComposite

在 Composite 规格中设置一个规格。

squeeze(dim: int | None = None)

返回一个新 Spec，其中所有大小为 1 的维度都已删除。

当给定 dim 时，仅在该维度上执行挤压操作。

参数:

dim (int 或 None) – 应用挤压操作的维度

to(dest: torch.dtype | DEVICE_TYPING) → T

将 TensorSpec 转换为设备或 dtype。

如果未进行更改，则返回相同的规格。

to_numpy(val: TensorDict, safe: bool | None = None) → dict

返回输入张量的 np.ndarray 对应项。

This is intended to be the inverse operation of encode().

参数:

• val (torch.Tensor) – 要转换为 numpy 的张量。

• safe (bool) – boolean value indicating whether a check should be performed on the value against the domain of the spec. Defaults to the value of the CHECK_SPEC_ENCODE environment variable.

返回:

一个 np.ndarray。

type_check(value: torch.Tensor | TensorDictBase, selected_keys: NestedKey | Sequence[NestedKey] | None = None)

Checks the input value dtype against the TensorSpec dtype and raises an exception if they don’t match.

参数:

• value (torch.Tensor) – 必须检查其 dtype 的张量。

• key (str, optional) – 如果 TensorSpec 具有键，则将检查值 dtype 是否与指定键指向的规格匹配。

unflatten(dim: int, sizes: tuple[int]) → T

解展一个 TensorSpec。

Check unflatten() for more information on this method.

unlock_(recurse: bool | None = None) → T

解锁 Composite 并允许修改其内容。

除非通过 recurse 参数另有指定，否则这仅是第一级的锁修改。

unsqueeze(dim: int)

返回一个新 Spec，其中在 dim 指定的位置增加了一个单例维度。

参数:

dim (int 或 None) – 应用 unsqueeze 操作的维度。

values(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Callable[[type], bool] | None = None, step_mdp_static_only: bool = False) → _CompositeSpecValuesView

Composite 的值。

参数:

• include_nested (bool, optional) – if False, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next"]. Default is ``False`, i.e. nested keys will not be returned.

• leaves_only (bool, optional) – if False, the values returned will contain every level of nesting, i.e. a Composite(next=Composite(obs=None)) will lead to the keys ["next", ("next", "obs")]. Default is False.

关键字参数:

• is_leaf (callable, optional) – 读取一个类型并返回一个布尔值，指示该类型是否应被视为叶子。默认情况下，所有非 Composite 节点都被视为叶子。

• step_mdp_static_only (bool, optional) – if True, only keys that are static under step_mdp will be returned. Default is False.

view(*shape) → T

重塑一个 TensorSpec。

Check reshape() for more information on this method.

zero(shape: Optional[Size] = None) → TensorDictBase

返回盒中的零填充张量。

注意

Even though there is no guarantee that 0 belongs to the spec domain, this method will not raise an exception when this condition is violated. The primary use case of zero is to generate empty data buffers, not meaningful data.

参数:

shape (torch.Size) – 零张量的形状

返回:

在 TensorSpec 框中采样的零填充张量。

zeros(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Proxy to zero().