Skip to content

lit_sb3.wrappers.atari.clip_reward

[docs] module lit_sb3.wrappers.atari.clip_reward 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from typing import SupportsFloat

import gymnasium as gym
import numpy as np


class ClipRewardEnv(gym.RewardWrapper):
    """Clip the reward to {+1, 0, -1} by its sign.

    Args:
        env: Environment to wrap
    """

    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)

    def reward(self, reward: SupportsFloat) -> float:
        """Bin reward to {+1, 0, -1} by its sign.

        Args:
            reward:
            reward: The reward to process.
        """
        return np.sign(float(reward))