用飞桨框架2.0造一个会下五子棋的AI模型

class Board(object):
    """棋盘游戏逻辑控制"""

    def __init__(self, **kwargs):
        self.width = int(kwargs.get("width", 15))  # 棋盘宽度
        self.height = int(kwargs.get("height", 15))  # 棋盘高度
        self.states = {}  # 棋盘状态为一个字典,键: 移动步数,值: 玩家的棋子类型
        self.n_in_row = int(kwargs.get("n_in_row", 5))  # 5个棋子一条线则获胜
        self.players = [1, 2]  # 玩家1,2

    def init_board(self, start_player=0):
        # 初始化棋盘

        # 当前棋盘的宽高小于5时,抛出异常(因为是五子棋)
        if self.width < self.n_in_row or self.height < self.n_in_row:
            raise Exception("棋盘的长宽不能少于{}".format(self.n_in_row))
        self.current_player = self.players[start_player]  # 先手玩家
        self.availables = list(
            range(self.width * self.height)
        )  # 初始化可用的位置列表
        self.states = {}  # 初始化棋盘状态
        self.last_move = -1  # 初始化最后一次的移动位置

    def current_state(self):
        """
        从当前玩家的角度返回棋盘状态。
        状态形式: 4 * 宽 * 高
        """
        # 使用4个15x15的二值特征平面来描述当前的局面
        # 前两个平面分别表示当前player的棋子位置和对手player的棋子位置，有棋子的位置是1，没棋子的位置是0
        # 第三个平面表示对手player最近一步的落子位置，也就是整个平面只有一个位置是1，其余全部是0
        # 第四个平面表示的是当前player是不是先手player，如果是先手player则整个平面全部为1，否则全部为0
        square_state = np.zeros((4, self.width, self.height))
        if self.states:
            moves, players = np.array(list(zip(*self.states.items())))
            move_curr = moves[
                players == self.current_player
            ]  # 获取棋盘状态上属于当前玩家的所有移动值
            move_oppo = moves[
                players != self.current_player
            ]  # 获取棋盘状态上属于对方玩家的所有移动值
            square_state[0][
                move_curr // self.width,  # 对第一个特征平面填充值(当前玩家)
                move_curr % self.height,
            ] = 1.0
            square_state[1][
                move_oppo // self.width,  # 对第二个特征平面填充值(对方玩家)
                move_oppo % self.height,
            ] = 1.0
            # 指出最后一个移动位置
            square_state[2][
                self.last_move
                // self.width,  # 对第三个特征平面填充值(对手最近一次的落子位置)
                self.last_move % self.height,
            ] = 1.0
        if (
            len(self.states) % 2 == 0
        ):  # 对第四个特征平面填充值,当前玩家是先手,则填充全1,否则为全0
            square_state[3][:, :] = 1.0
        # 将每个平面棋盘状态按行逆序转换(第一行换到最后一行,第二行换到倒数第二行..)
        return square_state[:, ::-1, :]

    def do_move(self, move):
        # 根据移动的数据更新各参数
        self.states[move] = self.current_player  # 将当前的参数存入棋盘状态中
        self.availables.remove(move)  # 从可用的棋盘列表移除当前移动的位置
        self.current_player = (
            self.players[0]
            if self.current_player == self.players[1]
            else self.players[1]
        )  # 改变当前玩家
        self.last_move = move  # 记录最后一次的移动位置

    def has_a_winner(self):
        # 是否产生赢家
        width = self.width  # 棋盘宽度
        height = self.height  # 棋盘高度
        states = self.states  # 状态
        n = self.n_in_row  # 获胜需要的棋子数量

        # 当前棋盘上所有的落子位置
        moved = list(set(range(width * height)) - set(self.availables))
        if len(moved) < self.n_in_row + 2:
            # 当前棋盘落子数在7个以上时会产生赢家,落子数低于7个时,直接返回没有赢家
            return False, -1

        # 遍历落子数
        for m in moved:
            h = m // width
            w = m % width  # 获得棋子的坐标
            player = states[m]  # 根据移动的点确认玩家

            # 判断各种赢棋的情况
            # 横向5个
            if (
                w in range(width - n + 1)
                and len(set(states.get(i, -1) for i in range(m, m + n))) == 1
            ):
                return True, player

            # 纵向5个
            if (
                h in range(height - n + 1)
                and len(
                    set(
                        states.get(i, -1)
                        for i in range(m, m + n * width, width)
                    )
                )
                == 1
            ):
                return True, player

            # 左上到右下斜向5个
            if (
                w in range(width - n + 1)
                and h in range(height - n + 1)
                and len(
                    set(
                        states.get(i, -1)
                        for i in range(m, m + n * (width + 1), width + 1)
                    )
                )
                == 1
            ):
                return True, player

            # 右上到左下斜向5个
            if (
                w in range(n - 1, width)
                and h in range(height - n + 1)
                and len(
                    set(
                        states.get(i, -1)
                        for i in range(m, m + n * (width - 1), width - 1)
                    )
                )
                == 1
            ):
                return True, player

        # 当前都没有赢家,返回False
        return False, -1

    def game_end(self):
        """检查当前棋局是否结束"""
        win, winner = self.has_a_winner()
        if win:
            return True, winner
        elif not len(self.availables):
            # 棋局布满,没有赢家
            return True, -1
        return False, -1

    def get_current_player(self):
        return self.current_player


