reference/generated/torchrl.envs.transforms.LastAction

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

LastAction

class torchrl.envs.transforms.LastAction(in_keys: Sequence[NestedKey] | NestedKey | None = None, out_keys: Sequence[NestedKey] | NestedKey | None = None, *, default: Literal['zeros', 'nan'] | float | int | Tensor = 'zeros', reset_key: NestedKey | None = None)

Copies the last action into the next observation.

This is the action analogue of CatFrames: the policy can condition on the action taken at the previous step (delayed control, recurrent policies, residual action heads). On each step() the action at time t is written under out_keys in the "next" tensordict; on reset() the same keys are filled with a default value (zeros, NaN, or a user-provided fill). On a batch-unlocked parent the default is expanded by the runtime reset batch, preserving the action feature shape.

out_keys are registered as Unbounded observation specs with the action’s shape, dtype and device, so reset fills (zeros on a one-hot action, NaN on a bounded action) remain in-spec.

Parameters:

• in_keys (NestedKey or sequence of NestedKey, optional) – keys pointing to the actions to remember. Defaults to the parent environment’s action_keys when the transform is attached, or ["action"] otherwise.

• out_keys (NestedKey or sequence of NestedKey, optional) – destination keys written into the observation. Defaults to each in_keys entry with its last component replaced by "last_action" (e.g. "action" -> "last_action", ("agents", "action") -> ("agents", "last_action")).

Keyword Arguments:

• default (str, number or torch.Tensor, optional) – value used to fill out_keys on reset(). "zeros" (default) writes zeros matching the action spec; "nan" writes NaNs (floating-point action specs only); a scalar is broadcast with fill_(); a tensor is broadcast to the action spec shape on the spec’s device and dtype. Defaults to "zeros".

• reset_key (NestedKey, optional) – the reset key to be used as a partial-reset indicator. Must be unique. If not provided, defaults to the only reset key of the parent environment (if it has only one) and raises an exception otherwise.

Examples

>>> from torchrl.envs import GymEnv, TransformedEnv
>>> from torchrl.envs.transforms import LastAction
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), LastAction())
>>> td = env.reset()
>>> td["last_action"]
tensor([0.])
>>> rollout = env.rollout(3)
>>> (rollout["next", "last_action"] == rollout["action"]).all()
tensor(True)

See also

CatFrames for stacking past observations, InitTracker for marking episode starts, and LastActionConfig for the Hydra configuration.

forward(tensordict: TensorDictBase) → TensorDictBase

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

By default, this method:

• calls directly _apply_transform().

• does not call _step() or _call().

This method is not called within env.step at any point. However, is is called within sample().

Note

forward also works with regular keyword arguments using dispatch to cast the args names to the keys.

Examples

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