| from scipy.misc import imresize |
| import gym |
| import random |
| import numpy as np |
| from queue import Queue |
|
|
| from matplotlib import pyplot as plt |
| from PIL import Image |
|
|
| class CatchEnv: |
| def __init__(self): |
| self.size = 21 |
| self.image = np.zeros((self.size, self.size)) |
| self.state = [] |
| self.fps = 4 |
| self.output_shape = (84, 84) |
|
|
| def reset_random(self): |
| self.image.fill(0) |
| self.pos = np.random.randint(2, self.size-2) |
| self.vx = np.random.randint(5) - 2 |
| self.vy = 1 |
| self.ballx, self.bally = np.random.randint(self.size), 4 |
| self.image[self.bally, self.ballx] = 1 |
| self.image[-5, self.pos - 2:self.pos + 3] = np.ones(5) |
|
|
| return self.step(2)[0] |
|
|
|
|
| def step(self, action): |
| def left(): |
| if self.pos > 3: |
| self.pos -= 2 |
| def right(): |
| if self.pos < 17: |
| self.pos += 2 |
| def noop(): |
| pass |
| {0: left, 1: right, 2: noop}[action]() |
|
|
| |
| self.image[self.bally, self.ballx] = 0 |
| self.ballx += self.vx |
| self.bally += self.vy |
| if self.ballx > self.size - 1: |
| self.ballx -= 2 * (self.ballx - (self.size-1)) |
| self.vx *= -1 |
| elif self.ballx < 0: |
| self.ballx += 2 * (0 - self.ballx) |
| self.vx *= -1 |
| self.image[self.bally, self.ballx] = 1 |
|
|
| self.image[-5].fill(0) |
| self.image[-5, self.pos-2:self.pos+3] = np.ones(5) |
| |
| terminal = self.bally == self.size - 1 - 4 |
| reward = int(self.pos - 2 <= self.ballx <= self.pos + 2) if terminal else 0 |
|
|
| [self.state.append(imresize(self.image, (84, 84))) for _ in range(self.fps - len(self.state) + 1)] |
| self.state = self.state[-self.fps:] |
|
|
| return np.transpose(self.state, [1, 2, 0]), reward, terminal |
|
|
| def get_num_actions(self): |
| return 3 |
|
|
| def reset(self): |
| return self.reset_random() |
|
|
| def state_shape(self): |
| return (self.fps,) + self.output_shape |
|
|
|
|
| def test(): |
| env = CatchEnv() |
| i = 0 |
|
|
| for ep in range(1): |
| env.reset() |
| |
| state, reward, terminal = env.step(1) |
|
|
| while not terminal: |
| state, reward, terminal = env.step(random.randint(0,2)) |
|
|
| state = np.squeeze(state) |
| |
| |
| |
| i += 1 |
|
|
| if __name__ == "__main__": |
| test() |
|
|
