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_action_spec(full_action_spec: Composite) → Composite

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

引數：

action_spec (TensorSpec) – 轉換前的 spec

返回值：

轉換後預期的 spec

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

轉換 input 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() 和 transform_full_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

transform_state_spec(full_state_spec: Composite) → Composite

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

引數：

state_spec (TensorSpec) – 轉換前的 spec

返回值：

轉換後預期的 spec