Nagi-ovo

def epsilon_greedy(restaurants, ratings, epsilon=0.1):
    if random.random() < epsilon:
        return random.choice(restaurants)  # 探索
    else:
        return restaurants[np.argmax(ratings)]  # 利用

class Restaurant:
    def __init__(self, name):
        self.name = name
        self.total_rating = 0
        self.visits = 0
        
def ucb_choice(restaurants, total_days, c=2):
    # 确保每家餐厅至少访问一次
    for r in restaurants:
        if r.visits == 0:
            return r
            
    # 使用UCB公式选择餐厅
    scores = []
    for r in restaurants:
        avg_rating = r.total_rating / r.visits
        exploration_term = c * math.sqrt(math.log(total_days) / r.visits)
        ucb_score = avg_rating + exploration_term
        scores.append(ucb_score)
        
    return restaurants[np.argmax(scores)]

class GomokuGame:
    def __init__(self, n=15):
        self.n = n  # 棋盘大小，默认15x15

    def getInitBoard(self):
        """
        返回初始的棋盘状态，所有位置都为空。
        """
        b = Board(self.n)
        return np.array(b.pieces)

    def getBoardSize(self):
        """
        返回棋盘的尺寸，即 (n, n)。
        """
        return (self.n, self.n)

    def getActionSize(self):
        """
        返回动作的总数，这里是 n * n，因为每个格子都可能是一个动作。
        """
        return self.n * self.n

    def getNextState(self, board, player, action):
        """
        执行动作，返回下一个棋盘状态和下一个玩家。

        参数：
        - board: 当前棋盘状态
        - player: 当前玩家（1 或 -1）
        - action: 当前动作，0 ~ n*n-1

        返回：
        - (next_board, next_player): 执行动作后的棋盘和下一玩家
        """
        b = Board(self.n)
        b.pieces = np.copy(board)
        # 将动作转换为坐标 (x, y)
        move = (action // self.n, action % self.n)
        b.execute_move(move, player)
        return (b.pieces, -player)

import torch
import torch.nn as nn
import torch.nn.functional as F

class AlphaZeroNNet(nn.Module):
    def __init__(self, game, args):
        """
        参数：
        - game: 游戏对象，提供棋盘大小和动作空间大小等信息。
        - args: 包含网络结构的参数，例如通道数、dropout 概率等。
        """
        super(AlphaZeroNNet, self).__init__()
        self.board_x, self.board_y = game.getBoardSize()  # 棋盘尺寸
        self.action_size = game.getActionSize()           # 动作空间大小
        self.args = args

        # 卷积层块
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, args.num_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(args.num_channels),
            nn.ReLU(),
            nn.Conv2d(args.num_channels, args.num_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(args.num_channels),
            nn.ReLU(),
            nn.Conv2d(args.num_channels, args.num_channels, kernel_size=3),
            nn.BatchNorm2d(args.num_channels),
            nn.ReLU(),
            nn.Conv2d(args.num_channels, args.num_channels, kernel_size=3),
            nn.BatchNorm2d(args.num_channels),
            nn.ReLU(),
        )

        # 计算卷积层输出的尺寸
        conv_output_size = self._get_conv_output_size()

        # 全连接层块
        self.fc_layers = nn.Sequential(
            nn.Linear(conv_output_size, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Dropout(args.dropout),
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(args.dropout),
        )

        # 策略头：输出每个动作的对数概率
        self.policy_head = nn.Linear(512, self.action_size)

        # 价值头：输出当前状态的价值估计
        self.value_head = nn.Linear(512, 1)

    def _get_conv_output_size(self):
        """
        计算卷积层输出的特征数量。
        """
        dummy_input = torch.zeros(1, 1, self.board_x, self.board_y)
        output_feat = self.conv_layers(dummy_input)
        n_size = output_feat.view(1, -1).size(1)
        return n_size

    def forward(self, s):
        """
        前向传播函数。

        参数：
        - s: 输入的棋盘状态，形状为 (batch_size, board_x, board_y)。

        返回：
        - log_policies: 策略的对数概率，形状为 (batch_size, action_size)。
        - values: 价值估计，形状为 (batch_size, 1)。
        """
        # 输入形状调整
        s = s.view(-1, 1, self.board_x, self.board_y)  # (batch_size, 1, board_x, board_y)

        # 卷积层提取特征
        s = self.conv_layers(s)  # (batch_size, num_channels, new_board_x, new_board_y)

        # 展平卷积层输出
        s = s.view(s.size(0), -1)  # (batch_size, conv_output_size)

        # 全连接层提取高级特征
        s = self.fc_layers(s)  # (batch_size, 512)

        # 策略头输出
        policies = self.policy_head(s)  # (batch_size, action_size)
        log_policies = F.log_softmax(policies, dim=1)

        # 价值头输出
        values = self.value_head(s)  # (batch_size, 1)
        values = torch.tanh(values)  # 将价值限定在 [-1, 1]

        return log_policies, values

class MCTS:
    def __init__(self, game, nnet, args):
        self.game = game         # 游戏环境
        self.nnet = nnet         # 神经网络
        self.args = args         # 参数
        self.Qsa = {}            # 存储 Q 值：Q(s,a)
        self.Nsa = {}            # 存储边的访问次数：N(s,a)
        self.Ns = {}             # 存储节点的访问次数：N(s)
        self.Ps = {}             # 存储策略先验：P(s,a)

        self.Es = {}             # 存储游戏结束信息：E(s)
        self.Vs = {}             # 存储合法动作：V(s)

if s not in self.Ps: # 不在策略先验存储 P(s,a) 中
    # 叶节点，根据神经网络输出策略进行扩展
    self.Ps[s], v = self.nnet.predict(canonicalBoard)
    valids = self.game.getValidMoves(canonicalBoard, 1)  # 获取合法动作的掩码，1 为合法
    self.Ps[s] = self.Ps[s] * valids  # 掩码非法动作
    sum_Ps_s = np.sum(self.Ps[s])
    
    if sum_Ps_s > 0:
        self.Ps[s] /= sum_Ps_s  # 归一化策略概率
    else:
        # 所有动作被屏蔽，进行均匀分布
        self.Ps[s] = self.Ps[s] + valids
        self.Ps[s] /= np.sum(self.Ps[s])

    self.Vs[s] = valids  # 存储合法动作的掩码
    self.Ns[s] = 0       # 初始化状态访问次数
    return v

cur_best = -float('inf') # 记录当前最大的 U 值
best_act = -1            # 记录具有最大 U 值的动作

# 遍历所有可能的动作
for a in range(self.game.getActionSize()):
    if valids[a]: 
        # 对于合法的动作，计算 PUCT 值
        if (s, a) in self.Qsa:
            # 如果之前访问过 (s, a)，使用已计算的 Q 值和访问次数
            u = self.Qsa[(s, a)] + self.args.cpuct * self.Ps[s][a] * \
                math.sqrt(self.Ns[s]) / (1 + self.Nsa[(s, a)])
        else:
            # 如果是未访问过的节点，Q 值为 0，鼓励探索
            u = self.args.cpuct * self.Ps[s][a] * math.sqrt(self.Ns[s] + 1e-8) # 对未访问过的动作给予更高的 U 值，避免除零错误

        # 选择具有最高 U 值的动作
        if u > cur_best:
            cur_best = u
            best_act = a

v = -self.search(next_s)

def search(self, canonicalBoard):
    """
    执行一次 MCTS 搜索。

    参数：
    - canonicalBoard: 当前玩家的棋盘状态（规范形式）

    返回：
    - v: 对当前节点的价值评估
    """
    s = self.game.stringRepresentation(canonicalBoard)

    if s not in self.Es:
        self.Es[s] = self.game.getGameEnded(canonicalBoard, 1)
    if self.Es[s] is not None:
        # 如果游戏结束，返回结果
        return self.Es[s]

    '''
    Valid Mask & 叶节点拓展
    '''

    valids = self.Vs[s]

    '''
    使用 PUCT 公式选择最优动作
    '''

    a = best_act
    next_s, next_player = self.game.getNextState(canonicalBoard, 1, a)
    next_s = self.game.getCanonicalForm(next_s, next_player)

    # 下一状态的价值取反
    v = -self.search(next_s)

    # 更新 Q 值和访问次数
    if (s, a) in self.Qsa:
        self.Qsa[(s, a)] = (self.Nsa[(s, a)] * self.Qsa[(s, a)] + v) / (self.Nsa[(s, a)] + 1)
        self.Nsa[(s, a)] += 1
    else:
        self.Qsa[(s, a)] = v
        self.Nsa[(s, a)] = 1

    self.Ns[s] += 1
    return v

def executeEpisode(self):
    """
    执行一局自我对弈，返回训练样本列表。

    返回：
        trainExamples: 一个列表，包含多个 (canonicalBoard, pi, v) 元组。
    """
    trainExamples = []
    board = self.game.getInitBoard()  # 初始化棋盘
    self.curPlayer = 1                # 当前玩家（1 或 -1）
    episodeStep = 0                   # 记录步数

    while True: # 直到游戏结束
        episodeStep += 1
        canonicalBoard = self.game.getCanonicalForm(board, self.curPlayer)

        temp = int(episodeStep < self.args.tempThreshold)  # 温度参数

        pi = self.mcts.getActionProb(canonicalBoard, temp=temp)  # 使用 MCTS 获取动作的概率分布（策略）
        sym = self.game.getSymmetries(canonicalBoard, pi)        # 数据增强（对称性）
        for b, p in sym:
            trainExamples.append([b, self.curPlayer, p, None])

        action = np.random.choice(len(pi), p=pi)  # 按照策略概率选择动作
        board, self.curPlayer = self.game.getNextState(board, self.curPlayer, action) # 执行动作，更新棋盘和当前玩家。

        r = self.game.getGameEnded(board, self.curPlayer)  # 检查游戏是否结束
        if r is not None:
            # 为每个训练样本赋予最终的价值，对于获胜的玩家为 +1，失败的玩家为 -1，平局为 0
            return [(x[0], x[2], r * ((-1)  (x[1] != self.curPlayer))) for x in trainExamples]

def learn(self):
    """
    进行多次迭代的自我对弈和模型更新。
    """
    for i in range(1, self.args.numIters + 1):
        log.info(f"Starting Iter #{i} ...")
        # 存储本次迭代的训练样本
        iterationTrainExamples = deque([], maxlen=self.args.maxlenOfQueue)

        # 进行指定次数的自我对弈
        for _ in tqdm(range(self.args.numEps), desc="Self Play"):
            self.mcts = MCTS(self.game, self.nnet, self.args)  # 重置 MCTS
            iterationTrainExamples += self.executeEpisode()

        # 保存训练样本
        self.trainExamplesHistory.append(iterationTrainExamples)

        # 维护固定长度的训练样本历史
        if len(self.trainExamplesHistory) > self.args.numItersForTrainExamplesHistory:
            log.warning(f"Removing the oldest entry in trainExamples.")
            self.trainExamplesHistory.pop(0)

        # 合并并打乱所有的训练样本
        trainExamples = []
        for e in self.trainExamplesHistory:
            trainExamples.extend(e)
        shuffle(trainExamples)

        # 训练神经网络
        self.nnet.save_checkpoint(folder=self.args.checkpoint, filename="temp.pth.tar")
        self.pnet.load_checkpoint(folder=self.args.checkpoint, filename="temp.pth.tar")
        pmcts = MCTS(self.game, self.pnet, self.args)

        self.nnet.train(trainExamples)
        nmcts = MCTS(self.game, self.nnet, self.args)

        # 评估新模型
        log.info("PITTING AGAINST PREVIOUS VERSION")
        arena = game.Arena(lambda x: np.argmax(pmcts.getActionProb(x, temp=0)),
                           lambda x: np.argmax(nmcts.getActionProb(x, temp=0)),
                           self.game)
        pwins, nwins, draws = arena.playGames(self.args.arenaCompare)

        log.info(f"NEW/PREV WINS : {nwins} / {pwins} ; DRAWS : {draws}")
        if pwins + nwins == 0 or float(nwins) / (pwins + nwins) < self.args.updateThreshold:
            log.info("REJECTING NEW MODEL")
            self.nnet.load_checkpoint(folder=self.args.checkpoint, filename="temp.pth.tar")
        else:
            log.info("ACCEPTING NEW MODEL")
            self.nnet.save_checkpoint(folder=self.args.checkpoint, filename="best.pth.tar")

def getActionProb(self, canonicalBoard, temp=1):
    """
    执行多次 MCTS 搜索，返回动作的概率分布。

    参数：
        canonicalBoard: 当前规范化的棋盘状态
        temp: 温度参数

    返回：
        probs: 动作的概率分布
    """
    # 进行指定次数的 MCTS 搜索
    for _ in range(self.args.numMCTSSims):
        self.search(canonicalBoard) 

    s = self.game.stringRepresentation(canonicalBoard)
    # 记录每个动作被访问的次数 N(s, a)
    counts = [self.Nsa[(s, a)] if (s, a) in self.Nsa else 0 
              for a in range(self.game.getActionSize())]

    if temp == 0:
        bestAs = np.array(np.argwhere(counts == np.max(counts))).flatten()
        bestA = np.random.choice(bestAs)
        probs = [0] * len(counts)
        probs[bestA] = 1
        return probs

    counts = [x  (1. / temp) for x in counts]
    counts_sum = float(sum(counts))
    probs = [x / counts_sum for x in counts]
    return probs

class NNetWrapper:
    def __init__(self, game, args):
        ...
        self.optimizer = optim.Adam(self.nnet.parameters(), lr=args.max_lr)
        
        # 1cycle 学习率参数
        self.total_steps = args.numIters * args.epochs * (args.maxlenOfQueue // args.batch_size)
        self.current_step = 0

    def get_learning_rate(self):
        """实现 1cycle 学习率策略"""
        if self.current_step >= self.total_steps:
            return self.args.min_lr
        
        half_cycle = self.total_steps // 2
        
        if self.current_step <= half_cycle:
            # 第一阶段：从 min_lr 增加到 max_lr
            phase = self.current_step / half_cycle
            lr = self.args.min_lr + (self.args.max_lr - self.args.min_lr) * phase
        else:
            # 第二阶段：从 max_lr 减少到 min_lr
            phase = (self.current_step - half_cycle) / half_cycle
            lr = self.args.max_lr - (self.args.max_lr - self.args.min_lr) * phase
        
        return lr

def train(self, examples):
    ...
    for epoch in range(self.args.epochs):
        ...
        for _ in t:
            ...
            # 计算损失并反向传播
            total_loss.backward()

            # 梯度裁剪
            if self.args.grad_clip:
                torch.nn.utils.clip_grad_norm_(self.nnet.parameters(), self.args.grad_clip)

            # 更新参数
            self.optimizer.step()
            ...