Normalization

src.jaxnasium.algorithms.utils.Normalizer

A container for running statistics on Observations and Rewards.

normalize_obs(obs: PyTree) -> PyTree

Normalizes the given batch of observations if normalization of observations is enabled.

normalize_reward(reward: Array) -> Array

Normalizes the given batch of rewards if normalization of rewards is enabled.

update(batch: Transition) -> Normalizer

Updates the normalization state with a new Transition containing both rewards and observations.