InitTracker

class torchrl.envs.transforms.InitTracker(init_key: str = 'is_init')

Reset tracker.

This transform populates the step/reset tensordict with a reset tracker entry that is set to True whenever reset() is called.

Parameters:

init_key (NestedKey, optional) – the key to be used for the tracker entry. In case of multiple _reset flags, this key is used as the leaf replacement for each.

Examples

>>> from torchrl.envs.libs.gym import GymEnv
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
>>> td = env.reset()
>>> print(td["is_init"])
tensor(True)
>>> td = env.rand_step(td)
>>> print(td["next", "is_init"])
tensor(False)

forward(tensordict: TensorDictBase) → TensorDictBase

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

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