import math
import random
import copy
import numpy as np

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

import chess
import chess.pgn

from tqdm import tqdm # 可视化进度

###############################################################################

A. 棋盘环境 (python-chess) & 基础辅助函数

###############################################################################

class ChessEnv:
"""
基于 python-chess 管理国际象棋对局:
- reset() 重置
- step(move_uci) 走一步
- get_legal_moves() 获取所有合法走子(uci字符串)
- current_player() 返回 1(白)或 -1(黑)
- copy() / fen() / is_game_over() / get_result() ...
"""
def init(self):
self.board = chess.Board()

text
def reset(self):
    self.board.reset()

def step(self, move_uci):
    """
    执行走子; 返回 (done, winner).
    winner=1表示白胜, -1黑胜,0和棋(或对局未结束).
    """
    move = chess.Move.from_uci(move_uci)
    if move in self.board.legal_moves:
        self.board.push(move)
    done = self.board.is_game_over()
    winner = 0
    if done:
        # '1-0' '0-1' '1/2-1/2'
        result = self.board.result()
        if result == '1-0':
            winner = 1
        elif result == '0-1':
            winner = -1
        else:
            winner = 0
    return done, winner

def get_legal_moves(self):
    return [m.uci() for m in self.board.legal_moves]

def copy(self):
    new_env = ChessEnv()
    new_env.board = self.board.copy()
    return new_env

def current_player(self):
    return 1 if self.board.turn else -1

def is_game_over(self):
    return self.board.is_game_over()

def get_result(self):
    """
    对局结果: 1 白胜, -1 黑胜, 0 和棋/未结束
    """
    if not self.board.is_game_over():
        return 0
    result = self.board.result()
    if result == '1-0':
        return 1
    elif result == '0-1':
        return -1
    else:
        return 0

def fen(self):
    return self.board.fen()

def get_repetition_count(board, history_boards):
"""
返回当前 board 在 history_boards 中出现过几次(>=2时视为2).
"""
count = 0
for hb in history_boards:
if hb == board:
count += 1
if count >= 2:
break
return count

###############################################################################

B. 输入特征编码 => [119, 8, 8]

###############################################################################

def encode_chess_position_119planes(env, history_boards, max_history=8):
"""
返回 [119,8,8] 的 numpy.float32, 类似AlphaZero风格的特征:
- 最多 8 步历史(每步14平面=12子平面+2重复次数)
- 1平面(当前执方)
- 4平面(王车易位权限)
- 1平面(fullmove_number 缩放)
- 1平面(halfmove_clock 缩放)
"""
board = env.board
planes = np.zeros((119, 8, 8), dtype=np.float32)

text
# 收集所有历史位面(含当前局面)
all_positions = history_boards + [board]
all_positions = all_positions[-max_history:]

