"""Utilities and definitions shared by reward-related code."""

from typing import Callable

import numpy as np
from stable_baselines3.common import vec_env

RewardFn = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]


def build_norm_reward_fn(
    *,
    reward_fn: RewardFn,
    vec_normalize: vec_env.VecNormalize,
    norm_reward: bool = True,
) -> RewardFn:
    """Wraps `reward_fn` to automatically normalize input.

    Args:
        reward_fn: The reward function that normalized inputs are evaluated on.
        vec_normalize: Instance of VecNormalize used to normalize inputs and
            rewards.
        norm_reward: If True, then also normalize reward before returning.

    Returns:
        A reward function that normalizes the inputs using `vec_normalize`,
        calls `reward_fn` and if `norm_reward` then normalizes the reward.
    """

    def inner(
        obs: np.ndarray,
        acts: np.ndarray,
        next_obs: np.ndarray,
        dones: np.ndarray,
    ) -> np.ndarray:
        """Normalizes `obs` and `next_obs` and computes reward from `reward_fn`.

        Args:
            obs: Observations before transition.
            acts: Actions.
            next_obs: Observations after transition.
            dones: Is the transition into terminal state at end of episode?

        Returns:
            The reward, normalized if `norm_reward` is true.
        """
        norm_obs = vec_normalize.normalize_obs(obs)
        norm_next_obs = vec_normalize.normalize_obs(next_obs)
        rew = reward_fn(norm_obs, acts, norm_next_obs, dones)
        if norm_reward:
            rew = vec_normalize.normalize_reward(rew)
        return rew

    return inner

@@ -1,31 +1,16 @@
 %22%22%22
-Utilities and definition
+Type alia
 s sh
@@ -93,53 +93,8 @@
 s np
-%0Afrom stable_baselines3.common import vec_env
 %0A%0ARe
@@ -177,1474 +177,4 @@
 ay%5D%0A
-%0A%0Adef build_norm_reward_fn(%0A    *,%0A    reward_fn: RewardFn,%0A    vec_normalize: vec_env.VecNormalize,%0A    norm_reward: bool = True,%0A) -%3E RewardFn:%0A    %22%22%22Wraps %60reward_fn%60 to automatically normalize input.%0A%0A    Args:%0A        reward_fn: The reward function that normalized inputs are evaluated on.%0A        vec_normalize: Instance of VecNormalize used to normalize inputs and%0A            rewards.%0A        norm_reward: If True, then also normalize reward before returning.%0A%0A    Returns:%0A        A reward function that normalizes the inputs using %60vec_normalize%60,%0A        calls %60reward_fn%60 and if %60norm_reward%60 then normalizes the reward.%0A    %22%22%22%0A%0A    def inner(%0A        obs: np.ndarray,%0A        acts: np.ndarray,%0A        next_obs: np.ndarray,%0A        dones: np.ndarray,%0A    ) -%3E np.ndarray:%0A        %22%22%22Normalizes %60obs%60 and %60next_obs%60 and computes reward from %60reward_fn%60.%0A%0A        Args:%0A            obs: Observations before transition.%0A            acts: Actions.%0A            next_obs: Observations after transition.%0A            dones: Is the transition into terminal state at end of episode?%0A%0A        Returns:%0A            The reward, normalized if %60norm_reward%60 is true.%0A        %22%22%22%0A        norm_obs = vec_normalize.normalize_obs(obs)%0A        norm_next_obs = vec_normalize.normalize_obs(next_obs)%0A        rew = reward_fn(norm_obs, acts, norm_next_obs, dones)%0A        if norm_reward:%0A            rew = vec_normalize.normalize_reward(rew)%0A        return rew%0A%0A    return inner%0A

Remove dead reward input normalization code (#426)