TimeMaxPool

class torchrl.envs.transforms.TimeMaxPool(in_keys: Optional[Sequence[Union[str, Tuple[str, ...]]]] = None, out_keys: Optional[Sequence[Union[str, Tuple[str, ...]]]] = None, T: int = 1)

Take the maximum value in each position over the last T observations.

This transform take the maximum value in each position for all in_keys tensors over the last T time steps.

Parameters:

• in_keys (sequence of NestedKey, optional) – input keys on which the max pool will be applied. Defaults to “observation” if left empty.

• out_keys (sequence of NestedKey, optional) – output keys where the output will be written. Defaults to in_keys if left empty.

• T (int, optional) – Number of time steps over which to apply max pooling.

forward(tensordict: TensorDictBase) → TensorDictBase

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

reset(tensordict: TensorDictBase) → TensorDictBase

Resets _buffers.

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