DTypeCastTransform

class torchrl.envs.transforms.DTypeCastTransform(dtype_in: torch.dtype, dtype_out: torch.dtype, in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = None, in_keys_inv: Sequence[NestedKey] | None = None, out_keys_inv: Sequence[NestedKey] | None = None)

將選定鍵的資料型別從一種轉換為另一種。

取決於在構造時是否提供了 in_keys 或 in_keys_inv，類的行為會有所不同

• 如果提供了鍵，則只有這些條目會從 dtype_in 轉換為 dtype_out 條目；

• 如果未提供鍵，並且物件位於 transforms 的環境登錄檔中，則將 dtype 設定為 dtype_in 的輸入和輸出 spec 將分別用作 in_keys_inv / in_keys。

• 如果未提供鍵且物件在沒有環境的情況下使用，則 forward / inverse 過程將掃描輸入 tensordict 以查詢所有 dtype_in 值，並將它們對映到 dtype_out tensor。對於大型資料結構，這可能會影響效能，因為掃描不是免費的。要轉換的鍵不會被快取。請注意，在這種情況下，不能傳遞 out_keys (或 out_keys_inv)，因為無法精確預測鍵的處理順序。

引數：

• dtype_in (torch.dtype) – 輸入 dtype (來自環境)。

• dtype_out (torch.dtype) – 輸出 dtype (用於模型訓練)。

• in_keys (Sequence of NestedKey, 可選) – 在暴露給外部物件和函式之前，要從 dtype_in 轉換為 dtype_out 的鍵列表。

• out_keys (Sequence of NestedKey, 可選) – 目標鍵列表。如果未提供，則預設為 in_keys。

• in_keys_inv (Sequence of NestedKey, 可選) – 在傳遞給包含的 base_env 或 storage 之前，要從 dtype_out 轉換為 dtype_in 的鍵列表。

• out_keys_inv (Sequence of NestedKey, 可選) – 逆向 transform 的目標鍵列表。如果未提供，則預設為 in_keys_inv。

示例

>>> td = TensorDict(
...     {'obs': torch.ones(1, dtype=torch.double),
...     'not_transformed': torch.ones(1, dtype=torch.double),
... }, [])
>>> transform = DTypeCastTransform(torch.double, torch.float, in_keys=["obs"])
>>> _ = transform(td)
>>> print(td.get("obs").dtype)
torch.float32
>>> print(td.get("not_transformed").dtype)
torch.float64

在“自動”模式下，所有 float64 條目都會被轉換

示例

>>> td = TensorDict(
...     {'obs': torch.ones(1, dtype=torch.double),
...     'not_transformed': torch.ones(1, dtype=torch.double),
... }, [])
>>> transform = DTypeCastTransform(torch.double, torch.float)
>>> _ = transform(td)
>>> print(td.get("obs").dtype)
torch.float32
>>> print(td.get("not_transformed").dtype)
torch.float32

在構造環境時未指定 transform 鍵的情況下，規則也是相同的行為

示例

>>> class MyEnv(EnvBase):
...     def __init__(self):
...         super().__init__()
...         self.observation_spec = Composite(obs=Unbounded((), dtype=torch.float64))
...         self.action_spec = Unbounded((), dtype=torch.float64)
...         self.reward_spec = Unbounded((1,), dtype=torch.float64)
...         self.done_spec = Unbounded((1,), dtype=torch.bool)
...     def _reset(self, data=None):
...         return TensorDict({"done": torch.zeros((1,), dtype=torch.bool), **self.observation_spec.rand()}, [])
...     def _step(self, data):
...         assert data["action"].dtype == torch.float64
...         reward = self.reward_spec.rand()
...         done = torch.zeros((1,), dtype=torch.bool)
...         obs = self.observation_spec.rand()
...         assert reward.dtype == torch.float64
...         assert obs["obs"].dtype == torch.float64
...         return obs.empty().set("next", obs.update({"reward": reward, "done": done}))
...     def _set_seed(self, seed):
...         pass
>>> env = TransformedEnv(MyEnv(), DTypeCastTransform(torch.double, torch.float))
>>> assert env.action_spec.dtype == torch.float32
>>> assert env.observation_spec["obs"].dtype == torch.float32
>>> assert env.reward_spec.dtype == torch.float32, env.reward_spec.dtype
>>> print(env.rollout(2))
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                obs: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([2, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([2]),
            device=cpu,
            is_shared=False),
        obs: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=cpu,
    is_shared=False)
>>> assert env.transform.in_keys == ["obs", "reward"]
>>> assert env.transform.in_keys_inv == ["action"]

forward(tensordict: TensorDictBase = None) → TensorDictBase

讀取輸入的 tensordict，並對選定的鍵應用 transform。

transform_input_spec(input_spec: TensorSpec) → TensorSpec

轉換輸入 spec，使結果 spec 與 transform 對映匹配。

引數：

input_spec (TensorSpec) – transform 前的 spec

返回：

transform 後的預期 spec

transform_observation_spec(observation_spec)

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

引數：

observation_spec (TensorSpec) – transform 前的 spec

返回：

transform 後的預期 spec

transform_output_spec(output_spec: Composite) → Composite

轉換輸出 spec，使結果 spec 與 transform 對映匹配。

通常應保持此方法不變。更改應使用 transform_observation_spec()、transform_reward_spec() 和 transform_full_done_spec() 實現。 :param output_spec: transform 前的 spec :type output_spec: TensorSpec

返回：

transform 後的預期 spec