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 是從轉換後的 env 外部看到的鍵的名稱。預設為 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，並對選定的鍵應用轉換。

預設情況下，此方法

• 直接呼叫 _apply_transform()。

• 不呼叫 _step() 或 _call()。

此方法不會在任何時候在 env.step 中呼叫。但是，它會在 sample() 中呼叫。

注意

forward 也可以使用 dispatch 將引數名稱轉換為鍵，並使用常規關鍵字引數。

示例

>>> class TransformThatMeasuresBytes(Transform):
...     '''Measures the number of bytes in the tensordict, and writes it under `"bytes"`.'''
...     def __init__(self):
...         super().__init__(in_keys=[], out_keys=["bytes"])
...
...     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
...         bytes_in_td = tensordict.bytes()
...         tensordict["bytes"] = bytes
...         return tensordict
>>> t = TransformThatMeasuresBytes()
>>> env = env.append_transform(t) # works within envs
>>> t(TensorDict(a=0))  # Works offline too.

transform_action_spec(full_action_spec: Composite) → Composite

轉換動作規範，使結果規範與變換對映匹配。

引數:

action_spec (TensorSpec) – 變換前的規範

返回:

轉換後的預期規範

transform_done_spec(full_done_spec: Composite) → Composite

變換 done spec，使結果 spec 與變換對映匹配。

引數:

done_spec (TensorSpec) – 變換前的 spec

返回:

轉換後的預期規範

transform_env_device(device)

轉換父環境的 device。

transform_input_spec(input_spec: Composite) → Composite

轉換輸入規範，使結果規範與轉換對映匹配。

引數:

input_spec (TensorSpec) – 轉換前的規範

返回:

轉換後的預期規範

transform_observation_spec(observation_spec: Composite) → Composite

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

引數:

observation_spec (TensorSpec) – 轉換前的規範

返回:

轉換後的預期規範

transform_output_spec(output_spec: Composite) → Composite

轉換輸出規範，使結果規範與轉換對映匹配。

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

返回:

轉換後的預期規範

transform_reward_spec(full_reward_spec: Composite) → Composite

轉換獎勵的 spec，使其與變換對映匹配。

引數:

reward_spec (TensorSpec) – 變換前的 spec

返回:

轉換後的預期規範

transform_state_spec(full_state_spec: Composite) → Composite

轉換狀態規範，使結果規範與變換對映匹配。

引數:

state_spec (TensorSpec) – 變換前的規範

返回:

轉換後的預期規範