如何实现使用图像作为观察的 2D OpenAI-Gym 环境？

Question

I need to create a 2D environment with a basic model of a robot arm and a target point. I want the arm to reach the target through a series of discrete actions (eg go right, left, up and down) and I need the observation space to be an RGB image of the screen that I will then use as input to DQN.

My problem is that I don't understand how can I use the image observation to move the robot and to get the positions of its arm and of the target (eg to create a reward function based on the their distance).

It looks like the only Gym environments I can take inspiration from are the Atari ones, but I cannot find specific codes of the single games probably because they are embedded in their ROMs.

So, if I want to follow the examples of Atari environments in Gym and use atari_env.py, does it mean that I also need to create a game and its ROM and then integrate it in Stella? Or is there another way? In general, are there other types of environments that use images as observations from which I can get inspiration?

Many thanks

Answer 1

You can create a custom Gym environment by simply implementing a class with the appropriate methods:

import gym
from gym.spaces import Discrete, Box

class RobotEnv(gym.Env):
  metadata = {'render.modes': ['human', 'rgb_array']}

  def __init__(self, arg1, arg2, ...):
    super().__init__()
    # The action and observation spaces need to be gym.spaces objects:
    self.action_space = Discrete(4)  # up, left, right, down
    # Here's an observation space for 200 wide x 100 high RGB image inputs:
    self.observation_space = Box(
        low=0, high=255, shape=(100, 200, 3), dtype=np.uint8)

  def step(self, action):
    # Execute one time step in the environment
    ...

  def reset(self):
    # Reset the state of the environment
    ...

  def render(self, mode='human', close=False):
    if mode == 'human':
        # render to screen
        ...
    elif mode == 'rgb_array':
        # render to a NumPy 100x200x3 array
        ...
        return array
    else:
        # raise an error, unsupported mode

No need to implement your own Atari ROM or anything like that. It can be pure Python (+ NumPy arrays). To produce NumPy arrays representing images, there are many options. For example, you could use Matplotlib, PIL, or other libraries. Or, if the images are very simple, you could even create the NumPy arrays manually. For example, the following code creates a random 200x100 RGB image:

import numpy as np

random_image = np.random.rand(100, 200, 3)