# 加上UI的布局的训练方式
class Game_UI(object):
    """游戏控制区域"""

    def __init__(self, board, **kwargs):
        self.board = board  # 加载棋盘控制类

        # 初始化 pygame
        pygame.init()

    def start_play_evaluate(self, player1, player2, start_player=0):
        """开始一局游戏，评估当前的价值策略网络的胜率"""
        if start_player not in (0, 1):
            # 如果玩家不在玩家1,玩家2之间,抛出异常
            raise Exception("开始的玩家必须为0(玩家1)或1(玩家2)")
        self.board.init_board(start_player)  # 初始化棋盘
        p1, p2 = self.board.players  # 加载玩家1,玩家2
        player1.set_player_ind(p1)  # 设置玩家1
        player2.set_player_ind(p2)  # 设置玩家2
        players = {p1: player1, p2: player2}

        while True:
            current_player = self.board.current_player  # 获取当前玩家
            player_in_turn = players[current_player]  # 当前玩家的信息
            move = player_in_turn.get_action(
                self.board
            )  # 基于MCTS的AI下一步落子
            self.board.do_move(move)  # 根据下一步落子的状态更新棋盘各参数

            # 判断当前棋局是否结束
            end, winner = self.board.game_end()
            # 结束
            if end:
                win = winner
                break

        return win

    def start_play_train(self, player, temp=1e-3):
        """
        开始自我博弈，使用MCTS玩家开始自己玩游戏,重新使用搜索树并存储自己玩游戏的数据
        (state, mcts_probs, z) 提供训练
        """
        self.board.init_board()  # 初始化棋盘
        states, mcts_probs, current_players = (
            [],
            [],
            [],
        )  # 状态,mcts的行为概率,当前玩家

        while True:
            # 根据当前棋盘状态返回可能得行为,及行为对应的概率
            move, move_probs = player.get_action(
                self.board, temp=temp, return_prob=1
            )
            # 存储数据
            states.append(self.board.current_state())  # 存储状态数据
            mcts_probs.append(move_probs)  # 存储行为概率数据
            current_players.append(self.board.current_player)  # 存储当前玩家
            # 执行一个移动
            self.board.do_move(move)

            # 判断该局游戏是否终止
            end, winner = self.board.game_end()
            if end:
                # 从每个状态的当时的玩家的角度看待赢家
                winners_z = np.zeros(len(current_players))
                if winner != -1:
                    # 没有赢家时
                    winners_z[np.array(current_players) == winner] = 1.0
                    winners_z[np.array(current_players) != winner] = -1.0
                # 重置MSCT的根节点
                player.reset_player()
                return winner, zip(states, mcts_probs, winners_z)

def softmax(x):
    probs = np.exp(x - np.max(x))
    probs /= np.sum(probs)
    return probs


def policy_value_fn(board):
    """
    接受状态并输出（动作，概率）列表的函数元组和状态的分数"""
    # 返回统一概率和0分的纯MCTS
    action_probs = np.ones(len(board.availables)) / len(board.availables)
    return zip(board.availables, action_probs), 0


class TreeNode(object):
    """MCTS树中的节点。

    每个节点跟踪其自身的值Q，先验概率P及其访问次数调整的先前得分u。
    """

    def __init__(self, parent, prior_p):
        self._parent = parent
        self._children = {}  # 从动作到TreeNode的映射
        self._n_visits = 0
        self._Q = 0
        self._u = 0
        self._P = prior_p

    def expand(self, action_priors):
        """通过创建新子项来展开树。
        action_priors：一系列动作元组及其先验概率根据策略函数.
        """
        for action, prob in action_priors:
            if action not in self._children:
                self._children[action] = TreeNode(self, prob)

    def select(self, c_puct):
        """在子节点中选择能够提供最大行动价值Q的行动加上奖金u（P）。
        return：（action，next_node）的元组
        """
        return max(
            self._children.items(),
            key=lambda act_node: act_node[1].get_value(c_puct),
        )

    def update(self, leaf_value):
        """从叶节点评估中更新节点值
        leaf_value: 这个子树的评估值来自从当前玩家的视角
        """
        # 统计访问次数
        self._n_visits += 1
        # 更新Q值,取对于所有访问次数的平均数
        self._Q += 1.0 * (leaf_value - self._Q) / self._n_visits

    def update_recursive(self, leaf_value):
        """就像调用update（）一样，但是对所有祖先进行递归应用。"""
        # 如果它不是根节点，则应首先更新此节点的父节点。
        if self._parent:
            self._parent.update_recursive(-leaf_value)
        self.update(leaf_value)

    def get_value(self, c_puct):
        """计算并返回此节点的值。它是叶评估Q和此节点的先验的组合
        调整了访问次数，u。
        c_puct：控制相对影响的（0，inf）中的数字，该节点得分的值Q和先验概率P.
        """
        self._u = (
            c_puct
            * self._P
            * np.sqrt(self._parent._n_visits)
            / (1 + self._n_visits)
        )
        return self._Q + self._u

    def is_leaf(self):
        """检查叶节点（即没有扩展的节点）。"""
        return self._children == {}

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


