ToTensorImage

class torchrl.envs.transforms.ToTensorImage(from_int: Optional[bool] = None, unsqueeze: bool = False, dtype: Optional[device] = None, *, in_keys: Optional[Sequence[NestedKey]] = None, out_keys: Optional[Sequence[NestedKey]] = None, shape_tolerant: bool = False)

将类似 numpy 的图像 (W x H x C) 转换为 pytorch 图像 (C x W x H)。

将观测图像从 (… x W x H x C) 张量转换为 (… x C x W x H) 张量。 可选地，将输入张量从 [0, 255] 范围缩放到 [0.0, 1.0] 范围（有关更多详细信息，请参阅 `from_int`）。

在其他情况下，张量在不缩放的情况下返回。

参数:

• from_int (bool, optional) – 如果为 `True`，张量将从 [0, 255] 范围缩放到 [0.0, 1.0] 范围。如果为 `False`，张量将不会缩放。如果为 `None`，如果张量不是浮点张量，则会进行缩放。默认值=None。

• unsqueeze (bool) – 如果为 `True`，则观测张量将沿第一个维度进行解压缩。默认值=False。

• dtype (torch.dtype, optional) – 用于结果观测值的数据类型。

关键字参数:

• in_keys (NestedKeys 列表) – 要处理的键。

• out_keys (NestedKeys 列表) – 要写入的键。

• shape_tolerant (bool, optional) – 如果为 `True`，将检查输入图像的形状。如果最后一个通道不是 `3`，则会忽略排列。默认为 `False`。

示例

>>> transform = ToTensorImage(in_keys=["pixels"])
>>> ri = torch.randint(0, 255, (1 , 1, 10, 11, 3), dtype=torch.uint8)
>>> td = TensorDict(
...     {"pixels": ri},
...     [1, 1])
>>> _ = transform(td)
>>> obs = td.get("pixels")
>>> print(obs.shape, obs.dtype)
torch.Size([1, 1, 3, 10, 11]) torch.float32

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

转换观测规范，使结果规范与转换映射匹配。

参数:

observation_spec (TensorSpec) – 转换前的规范

返回:

转换后的预期规范