diff --git a/gymnasium/wrappers/stateful_reward.py b/gymnasium/wrappers/stateful_reward.py index 053050d2c..002b3638b 100644 --- a/gymnasium/wrappers/stateful_reward.py +++ b/gymnasium/wrappers/stateful_reward.py @@ -19,8 +19,9 @@ class NormalizeReward( gym.Wrapper[ObsType, ActType, ObsType, ActType], gym.utils.RecordConstructorArgs ): - r"""Normalizes immediate rewards such that their exponential moving average has a fixed variance. + r"""This wrapper will scale rewards s.t. the discounted returns have a mean of 0 and std of 1. + In a nutshell, the rewards are divided through by the standard deviation of a rolling discounted sum of the reward. The exponential moving average will have variance :math:`(1 - \gamma)^2`. The property `_update_running_mean` allows to freeze/continue the running mean calculation of the reward @@ -29,6 +30,11 @@ class NormalizeReward( A vector version of the wrapper exists :class:`gymnasium.wrappers.vector.NormalizeReward`. + Important note: + Contrary to what the name suggests, this wrapper does not normalize the rewards to have a mean of 0 and a standard + deviation of 1. Instead, it scales the rewards such that **discounted returns** have approximately unit variance. + See [Engstrom et al.](https://openreview.net/forum?id=r1etN1rtPB) on "reward scaling" for more information. + Note: In v0.27, NormalizeReward was updated as the forward discounted reward estimate was incorrectly computed in Gym v0.25+. For more detail, read [#3154](https://github.com/openai/gym/pull/3152). diff --git a/gymnasium/wrappers/vector/stateful_reward.py b/gymnasium/wrappers/vector/stateful_reward.py index 8b96ae3ad..b086825f3 100644 --- a/gymnasium/wrappers/vector/stateful_reward.py +++ b/gymnasium/wrappers/vector/stateful_reward.py @@ -18,14 +18,20 @@ class NormalizeReward(VectorWrapper, gym.utils.RecordConstructorArgs): - r"""This wrapper will normalize immediate rewards s.t. their exponential moving average has a fixed variance. + r"""This wrapper will scale rewards s.t. the discounted returns have a mean of 0 and std of 1. + In a nutshell, the rewards are divided through by the standard deviation of a rolling discounted sum of the reward. The exponential moving average will have variance :math:`(1 - \gamma)^2`. The property `_update_running_mean` allows to freeze/continue the running mean calculation of the reward statistics. If `True` (default), the `RunningMeanStd` will get updated every time `self.normalize()` is called. If False, the calculated statistics are used but not updated anymore; this may be used during evaluation. + Important note: + Contrary to what the name suggests, this wrapper does not normalize the rewards to have a mean of 0 and a standard + deviation of 1. Instead, it scales the rewards such that **discounted returns** have approximately unit variance. + See [Engstrom et al.](https://openreview.net/forum?id=r1etN1rtPB) on "reward scaling" for more information. + Note: The scaling depends on past trajectories and rewards will not be scaled correctly if the wrapper was newly instantiated or the policy was changed recently.