class MCTS(object):
    """对蒙特卡罗树搜索的一个简单实现"""

    def __init__(
        self, policy_value_fn, c_puct=5, n_playout=10000, mode="train"
    ):
        """
        policy_value_fn：一个接收板状态和输出的函数（动作，概率）元组列表以及[-1,1]中的分数
             （即来自当前的最终比赛得分的预期值玩家的观点）对于当前的玩家。
        c_puct：（0，inf）中的数字，用于控制探索的速度收敛于最大值政策。 更高的价值意味着
                 依靠先前的更多。
        """
        self._root = TreeNode(None, 1.0)
        self._policy = policy_value_fn
        self._c_puct = c_puct
        self._n_playout = n_playout
        self.mode = mode

    def _playout(self, state):
        """从根到叶子运行单个播出，获取值
        叶子并通过它的父母传播回来。
        State已就地修改，因此必须提供副本。
        """
        node = self._root
        while True:
            if node.is_leaf():
                break
            # 贪心算法选择下一步行动
            action, node = node.select(self._c_puct)
            state.do_move(action)

        # 使用网络评估叶子，该网络输出（动作，概率）元组p的列表以及当前玩家的[-1,1]中的分数v。
        action_probs, leaf_value = self._policy(state)
        # 查看游戏是否结束
        end, winner = state.game_end()
        if not end:
            node.expand(action_probs)
        if self.mode == "train":
            if end:
                # 对于结束状态,将叶子节点的值换成"true"
                if winner == -1:  # tie
                    leaf_value = 0.0
                else:
                    leaf_value = (
                        1.0 if winner == state.get_current_player() else -1.0
                    )
        else:
            # 通过随机的rollout评估叶子结点
            leaf_value = self._evaluate_rollout(state)
        # 在本次遍历中更新节点的值和访问次数
        node.update_recursive(-leaf_value)

    @staticmethod
    def _evaluate_rollout(state, limit=1000):
        """使用推出策略直到游戏结束，
        如果当前玩家获胜则返回+1，如果对手获胜则返回-1，
        如果是平局则为0。
        """
        player = state.get_current_player()
        winner = -1
        for i in range(limit):
            end, winner = state.game_end()
            if end:
                break
            max_action = np.random.choice(board.availables)
            state.do_move(max_action)
        else:
            # 如果没有从循环中断，请发出警告。
            print("WARNING: rollout reached move limit")
        if winner == -1:  # tie
            return 0
        else:
            return 1 if winner == player else -1

    def get_move(self, state, temp=1e-3):
        """
        如果 prob 为 True，则按顺序运行所有播出并返回可用的操作及其相应的概率。
        否则按顺序运行所有播出并返回访问量最大的操作。
        """
        for n in range(self._n_playout):
            state_copy = copy.deepcopy(state)
            self._playout(state_copy)
        if self.mode == "train":
            # 根据根节点处的访问计数来计算移动概率
            act_visits = [
                (act, node._n_visits)
                for act, node in self._root._children.items()
            ]
            acts, visits = zip(*act_visits)
            act_probs = softmax(1.0 / temp * np.log(np.array(visits) + 1e-10))

            return acts, act_probs

        return max(
            self._root._children.items(),
            key=lambda act_node: act_node[1]._n_visits,
        )[0]

    def update_with_move(self, last_move):
        """保留我们已经知道的关于子树的信息"""
        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)

    def __str__(self):
        return "MCTS"


class MCTSPlayer(object):
    """基于MCTS的AI玩家"""

    def __init__(
        self,
        policy_value_function=policy_value_fn,
        c_puct=5,
        n_playout=2000,
        is_selfplay=0,
        mode="train",
    ):
        self.mcts = MCTS(policy_value_function, c_puct, n_playout, mode)
        self._is_selfplay = is_selfplay

    def set_player_ind(self, p):
        self.player = p

    def reset_player(self):
        self.mcts.update_with_move(-1)

    def get_action(self, board, temp=1e-3, return_prob=0):
        sensible_moves = board.availables
        # 像alphaGo Zero论文一样使用MCTS算法返回的pi向量
        move_probs = np.zeros(board.width * board.height)
        if len(sensible_moves) > 0:
            if self.mcts.mode == "train":
                acts, probs = self.mcts.get_move(board, temp)
                move_probs[list(acts)] = probs
                if self._is_selfplay:
                    # 添加Dirichlet Noise进行探索（自我训练所需）
                    move = np.random.choice(
                        acts,
                        p=0.75 * probs
                        + 0.25 * np.random.dirichlet(0.3 * np.ones(len(probs))),
                    )
                    # 更新根节点并重用搜索树
                    self.mcts.update_with_move(move)
                else:
                    # 使用默认的temp = 1e-3，它几乎相当于选择具有最高概率的移动
                    move = np.random.choice(acts, p=probs)
                    # 重置根节点
                    self.mcts.update_with_move(-1)

                if return_prob:
                    return move, move_probs
                else:
                    return move
            else:
                move = self.mcts.get_move(board)
                self.mcts.update_with_move(-1)
                return move
        else:
            print("棋盘已满")

    def __str__(self):
        return "MCTS {}".format(self.player)