OneHotDiscreteTensorSpec

class torchrl.data.OneHotDiscreteTensorSpec(n: int, shape: Optional[torch.Size] = None, device: Optional[DEVICE_TYPING] = None, dtype: Optional[Union[str, torch.dtype]] = torch.bool, use_register: bool = False, mask: torch.Tensor | None = None)

一维独热离散张量规范。

默认情况下，TorchRL 假设分类变量被编码为变量的独热编码。这允许对张量进行简单的索引，例如：

>>> batch, size = 3, 4
>>> action_value = torch.arange(batch*size)
>>> action_value = action_value.view(batch, size).to(torch.float)
>>> action = (action_value == action_value.max(-1,
...    keepdim=True)[0]).to(torch.long)
>>> chosen_action_value = (action * action_value).sum(-1)
>>> print(chosen_action_value)
tensor([ 3.,  7., 11.])

形状的最后一维（变量域）不能被索引。

参数：

• n (int) – 可能结果的数量。

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

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

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

• user_register (bool) – 实验性功能。如果为 True，则每个整数都将按其出现的顺序映射到一个二进制向量。此功能专为没有先验结果数量定义的环境而设计（例如，离散结果从任意集合中采样，其元素将在寄存器中映射到一系列唯一的独热二进制向量）。

• mask (torch.Tensor 或 None) – 在采样时屏蔽一些可能的结果。有关更多信息，请参见 update_mask()。

assert_is_in(value: Tensor) → None

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

参数：

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

clear_device_()

对于所有叶子规范（必须具有设备）而言，这是一个无操作。

contains(item)

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

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

encode(val: Union[ndarray, Tensor], space: Optional[DiscreteBox] = None, *, ignore_device: bool = False) → Tensor

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

参数：

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

关键字参数：

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

返回值：

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

expand(*shape)

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

参数：

*shape (tuple 或 int 的可迭代对象) – Spec 的新形状。必须符合当前形状：其长度必须至少与当前形状长度相同，并且其最后一个值也必须符合；即，只有当当前维度为单例时，它们才能与之不同。

flatten(start_dim, end_dim)

展平一个 TensorSpec。

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

classmethod implements_for_spec(torch_function: Callable) → Callable

为 TensorSpec 注册一个 torch 函数重写。

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

索引输入张量。

参数：

• index (int, torch.Tensor, slice 或 list) – 张量的索引

• 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=None) → Tensor

返回 spec 定义的空间中的随机张量。除非范围是无界的，否则采样将是均匀的。

参数：

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

返回值：

在 TensorSpec 范围内采样的随机张量。

reshape(*shape)

重新整形一个 TensorSpec。

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

property sample

返回 spec 定义的空间中的随机张量。

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

squeeze(dim=None)

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

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

参数：

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

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

将给定的独热张量转换为分类格式。

参数：

• val (torch.Tensor, 可选) – 要转换为分类格式的独热张量。

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

返回值：

分类张量。

to_categorical_spec() → DiscreteTensorSpec

将 spec 转换为等效的分类 spec。

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

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

参数：

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

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

返回值：

一个 np.ndarray

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

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

参数：

• 值 (torch.Tensor) – 需要检查数据类型的张量。

• 键 (str, 可选) – 如果 TensorSpec 具有键，则将根据指示键指向的规范检查值的 dtype。

unflatten(dim, sizes)

展平张量规范。

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

update_mask(mask)

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

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

参数：

mask (torch.Tensor 或 None) – 布尔掩码。如果为 None，则禁用掩码。否则，掩码的形状必须可扩展到规范的形状。False 掩盖结果，True 保持结果未掩盖。如果所有可能的结果都被掩盖，则在采样时会引发错误。

示例

>>> mask = torch.tensor([True, False, False])
>>> ts = OneHotDiscreteTensorSpec(3, (2, 3,), dtype=torch.int64, mask=mask)
>>> # All but one of the three possible outcomes are masked
>>> ts.rand()
tensor([[1, 0, 0],
        [1, 0, 0]])

view(*shape)

重新整形一个 TensorSpec。

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

zero(shape=None) → Tensor

返回框中填充零的张量。

参数：

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

返回值：

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