Search code examples
machine-learningartificial-intelligenceopenai-gymopenai-api

Is it possible to modify an OpenAI gym state before and during training?


What I would like to do is modify an environment, for example take the Super Mario Bros gym environment, and blur the image that the agent trains on, and see whether a reinforcement learning agent is still capable of learning on these "blurred" states.

Does OpenAI make it possible to do something like this? How would I add a gym environment pre-processing step?


Solution

  • I recommend you to make a wrapper for your gym environment that adds a treatment in the step() and reset() functions

    Here is a bit of code to illustrate the idea :

    class EnvWrapper(gym.Env):
        def __init__(self, config):
            self.env = gym.make("Your-Env-Name")    # The wrapper encapsulates the gym env
        
        def step(self, action):
            obs, reward, done, info = self.env.step(action)   # calls the gym env methods
            obs = self._blur(obs)                             # applies your specific treatment
            return obs, reward, done, info
    
        def reset(self):
            obs = self.env.reset()    # same for reset
            return self._blur(obs)
    
        def _blur(self):
            do_whatever_you_need
    

    With this method your do not need to make any change to the original environment, which is generally a good idea