for i in range(max_history):
    plane_offset = i * 14
    real_index = len(all_positions) - max_history + i
    if real_index < 0:
        continue
    b = all_positions[real_index]
    if b is None:
        continue

    # 先记录重复次数
    rep_count = get_repetition_count(b, all_positions[:real_index])
    if rep_count >= 1:
        planes[plane_offset + 12, :, :] = 1.0
    if rep_count >= 2:
        planes[plane_offset + 13, :, :] = 1.0

    # 记录棋子分布
    piece_map = b.piece_map()
    for sq, piece in piece_map.items():
        row = 7 - (sq // 8)
        col = sq % 8
        symbol = piece.symbol()  # 大写=白, 小写=黑
        if symbol == 'P':
            planes[plane_offset + 0, row, col] = 1
        elif symbol == 'N':
            planes[plane_offset + 1, row, col] = 1
        elif symbol == 'B':
            planes[plane_offset + 2, row, col] = 1
        elif symbol == 'R':
            planes[plane_offset + 3, row, col] = 1
        elif symbol == 'Q':
            planes[plane_offset + 4, row, col] = 1
        elif symbol == 'K':
            planes[plane_offset + 5, row, col] = 1
        elif symbol == 'p':
            planes[plane_offset + 6, row, col] = 1
        elif symbol == 'n':
            planes[plane_offset + 7, row, col] = 1
        elif symbol == 'b':
            planes[plane_offset + 8, row, col] = 1
        elif symbol == 'r':
            planes[plane_offset + 9, row, col] = 1
        elif symbol == 'q':
            planes[plane_offset + 10, row, col] = 1
        elif symbol == 'k':
            planes[plane_offset + 11, row, col] = 1

plane_offset = max_history * 14  # 112

# 当前执方 (白:1, 黑:0)
if board.turn:
    planes[plane_offset, :, :] = 1.0
else:
    planes[plane_offset, :, :] = 0.0

import chess
# 王车易位权限
if board.has_kingside_castling_rights(chess.WHITE):
    planes[plane_offset+1, :, :] = 1.0
if board.has_queenside_castling_rights(chess.WHITE):
    planes[plane_offset+2, :, :] = 1.0
if board.has_kingside_castling_rights(chess.BLACK):
    planes[plane_offset+3, :, :] = 1.0
if board.has_queenside_castling_rights(chess.BLACK):
    planes[plane_offset+4, :, :] = 1.0

# fullmove_number & halfmove_clock (简单剪裁+归一化)
tmc = min(board.fullmove_number, 100)/100.0
hmc = min(board.halfmove_clock, 100)/100.0
planes[plane_offset+5, :, :] = tmc
planes[plane_offset+6, :, :] = hmc

return planes

###############################################################################

C. 全局动作表构建 (为 policy_head 输出做“扁平化”映射)

###############################################################################

import chess

def build_all_moves():
"""
先添加普通走子，再对所有满足兵升变条件的 from->to 增加4种升变.
"""
all_moves_set = set()

text
for from_sq in range(64):
    for to_sq in range(64):
        if from_sq == to_sq:
            continue

        # 先添加普通走子 (不带promotion)
        mv = chess.Move(from_sq, to_sq)
        all_moves_set.add(mv)

        # 如果是兵升变，则额外添加4种升变走子
        from_rank = from_sq // 8
        to_rank   = to_sq // 8
        is_white_promo = (from_rank == 6 and to_rank == 7)
        is_black_promo = (from_rank == 1 and to_rank == 0)
        if is_white_promo or is_black_promo:
            for promo_piece in [chess.QUEEN, chess.ROOK, chess.BISHOP, chess.KNIGHT]:
                mv_promo = chess.Move(from_sq, to_sq, promotion=promo_piece)
                all_moves_set.add(mv_promo)

all_moves_list = sorted(list(all_moves_set), key=lambda m: m.uci())
move2idx = {}
idx2move = {}
for i, mv in enumerate(all_moves_list):
    uci_str = mv.uci()
    move2idx[uci_str] = i
    idx2move[i] = uci_str
return all_moves_list, move2idx, idx2move

ALL_MOVES, move2idx, idx2move = build_all_moves()
print("Total distinct moves:", len(ALL_MOVES))

###############################################################################

D. 定义 AlphaZero 风格网络（ResNet）

###############################################################################

class BasicBlock(nn.Module):
"""简化残差块"""
def init(self, channels=128):
super(BasicBlock, self).init()
self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(channels)
self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(channels)

text
def forward(self, x):
    identity = x
    out = F.relu(self.bn1(self.conv1(x)))
    out = self.bn2(self.conv2(out))
    out = out + identity
    out = F.relu(out)
    return out

class AlphaZeroChessNet(nn.Module):
"""
单一网络，输出 (policy_logits, value).
policy_logits 大小 = len(ALL_MOVES).
"""
def init(self, input_channels=119, policy_output_size=None,
num_blocks=4, base_channels=64):
super(AlphaZeroChessNet, self).init()
if policy_output_size is None:
policy_output_size = len(ALL_MOVES)
self.policy_output_size = policy_output_size

text
    self.conv_in = nn.Conv2d(input_channels, base_channels, kernel_size=3, padding=1)
    self.bn_in   = nn.BatchNorm2d(base_channels)

    self.res_blocks = nn.ModuleList([BasicBlock(base_channels) for _ in range(num_blocks)])

    # 策略头
    self.policy_conv = nn.Conv2d(base_channels, 2, kernel_size=1)
    self.policy_bn   = nn.BatchNorm2d(2)
    self.policy_fc   = nn.Linear(2*8*8, policy_output_size)

    # 价值头
    self.value_conv = nn.Conv2d(base_channels, 1, kernel_size=1)
    self.value_bn   = nn.BatchNorm2d(1)
    self.value_fc1  = nn.Linear(8*8*1, 128)
    self.value_fc2  = nn.Linear(128, 1)

def forward(self, x):
    # x: [batch_size,119,8,8]
    x = F.relu(self.bn_in(self.conv_in(x)))
    for block in self.res_blocks:
        x = block(x)

    # policy
    p = F.relu(self.policy_bn(self.policy_conv(x)))
    p = p.view(p.size(0), -1)  # [batch_size, 2*8*8]
    p = self.policy_fc(p)

    # value
    v = F.relu(self.value_bn(self.value_conv(x)))
    v = v.view(v.size(0), -1)
    v = F.relu(self.value_fc1(v))
    v = torch.tanh(self.value_fc2(v))

    return p, v

###############################################################################

E. MCTS 实现: 增加 root 节点的 Dirichlet 噪声

###############################################################################

class TreeNode:
"""
MCTS节点
"""
def init(self, parent, prior_p):
self.parent = parent
self.children = {}
self.Q = 0.0
self.P = prior_p
self.N = 0

text
def is_leaf(self):
    return (len(self.children) == 0)

def is_root(self):
    return (self.parent is None)

def expand(self, action_priors):
    """
    action_priors: [(move_uci, prob), ...]
    """
    for (move_uci, prob) in action_priors:
        if move_uci not in self.children:
            self.children[move_uci] = TreeNode(self, prob)

def update(self, leaf_value):
    """
    leaf_value: 从子节点回传的价值（以当前节点视角）
    """
    self.N += 1
    self.Q += (leaf_value - self.Q)/self.N

def update_recursive(self, leaf_value):
    """
    递归地向上更新祖先节点
    """
    if self.parent is not None:
        self.parent.update_recursive(-leaf_value)
    self.update(leaf_value)

def get_value(self, c_puct):
    """
    UCB：Q + U
    U = c_puct * P * sqrt(parent.N)/(1+N)
    """
    U = c_puct * self.P * math.sqrt(self.parent.N)/(1+self.N)
    return self.Q + U

def select(self, c_puct):
    """
    在children里选 get_value() 最大的action
    返回 (action_uci, child_node)
    """
    return max(self.children.items(), key=lambda act_node: act_node[1].get_value(c_puct))

class MCTS:
def init(self, model, c_puct=5, n_playout=1000, device='cpu',
dirichlet_alpha=0.3, dirichlet_epsilon=0.25):
"""
model: AlphaZeroChessNet
"""
self.model = model
self.c_puct = c_puct
self.n_playout = n_playout
self.device = device

text
    # Dirichlet 噪声相关超参(仅在 root 上添加一次)
    self.dirichlet_alpha = dirichlet_alpha
    self.dirichlet_epsilon = dirichlet_epsilon

    self.root = TreeNode(None, 1.0)

def _playout(self, env, history_boards):
    """
    执行一次从根到叶的模拟
    """
    node = self.root
    # 1) selection
    while not node.is_leaf():
        action_uci, node = node.select(self.c_puct)
        done, _ = env.step(action_uci)
        if done:
            break

    # 2) expansion & evaluation
    if not env.is_game_over():
        # 用model评估
        planes = encode_chess_position_119planes(env, history_boards, 8)
        planes_t = torch.from_numpy(planes).unsqueeze(0).to(self.device)
        with torch.no_grad():
            policy_logits, value = self.model(planes_t)
            policy = F.softmax(policy_logits, dim=1)[0].cpu().numpy()
            leaf_value = value.item()

        # 只保留当前合规则子，对其概率做归一化
        legal_uci = env.get_legal_moves()
        action_priors = []
        sum_p = 0.0
        for m_uci in legal_uci:
            idx = move2idx[m_uci]
            p = policy[idx]
            sum_p += p
        if sum_p < 1e-8:
            # 若网络给的概率全为0(或极小), 均分
            prob = 1.0/len(legal_uci)
            for m_uci in legal_uci:
                action_priors.append((m_uci, prob))
        else:
            for m_uci in legal_uci:
                idx = move2idx[m_uci]
                p = policy[idx]/sum_p
                action_priors.append((m_uci, p))

        node.expand(action_priors)
    else:
        # 若已经终局，leaf_value 就是最终胜负(对当前节点玩家)
        leaf_value = env.get_result()

    # 3) backup
    node.update_recursive(-leaf_value)

def get_action_probs(self, env, history_boards, temp=1.0):
    """
    对当前root进行 n_playout 次模拟，然后根据 child.N 分布返回策略概率(AlphaZero风格).
    返回大小与 model.policy_output_size 相同，仅合法招有非零概率.
    """
    # ---- 如果root未展开, 先让它展开一次 (以便后续添加Dirichlet噪声) ----
    if self.root.is_leaf():
        planes = encode_chess_position_119planes(env, history_boards, 8)
        planes_t = torch.from_numpy(planes).unsqueeze(0).to(self.device)
        with torch.no_grad():
            policy_logits, value = self.model(planes_t)
            policy = F.softmax(policy_logits, dim=1)[0].cpu().numpy()

        legal_uci = env.get_legal_moves()
        sum_p = 0.0
        action_priors = []
        for m_uci in legal_uci:
            idx = move2idx[m_uci]
            sum_p += policy[idx]
        if sum_p < 1e-8:
            prob = 1.0/len(legal_uci)
            for m_uci in legal_uci:
                action_priors.append((m_uci, prob))
        else:
            for m_uci in legal_uci:
                idx = move2idx[m_uci]
                p = policy[idx]/sum_p
                action_priors.append((m_uci, p))
        self.root.expand(action_priors)

    # ---- 在 root 的先验上添加 Dirichlet 噪声 (仅自对弈时使用) ----
    self._add_dirichlet_noise()

    # ---- 进行 n_playout 次模拟 ----
    for _ in range(self.n_playout):
        env_copy = env.copy()
        self._playout(env_copy, copy.deepcopy(history_boards))

    # ---- 统计并返回基于 N 的最终策略分布 ----
    children_items = list(self.root.children.items())
    if not children_items:
        # 无子节点(终局)
        pi = np.zeros(self.model.policy_output_size, dtype=np.float32)
        pi[0] = 1.0
        return pi

    actions, child_nodes = zip(*children_items)
    counts = np.array([c.N for c in child_nodes], dtype=np.float32)

    # 如果 temp <= 1e-3，直接近似贪心
    if temp <= 1e-3:
        best_id = np.argmax(counts)
        probs = np.zeros_like(counts)
        probs[best_id] = 1.0
    else:
        # 温度因子
        counts_exp = counts ** (1.0 / temp)
        # 若出现inf，做个小clip
        counts_exp = np.clip(counts_exp, 1e-30, 1e30)
        probs = counts_exp / np.sum(counts_exp)

    # 组装为全局向量 pi
    pi = np.zeros(self.model.policy_output_size, dtype=np.float32)
    for a, p in zip(actions, probs):
        idx = move2idx[a]
        pi[idx] = p
    return pi

def _add_dirichlet_noise(self):
    """在 root 的 children 概率上添加一次狄利克雷噪声, 仅用于自对弈时."""
    if not self.root.is_leaf():  # root 已展开
        children_items = list(self.root.children.items())
        actions, child_nodes = zip(*children_items)
        # 生成 dirichlet
        alpha_vec = [self.dirichlet_alpha]*len(child_nodes)
        noise = np.random.dirichlet(alpha_vec)
        for a, n in zip(actions, noise):
            self.root.children[a].P = \
                (1 - self.dirichlet_epsilon)*self.root.children[a].P \
                + self.dirichlet_epsilon * n

def update_with_move(self, last_move):
    """
    对外部执行一步动作后, 将MCTS的root切换到对应子节点(加速搜索).
    如果对应分支不存在, 则重置。
    """
    if last_move in self.root.children:
        self.root = self.root.children[last_move]
        self.root.parent = None
    else:
        self.root = TreeNode(None, 1.0)

###############################################################################

F. Elo评分

###############################################################################

def probability(rA, rB):
return 1.0/(1+10**((rB-rA)/400.0))

def elo_update(rA, rB, outcome, K=30):
pA = probability(rA, rB)
newRa = rA + K*(outcome - pA)
newRb = rB + K*((1.0-outcome) - (1.0-pA))
return newRa, newRb

def evaluate_models_elo(modelA, modelB, n_games=10, device='cpu'):
"""
A vs B 共下 n_games 局，每局更新 Elo. 初始各 2000 分.
可以交替先手, 这里简单让 A 始终先手也行.
"""
RA, RB = 2000.0, 2000.0
for _ in range(n_games):
env = ChessEnv()
history_boards = []
mctsA = MCTS(modelA, c_puct=5, n_playout=16, device=device,
dirichlet_alpha=0.3, dirichlet_epsilon=0.25)
mctsB = MCTS(modelB, c_puct=5, n_playout=16, device=device,
dirichlet_alpha=0.3, dirichlet_epsilon=0.25)

text
    done = False
    while not done:
        cp = env.current_player()
        if cp == 1:  # A走(白)
            pi = mctsA.get_action_probs(env, history_boards, temp=1e-3)
            move_idx = np.argmax(pi)
            move_uci = idx2move[move_idx]
            done, _ = env.step(move_uci)
            mctsA.update_with_move(move_uci)
        else:        # B走(黑)
            pi = mctsB.get_action_probs(env, history_boards, temp=1e-3)
            move_idx = np.argmax(pi)
            move_uci = idx2move[move_idx]
            done, _ = env.step(move_uci)
            mctsB.update_with_move(move_uci)

        history_boards.append(env.board.copy())

    res = env.get_result()  # +1/-1/0
    if res == 1:
        outcome = 1.0
    elif res == -1:
        outcome = 0.0
    else:
        outcome = 0.5

    RA, RB = elo_update(RA, RB, outcome)

return RA, RB

###############################################################################

G. 训练过程 (AlphaZero风格: 自对弈 + 训练 + 可选的新旧网络对弈验证)

###############################################################################

def generate_self_play_data(model, n_games=1, n_playout=32, device='cpu',
temp_threshold=30):
"""
使用当前 model 进行自对弈, 生成 (plane, pi, z).
- 为了模拟 AlphaZero 对温度的处理:
在对局前若干步(例如前30步)使用 temperature=1.0,
超过这个步数后使用近似 0 的温度(如 1e-3)，更趋向贪心。
"""
data = []
for _ in range(n_games):
env = ChessEnv()
history_boards = []
# MCTS 实例(自对弈时, root 会加 dirichlet noise)
mcts = MCTS(model, c_puct=5, n_playout=n_playout, device=device,
dirichlet_alpha=0.3, dirichlet_epsilon=0.25)

text
    states = []
    mcts_probs = []
    players = []

    move_count = 0

    while not env.is_game_over():
        move_count += 1
        # 根据 move_count 动态调节 temperature
        if move_count <= temp_threshold:
            temperature = 1.0
        else:
            temperature = 1e-3

        planes = encode_chess_position_119planes(env, history_boards, max_history=8)
        pi = mcts.get_action_probs(env, history_boards, temp=temperature)

        # 只在 合法走子 中做采样 (若 temperature=1.0)
        legal_moves = env.get_legal_moves()
        legal_pi_sum = 0.0
        legal_pairs = []
        for m_uci in legal_moves:
            idx = move2idx[m_uci]
            p = pi[idx]
            legal_pi_sum += p
            legal_pairs.append((m_uci, idx, p))

        if legal_pi_sum < 1e-8:
            # 若总概率极小，直接随机选
            chosen_move = random.choice(legal_moves)
            chosen_idx  = move2idx[chosen_move]
        else:
            # 归一化后做随机采样 (在 temperature=1 时有效果)
            for i, (m_uci, idx, p) in enumerate(legal_pairs):
                legal_pairs[i] = (m_uci, idx, p / legal_pi_sum)
            moves_uci  = [lp[0] for lp in legal_pairs]
            moves_idx  = [lp[1] for lp in legal_pairs]
            moves_prob = [lp[2] for lp in legal_pairs]

            chosen_id = np.random.choice(len(moves_prob), p=moves_prob)
            chosen_move = moves_uci[chosen_id]
            chosen_idx  = moves_idx[chosen_id]

        # 记录(s, pi, player)
        states.append(planes)
        mcts_probs.append(pi)  # 对应全局move的分布
        players.append(env.current_player())

        # 执行该动作
        done, winner = env.step(chosen_move)
        mcts.update_with_move(chosen_move)
        history_boards.append(env.board.copy())

    # 终局, 回填 z
    final_res = env.get_result()  # +1/-1/0
    for i, plyr in enumerate(players):
        if final_res == 0:
            z = 0.0
        else:
            z = 1.0 if (final_res == plyr) else -1.0
        data.append((states[i], mcts_probs[i], z))

return data

def train_alpha_zero_net(model, optimizer, dataset, batch_size=32, epochs=1, device='cpu'):
"""
对 dataset 训练 (policy-loss + value-loss).
dataset: [(plane[119,8,8], pi[policy_size], z), ...]
"""
model.train()
for e in range(epochs):
random.shuffle(dataset)
total_loss = 0
for i in range(0, len(dataset), batch_size):
mini_batch = dataset[i:i+batch_size]
planes_batch = []
pi_batch = []
z_batch = []
for (planes, pi_vec, z_val) in mini_batch:
planes_batch.append(planes)
pi_batch.append(pi_vec)
z_batch.append(z_val)

text
        planes_batch = torch.tensor(planes_batch, dtype=torch.float32, device=device)
        pi_batch     = torch.tensor(pi_batch, dtype=torch.float32, device=device)
        z_batch      = torch.tensor(z_batch, dtype=torch.float32, device=device).unsqueeze(-1)

        optimizer.zero_grad()
        policy_logits, value = model(planes_batch)
        # policy loss (交叉熵)
        logp = F.log_softmax(policy_logits, dim=1)
        policy_loss = -torch.mean(torch.sum(pi_batch * logp, dim=1))
        # value loss (均方误差)
        value_loss = F.mse_loss(value, z_batch)
        loss = policy_loss + value_loss
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    print(f"Train epoch {e}, loss={total_loss:.4f}")

###############################################################################

H. 主训练循环示例

###############################################################################

def update_best_model(current_model, new_model, n_test_games=10, device='cpu'):
"""
可选：对比 current_model / new_model 对弈结果，如 new_model 更强则替换。
这里演示用 Elo 来简单判断，也可统计胜率。
"""
RA, RB = evaluate_models_elo(current_model, new_model, n_games=n_test_games, device=device)
# 以最终 Elo 来粗略判定
if RB > RA: # 新网络更强
print("[Gating] new model is stronger, update best model.")
# 把 new_model 参数复制给 current_model
current_model.load_state_dict(new_model.state_dict())
else:
print("[Gating] new model is not stronger, keep old model.")

def main_training_loop():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

text
# 1) 初始化“最佳”网络 (best_model)
policy_size = len(ALL_MOVES)  # 与构建的动作表一致
best_model = AlphaZeroChessNet(
    input_channels=119,
    policy_output_size=policy_size,
    num_blocks=4,  # 可改大
    base_channels=64
).to(device)

# 优化器
best_model_optimizer = optim.Adam(best_model.parameters(), lr=1e-3)

# 2) 经验缓存
replay_buffer = []
max_buffer_size = 20000

# 3) 多轮迭代训练
num_iterations = 30
evaluate_interval = 5

for it in range(1, num_iterations+1):
    # 1) 用“当前最佳”模型自对弈生成数据
    new_data = generate_self_play_data(best_model, n_games=2, n_playout=32,
                                       device=device, temp_threshold=30)
    # 2) 放入 replay_buffer
    replay_buffer.extend(new_data)
    if len(replay_buffer) > max_buffer_size:
        replay_buffer = replay_buffer[-max_buffer_size:]

    # 3) 用 replay_buffer 训练“新”模型(先复制一份 best_model)
    new_model = AlphaZeroChessNet(
        input_channels=119,
        policy_output_size=policy_size,
        num_blocks=4,
        base_channels=64
    ).to(device)
    new_model.load_state_dict(best_model.state_dict())  # 初始化参数相同
    new_model_optimizer = optim.Adam(new_model.parameters(), lr=1e-3)

    # 这里仅训练 1 epoch 演示，实际可多轮
    train_alpha_zero_net(new_model, new_model_optimizer, replay_buffer,
                         batch_size=32, epochs=1, device=device)

    # 4) 可选：对新旧模型做一次对弈评估，若新模型更强则替换
    update_best_model(best_model, new_model, n_test_games=4, device=device)

    # 5) 也可定期输出 self-play Elo 测试等
    if it % evaluate_interval == 0:
        RA, RB = evaluate_models_elo(best_model, best_model, n_games=4, device=device)
        print(f"[Iter {it}] Elo self-play test: RA={RA:.1f}, RB={RB:.1f}")

torch.save(best_model.state_dict(), "best_model.pth")
print("Training finished.")

if name == "main":
main_training_loop()

帮我优化：
1）Self-play的数据产生速度，考虑用C++或cython优化这一环节；
2）self-play中的Elo评估太慢了，评价值忽高忽低，请给出更稳健的elo rate评估策略；
3）增大batch_size和训练速度，我现在8GB显存才占了不到1GB，把所有显存都用上！