class torchrl.envs.transforms.RewardSum(in_keys: Optional[Sequence[Union[str, Tuple[str, ...]]]] = None, out_keys: Optional[Sequence[Union[str, Tuple[str, ...]]]] = None)[source]

Tracks episode cumulative rewards.

This transform accepts a list of tensordict reward keys (i.e. ´in_keys´) and tracks their cumulative value along each episode. When called, the transform creates a new tensordict key for each in_key named ´episode_{in_key}´ where the cumulative values are written. All ´in_keys´ should be part of the env reward and be present in the env reward_spec.

If no in_keys are specified, this transform assumes ´reward´ to be the input key. However, multiple rewards (e.g. reward1 and reward2) can also be specified. If ´in_keys´ are not present in the provided tensordict, this transform hos no effect.

forward(tensordict: TensorDictBase) TensorDictBase[source]

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

reset(tensordict: TensorDictBase) TensorDictBase[source]

Resets episode rewards.

transform_input_spec(input_spec: TensorSpec) TensorSpec[source]

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


input_spec (TensorSpec) – spec before the transform


expected spec after the transform

transform_observation_spec(observation_spec: TensorSpec) TensorSpec[source]

Transforms the observation spec, adding the new keys generated by RewardSum.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources