def test_benchmark():
env = ChessEnvV1(opponent="none", log=False)
num_episodes = 10
num_steps = 50
total_steps = 0
start = time.time()
for e in range(num_episodes):
env.reset()
for step in range(num_steps):
total_steps += 1
actions = env.get_possible_actions()
if not actions:
break
idx = np.random.choice(np.arange(len(actions)))
action = actions[idx]
state, reward, done, info = env.step(action)
if done:
break
end = time.time()
diff = end - start
print("Total time (s)", diff)
print("Total episodes", num_episodes)
print("Total steps", total_steps)
print("Time per episode (s)", diff / num_episodes)
print("Time per step (s)", diff / total_steps)
# assert that it's less than 100 seconds
assert diff < 100
def test_queen_moves():
BOARD = copy(BASIC_BOARD)
BOARD[4, 4] = QUEEN_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves(attack=True, skip_pawns=True)
env.render_moves(moves)
expected_attacks = set([
(6, 2),
(4, 0),
(5, 5),
(3, 4),
(4, 3),
(5, 4),
(4, 6),
(6, 6),
(6, 4),
(4, 2),
(4, 5),
(3, 3),
(5, 3),
(4, 1),
(4, 7),
(3, 5),
])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
def test_queen_side_castle():
BOARD = copy(BASIC_BOARD)
BOARD[7, 0] = ROOK_ID
BOARD[7, 4] = KING_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.castle_moves(env.current_player)
env.render_moves(moves)
assert moves == [CASTLE_QUEEN_SIDE_WHITE]
def test_pawn_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[6, 0] = PAWN_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
expected_attacks = set([(4, 0), (5, 0)])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
def test_knight_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[4, 4] = KNIGHT_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
expected_attacks = set([(6, 5), (2, 3), (6, 3), (5, 6), (3, 6), (3, 2),
(2, 5), (5, 2)])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
def test_king_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[7, 7] = 0
BOARD[4, 4] = KING_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
expected_attacks = set([(5, 5), (3, 4), (4, 3), (5, 4), (4, 5), (3, 3),
(5, 3), (3, 5)])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
def test_king_moves_2():
BOARD = copy(BASIC_BOARD)
BOARD[3, 4] = KING_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
king_is_checked = env.king_is_checked()
expected_attacks = set([(2, 4), (4, 3), (2, 3), (4, 5), (2, 5)])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
assert not king_is_checked
def test_king_has_moved_castling():
BOARD = copy(BASIC_BOARD)
BOARD[7, 0] = ROOK_ID
BOARD[7, 4] = KING_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
king_moves = env.king_moves(env.current_player, np.array([7, 4]))
action = env.move_to_action(king_moves[0])
env.step(action)
env.render()
env.current_player = "white"
moves = env.castle_moves(env.current_player)
env.render_moves(moves)
assert moves == []
def test_attacked_square_castling_path():
BOARD = copy(BASIC_BOARD)
BOARD[0, 2] = -ROOK_ID
BOARD[6, 2] = 0
BOARD[7, 0] = ROOK_ID
BOARD[7, 4] = KING_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
opponent = env.get_other_player(env.current_player)
moves = env.get_possible_moves(player=opponent, attack=True)
env.render_moves(moves)
squares_under_attack = env.get_squares_attacked_by_player(
env.state, opponent)
squares_under_attack_hashmap = defaultdict(lambda: None)
for sq in squares_under_attack:
squares_under_attack_hashmap[tuple(sq)] = True
moves = env.castle_moves(
env.current_player,
squares_under_attack_hashmap=squares_under_attack_hashmap)
env.render_moves(moves)
assert moves == []
def test_pawn_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[6, 0] = PAWN_ID
BOARD[1, 0] = -PAWN_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
# player_1
actions = env.get_possible_actions()
env.step(actions[0])
# player_2
actions = env.get_possible_actions()
env.step(actions[0])
# player_3
actions = env.get_possible_actions()
env.step(actions[0])
# player_4
actions = env.get_possible_actions()
env.step(actions[0])
EXPECTED_BOARD = copy(BASIC_BOARD)
EXPECTED_BOARD[4, 0] = PAWN_ID
EXPECTED_BOARD[3, 0] = -PAWN_ID
assert (env.state == EXPECTED_BOARD).all()
def test_bishop_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[4, 4] = BISHOP_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
expected_attacks = set([
(6, 2),
(5, 5),
(7, 1),
(7, 7),
(0, 0),
(1, 1),
(6, 6),
(1, 7),
(3, 3),
(2, 6),
(2, 2),
(5, 3),
(3, 5),
])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
def test_rook_basic_moves():
BOARD = copy(BASIC_BOARD)
BOARD[4, 4] = ROOK_ID
env = ChessEnvV1(opponent="none", initial_state=BOARD)
moves = env.get_possible_moves()
env.render_moves(moves)
expected_attacks = set([
(7, 4),
(2, 4),
(4, 0),
(0, 4),
(3, 4),
(4, 3),
(5, 4),
(4, 6),
(6, 4),
(1, 4),
(4, 2),
(4, 5),
(4, 1),
(4, 7),
])
squares_attacked = set([tuple(move[1]) for move in moves])
assert squares_attacked == expected_attacks
