File size: 3,850 Bytes
c80bb82 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | 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()
|