DeviceCastTransform

class torchrl.envs.transforms.DeviceCastTransform(device, orig_device=None, *, in_keys=None, out_keys=None, in_keys_inv=None, out_keys_inv=None)

將資料從一個裝置移動到另一個裝置。

參數:

• device (torch.device 或 等效) – 目的地裝置。

• orig_device (torch.device 或 等效) – 原始裝置。如果未指定，且存在父環境，則會從中擷取。在所有其他情況下，它保持未指定狀態。

關鍵字引數:

• in_keys (NestedKey 的列表) – 要對應到不同裝置的項目列表。預設為 None。

• out_keys (NestedKey 的列表) – 對應到裝置上的項目的輸出名稱。預設為 in_keys 的值。

• in_keys_inv (NestedKey 的列表) – 要對應到不同裝置的項目列表。in_keys_inv 是基礎環境預期的名稱。預設為 None。

• out_keys_inv (NestedKey 的列表) – 映射到裝置上的條目的輸出名稱。out_keys_inv 是從轉換後的環境外部看到的鍵的名稱。預設為 in_keys_inv 的值。

範例

>>> td = TensorDict(
...     {'obs': torch.ones(1, dtype=torch.double),
... }, [], device="cpu:0")
>>> transform = DeviceCastTransform(device=torch.device("cpu:2"))
>>> td = transform(td)
>>> print(td.device)
cpu:2

forward(tensordict: TensorDictBase = None) → TensorDictBase

讀取輸入的 tensordict，並對所選的鍵應用轉換。

transform_done_spec(full_done_spec: Composite) → Composite

轉換 done spec，使得到的 spec 與轉換映射相符。

參數:

done_spec (TensorSpec) – 轉換前的 spec

返回值:

轉換後預期的 spec

transform_env_device(device)

轉換父環境的裝置。

transform_input_spec(input_spec: Composite) → Composite

轉換輸入 spec，使得到的 spec 與轉換映射相符。

參數:

input_spec (TensorSpec) – 轉換前的 spec

返回值:

轉換後預期的 spec

transform_observation_spec(observation_spec: Composite) → Composite

轉換 observation spec，使得到的 spec 與轉換映射相符。

參數:

observation_spec (TensorSpec) – 轉換前的 spec

返回值:

轉換後預期的 spec

transform_output_spec(output_spec: Composite) → Composite

轉換 output spec，使得到的 spec 與轉換映射相符。

通常這個方法不應該被修改。應該使用 transform_observation_spec()、 transform_reward_spec() 和 transformfull_done_spec() 來實作變更。 :param output_spec: 轉換前的 spec :type output_spec: TensorSpec

返回值:

轉換後預期的 spec

transform_reward_spec(full_reward_spec: Composite) → Composite

轉換 reward spec，使得到的 spec 與轉換映射相符。

參數:

reward_spec (TensorSpec) – 轉換前的 spec

返回值:

轉換後預期的 spec