DeviceCastTransform

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

Moves data from one device to another.

Parameters:

• device (torch.device or equivalent) – the destination device.

• orig_device (torch.device or equivalent) – the origin device. If not specified and a parent environment exists, it it retrieved from it. In all other cases, it remains unspecified.

Examples

>>> 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

Reads the input tensordict, and for the selected keys, applies the transform.

transform_done_spec(full_done_spec: Composite) → Composite

Transforms the done spec such that the resulting spec matches transform mapping.

Parameters:

done_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_env_device(device)

Transforms the device of the parent env.

transform_input_spec(input_spec: Composite) → Composite

Transforms the input spec such that the resulting spec matches transform mapping.

Parameters:

input_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_observation_spec(observation_spec: Composite) → Composite

Transforms the observation spec such that the resulting spec matches transform mapping.

Parameters:

observation_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_output_spec(output_spec: Composite) → Composite

Transforms the output spec such that the resulting spec matches transform mapping.

This method should generally be left untouched. Changes should be implemented using transform_observation_spec(), transform_reward_spec() and transformfull_done_spec(). :param output_spec: spec before the transform :type output_spec: TensorSpec

Returns:

expected spec after the transform

transform_reward_spec(full_reward_spec: Composite) → Composite

Transforms the reward spec such that the resulting spec matches transform mapping.

Parameters:

reward_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform