MultiDiscreteTensorSpec

class torchrl.data.MultiDiscreteTensorSpec(nvec: Union[Sequence[int], torch.Tensor, int], shape: Optional[torch.Size] = None, device: Optional[DEVICE_TYPING] = None, dtype: Optional[Union[str, torch.dtype]] = torch.int64, mask: torch.Tensor | None = None, remove_singleton: bool = True)

离散张量规范的串联。

参数:

• nvec (整数的可迭代对象 或 torch.Tensor) – 张量中每个元素的基数。可以有多个轴。

• shape (torch.Size, 可选) – 采样张量的总形状。如果提供，最后 m 个维度必须与 nvec.shape 匹配。

• device (str, int 或 torch.device, 可选) – 张量的设备。

• dtype (str 或 torch.dtype, 可选) – 张量的 dtype。

• remove_singleton (bool, 可选) – 如果为 True，则将单例样本（大小为 [1]）压缩。默认为 True。

• mask (torch.Tensor 或 None) – 在进行采样时，掩盖一些可能的输出结果。有关更多信息，请参阅 update_mask()。

示例

>>> ts = MultiDiscreteTensorSpec((3, 2, 3))
>>> ts.is_in(torch.tensor([2, 0, 1]))
True
>>> ts.is_in(torch.tensor([2, 2, 1]))
False

assert_is_in(value: Tensor) → None

断言张量是否属于该框，否则引发异常。

参数:

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

clear_device_()

所有叶子规范的无操作（必须具有设备）。

contains(item)

返回样本是否包含在 TensorSpec 定义的空间内。

有关更多信息，请参阅 is_in()。

encode(val: Union[ndarray, Tensor], *, ignore_device=False) → Tensor

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

参数:

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

关键字参数:

ignore_device (bool, 可选) – 如果 True，则忽略规格设备。这用于在调用 TensorDict(..., device="cuda") 时对张量进行分组转换，这速度更快。

返回值:

与所需张量规格匹配的 torch.Tensor。

expand(*shape)

返回具有扩展形状的新规格。

参数:

*shape (元组 或 int 可迭代对象) – 规格的新形状。必须符合当前形状：其长度必须至少与当前形状长度一样长，并且其最后的值也必须符合；即它们只能在当前维度是单例时才与它不同。

flatten(start_dim, end_dim)

将张量规格展平。

查看 flatten() 以获取有关此方法的更多信息。

classmethod implements_for_spec(torch_function: Callable) → Callable

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

abstract index(index: Union[int, Tensor, ndarray, slice, List], tensor_to_index: Tensor) → Tensor

索引输入张量。

参数:

• index (int, torch.Tensor, 切片 或 列表) – 张量的索引

• tensor_to_index – 要索引的张量

返回值:

索引的张量

is_in(val: Tensor) → bool

如果值 val 位于 TensorSpec 定义的框中，则返回 True，否则返回 False。

参数:

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

返回值:

一个布尔值，指示值是否属于 TensorSpec 框

project(val: Tensor) → Tensor

如果输入张量不在 TensorSpec 框中，它会根据一些启发式方法将其映射回框中。

参数:

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

返回值:

一个属于 TensorSpec 框的 torch.Tensor。

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

在规格定义的空间中返回一个随机张量。除非框是无界的，否则采样将是均匀的。

参数:

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

返回值:

在 TensorSpec 框中采样的一个随机张量。

reshape(*shape)

重新整形张量规格。

查看 reshape() 以获取有关此方法的更多信息。

property sample

在规格定义的空间中返回一个随机张量。

有关详细信息，请参见 rand()。

squeeze(dim: int | None = None)

返回一个新的规格，其中所有大小为 1 的维度都被移除。

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

参数:

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

to_numpy(val: Tensor, safe: Optional[bool] = None) → dict

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

参数:

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

• safe (bool) – 布尔值，指示是否应根据规格的域对值执行检查。默认为 CHECK_SPEC_ENCODE 环境变量的值。

返回值:

一个 np.ndarray

to_one_hot(val: Tensor, safe: Optional[bool] = None) → Union[MultiOneHotDiscreteTensorSpec, Tensor]

将规格域中的离散张量编码为其对应的独热编码。

参数:

• val (torch.Tensor, 可选) – 要进行独热编码的张量。

• safe (bool) – 布尔值，指示是否应根据规格的域对值执行检查。默认为 CHECK_SPEC_ENCODE 环境变量的值。

返回值:

独热编码后的张量。

to_one_hot_spec() → MultiOneHotDiscreteTensorSpec

将规格转换为等效的独热规格。

type_check(value: Tensor, key: Optional[NestedKey] = None) → None

检查输入值的 dtype 是否与 TensorSpec 的 dtype 相匹配，如果不匹配，则抛出异常。

参数:

• value (torch.Tensor) – 需要检查 dtype 的张量

• key (str, optional) – 如果 TensorSpec 有键，则会根据指定的键检查 value 的 dtype 是否与 spec 相匹配。

unflatten(dim, sizes)

将一个 TensorSpec 展开。

有关此方法的更多信息，请查看 unflatten()。

update_mask(mask)

设置一个掩码，以防止在采样时出现一些可能的结果。

也可以在 spec 初始化期间设置掩码。

参数:

mask (torch.Tensor or None) – 布尔掩码。如果为 None，则禁用掩码。否则，掩码的形状必须可扩展到等效 one-hot spec 的形状。 False 掩盖一个结果，True 保留结果不被掩盖。如果所有可能的结果都被掩盖，则在采样时会抛出错误。

示例

>>> mask = torch.tensor([False, False, True,
...                      True, True])
>>> ts = MultiDiscreteTensorSpec((3, 2), (5, 2,), dtype=torch.int64, mask=mask)
>>> # All but one of the three possible outcomes for the first
>>> # group are masked, but neither of the two possible
>>> # outcomes for the second group are masked.
>>> ts.rand()
tensor([[2, 1],
        [2, 0],
        [2, 1],
        [2, 1],
        [2, 0]])

view(*shape)

重新整形张量规格。

查看 reshape() 以获取有关此方法的更多信息。

zero(shape=None) → Tensor

返回一个在 box 内的填充为零的张量。

参数:

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

返回值:

在 TensorSpec box 内采样的一个填充为零的张量。