| 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] = {} |
|
|
| @staticmethod |
| def key(board: np.ndarray, player: int) -> tuple: |
| return (*board.tolist(), player) |
|
|
| @torch.inference_mode() |
| 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() |
|
|
