| import random |
| import numpy as np |
| from scipy.misc import imresize |
|
|
| import random |
| import numpy as np |
| from scipy.ndimage import rotate |
| from scipy.misc import imresize |
| from matplotlib import pyplot as plt |
|
|
| class CatchEnv2: |
| 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) |
| |
| for i in range(0, self.size): |
| for j in range(0, self.size): |
| self.image[i][j] = random.randint(2,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(random.randint(2,5)) |
| 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 show_state(self, i): |
| plt.imshow(self.image) |
| plt.imsave('image_'+str(i)+'.jpg', self.image) |
|
|
| def test(): |
| env = CatchEnv2() |
| i = 0 |
| for ep in range(1): |
| env.reset() |
| env.show_state(i) |
| |
| state, reward, terminal = env.step(1) |
| while not terminal: |
| env.show_state(i) |
| state, reward, terminal = env.step(np.random.randint(0,2)) |
| i += 1 |
| |
|
|
| if __name__ == "main": |
| test() |
|
|
