ToTensorImage

class torchrl.envs.transforms.ToTensorImage(from_int: bool | None = None, unsqueeze: bool = False, dtype: torch.device | None = 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，則觀察張量沿第一個維度進行擴充套件。預設值為 False。

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

關鍵字引數:

• in_keys (list of NestedKeys) – 要處理的鍵。

• out_keys (list of 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) – 轉換前的規範

返回:

轉換後的預期規範