RewardSum

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

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

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

reset(tensordict: TensorDictBase) → TensorDictBase

Resets episode rewards.

transform_input_spec(input_spec: TensorSpec) → TensorSpec

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

Parameters:

input_spec (TensorSpec) – spec before the transform

Returns:

expected spec after the transform

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

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