TensorSpec

class torchrl.data.TensorSpec(shape: torch.Size, space: Union[None, Box], device: torch.device | None = <property object>, dtype: torch.dtype = torch.float32, domain: str = '')

张量元数据容器（用于观察、动作和奖励）的父类。

参数：

• shape (torch.Size) – 张量的尺寸

• space (Box) – 描述可以预期哪些值的 Box 实例

• device (torch.device) – 张量的设备

• dtype (torch.dtype) – 张量的 dtype

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。

abstract expand(*shape)

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

参数：

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

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

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

参数：

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

返回值：

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

reshape(*shape)

重塑张量规范。

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

属性 sample

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

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

squeeze(dim: int | None = None)

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

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

参数：

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

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

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

参数：

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

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

返回值：

一个np.ndarray

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

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

参数：

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

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

unflatten(dim, sizes)

展开张量规范。

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

view(*shape)

重塑张量规范。

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

zero(shape=None) → Tensor

返回范围内的零填充张量。

参数：

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

返回值：

在TensorSpec范围内采样的零填充张量。