import gym

class RewardShapingWrapper(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.env = env

    def step(self, action):
        next_state, reward, done, info = self.env.step(action)
        reward += ((abs(next_state[1]))) * 10
        return next_state, reward, done, info
