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: ネットワーク構造のパラメータ（チャンネル数、ドロップアウト確率など）を含みます。
        """
        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]

    '''
    有効マスク & 葉ノード拡張
    '''

    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()
            ...