SelectTransform

class torchrl.envs.transforms.SelectTransform(*selected_keys)

Select keys from the input tensordict.

In general, the ExcludeTransform should be preferred: this transforms also

selects the “action” (or other keys from input_spec), “done” and “reward” keys but other may be necessary.

Parameters:

*selected_keys (iterable of NestedKey) – The name of the keys to select. If the key is not present, it is simply ignored.

forward(tensordict: TensorDictBase) → TensorDictBase

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

reset(tensordict: TensorDictBase) → TensorDictBase

Resets a tranform if it is stateful.

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

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