from gridmind.wrappers.env_wrappers.base_gym_wrapper import BaseGymWrapper
[docs]class NormalizedObservationWrapper(BaseGymWrapper):
def __init__(self, env):
super().__init__(env)
[docs] self.observation_space_low = self.env.observation_space.low
[docs] self.observation_space_high = self.env.observation_space.high
[docs] def normalize_observation(self, observation):
normalized_observation = (observation - self.observation_space_low) / (
self.observation_space_high - self.observation_space_low
)
return normalized_observation
[docs] def reset(self):
observation, info = self.env.reset()
normalized_observation = self.normalize_observation(observation)
return normalized_observation, info
[docs] def step(self, action):
observation, reward, terminated, truncated, info = self.env.step(action)
normalized_observation = self.normalize_observation(observation)
return normalized_observation, reward, terminated, truncated, info