|
import gym |
|
from gym import spaces |
|
from tetris_gym.utils.board_utils import get_heights, get_bumps_from_heights |
|
from agent.utils import calc_holes_array |
|
|
|
import numpy as np |
|
|
|
|
|
class CustomObsWrapper(gym.ObservationWrapper): |
|
def __init__(self, env): |
|
super().__init__(env) |
|
self.observation_space = spaces.Dict({ |
|
"board": env.observation_space["board"], |
|
"piece": env.observation_space["piece"], |
|
"holes_list": spaces.Box( |
|
low=1, |
|
high=env.height, |
|
shape=(env.width,), |
|
dtype=np.uint8, |
|
), |
|
"x": spaces.Discrete(env.width), |
|
"y": spaces.Discrete(env.width), |
|
"piece_shape": spaces.Box( |
|
low=0, |
|
high=1, |
|
shape=(4, 4), |
|
dtype=np.uint8, |
|
), |
|
"empty_above": spaces.Box( |
|
low=0, |
|
high=env.height, |
|
shape=(env.width,), |
|
dtype=np.uint8, |
|
), |
|
"heights": spaces.Box( |
|
low=0, |
|
high=env.height, |
|
shape=(env.width,), |
|
dtype=np.uint8, |
|
), |
|
"bumps": spaces.Box( |
|
low=0, |
|
high=env.height, |
|
shape=(env.width - 1,), |
|
dtype=int, |
|
) |
|
}) |
|
|
|
def observation(self, obs): |
|
board = obs["board"] |
|
piece = obs["piece"] |
|
|
|
heights = get_heights(board) |
|
bumps = get_bumps_from_heights(heights) |
|
holes_array = calc_holes_array(self, board, heights) |
|
empty_above = np.max(heights) - heights[:] |
|
piece_shape = np.zeros((4, 4), dtype=np.uint8) |
|
piece_shape[:len(self.piece), :len(self.piece[0])] = self.piece[:] |
|
|
|
|
|
|
|
obs = { |
|
"board": board, |
|
"x": self.current_pos["x"], |
|
"y": self.current_pos["y"], |
|
"piece_shape": piece_shape, |
|
"piece": piece, |
|
"empty_above": empty_above, |
|
"holes_list": holes_array, |
|
"heights": heights, |
|
"bumps": bumps |
|
} |
|
|
|
return obs |
|
|