mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-26 16:27:11 +00:00
195 lines
4.5 KiB
Python
195 lines
4.5 KiB
Python
import numpy as np
|
|
import gym
|
|
from gym import spaces
|
|
from gym.utils import seeding
|
|
|
|
# Unit test environment for CNNs.
|
|
# Looks like this (RGB observations):
|
|
#
|
|
# ---------------------------
|
|
# | |
|
|
# | ****** |
|
|
# | ****** |
|
|
# | ** ** |
|
|
# | ** ** |
|
|
# | ** |
|
|
# | ** |
|
|
# | **** |
|
|
# | **** |
|
|
# | **** |
|
|
# | **** |
|
|
# | ********** |
|
|
# | ********** |
|
|
# | |
|
|
# ---------------------------
|
|
#
|
|
# Agent should hit action 2 to gain reward. Catches off-by-one errors in your agent.
|
|
#
|
|
# To see how it works, run:
|
|
#
|
|
# python examples/agents/keyboard_agent.py MemorizeDigits-v0
|
|
|
|
FIELD_W = 32
|
|
FIELD_H = 24
|
|
|
|
bogus_mnist = \
|
|
[[
|
|
" **** ",
|
|
"* *",
|
|
"* *",
|
|
"* *",
|
|
"* *",
|
|
" **** "
|
|
], [
|
|
" ** ",
|
|
" * * ",
|
|
" * ",
|
|
" * ",
|
|
" * ",
|
|
" *** "
|
|
], [
|
|
" **** ",
|
|
"* *",
|
|
" *",
|
|
" *** ",
|
|
"** ",
|
|
"******"
|
|
], [
|
|
" **** ",
|
|
"* *",
|
|
" ** ",
|
|
" *",
|
|
"* *",
|
|
" **** "
|
|
], [
|
|
" * * ",
|
|
" * * ",
|
|
" * * ",
|
|
" **** ",
|
|
" * ",
|
|
" * "
|
|
], [
|
|
" **** ",
|
|
" * ",
|
|
" **** ",
|
|
" * ",
|
|
" * ",
|
|
" **** "
|
|
], [
|
|
" *** ",
|
|
" * ",
|
|
" **** ",
|
|
" * * ",
|
|
" * * ",
|
|
" **** "
|
|
], [
|
|
" **** ",
|
|
" * ",
|
|
" * ",
|
|
" * ",
|
|
" * ",
|
|
" * "
|
|
], [
|
|
" **** ",
|
|
"* *",
|
|
" **** ",
|
|
"* *",
|
|
"* *",
|
|
" **** "
|
|
], [
|
|
" **** ",
|
|
"* *",
|
|
"* *",
|
|
" *****",
|
|
" *",
|
|
" **** "
|
|
]]
|
|
|
|
color_black = np.array((0,0,0)).astype('float32')
|
|
color_white = np.array((255,255,255)).astype('float32')
|
|
|
|
class MemorizeDigits(gym.Env):
|
|
metadata = {
|
|
'render.modes': ['human', 'rgb_array'],
|
|
'video.frames_per_second' : 60,
|
|
'video.res_w' : FIELD_W,
|
|
'video.res_h' : FIELD_H,
|
|
}
|
|
|
|
use_random_colors = False
|
|
|
|
def __init__(self):
|
|
self.seed()
|
|
self.viewer = None
|
|
self.observation_space = spaces.Box(0, 255, (FIELD_H,FIELD_W,3), dtype=np.uint8)
|
|
self.action_space = spaces.Discrete(10)
|
|
self.bogus_mnist = np.zeros( (10,6,6), dtype=np.uint8 )
|
|
for digit in range(10):
|
|
for y in range(6):
|
|
self.bogus_mnist[digit,y,:] = [ord(char) for char in bogus_mnist[digit][y]]
|
|
self.reset()
|
|
|
|
def seed(self, seed=None):
|
|
self.np_random, seed = seeding.np_random(seed)
|
|
return [seed]
|
|
|
|
def random_color(self):
|
|
return np.array([
|
|
self.np_random.randint(low=0, high=255),
|
|
self.np_random.randint(low=0, high=255),
|
|
self.np_random.randint(low=0, high=255),
|
|
]).astype('uint8')
|
|
|
|
def reset(self):
|
|
self.digit_x = self.np_random.randint(low=FIELD_W//5, high=FIELD_W//5*4)
|
|
self.digit_y = self.np_random.randint(low=FIELD_H//5, high=FIELD_H//5*4)
|
|
self.color_bg = self.random_color() if self.use_random_colors else color_black
|
|
self.step_n = 0
|
|
while 1:
|
|
self.color_digit = self.random_color() if self.use_random_colors else color_white
|
|
if np.linalg.norm(self.color_digit - self.color_bg) < 50: continue
|
|
break
|
|
self.digit = -1
|
|
return self.step(0)[0]
|
|
|
|
def step(self, action):
|
|
reward = -1
|
|
done = False
|
|
self.step_n += 1
|
|
if self.digit==-1:
|
|
pass
|
|
else:
|
|
if self.digit==action:
|
|
reward = +1
|
|
done = self.step_n > 20 and 0==self.np_random.randint(low=0, high=5)
|
|
self.digit = self.np_random.randint(low=0, high=10)
|
|
obs = np.zeros( (FIELD_H,FIELD_W,3), dtype=np.uint8 )
|
|
obs[:,:,:] = self.color_bg
|
|
digit_img = np.zeros( (6,6,3), dtype=np.uint8 )
|
|
digit_img[:] = self.color_bg
|
|
xxx = self.bogus_mnist[self.digit]==42
|
|
digit_img[xxx] = self.color_digit
|
|
obs[self.digit_y-3:self.digit_y+3, self.digit_x-3:self.digit_x+3] = digit_img
|
|
self.last_obs = obs
|
|
return obs, reward, done, {}
|
|
|
|
def render(self, mode='human'):
|
|
if mode == 'rgb_array':
|
|
return self.last_obs
|
|
|
|
elif mode == 'human':
|
|
from gym.envs.classic_control import rendering
|
|
if self.viewer is None:
|
|
self.viewer = rendering.SimpleImageViewer()
|
|
self.viewer.imshow(self.last_obs)
|
|
return self.viewer.isopen
|
|
|
|
else:
|
|
assert 0, "Render mode '%s' is not supported" % mode
|
|
|
|
def close(self):
|
|
if self.viewer is not None:
|
|
self.viewer.close()
|
|
self.viewer = None
|
|
|