| import gym |
| import math |
| import numpy as np |
|
|
| class MyWrapper(gym.Wrapper): |
| def __init__(self, env): |
| super().__init__(env) |
| self.prevBoardState = np.zeros((self.env.height, self.env.width)) |
|
|
| def check_cleared_rows(self, board): |
| to_delete = [] |
| for i, row in enumerate(board[::-1]): |
| if 0 not in row: |
| to_delete.append(len(board) - 1 - i) |
| if len(to_delete) > 0: |
| board = self.remove_row(board, to_delete) |
| return len(to_delete) |
|
|
|
|
| def get_bumpiness_and_height(self, board): |
| board = np.array(board) |
| mask = board != 0 |
| invert_heights = np.where(mask.any(axis=0), np.argmax(mask, axis=0), self.height) |
| heights = self.height - invert_heights |
| total_height = np.sum(heights) |
| currs = heights[:-1] |
| nexts = heights[1:] |
| diffs = np.abs(currs - nexts) |
| total_bumpiness = np.sum(diffs) |
| return total_bumpiness, total_height, np.max(heights) |
|
|
| def get_holes(self, board): |
| num_holes = 0 |
| for col in zip(*board): |
| row = 0 |
| while row < self.height and col[row] == 0: |
| row += 1 |
| num_holes += len([x for x in col[row + 1:] if x == 0]) |
| return num_holes |
|
|
|
|
| def step(self, action): |
| |
| |
|
|
| obs, reward, done, info = self.env.step(action) |
| |
|
|
| currBoard = obs['board'].copy() |
| cleared_lines = self.check_cleared_rows(currBoard) |
| bumpiness, total_height, max_height = self.get_bumpiness_and_height(currBoard) |
| num_holes = self.get_holes(currBoard) |
|
|
| reward = 1 + math.pow(cleared_lines, 4) * obs['board'].shape[1] - 0.28 * bumpiness - 0.36 * num_holes |
|
|
| if done: |
| reward -= 1000 |
|
|
| self.prevBoardState = obs['board'].copy() |
|
|
| return obs, reward, done, info |
| |
|
|
|
|
