Как создать свой open-ai gym на основе изображения?

Я использую алгоритм reinforcement learning A3C. Среда в которой учатся агенты должна быть простой формой регистрации. То есть мне нужно создать свой собственный open-ai gym. Отличие состоит в том, что я хочу создать среду на основе скриншота. Я пытаюсь создавать среду, используя это описание https://github.com/andriusbern/RL. Если инициализировать observation_space вот так self.observation_space = gym.spaces.Box(low=0, high=255, shape=(128, 128, 3), dtype=np.uint8), то я получаю просто картинку. Вопрос, как сделать так, чтобы создавать среду на основе скриншота. Оригинальный алгоритм того, что я пытаюсь сделать описан здесь https://drive.google.com/file/d/1oenbo13OkGtKX2nTfRi7ggyb520tAu2j/view?usp=sharing. Возможно я неправильно поняла алгоритм и то, как это должно работать. Возможно среда должна быть вне зависимости от изображения, а скриншоты это просто состояния среды.

Ссылка на алгоритм A3C: https://github.com/dgriff777/rl_a3c_pytorch