Spaces:
Running
Running
| from __future__ import annotations | |
| import math | |
| import numpy as np | |
| import torch | |
| from game import canonical, is_terminal, legal_actions, play, winner | |
| from model import MicroZeroNet | |
| class NeuralMCTS: | |
| def __init__( | |
| self, | |
| model: MicroZeroNet, | |
| simulations: int, | |
| exploration: float = 1.35, | |
| ) -> None: | |
| self.model = model | |
| self.simulations = simulations | |
| self.exploration = exploration | |
| self.priors: dict[tuple, np.ndarray] = {} | |
| self.state_visits: dict[tuple, int] = {} | |
| self.edge_visits: dict[tuple, int] = {} | |
| self.edge_value: dict[tuple, float] = {} | |
| def key(board: np.ndarray, player: int) -> tuple: | |
| return (*board.tolist(), player) | |
| def expand(self, board: np.ndarray, player: int) -> float: | |
| key = self.key(board, player) | |
| logits, value = self.model(torch.from_numpy(canonical(board, player))[None]) | |
| logits = logits[0].numpy() | |
| legal = legal_actions(board) | |
| masked = np.full(9, -1e9, dtype=np.float64) | |
| masked[legal] = logits[legal] | |
| probabilities = np.exp(masked - np.max(masked)) | |
| probabilities /= probabilities.sum() | |
| self.priors[key] = probabilities | |
| self.state_visits[key] = 0 | |
| return float(value[0]) | |
| def search(self, board: np.ndarray, player: int) -> float: | |
| result = winner(board) | |
| if result: | |
| return 1.0 if result == player else -1.0 | |
| if is_terminal(board): | |
| return 0.0 | |
| state_key = self.key(board, player) | |
| if state_key not in self.priors: | |
| return self.expand(board, player) | |
| best_score = -float("inf") | |
| best_action = -1 | |
| root_visits = max(1, self.state_visits[state_key]) | |
| for action in legal_actions(board): | |
| edge = (*state_key, int(action)) | |
| visits = self.edge_visits.get(edge, 0) | |
| value = self.edge_value.get(edge, 0.0) | |
| prior = self.priors[state_key][action] | |
| score = value + self.exploration * prior * math.sqrt(root_visits) / ( | |
| 1 + visits | |
| ) | |
| if score > best_score: | |
| best_score = score | |
| best_action = int(action) | |
| child = play(board, best_action, player) | |
| value = -self.search(child, -player) | |
| edge = (*state_key, best_action) | |
| visits = self.edge_visits.get(edge, 0) + 1 | |
| previous = self.edge_value.get(edge, 0.0) | |
| self.edge_visits[edge] = visits | |
| self.edge_value[edge] = previous + (value - previous) / visits | |
| self.state_visits[state_key] += 1 | |
| return value | |
| def policy( | |
| self, | |
| board: np.ndarray, | |
| player: int, | |
| *, | |
| temperature: float, | |
| add_noise: bool, | |
| rng: np.random.Generator, | |
| ) -> np.ndarray: | |
| key = self.key(board, player) | |
| if key not in self.priors: | |
| self.expand(board, player) | |
| if add_noise: | |
| legal = legal_actions(board) | |
| noise = rng.dirichlet(np.full(len(legal), 0.30)) | |
| priors = self.priors[key].copy() | |
| priors[legal] = 0.75 * priors[legal] + 0.25 * noise | |
| self.priors[key] = priors | |
| for _ in range(self.simulations): | |
| self.search(board, player) | |
| visits = np.zeros(9, dtype=np.float64) | |
| for action in legal_actions(board): | |
| visits[action] = self.edge_visits.get((*key, int(action)), 0) | |
| if temperature <= 1e-6: | |
| policy = np.zeros(9, dtype=np.float64) | |
| policy[int(np.argmax(visits))] = 1.0 | |
| return policy | |
| adjusted = visits ** (1 / temperature) | |
| if adjusted.sum() == 0: | |
| adjusted[legal_actions(board)] = 1 | |
| return adjusted / adjusted.sum() | |