Cecil’s B-Spline Gomoku Training

cecils_bspline_gomoku.py

# cecils_bspline_gomoku

# Large scalability issues fixed & prevents gradient explosions

import os
import math
import time
import json
import random
from dataclasses import dataclass
from typing import Tuple, Dict, Any, List, Optional, Deque
from collections import deque
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

BOARD_SIZE = 15
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
GRID_RESOLUTION = 8  
COMPLEX_CHANNELS = 32  
SVD_ENERGY_THRESH = 0.92
SVD_MIN_RANK = 4
SVD_MAX_RANK = 16
SELF_PLAY_GAMES = 2
SELF_PLAY_SIMULATIONS = 40  
REPLAY_BUFFER_SIZE = 2000
BATCH_SIZE = 32
TRAIN_ITERS = 10
LR = 3e-4
CHECKPOINT_DIR = "./checkpoints_bspline"
META_STATE_FILE = os.path.join(CHECKPOINT_DIR, "meta_state.json")

class Gomoku:
    EMPTY = 0
    BLACK = 1
    WHITE = -1

    def __init__(self, size=BOARD_SIZE):
        self.size = size
        self.board = np.zeros((size, size), dtype=np.int8)
        self.to_move = Gomoku.BLACK
        self.moves = []  
        self.winner = None  

    def copy(self):
        g = Gomoku(self.size)
        g.board = self.board.copy()
        g.to_move = self.to_move
        g.moves = self.moves.copy()
        g.winner = self.winner
        return g

    def legal_moves(self) -> List[Tuple[int, int]]:
        return list(zip(*np.where(self.board == Gomoku.EMPTY)))

    def play(self, r: int, c: int):
        if self.board[r, c] != Gomoku.EMPTY:
            raise ValueError(f"Illegal move at ({r},{c})")
        self.board[r, c] = self.to_move
        self.moves.append((r, c))
        self._update_winner(r, c)
        self.to_move = -self.to_move

    def _update_winner(self, r: int, c: int):
        player = self.board[r, c]
        directions = [(1,0), (0,1), (1,1), (1,-1)]
        for dr, dc in directions:
            cnt = 1
            for dirn in (1, -1):
                rr, cc = r, c
                while True:
                    rr += dr * dirn
                    cc += dc * dirn
                    if 0 <= rr < self.size and 0 <= cc < self.size and self.board[rr, cc] == player:
                        cnt += 1
                    else:
                        break
            if cnt >= 5:
                self.winner = int(player)
                return
        if np.all(self.board != Gomoku.EMPTY):
            self.winner = 0

    def result(self) -> Optional[int]:
        return self.winner

    def to_tensor(self) -> torch.Tensor:
        cur = (self.board == self.to_move).astype(np.float32)
        opp = (self.board == -self.to_move).astype(np.float32)
        return torch.from_numpy(np.stack([cur, opp], axis=0))


def stabilized_svd_compress(weight: torch.Tensor, energy_thresh: float) -> torch.Tensor:
    with torch.no_grad():
        shape = weight.shape
        flat_w = weight.view(shape[0], -1)
        U, S, Vh = torch.linalg.svd(flat_w, full_matrices=False)
        cum_energy = torch.cumsum(S**2, dim=-1)
        total_energy = cum_energy[-1] + 1e-12
        ranks = torch.where(cum_energy / total_energy >= energy_thresh)[0]
        k = int(ranks[0].item()) + 1 if len(ranks) > 0 else SVD_MIN_RANK
        k = max(min(k, SVD_MAX_RANK), SVD_MIN_RANK)
        k = min(k, S.shape[0])
        w_reconstructed = U[:, :k] @ torch.diag(S[:k]) @ Vh[:k, :]
        return w_reconstructed.view(shape)

