delete breakthrough_mcvs.py

breakthrough_mcvs.py

@@ -1,1054 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.nn.functional as F
-from torch.utils.data import DataLoader, TensorDataset
-from typing import List, Dict, Tuple
-import random
-import time
-import os
-import bisect
-
-from abc_model import ABCModelDynamic
-from matrix_model import WeightedMatrixABC
-
-# ============================================================================
-# HILBERT CURVE UTILITIES
-# ============================================================================
-
-def xy2d(n: int, x: int, y: int) -> int:
-    """Convert (x,y) to Hilbert curve distance."""
-    d = 0
-    s = n // 2
-    while s > 0:
-        rx = 1 if (x & s) > 0 else 0
-        ry = 1 if (y & s) > 0 else 0
-        d += s * s * ((3 * rx) ^ ry)
-        if ry == 0:
-            if rx == 1:
-                x = n - 1 - x
-                y = n - 1 - y
-            x, y = y, x
-        s //= 2
-    return d
-
-def matrix_to_hilbert_index(W: np.ndarray) -> int:
-    """
-    Convert 64x64 matrix to Hilbert curve index.
-    """
-    if W.size == 0:
-        return 0
-    
-    # Find location of maximum value
-    W_flat = W.flatten()
-    max_idx = np.argmax(W_flat)
-    
-    x = max_idx % 64
-    y = max_idx // 64
-    idx = xy2d(64, x, y)
-    # Sanity check: index must be in [0, 64*64-1]
-    assert 0 <= idx < 64 * 64, f"Hilbert index out of range: {idx}"
-    return idx
-
-
-def move_to_index(fr: int, fc: int, tr: int, tc: int) -> int:
-    """Convert a move tuple (fr,fc,tr,tc) to a policy index in [0,4095].
-
-    Mapping: idx = (from_sq) * 64 + to_sq, where from_sq = fr*8 + fc, to_sq = tr*8 + tc
-    """
-    from_sq = fr * 8 + fc
-    to_sq = tr * 8 + tc
-    return from_sq * 64 + to_sq
-
-# ============================================================================
-# BREAKTHROUGH GAME LOGIC 
-# ============================================================================
-
-class Breakthrough:
-    ROWS, COLS = 8, 8
-    EMPTY, PLAYER1, PLAYER2 = 0, 1, 2
-
-    def __init__(self):
-        self.board = np.zeros((self.ROWS, self.COLS), dtype=np.int32)
-        self.board[0:2, :] = self.PLAYER1
-        self.board[6:8, :] = self.PLAYER2
-        self.move_count = 0
-        self.abc_model = ABCModelDynamic(n=2, t=1.0, T=1.41)
-        self._cached_matrix = None
-        self._setup_abc_initial()
-
-    def _setup_abc_initial(self):
-        """Initialize ABC model with starting pieces."""
-        for row in range(2):
-            for col in range(self.COLS):
-                x = col - 3.5
-                y = 3.5 - row
-                self.abc_model.add_piece((x, y), (1.0, 1.0, 1.0))
-
-        for row in range(6, 8):
-            for col in range(self.COLS):
-                x = col - 3.5
-                y = 3.5 - row
-                self.abc_model.add_piece((x, y), (1.0, 1.1, 1.0))
-
-    def copy(self):
-        """Create a deep copy of the game state."""
-        new_game = Breakthrough()
-        new_game.board = self.board.copy()
-        new_game.move_count = self.move_count
-        new_game.abc_model = ABCModelDynamic(n=2, t=1.0, T=1.41)
-        
-        for entry in self.abc_model.history:
-            pos = entry['SB_position']
-            # Preserve kappa (tokenized vector) if present
-            kappa_vec = entry.get('kappa', None)
-            new_game.abc_model.add_piece(pos, entry['delta'], kappa_vector=kappa_vec)
-        
-        new_game._cached_matrix = self._cached_matrix.copy() if self._cached_matrix is not None else None
-        return new_game
-
-    def get_legal_moves(self) -> List[Tuple[int, int, int, int]]:
-        """Get all legal moves for the current player."""
-        moves = []
-        player = self.PLAYER1 if self.move_count % 2 == 0 else self.PLAYER2
-        direction = 1 if player == self.PLAYER1 else -1
-
-        for r in range(self.ROWS):
-            for c in range(self.COLS):
-                if self.board[r, c] == player:
-                    # Forward move
-                    nr = r + direction
-                    if 0 <= nr < self.ROWS and self.board[nr, c] == self.EMPTY:
-                        moves.append((r, c, nr, c))
-                    
-                    # Diagonal captures
-                    for dc in [-1, 1]:
-                        nc = c + dc
-                        nr = r + direction
-                        opponent = 3 - player
-                        if 0 <= nr < self.ROWS and 0 <= nc < self.COLS and self.board[nr, nc] == opponent:
-                            moves.append((r, c, nr, nc))
-        
-        return moves
-
-    def apply_move(self, move: Tuple[int, int, int, int]) -> None:
-        """Apply a move and update the ABC model."""
-        fr, fc, tr, tc = move
-        player = self.board[fr, fc]
-        self.board[tr, tc] = player
-        self.board[fr, fc] = self.EMPTY
-        self.move_count += 1
-        
-        x = tc - 3.5
-        y = 3.5 - tr 
-        delta_2 = 1.0 if player == self.PLAYER1 else 1.1
-        
-        self.abc_model.add_piece((x, y), (1.0, delta_2, 1.0))
-        self._cached_matrix = None
-
-    def check_winner(self) -> int:
-        """Check if there's a winner."""
-        if np.any(self.board[0, :] == self.PLAYER2):
-            return self.PLAYER2
-        if np.any(self.board[7, :] == self.PLAYER1):
-            return self.PLAYER1
-        return self.EMPTY
-
-    def is_terminal(self) -> bool:
-        """Check if the game has ended."""
-        return self.check_winner() != self.EMPTY or len(self.get_legal_moves()) == 0
-
-    def get_weighted_adjacency_matrix(self) -> np.ndarray:
-        if self._cached_matrix is not None:
-            return self._cached_matrix
-
-        occupied = []
-        for r in range(8):
-            for c in range(8):
-                piece = self.board[r, c]
-                if piece != self.EMPTY:
-                    x = c - 3.5
-                    y = 3.5 - r
-                    delta_2 = 1.0 if piece == self.PLAYER1 else 1.1
-                    kappa = (1.0, delta_2, 1.0)  # δ1=1.0 fixed, δ2=color, δ3=1.0 (no piece-type variation)
-                    hilbert_d = xy2d(8, c, r)  # file=c, rank=r
-                    occupied.append((hilbert_d, (x, y), kappa))
-
-        occupied.sort(key=lambda t: t[0])  # Hilbert order
-
-        abc = ABCModelDynamic(n=2, t=1.0, T=1.41)
-        for _, pos, kappa_vec in occupied:
-            abc.add_piece(pos, delta_values=kappa_vec)
-
-        builder = WeightedMatrixABC(abc, sigma=1.0)
-        W_small = builder.compute_weighted_matrix()
-
-        n = len(occupied)
-        W = np.zeros((64, 64), dtype=np.float32)
-        if n > 0:
-            W[:n, :n] = W_small
-
-        self._cached_matrix = W
-        return W
-
-# ============================================================================
-# NEURAL NETWORKS
-# ============================================================================
-
-class PolicyNetworkCNN(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.relu = nn.ReLU()
-        self.pool = nn.AdaptiveAvgPool2d(8)
-        self.fc = nn.Linear(64 * 8 * 8, 4096)
-
-    def forward(self, Wt):
-        x = self.relu(self.conv1(Wt))
-        x = self.relu(self.conv2(x))
-        x = self.pool(x)
-        x = x.view(x.size(0), -1)
-        logits = self.fc(x)
-        return logits
-
-
-class ValueNetworkCNN(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.relu = nn.ReLU()
-        self.pool = nn.AdaptiveAvgPool2d(4)
-        self.fc1 = nn.Linear(64 * 4 * 4, 64)
-        self.fc2 = nn.Linear(64, 1)
-        self.tanh = nn.Tanh()
-
-    def forward(self, Wt):
-        x = self.relu(self.conv1(Wt))
-        x = self.relu(self.conv2(x))
-        x = self.pool(x)
-        x = x.view(x.size(0), -1)
-        x = self.relu(self.fc1(x))
-        value = self.tanh(self.fc2(x))
-        return value
-
-
-# ============================================================================
-# ZONE DATABASE
-# ============================================================================
-
-class HilbertOrderedZoneDatabase:
-    def __init__(self, filepath: str = "breakthrough_zone_db.npz", max_size: int = 10000):
-        self.filepath = filepath
-        self.max_size = max_size  # Maximum matrices per zone
-        self.winning_matrices = []
-        self.winning_indices = []
-        self.losing_matrices = []
-        self.losing_indices = []
-        self.draw_matrices = []
-        self.draw_indices = []
-        self.load()
-    
-    def add_winning_matrix(self, W: np.ndarray):
-        if W.shape != (64, 64):
-            return
-        
-        hilbert_idx = matrix_to_hilbert_index(W)
-        pos = bisect.bisect_left(self.winning_indices, hilbert_idx)
-        self.winning_indices.insert(pos, hilbert_idx)
-        self.winning_matrices.insert(pos, W.copy())
-    
-    def add_losing_matrix(self, W: np.ndarray):
-        if W.shape != (64, 64):
-            return
-        
-        hilbert_idx = matrix_to_hilbert_index(W)
-        pos = bisect.bisect_left(self.losing_indices, hilbert_idx)
-        self.losing_indices.insert(pos, hilbert_idx)
-        self.losing_matrices.insert(pos, W.copy()) 
-    
-    def add_draw_matrix(self, W: np.ndarray):
-        if W.shape != (64, 64):
-            return
-        
-        hilbert_idx = matrix_to_hilbert_index(W)
-        pos = bisect.bisect_left(self.draw_indices, hilbert_idx)
-        self.draw_indices.insert(pos, hilbert_idx)
-        self.draw_matrices.insert(pos, W.copy())
-    
-    def add_game_record(self, trajectory: List[Breakthrough], result: int, sample_rate: float = 0.3):
-        """Add game record with sampling to control database size."""
-        # Sample states: always keep first and last, random subset of middle
-        sampled_states = []
-        
-        for i, state in enumerate(trajectory):
-            if i == 0 or i == len(trajectory) - 1:
-                # Always keep start and end positions
-                sampled_states.append(state)
-            elif random.random() < sample_rate:
-                # Randomly sample middle positions
-                sampled_states.append(state)
-        
-        for state in sampled_states:
-            W = state.get_weighted_adjacency_matrix()
-            
-            if result == 1:
-                self.add_winning_matrix(W)
-            elif result == -1:
-                self.add_losing_matrix(W)
-            else:
-                self.add_draw_matrix(W)
-    
-    def save(self):
-        """Save database with pruning if too large."""
-        # Prune if any zone exceeds 80% of max_size
-        threshold = int(self.max_size * 0.8)
-        
-        if (len(self.winning_matrices) > threshold or 
-            len(self.losing_matrices) > threshold or 
-            len(self.draw_matrices) > threshold):
-            print(f"Pruning database before save...")
-            self.prune_database(target_size=int(self.max_size * 0.7))
-        
-        try:
-            np.savez_compressed(
-                self.filepath,
-                winning=np.array(self.winning_matrices, dtype=object),
-                winning_indices=np.array(self.winning_indices),
-                losing=np.array(self.losing_matrices, dtype=object),
-                losing_indices=np.array(self.losing_indices),
-                draw=np.array(self.draw_matrices, dtype=object),
-                draw_indices=np.array(self.draw_indices),
-            )
-            print(f"Zone DB saved: W={len(self.winning_matrices)}, L={len(self.losing_matrices)}, D={len(self.draw_matrices)}")
-        except MemoryError:
-            print(f"MemoryError during save! Pruning more aggressively...")
-            self.prune_database(target_size=1000)
-            # Try again with much smaller size
-            np.savez_compressed(
-                self.filepath,
-                winning=np.array(self.winning_matrices, dtype=object),
-                winning_indices=np.array(self.winning_indices),
-                losing=np.array(self.losing_matrices, dtype=object),
-                losing_indices=np.array(self.losing_indices),
-                draw=np.array(self.draw_matrices, dtype=object),
-                draw_indices=np.array(self.draw_indices),
-            )
-    
-    def load(self):
-        if os.path.exists(self.filepath):
-            try:
-                data = np.load(self.filepath, allow_pickle=True)
-                self.winning_matrices = list(data['winning'])
-                self.winning_indices = list(data['winning_indices'])
-                self.losing_matrices = list(data['losing'])
-                self.losing_indices = list(data['losing_indices'])
-                self.draw_matrices = list(data['draw'])
-                self.draw_indices = list(data['draw_indices'])
-                print(f"Loaded zone DB: W={len(self.winning_matrices)}, L={len(self.losing_matrices)}, D={len(self.draw_matrices)}")
-            except Exception as e:
-                print(f"Failed to load zone database: {e}")
-    
-    def prune_database(self, target_size: int = 5000):
-        """Keep only evenly-spaced diverse samples along Hilbert curve."""
-        def prune_zone(matrices, indices, target):
-            if len(matrices) <= target:
-                return matrices, indices
-            
-            # Keep evenly spaced samples
-            step = max(1, len(matrices) // target)
-            keep_idx = list(range(0, len(matrices), step))[:target]
-            
-            new_matrices = [matrices[i] for i in keep_idx]
-            new_indices = [indices[i] for i in keep_idx]
-            
-            return new_matrices, new_indices
-        
-        old_sizes = (len(self.winning_matrices), len(self.losing_matrices), len(self.draw_matrices))
-        
-        self.winning_matrices, self.winning_indices = prune_zone(
-            self.winning_matrices, self.winning_indices, target_size
-        )
-        self.losing_matrices, self.losing_indices = prune_zone(
-            self.losing_matrices, self.losing_indices, target_size
-        )
-        self.draw_matrices, self.draw_indices = prune_zone(
-            self.draw_matrices, self.draw_indices, target_size
-        )
-        
-        new_sizes = (len(self.winning_matrices), len(self.losing_matrices), len(self.draw_matrices))
-        print(f"Pruned: W {old_sizes}→{new_sizes}, L {old_sizes}→{new_sizes}, D {old_sizes}→{new_sizes}")
-    
-    def compute_zone_score(self, W: np.ndarray, k: int = 5, beta: float = 0.5) -> float:
-        """Compute zone guidance score Z(x(t), a) ∈ [-1, 1]."""
-        if W.shape != (64, 64):
-            return 0.0
-        
-        def knn_similarity(matrices, k_val):
-            if len(matrices) == 0:
-                return 0.0
-            
-            k_actual = min(k_val, len(matrices))
-            distances = []
-            
-            for mat in matrices[:k_actual]:
-                dist = np.sum(np.abs(W - mat)) / (64.0 * 64.0)
-                distances.append(dist)
-            
-            similarities = [1.0 - d for d in distances]
-            return np.mean(similarities)
-        
-        zone_win = knn_similarity(self.winning_matrices, k)
-        zone_loss = knn_similarity(self.losing_matrices, k)
-        zone_draw = knn_similarity(self.draw_matrices, k)
-        
-        Z = zone_win - zone_loss + beta * zone_draw
-        return float(np.clip(Z, -1.0, 1.0))
-
-
-
-# ============================================================================
-# MCVS
-# ============================================================================
-
-class MCVSNode:
-    def __init__(self, game: Breakthrough, parent=None, move=None):
-        self.game = game.copy()
-        self.parent = parent
-        self.move = move
-        self.children = {}
-        self.N = 0
-        self.Q = 0.0
-        self.P = 0.0
-        self.Z = 0.0  # Zone guidance score Z(x(t), a) ∈ [-1,1]
-
-    def value(self) -> float:
-        return self.Q / self.N if self.N > 0 else 0.0
-
-class MCVSSearcher:
-    class Node:
-        def __init__(self, prior: float = 0.0, zone: float = 0.0):
-            self.prior = prior
-            self.zone = zone
-            self.visit_count = 0
-            self.value_sum = 0.0
-            self.children = {}
-
-    def __init__(self, policy_net=None, value_net=None, zone_db=None, device='cpu', 
-                 cpuct=1.41, lambda_zone=0.0, k_zone=5, use_nets=True):
-        self.policy_net = policy_net.to(device) if policy_net and use_nets else None
-        self.value_net = value_net.to(device) if value_net and use_nets else None
-        self.zone_db = zone_db
-        self.device = device
-        self.cpuct = cpuct
-        self.lambda_zone = lambda_zone
-        self.k_zone = k_zone
-        self.use_nets = use_nets
-
-    def _quick_rollout(self, game):
-        current = game.copy()
-        for _ in range(100):
-            if current.is_terminal():
-                break
-            moves = current.get_legal_moves()
-            if not moves:
-                break
-            current.apply_move(random.choice(moves))
-        winner = current.check_winner()
-        if winner == Breakthrough.PLAYER1:
-            return 1.0
-        elif winner == Breakthrough.PLAYER2:
-            return -1.0
-        else:
-            return 0.0
-
-    def search_with_time_budget(self, game, time_budget: float):
-        root = self.Node()
-        start_time = time.time()
-        simulations = 0
-        eps = 1e-6
-
-        # If there are no legal moves at the root, return empty visits and zero sims
-        root_legal = game.get_legal_moves()
-        if not root_legal:
-            return {}, 0
-
-        while time.time() - start_time < time_budget:
-            current_game = game.copy()
-            node = root
-            path = []
-
-            # SELECTION
-            while node.children:
-                path.append(node)
-                total_n = sum(c.visit_count for c in node.children.values()) + 1
-                best_score = -float('inf')
-                best_move = None
-                best_child = None
-                for move, child in node.children.items():
-                    q = child.value_sum / child.visit_count if child.visit_count > 0 else 0.0
-                    p_lambda = max(child.prior + self.lambda_zone * child.zone, eps)
-                    u = self.cpuct * p_lambda * np.sqrt(total_n) / (1 + child.visit_count)
-                    score = q + u
-                    if score > best_score:
-                        best_score = score
-                        best_move = move
-                        best_child = child
-                # Advance game and node
-                if best_move is None:
-                    break
-                current_game.apply_move(best_move)
-                node = best_child
-
-            path.append(node)
-
-            # EXPANSION / EVALUATION
-            if current_game.is_terminal():
-                w = current_game.check_winner()
-                if w == Breakthrough.PLAYER1:
-                    value = 1.0
-                elif w == Breakthrough.PLAYER2:
-                    value = -1.0
-                else:
-                    value = 0.0
-                legal_moves = []
-                prior_dict = {}
-            else:
-                legal_moves = current_game.get_legal_moves()
-                if not legal_moves:
-                    prior_dict = {}
-                    value = 0.0
-                else:
-                    if self.use_nets and self.policy_net is not None and self.value_net is not None:
-                        W = current_game.get_weighted_adjacency_matrix()
-                        W_tensor = torch.tensor(W, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(self.device)
-                        with torch.no_grad():
-                            logits = self.policy_net(W_tensor)
-                            probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy()
-                            val = float(self.value_net(W_tensor).squeeze().cpu().numpy())
-                            value = float(np.clip(val, -1.0, 1.0))
-
-                        prior_dict = {}
-                        total_p = 0.0
-                        for m in legal_moves:
-                            fr, fc, tr, tc = m
-                            from_sq = fr * 8 + fc
-                            to_sq = tr * 8 + tc
-                            idx = from_sq * 64 + to_sq
-                            p = float(probs[idx]) if idx < probs.size else 0.0
-                            prior_dict[m] = p
-                            total_p += p
-                        if total_p <= 0.0:
-                            prior_dict = {m: 1.0 / len(legal_moves) for m in legal_moves}
-                        else:
-                            for m in prior_dict:
-                                prior_dict[m] /= total_p
-                    else:
-                        prior_dict = {m: 1.0 / len(legal_moves) for m in legal_moves}
-                        value = self._quick_rollout(current_game)
-
-            # EXPAND CHILDREN
-            for m in legal_moves:
-                child_game = current_game.copy()
-                child_game.apply_move(m)
-                child_W = child_game.get_weighted_adjacency_matrix()
-                child_Z = self.zone_db.compute_zone_score(child_W, k=self.k_zone) if self.zone_db else 0.0
-                child_prior = prior_dict.get(m, 1.0 / len(legal_moves))
-                node_prior = max(child_prior + self.lambda_zone * child_Z, eps)
-                node.children[m] = self.Node(prior=node_prior, zone=child_Z)
-
-            # BACKPROPAGATION
-            v = value
-            for n in reversed(path):
-                n.visit_count += 1
-                n.value_sum += v
-                v = -v
-
-            simulations += 1
-
-        visits = {move: child.visit_count for move, child in root.children.items()}
-        if not visits:
-            legal = game.get_legal_moves()
-            visits = {random.choice(legal): 1} if legal else {}
-        return visits, simulations
-
-# ============================================================================
-# UCT-MCTS (Baseline without guidance)
-# ============================================================================
-
-class UCTNode:
-    def __init__(self, game: Breakthrough, parent=None, move=None):
-        self.game = game.copy()
-        self.parent = parent
-        self.move = move
-        self.children: Dict[Tuple[int,int,int,int], "UCTNode"] = {}
-        self.N = 0          # visit count
-        self.Q = 0.0        # total value
-
-    def value(self) -> float:
-        return self.Q / self.N if self.N > 0 else 0.0
-
-class UCTSearcher:
-    def __init__(self, cpuct: float = np.sqrt(2.0)):
-        self.cpuct = cpuct
-
-    def uct_score(self, parent: UCTNode, child: UCTNode) -> float:
-        if child.N == 0:
-            # Encourage exploration of unvisited nodes
-            return float('inf')
-        Q = child.value()
-        U = self.cpuct * np.sqrt(parent.N + 1) / child.N
-        return Q + U
-
-    def simulate(self, game: Breakthrough) -> float:
-        """
-        Default policy: random playout to terminal state.
-        Return +1 if PLAYER1 wins, -1 if PLAYER2 wins, 0 for draw.
-        Value is from the perspective of the player to move at root,
-        but since Breakthrough is alternating and we flip sign in backprop,
-        we can just return game outcome from fixed perspective.
-        """
-        sim_game = game.copy()
-        while not sim_game.is_terminal():
-            moves = sim_game.get_legal_moves()
-            if not moves:
-                break
-            move = random.choice(moves)
-            sim_game.apply_move(move)
-        winner = sim_game.check_winner()
-        if winner == Breakthrough.PLAYER1:
-            return 1.0
-        elif winner == Breakthrough.PLAYER2:
-            return -1.0
-        else:
-            return 0.0
-
-    def search_with_time_budget(self, game: Breakthrough, time_budget: float) -> Dict[Tuple[int,int,int,int], int]:
-        """
-        Standard UCT search with a time budget (seconds).
-        Returns a tuple: (move->visit count dict, simulations)
-        """
-        root = UCTNode(game)
-        start_time = time.time()
-        simulations = 0
-
-        legal_moves = game.get_legal_moves()
-        if not legal_moves:
-            return {}, 0
-
-        # Run simulations until time budget exhausted
-        while time.time() - start_time < time_budget:
-            node = root
-
-            # SELECTION & EXPANSION
-            while not node.game.is_terminal():
-                legal = node.game.get_legal_moves()
-                if not legal:
-                    break
-
-                # If any unexpanded moves exist, expand one
-                if len(node.children) < len(legal):
-                    untried = [m for m in legal if m not in node.children]
-                    if untried:
-                        move = random.choice(untried)
-                        child_game = node.game.copy()
-                        child_game.apply_move(move)
-                        child = UCTNode(child_game, parent=node, move=move)
-                        node.children[move] = child
-                        node = child
-                        break
-
-                # Otherwise, select best child by UCT score
-                best_move = max(node.children.keys(),
-                                key=lambda m: self.uct_score(node, node.children[m]))
-                node = node.children[best_move]
-
-            # SIMULATION
-            value = self.simulate(node.game)
-
-            # BACKPROPAGATION
-            current = node
-            while current is not None:
-                current.N += 1
-                current.Q += value
-                value = -value  # switch perspective
-                current = current.parent
-            simulations += 1
-        return {move: child.N for move, child in root.children.items()}, simulations
-
-
-# ============================================================================
-# Self-play Data Generation
-# ============================================================================
-
-def generate_self_play_data(policy_net, value_net, zone_db, max_time_seconds=60.0,
-                            lambda_zone=0.5, device='cpu'):
-    """Generate self-play games using MCTS with zone guidance."""
-    searcher = MCVSSearcher(policy_net, value_net, zone_db, device=device,
-                           cpuct=1.41, lambda_zone=lambda_zone, k_zone=5)
-    
-    W_states = []
-    policies = []
-    values = []
-    trajectories = []
-    
-    start_time = time.time()
-    game_idx = 0
-    
-    while time.time() - start_time < max_time_seconds:
-        game_start_time = time.time()
-        game = Breakthrough()
-        trajectory = [game.copy()]
-        W_state_list = []
-        policy_list = []
-        
-        while not game.is_terminal():
-            time_per_move = 0.3  # Increased from 0.05 for more exploration/variety
-            W_t = game.get_weighted_adjacency_matrix()
-            visits, _ = searcher.search_with_time_budget(game, time_per_move)
-            
-            if not visits:
-                break
-            
-            moves = list(visits.keys())
-            visit_counts = np.array([visits[m] for m in moves], dtype=float)
-            total_visits = visit_counts.sum()
-            probs = visit_counts / total_visits
-            
-            # Add temperature for more variety in selection
-            temperature = 1.0 if game.move_count < 8 else 0.25  # Higher early, lower late
-            if temperature != 1.0:
-                probs = probs ** (1.0 / temperature)
-                probs = probs / probs.sum()
-            
-            policy_target = np.zeros(4096)
-            for move, count in visits.items():
-                fr, fc, tr, tc = move
-                idx = move_to_index(fr, fc, tr, tc)
-                if 0 <= idx < 4096:
-                    policy_target[idx] = count / total_visits
-            
-            W_state_list.append(W_t)
-            policy_list.append(policy_target)
-            
-            if len(moves) > 1:
-                chosen_idx = np.random.choice(len(moves), p=probs)
-                chosen_move = moves[chosen_idx]
-            else:
-                chosen_move = moves[0]
-            
-            game.apply_move(chosen_move)
-            trajectory.append(game.copy())
-        
-        winner = game.check_winner()
-        result = 1 if winner == Breakthrough.PLAYER1 else -1 if winner == Breakthrough.PLAYER2 else 0
-        
-        for i in range(len(W_state_list)):
-            player_to_move = 1 if i % 2 == 0 else -1
-            value = result if result == 0 else (result if result == player_to_move else -result)
-            
-            W_states.append(W_state_list[i])
-            policies.append(policy_list[i])
-            values.append(value)
-        
-        zone_db.add_game_record(trajectory, result)
-        trajectories.append(trajectory)
-        
-        game_idx += 1
-        game_duration = time.time() - game_start_time
-        total_elapsed = time.time() - start_time
-        
-        print(f"Game {game_idx}: Winner={winner}, Moves={game.move_count}, Duration={game_duration:.1f}s, Total={total_elapsed:.1f}s")
-    
-    elapsed = time.time() - start_time
-    print(f"\n✅ Self-play finished: {game_idx} games in {elapsed:.1f}s, {len(W_states)} training examples collected")
-    
-    return W_states, policies, values, trajectories
-
-
-# ============================================================================
-# Training
-# ============================================================================
-
-def train_networks(policy_net, value_net, W_states, policies, values, epochs=5, batch_size=32, device='cpu', learning_rate=None, lr=0.001):
-    """Train policy and value networks.
-
-    Accepts either `learning_rate=` (preferred) or `lr=` for backward compatibility.
-    """
-    if learning_rate is not None:
-        lr = learning_rate
-
-    policy_net.train()
-    value_net.train()
-    opt_p = optim.Adam(policy_net.parameters(), lr=lr)
-    opt_v = optim.Adam(value_net.parameters(), lr=lr)
-
-    W_tensor = torch.tensor(np.array(W_states)[:, np.newaxis, :, :], dtype=torch.float32)
-    pi_tensor = torch.tensor(np.array(policies), dtype=torch.float32)
-    v_tensor = torch.tensor(np.array(values), dtype=torch.float32).unsqueeze(1)
-
-    dataset = TensorDataset(W_tensor, pi_tensor, v_tensor)
-    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
-
-    for epoch in range(epochs):
-        total_loss = 0.0
-        for batch_W, batch_pi, batch_v in loader:
-            batch_W = batch_W.to(device)
-            batch_pi = batch_pi.to(device)
-            batch_v = batch_v.to(device)
-
-            logits = policy_net(batch_W)
-            # Fixed policy loss: KL divergence between network policy and MCTS visit distribution
-            policy_loss = F.kl_div(F.log_softmax(logits, dim=1), batch_pi, reduction='batchmean')
-
-            pred_v = value_net(batch_W)
-            value_loss = F.mse_loss(pred_v, batch_v)
-
-            loss = policy_loss + value_loss
-            total_loss += loss.item()
-
-            opt_p.zero_grad()
-            opt_v.zero_grad()
-            loss.backward()
-
-            # Optional but recommended: gradient clipping
-            torch.nn.utils.clip_grad_norm_(policy_net.parameters(), max_norm=1.0)
-            torch.nn.utils.clip_grad_norm_(value_net.parameters(), max_norm=1.0)
-
-            opt_p.step()
-            opt_v.step()
-
-        # Optional: print epoch loss for monitoring
-        print(f"Epoch {epoch+1}/{epochs} - Loss: {total_loss/len(loader):.4f}")
-
-
-# ============================================================================
-# Diagnostic Test for Move Diversity
-# ============================================================================
-
-def test_move_diversity(policy_net, value_net, zone_db, num_trials=30, time_budget=2.0, device='cpu'):
-    """Test if the AI chooses different moves in the same starting position."""
-    searcher = MCVSSearcher(policy_net, value_net, zone_db, device=device)
-    game = Breakthrough()  # Starting position; change to any position you want to test
-    
-    move_counts = {}
-    for i in range(num_trials):
-        visits, _ = searcher.search_with_time_budget(game, time_budget)
-        if visits:
-            # Use argmax for the "best" move, but since there's noise/temp, it should vary
-            best_move = max(visits, key=visits.get)
-            move_counts[best_move] = move_counts.get(best_move, 0) + 1
-    
-    print("\nMove diversity test (starting position):")
-    total = sum(move_counts.values())
-    for m, cnt in sorted(move_counts.items(), key=lambda x: -x[1]):
-        print(f"{m}: {cnt:2d}×  ({cnt/total:.1%})")
-    print(f"Unique moves chosen: {len(move_counts)} out of {len(game.get_legal_moves())} legal\n")
-
-# ============================================================================
-# Guided MCVS vs UCT-MCTS Match
-# ============================================================================
-
-def play_match_mcvs_vs_uct(policy_net,
-                           value_net,
-                           zone_db,
-                           num_games: int = 200,
-                           time_per_move: float = 0.3,
-                           device: str = 'cpu'):
-    """
-    Play a tournament between guided MCVS and UCT.
-    Odd games: MCVS as PLAYER1 (first), UCT as PLAYER2.
-    Even games: UCT as PLAYER1, MCVS as PLAYER2.
-    """
-    mcvs_searcher = MCVSSearcher(policy_net, value_net, zone_db,
-                                 device=device, cpuct=1.41,
-                                 lambda_zone=1.0, k_zone=5)
-    uct_searcher = UCTSearcher(cpuct=np.sqrt(2.0))
-
-    mcvs_wins = 0
-    uct_wins = 0
-    draws = 0
-
-    for g in range(1, num_games + 1):
-        game = Breakthrough()
-        player_to_move = Breakthrough.PLAYER1  # always 1 then 2 then 1...
-
-        # Alternate roles
-        # True: MCVS plays PLAYER1 in this game
-        mcvs_as_player1 = (g % 2 == 1)
-
-        while not game.is_terminal():
-            moves = game.get_legal_moves()
-            if not moves:
-                break
-
-            # Decide which engine moves
-            if player_to_move == Breakthrough.PLAYER1:
-                use_mcvs = mcvs_as_player1
-            else:
-                use_mcvs = not mcvs_as_player1
-
-            if use_mcvs:
-                # Guided MCVS move
-                visits, _ = mcvs_searcher.search_with_time_budget(game, time_per_move)
-            else:
-                # UCT move
-                visits, _ = uct_searcher.search_with_time_budget(game, time_per_move)
-
-            if not visits:
-                break
-
-            # Choose move by max visit count
-            moves_list = list(visits.keys())
-            counts = np.array([visits[m] for m in moves_list], dtype=float)
-            best_index = int(np.argmax(counts))
-            chosen_move = moves_list[best_index]
-
-            game.apply_move(chosen_move)
-            player_to_move = Breakthrough.PLAYER2 if player_to_move == Breakthrough.PLAYER1 else Breakthrough.PLAYER1
-
-        winner = game.check_winner()
-        if winner == Breakthrough.PLAYER1:
-            if mcvs_as_player1:
-                mcvs_wins += 1
-            else:
-                uct_wins += 1
-        elif winner == Breakthrough.PLAYER2:
-            if mcvs_as_player1:
-                uct_wins += 1
-            else:
-                mcvs_wins += 1
-        else:
-            draws += 1
-
-        print(f"Game {g}/{num_games} finished. Winner = {winner} "
-              f"(MCVS as P1: {mcvs_as_player1})")
-
-    print("\n=== MCVS vs UCT Tournament Results ===")
-    print(f"Total games: {num_games}")
-    print(f"MCVS wins : {mcvs_wins}")
-    print(f"UCT wins  : {uct_wins}")
-    print(f"Draws     : {draws}")
-    if num_games > 0:
-        mcvs_rate = mcvs_wins / num_games * 100.0
-        uct_rate = uct_wins / num_games * 100.0
-        draw_rate = draws / num_games * 100.0
-        print(f"Win rates: MCVS={mcvs_rate:.1f}%, UCT={uct_rate:.1f}%, Draw={draw_rate:.1f}%\n")
-
-
-# ============================================================================
-# MAIN
-# ============================================================================
-
-if __name__ == "__main__":
-    print("="*80)
-    print("BREAKTHROUGH MCVS - INCREMENTAL TRAINING")
-    print("="*80)
-    
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Device: {device}")
-    
-    # Initialize networks
-    policy_net = PolicyNetworkCNN().to(device)
-    value_net = ValueNetworkCNN().to(device)
-    
-    # Initialize zone database (with fixed version - max_size limits growth)
-    zone_db = HilbertOrderedZoneDatabase("breakthrough_zone_db.npz", max_size=10000)
-    
-    # Try to load existing checkpoint
-    checkpoint_path = "breakthrough_checkpoint.pt"
-    iteration = 0
-    
-    if os.path.exists(checkpoint_path):
-        print(f"Loading checkpoint from {checkpoint_path}...")
-        checkpoint = torch.load(checkpoint_path, map_location=device)
-        policy_net.load_state_dict(checkpoint['policy'])
-        value_net.load_state_dict(checkpoint['value'])
-        iteration = checkpoint.get('iteration', 0)
-        print(f"Resumed from iteration {iteration}")
-    
-    # Accumulate training data across iterations
-    all_W_states = []
-    all_policies = []
-    all_values = []
-    
-    # Main training loop
-    while True:
-        print("="*80)
-        print(f"ITERATION {iteration}")
-        print("="*80)
-        
-        # 1. Generate self-play games (adds to zone_db automatically)
-        print("Generating self-play data (30 minutes)...")
-        lambda_zone = 0.0 if iteration == 0 else 1.0  # No guidance first iteration
-        
-        W_states, policies, values, trajectories = generate_self_play_data(
-            policy_net, value_net, zone_db,
-            max_time_seconds=1800.0,  # 30 minutes
-            lambda_zone=lambda_zone,
-            device=device
-        )
-        
-        print(f"Collected {len(W_states)} training examples from {len(trajectories)} games")
-        
-        # 2. Accumulate training data
-        all_W_states.extend(W_states)
-        all_policies.extend(policies)
-        all_values.extend(values)
-        
-        # Optional: Limit total training data size (keep most recent N examples)
-        max_training_examples = 50000
-        if len(all_W_states) > max_training_examples:
-            print(f"Trimming training data to most recent {max_training_examples} examples...")
-            all_W_states = all_W_states[-max_training_examples:]
-            all_policies = all_policies[-max_training_examples:]
-            all_values = all_values[-max_training_examples:]
-        
-        print(f"Total accumulated training examples: {len(all_W_states)}")
-        
-        # 3. Train networks on ALL accumulated data
-        print("Training networks...")
-        train_networks(
-            policy_net, value_net,
-            all_W_states, all_policies, all_values,
-            epochs=10,  # Fewer epochs since we train more frequently
-            device=device,
-            learning_rate=0.001
-        )
-        
-        # 4. Save checkpoint (ONE .pt file)
-        print("Saving checkpoint...")
-        torch.save({
-            'policy': policy_net.state_dict(),
-            'value': value_net.state_dict(),
-            'iteration': iteration,
-            'num_training_examples': len(all_W_states)
-        }, checkpoint_path)
-        print(f"✅ Checkpoint saved: {checkpoint_path}")
-        
-        # 5. Save zone database (ONE .npz file)
-        zone_db.save()
-        print(f"✅ Zone DB saved: W={len(zone_db.winning_matrices)}, "
-              f"L={len(zone_db.losing_matrices)}, "
-              f"D={len(zone_db.draw_matrices)}")
-        
-        iteration += 1
-        
-        print("\n" + "="*80)
-        print(f"Iteration {iteration} complete!")
-        print(f"Files updated:")
-        print(f"  - {checkpoint_path}")
-        print(f"  - {zone_db.filepath}")
-        print("="*80 + "\n")
-        
-        # Optional: Run tournament evaluation every N iterations
-        if iteration % 5 == 0:
-            print("\nRunning tournament evaluation...")
-            play_match_mcvs_vs_uct(
-                policy_net, value_net, zone_db,
-                num_games=20,
-                time_per_move=0.5,
-                device=device
-            )