import gym


class AngleWrapper(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)
        next_state = next_state[2:]
        return next_state, reward, done, info


class FlattenWrapper(gym.ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)

    def observation(self, obs):
        obs = obs.reshape(-1)
        return obs
