MCTS.py

import math
import numpy as np

EPS = 1e-8
DEPTHMAX = 30


class MCTS:
    """
    Monte-Carlo tree search
    """

    def __init__(self, game, nnet, args):
        self.game = game
        self.nnet = nnet
        self.args = args
        self.Qsa = {}  # stores Q values for s,a (as defined in the paper)
        self.Nsa = {}  # stores #times edge s,a was visited
        self.Ns = {}  # stores #times board s was visited
        self.Ps = {}  # stores initial policy (returned by neural net)
        self.Loop = {}
        self.Visited = []
        self.wrong_predictions = [0, 0]

        self.Es = {}  # stores game.getGameEnded ended for board s
        self.Vs = {}  # stores game.getValidMoves for board s

    def get_action_prob(self, board, player, best=False):
        """
        the main function of the tree search
        :param board:   current board
        :param player:  current player
        :param best:    boolean that denotes if best selection or varied selection is used
        :return:        probability vector
        """
        self.reset()

        s = self.game.stringRepresentation(board)

        for i in range(self.args.numMCTSSims):
            self.search(board, player, 0)

        counts = [self.Nsa[(s, a, player)] if (s, a, player) in self.Nsa else 0 for a in range(self.game.getActionSize())]

        if best:
            bestA = np.argmax(counts)
            probs = [0] * len(counts)
            probs[bestA] = 1
            return probs

        sum_counts = sum(counts)
        if sum_counts == 0:
            print(board)
            print("Random Move by " + str(player) + "!")
            counts = self.game.getValidMoves(board, player)
            sum_counts = sum(counts)

        probs = [x / float(sum_counts) for x in counts]
        return probs

    def search(self, board, player, depth):
        """
        recursive search function
        one search iteration expands the tree by a leaf node, except when a loop-cut is executed
        and updates tree values
        :param board:   current board
        :param player:  current player
        :param depth:   current depth
        :return:        scores that get propagated by the recursioun
        """
        s = self.game.stringRepresentation(board)

        scores = self.game.getGameEnded(board, True).astype('float16')

        canonical_board = self.game.getCanonicalForm(board, player)

        if s not in self.Es:
            self.Es[s] = np.copy(scores)
        if np.count_nonzero(self.Es[s]) > 1:
            # terminal node
            return scores

        if depth == 0:
            self.Loop = [(s, player)]
        else:
            if s in self.Visited:
                # print("VISITED")
                return scores
            if (s, player) in self.Loop:
                print("Prevented Loop! Depth: " + str(depth))
                return scores
            else:
                self.Loop.append((s, player))

        if (s, player) not in self.Ps or depth > DEPTHMAX:
            if depth > DEPTHMAX:

                print("CUT LEAF")
            # leaf node
            self.Ps[(s, player)], scores_nn = self.nnet.predict(canonical_board)
            if player == 2:
                scores_nn = np.array([scores_nn[2], scores_nn[0], scores_nn[1]])
            elif player == 3:
                scores_nn = np.array([scores_nn[1], scores_nn[2], scores_nn[0]])
            valids = self.game.getValidMoves(canonical_board, 1)
            self.Ps[(s, player)] = self.Ps[(s, player)] * valids  # masking invalid moves
            sum_Ps_s = np.sum(self.Ps[(s, player)])
            self.wrong_predictions = [self.wrong_predictions[0] + 1, self.wrong_predictions[1] + 1 - sum_Ps_s]
            if sum_Ps_s > 0:
                self.Ps[(s, player)] /= sum_Ps_s  # renormalize
            else:
                # if all valid moves were masked make all valid moves equally probable

                # NB! All valid moves may be masked if either your NNet architecture is insufficient or you've get overfitting or something else.
                # If you have got dozens or hundreds of these messages you should pay attention to your NNet and/or training process.
                print("All valid moves were masked, do workaround.")
                self.Ps[(s, player)] = self.Ps[(s, player)] + valids
                self.Ps[(s, player)] /= np.sum(self.Ps[(s, player)])

            self.Vs[(s, player)] = valids
            self.Ns[(s, player)] = 0

            for i in range(3):
                if scores[i] == 0:
                    scores[i] = scores_nn[i]

            return scores

        valids = self.Vs[(s, player)]
        cur_best = -float('inf')
        best_act = -1

        # pick the action with the highest upper confidence bound
        for a in range(self.game.getActionSize()):
            if valids[a]:
                if (s, a, player) in self.Qsa:
                    u = self.Qsa[(s, a, player)] + self.args.cpuct * self.Ps[(s, player)][a] * math.sqrt(self.Ns[s, player]) / (
                                1 + self.Nsa[(s, a, player)])
                else:
                    u = self.args.cpuct * self.Ps[(s, player)][a] * math.sqrt(self.Ns[(s, player)] + EPS)  # Q = 0 ?
                    u = u

                if u > cur_best:
                    cur_best = u
                    best_act = a

        a = best_act
        next_s, next_player = self.game.getNextState(board, player, a)

        scores = self.search(next_s, next_player, depth + 1)

        if (s, a, player) in self.Qsa:
            self.Qsa[(s, a, player)] = (self.Nsa[(s, a, player)] * self.Qsa[(s, a, player)] + scores[player-1]) / (self.Nsa[(s, a, player)] + 1)
            self.Nsa[(s, a, player)] += 1

        else:
            self.Qsa[(s, a, player)] = scores[player-1]
            self.Nsa[(s, a, player)] = 1

        self.Ns[(s, player)] += 1

        return scores

    def reset(self):
        """
        resets tree
        """
        self.Qsa = {}  # stores Q values for s,a (as defined in the paper)
        self.Nsa = {}  # stores #times edge s,a was visited
        self.Ns = {}  # stores #times board s was visited
        self.Ps = {}  # stores initial policy (returned by neural net)

        self.Es = {}  # stores game.getGameEnded ended for board s
        self.Vs = {}  # stores game.getValidMoves for board s

    def get_wrong_prediction_rate(self):
        """
        :return:    average pre-evaluation scores assigned to illegal moves
        """
        wrong_sum = self.wrong_predictions[1]
        counts = self.wrong_predictions[0]
        self.wrong_predictions = [0, 0]
        return wrong_sum / counts