ToTensorImage

class torchrl.envs.transforms.ToTensorImage(from_int: Optional[bool] = None, unsqueeze: bool = False, dtype: Optional[torch.device] = None, *, in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = 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，則沿第一維對觀察張量進行 unsqueeze 操作。預設值為 False。

• dtype (torch.dtype, optional) – 用於結果觀察值的 dtype。

關鍵字引數:

• 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

轉換觀察 spec，使結果 spec 與轉換對映匹配。

引數:

observation_spec (TensorSpec) – 轉換前的 spec

返回值:

轉換後預期的 spec