| import math |
| import torch |
| import numpy as np |
|
|
| class Node: |
|
|
| def __init__(self, game, args, state, parent = None, action = None, prob = 0, visits = 0): |
| self.game = game |
| self.args = args |
| self.state = state |
| self.parent = parent |
| self.action = action |
| self.prob = prob |
|
|
| self.children = [] |
| self.visits = visits |
| self.value = 0 |
|
|
| def leaf_or_not(self): |
| return len(self.children) > 0 |
|
|
| def search(self): |
| best_child = None |
| best_ucb = -np.inf |
| for child in self.children: |
| ucb = self.get_ucb(child) |
| if best_ucb < ucb: |
| best_ucb = ucb |
| best_child = child |
| return best_child |
|
|
| def get_ucb(self, child): |
| if child.visits == 0: |
| q_value = 0 |
| else: |
| q_value = 1 - (((child.value / child.visits) + 1) / 2) |
|
|
| return q_value + self.args['EXPLORATION_CONSTANT'] * (math.sqrt(self.visits) / (child.visits + 1)) * child.prob |
|
|
| def expand(self, policy): |
|
|
| for move, prob in enumerate(policy): |
| if prob > 0: |
| child = self.state.copy() |
| child = self.game.make_move(child, move, 1) |
| if self.args["ADVERSARIAL"]: |
| child = self.game.change_perspective(child, player = -1) |
|
|
| child = Node(self.game, self.args, child, self, move, prob) |
| self.children.append(child) |
|
|
| def backpropagate(self,state_value): |
| self.value += state_value |
| self.visits += 1 |
| if self.args["ADVERSARIAL"]: |
| state_value = self.game.get_opponent_value(state_value) |
| if self.parent is not None: |
| self.parent.backpropagate(state_value) |
|
|
|
|
| class Alpha_MCTS: |
|
|
| def __init__(self, game, args, model): |
| self.game = game |
| self.args = args |
| self.model = model |
|
|
| @torch.no_grad() |
| def search(self, state): |
| root = Node(self.game, self.args, state, visits = 1) |
|
|
| if self.args["ROOT_RANDOMNESS"]: |
| policy, _ = self.model( |
| torch.tensor(self.game.get_encoded_state(state), device = self.model.device |
| ).unsqueeze(0)) |
|
|
| policy = torch.softmax(policy, axis = 1).squeeze(0).cpu().numpy() |
|
|
| policy = (1 - self.args["DIRICHLET_EPSILON"]) * policy + self.args["DIRICHLET_EPSILON"] * np.random.dirichlet([self.args["DIRICHLET_ALPHA"]] * self.game.possible_state) |
|
|
| valid_state = self.game.get_valid_moves(state) |
| policy *= valid_state |
| policy /= np.sum(policy) |
| root.expand(policy) |
|
|
| for _ in range(self.args["NO_OF_SEARCHES"]): |
| node = root |
| no_moves = 0 |
| while node.leaf_or_not(): |
| node = node.search() |
| no_moves += 1 |
|
|
| is_terminal, value = self.game.know_terminal_value(node.state, node.action) |
|
|
| if self.args["ADVERSARIAL"]: |
| value = self.game.get_opponent_value(value) |
|
|
| if not is_terminal: |
|
|
| policy, value = self.model( |
| torch.tensor(self.game.get_encoded_state(node.state), device = self.model.device).unsqueeze(0) |
| ) |
|
|
| valid_state = self.game.get_valid_moves(node.state) |
| policy = torch.softmax(policy, axis = 1).squeeze(0).cpu().numpy().astype(np.float64) |
|
|
| policy *= valid_state |
| policy /= np.sum(policy) |
|
|
| value = value.item() |
|
|
| node.expand(policy) |
|
|
| node.backpropagate(value) |
|
|
| move_probability = np.zeros(self.game.possible_state) |
| for children in root.children: |
| move_probability[children.action] = children.visits |
| move_probability /= np.sum(move_probability) |
| return move_probability |
|
|