class ComplexBSpline2DConvLayer(nn.Module):
    def __init__(self, in_ch: int, out_ch_complex: int, grid_res: int = GRID_RESOLUTION):
        super().__init__()
        self.in_ch = in_ch
        self.out_ch_complex = out_ch_complex
        self.grid_res = grid_res
        self.weight_real = nn.Parameter(torch.randn(out_ch_complex, in_ch, 3, 3) * 0.05)
        self.weight_imag = nn.Parameter(torch.randn(out_ch_complex, in_ch, 3, 3) * 0.05)
        self.polarity_gates = nn.Parameter(torch.ones(out_ch_complex) * 0.5)

    def forward(self, x: torch.Tensor, meta_gates: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, _, H, W = x.shape
        w_r = stabilized_svd_compress(self.weight_real, SVD_ENERGY_THRESH)
        w_i = stabilized_svd_compress(self.weight_imag, SVD_ENERGY_THRESH)
        x_grid = F.interpolate(x, size=(self.grid_res, self.grid_res), mode='bilinear', align_corners=True)
        real_grid = F.conv2d(x_grid, w_r, padding=1)
        imag_grid = F.conv2d(x_grid, w_i, padding=1)
        real_part = F.interpolate(real_grid, size=(H, W), mode='bilinear', align_corners=True)
        imag_part = F.interpolate(imag_grid, size=(H, W), mode='bilinear', align_corners=True)
        gates = self.polarity_gates.view(1, -1, 1, 1)
        if meta_gates is not None:
            gates = gates * (1.0 + meta_gates.view(B, self.out_ch_complex, 1, 1))
        return torch.cat([real_part * gates, imag_part * gates], dim=1)
        
class RotationalNode(nn.Module):
    def __init__(self, num_bins: int = 24):
        super().__init__()
        self.num_bins = num_bins
        self.bin_bias = nn.Parameter(torch.zeros(num_bins))
        self.spatial_proj = nn.Linear(num_bins, BOARD_SIZE * BOARD_SIZE)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        # Preserve local phase information by flattening the spatial dimension first
        real_spatial = x[:, :C//2, :, :].reshape(B, C//2, -1)
        imag_spatial = x[:, C//2:, :, :].reshape(B, C//2, -1)
        real_avg = real_spatial.mean(dim=-1).mean(dim=-1)
        imag_avg = imag_spatial.mean(dim=-1).mean(dim=-1)
        angles = torch.atan2(imag_avg, real_avg + 1e-8).unsqueeze(-1)
        bin_centers = torch.linspace(-math.pi, math.pi, self.num_bins, device=x.device)
        d = torch.remainder(angles - bin_centers + math.pi, 2 * math.pi) - math.pi
        logits = - (d ** 2) + self.bin_bias
        return self.spatial_proj(F.softmax(logits, dim=-1))

class BSplineNet(nn.Module):
    def __init__(self, board_size=BOARD_SIZE, in_planes=2, complex_channels=COMPLEX_CHANNELS, grid_res=GRID_RESOLUTION):
        super().__init__()
        self.board_size = board_size
        self.input_conv = nn.Conv2d(in_planes, in_planes * 4, kernel_size=3, padding=1)
        self.bspline = ComplexBSpline2DConvLayer(in_ch=in_planes * 4, out_ch_complex=complex_channels, grid_res=grid_res)
        self.agg_conv = nn.Conv2d(complex_channels * 2, 64, kernel_size=3, padding=1)
        self.rot_node = RotationalNode(num_bins=24)
        self.policy_head = nn.Conv2d(64, 1, kernel_size=1)
        self.value_head = nn.Sequential(
            nn.Conv2d(64, 16, kernel_size=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(16 * board_size * board_size, 64),
            nn.ReLU(),
            nn.Linear(64, 1))

    def forward(self, x: torch.Tensor, meta_gates: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        h = F.relu(self.input_conv(x))
        h_spline = self.bspline(h, meta_gates=meta_gates)
        h_agg = F.relu(self.agg_conv(h_spline))
        p = self.policy_head(h_agg).view(x.shape[0], -1) 
        p = p + self.rot_node(h_spline)
        v = self.value_head(h_agg).squeeze(-1)
        return p, v

class MetaController:
    def __init__(self, dim: int, mc_id: Optional[str] = None):
        self.id = mc_id or f"mc_{int(time.time()*1000)}_{random.randint(0,9999)}"
        self.dim = dim
        self.state = torch.randn(dim) * 0.1
        self.angular_offset = random.uniform(0, 2 * math.pi)
        self.timing_interval = random.uniform(0.2, 0.8)

    def step(self):
        self.angular_offset = (self.angular_offset + self.timing_interval) % (2 * math.pi)
        with torch.no_grad():
            decay_factor = math.cos(self.angular_offset) * 0.05
            self.state = self.state * (0.95 + decay_factor)

    def gating_vector(self) -> torch.Tensor:
        return torch.tanh(self.state)

class SearchNode:
    def __init__(self, prior: float = 0.0):
        self.prior = prior
        self.visits = 0
        self.total_value = 0.0
        self.children: Dict[Tuple[int, int], SearchNode] = {}
        self.is_expanded = False

class PUCTSearcher:
    def __init__(self, net: BSplineNet, c_puct: float = 1.4, device: torch.device = DEVICE):
        self.net = net
        self.c_puct = c_puct
        self.device = device

    def search(self, root_game: Gomoku, sims: int, meta_gates: Optional[torch.Tensor] = None) -> Dict[Tuple[int, int], int]:
        root = SearchNode()
        self._expand_node(root, root_game, meta_gates)
        for _ in range(sims):
            node = root
            path = []
            game = root_game.copy()
            while node.is_expanded and node.children:
                total_N = sum(child.visits for child in node.children.values())
                best_mv, best_child, best_score = None, None, -1e9
                for mv, child in node.children.items():
                    Q = child.total_value / child.visits if child.visits > 0 else 0.0
                    U = self.c_puct * child.prior * math.sqrt(total_N + 1) / (1 + child.visits)
                    score = Q + U
                    if score > best_score:
                        best_score, best_mv, best_child = score, mv, child
                game.play(best_mv[0], best_mv[1])
                path.append((node, best_mv))
                node = best_child
                if game.result() is not None:
                    break
            if game.result() is not None:
                outcome = game.result()
                leaf_value = 0.0 if outcome == 0 else (1.0 if outcome == root_game.to_move else -1.0)
            else:
                if not node.is_expanded:
                    self._expand_node(node, game, meta_gates)
                board_t = game.to_tensor().unsqueeze(0).to(self.device)
                with torch.no_grad():
                    _, value = self.net(board_t, meta_gates=meta_gates)
                leaf_value = float(value.cpu().item())
            curr_val = leaf_value
            for parent_node, move in reversed(path):
                child = parent_node.children[move]
                child.visits += 1
                child.total_value += curr_val
                curr_val = -curr_val
        return {mv: child.visits for mv, child in root.children.items()}

    def _expand_node(self, node: SearchNode, game: Gomoku, meta_gates: Optional[torch.Tensor]):
        board_t = game.to_tensor().unsqueeze(0).to(self.device)
        with torch.no_grad():
            logits, _ = self.net(board_t, meta_gates=meta_gates)
        priors = F.softmax(logits, dim=-1).cpu().numpy()[0]
        for (r, c) in game.legal_moves():
            idx = r * game.size + c
            node.children[(r, c)] = SearchNode(prior=float(priors[idx]))
        node.is_expanded = True

class ReplayBuffer:
    def __init__(self, capacity=REPLAY_BUFFER_SIZE):
        self.buf = deque(maxlen=capacity)

    def push(self, transition):
        self.buf.append(transition)

    def sample(self, batch_size=BATCH_SIZE):
        batch = random.sample(self.buf, min(batch_size, len(self.buf)))
        states = torch.stack([t[0] for t in batch]).to(DEVICE)
        pis = torch.stack([t[1] for t in batch]).to(DEVICE)
        vals = torch.tensor([t[2] for t in batch], dtype=torch.float32, device=DEVICE)
        gates = torch.stack([t[3] for t in batch]).to(DEVICE)
        return states, pis, vals, gates

    def __len__(self):
        return len(self.buf)


def self_play_game(net: BSplineNet, searcher: PUCTSearcher, meta_controller: MetaController) -> List[Tuple[torch.Tensor, torch.Tensor, float, torch.Tensor]]:
    g = Gomoku()
    trajectory = []
    while g.result() is None and len(g.moves) < (BOARD_SIZE * BOARD_SIZE):
        mg = meta_controller.gating_vector().to(DEVICE).unsqueeze(0)
        visits = searcher.search(g, sims=SELF_PLAY_SIMULATIONS, meta_gates=mg)
        pi = np.zeros(g.size * g.size, dtype=np.float32)
        for (r, c), v in visits.items():
            pi[r * g.size + c] = v
        if pi.sum() == 0:
            for (r, c) in g.legal_moves():
                pi[r * g.size + c] = 1.0
        pi /= pi.sum()
        idx = np.random.choice(len(pi), p=pi)
        r, c = divmod(idx, g.size)
        trajectory.append((g.to_tensor(), torch.from_numpy(pi), g.to_move, mg.squeeze(0).cpu()))
        g.play(r, c)
        meta_controller.step()
    outcome = g.result()
    final_val = 0.0 if outcome == 0 else (1.0 if outcome == Gomoku.BLACK else -1.0)
    return [(s, p, (final_val if turn == Gomoku.BLACK else -final_val), g_v) for s, p, turn, g_v in trajectory]

def train_loop():
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
    net = BSplineNet().to(DEVICE)
    opt = torch.optim.AdamW(net.parameters(), lr=LR, weight_decay=1e-4)
    buffer = ReplayBuffer()
    mc = MetaController(dim=COMPLEX_CHANNELS)
    searcher = PUCTSearcher(net=net, device=DEVICE)
    print(f"Beginning pipeline optimization run on target execution unit: {DEVICE}...")
    for it in range(TRAIN_ITERS):
        for _ in range(SELF_PLAY_GAMES):
            traj = self_play_game(net, searcher, mc)
            for item in traj:
                buffer.push(item)
        if len(buffer) < BATCH_SIZE:
            continue
        states, pis, vals, gates = buffer.sample(BATCH_SIZE)
        opt.zero_grad()
        p_logits, pred_vals = net(states, meta_gates=gates)
        logp = F.log_softmax(p_logits, dim=-1)
        policy_loss = - (pis * logp).sum(dim=1).mean()
        value_loss = F.mse_loss(pred_vals, vals)
        loss = policy_loss + value_loss
        loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
        opt.step()
        print(f"Iteration {it+1}/{TRAIN_ITERS} | Loss: {loss.item():.4f} | Samples Gathered: {len(buffer)}")

if __name__ == "__main__":
    train_loop()
    print("Optimization workflow verified successfully.")