LightZero/lzero/mcts/ptree/ptree_az_sampled.py

"""
Overview:
    This code implements the Monte Carlo Tree Search (MCTS) algorithm with the integration of neural networks.
    The Node class represents a node in the Monte Carlo tree and implements the basic functionalities expected in a node.
    The MCTS class implements the specific search functionality and provides the optimal action through the ``get_next_action`` method.
    Compared to traditional MCTS, the introduction of value networks and policy networks brings several advantages.
    During the expansion of nodes, it is no longer necessary to explore every single child node, but instead,
    the child nodes are directly selected based on the prior probabilities provided by the neural network.
    This reduces the breadth of the search. When estimating the value of leaf nodes, there is no need for a rollout;
    instead, the value output by the neural network is used, which saves the depth of the search.
"""

import math
from typing import List, Tuple, Union, Callable, Type, Dict, Any

import numpy as np
import torch
from ding.envs import BaseEnv
from easydict import EasyDict
from numpy import ndarray

from lzero.mcts.ptree.ptree_sez import Action


class Node(object):
    """
    Overview:
        A class for a node in a Monte Carlo Tree. The properties of this class store basic information about the node,
        such as its parent node, child nodes, and the number of times the node has been visited.
        The methods of this class implement basic functionalities that a node should have, such as propagating the value back,
        checking if the node is the root node, and determining if it is a leaf node.
    """

    def __init__(self, parent: "Node" = None, prior_p: float = 1.0) -> None:
        """
        Overview:
            Initialize a Node object.
        Arguments:
            - parent (:obj:`Node`): The parent node of the current node.
            - prior_p (:obj:`Float`): The prior probability of selecting this node.
        """
        # The parent node.
        self._parent = parent
        # A dictionary representing the children of the current node. The keys are the actions, and the values are
        # the child nodes.
        self._children = {}
        # The number of times this node has been visited.
        self._visit_count = 0
        # The sum of the values of all child nodes of this node.
        self._value_sum = 0
        # The prior probability of selecting this node.
        self.prior_p = prior_p

    @property
    def value(self) -> float:
        """
        Overview:
            The value of the current node.
        Returns:
            - output (:obj:`Int`): Current value, used to compute ucb score.
        """
        # Computes the average value of the current node.
        if self._visit_count == 0:
            return 0
        return self._value_sum / self._visit_count

    def update(self, value: float) -> None:
        """
        Overview:
            Update the current node information, such as ``_visit_count`` and ``_value_sum``.
        Arguments:
            - value (:obj:`Float`): The value of the node.
        """
        # Updates the number of times this node has been visited.
        self._visit_count += 1
        # Updates the sum of the values of all child nodes of this node.
        self._value_sum += value

    def update_recursive(self, leaf_value: float, battle_mode_in_simulation_env: str) -> None:
        """
        Overview:
            Update node information recursively.
            The same game state has opposite values in the eyes of two players playing against each other.
            The value of a node is evaluated from the perspective of the player corresponding to its parent node.
            In ``self_play_mode``, because the player corresponding to a node changes every step during the backpropagation process, the value needs to be negated once.
            In ``play_with_bot_mode``, since all nodes correspond to the same player, the value does not need to be negated.

        Arguments:
            - leaf_value (:obj:`Float`): The value of the node.
            - battle_mode_in_simulation_env (:obj:`str`): The mode of MCTS, can be 'self_play_mode' or 'play_with_bot_mode'.
        """
        # Update the node information recursively based on the MCTS mode.
        if battle_mode_in_simulation_env == 'self_play_mode':
            # Update the current node's information.
            self.update(leaf_value)
            # If the current node is the root node, return.
            if self.is_root():
                return
            # Update the parent node's information recursively. When propagating the value back to the parent node,
            # the value needs to be negated once because the perspective of evaluation has changed.
            self._parent.update_recursive(-leaf_value, battle_mode_in_simulation_env)
        if battle_mode_in_simulation_env == 'play_with_bot_mode':
            # Update the current node's information.
            self.update(leaf_value)
            # If the current node is the root node, return.
            if self.is_root():
                return
            # Update the parent node's information recursively. In ``play_with_bot_mode``, since the nodes' values
            # are always evaluated from the perspective of the agent player, there is no need to negate the value
            # during value propagation.
            self._parent.update_recursive(leaf_value, battle_mode_in_simulation_env)

    def is_leaf(self) -> bool:
        """
        Overview:
            Check if the current node is a leaf node or not.
        Returns:
            - output (:obj:`Bool`): If self._children is empty, it means that the node has not
            been expanded yet, which indicates that the node is a leaf node.
        """
        # Returns True if the node is a leaf node (i.e., has no children), and False otherwise.
        return self._children == {}

    def is_root(self) -> bool:
        """
        Overview:
            Check if the current node is a root node or not.
        Returns:
            - output (:obj:`Bool`): If the node does not have a parent node,
            then it is a root node.
        """
        return self._parent is None

    @property
    def parent(self) -> None:
        """
        Overview:
            Get the parent node of the current node.
        Returns:
            - output (:obj:`Node`): The parent node of the current node.
        """
        return self._parent

    @property
    def children(self) -> None:
        """
        Overview:
            Get the dictionary of children nodes of the current node.
        Returns:
            - output (:obj:`dict`): A dictionary representing the children of the current node.
        """
        return self._children

    @property
    def visit_count(self) -> None:
        """
        Overview:
            Get the number of times the current node has been visited.
        Returns:
            - output (:obj:`Int`): The number of times the current node has been visited.
        """
        return self._visit_count


class MCTS(object):
    """
    Overview:
        A class for Monte Carlo Tree Search (MCTS). The methods in this class implement the steps involved in MCTS, such as selection and expansion.
        Based on this, the ``_simulate`` method is used to traverse from the root node to a leaf node.
        Finally, by repeatedly calling ``_simulate`` through ``get_next_action``, the optimal action is obtained.
    """

    def __init__(self, cfg: EasyDict, simulate_env: Type[BaseEnv]) -> None:
        """
        Overview:
            Initializes the MCTS process.
        Arguments:
            - cfg (:obj:`EasyDict`): A dictionary containing the configuration parameters for the MCTS process.
        """
        # Stores the configuration parameters for the MCTS search process.
        self._cfg = cfg

        # ==============================================================
        # sampled related core code
        # ==============================================================
        self.legal_actions = self._cfg.legal_actions
        self.action_space_size = self._cfg.action_space_size
        self.num_of_sampled_actions = self._cfg.num_of_sampled_actions
        print(f'num_of_sampled_actions: {self.num_of_sampled_actions}')
        self.continuous_action_space = self._cfg.continuous_action_space

        # The maximum number of moves allowed in a game.
        self._max_moves = self._cfg.get('max_moves', 512)  # for chess and shogi, 722 for Go.
        # The number of simulations to run for each move.
        self._num_simulations = self._cfg.get('num_simulations', 800)

        # UCB formula
        self._pb_c_base = self._cfg.get('pb_c_base', 19652)  # 19652
        self._pb_c_init = self._cfg.get('pb_c_init', 1.25)  # 1.25

        # Root prior exploration noise.
        self._root_dirichlet_alpha = self._cfg.get(
            'root_dirichlet_alpha', 0.3
        )  # 0.3  # for chess, 0.03 for Go and 0.15 for shogi.
        self._root_noise_weight = self._cfg.get('root_noise_weight', 0.25)  # 0.25
        self.mcts_search_cnt = 0
        self.simulate_env = simulate_env

    def get_next_action(
            self,
            state_config_for_env_reset: Dict[str, Any],
            policy_value_func: Callable,
            temperature: float = 1.0,
            sample: bool = True
    ) -> Tuple[int, List[float]]:
        """
        Overview:
            Get the next action to take based on the current state of the game.
        Arguments:
            - state_config_for_env_reset (:obj:`Dict`): The config of state when reset the env.
            - policy_value_func (:obj:`Function`): The Callable to compute the action probs and state value.
            - temperature (:obj:`Float`): The exploration temperature.
            - sample (:obj:`Bool`): Whether to sample an action from the probabilities or choose the most probable action.
        Returns:
            - action (:obj:`Int`): The selected action to take.
            - action_probs (:obj:`List`): The output probability of each action.
        """
        # Create a new root node for the MCTS search.
        self.root = Node()
        self.simulate_env.reset(
            start_player_index=state_config_for_env_reset.start_player_index,
            init_state=state_config_for_env_reset.init_state,
        )
        # self.simulate_env_root = copy.deepcopy(self.simulate_env)
        self._expand_leaf_node(self.root, self.simulate_env, policy_value_func)

        if sample:
            self._add_exploration_noise(self.root)

        for n in range(self._num_simulations):
            self.simulate_env.reset(
                start_player_index=state_config_for_env_reset.start_player_index,
                init_state=state_config_for_env_reset.init_state,
            )
            self.simulate_env.battle_mode = self.simulate_env.battle_mode_in_simulation_env
            self._simulate(self.root, self.simulate_env, policy_value_func)

        # sampled related code
        # Get the visit count for each possible action at the root node.
        action_visits = []
        for action in range(self.simulate_env.action_space.n):
            # Create an Action object for the current action
            current_action_object = Action(action)

            # Use the Action object to look up the child node in the dictionary
            if current_action_object in self.root.children:
                action_visits.append((action, self.root.children[current_action_object].visit_count))
            else:
                action_visits.append((action, 0))

        # Unpack the tuples in action_visits list into two separate tuples: actions and visits.
        actions, visits = zip(*action_visits)
        # print('action_visits= {}'.format(visits))

        visits_t = torch.as_tensor(visits, dtype=torch.float32)
        visits_t /= temperature
        action_probs = (visits_t / visits_t.sum()).numpy()

        if sample:
            action = np.random.choice(actions, p=action_probs)
        else:
            action = actions[np.argmax(action_probs)]
        self.mcts_search_cnt += 1

        # get the root sampled actions according to the action_probs
        self.root_sampled_actions = np.nonzero(action_probs)[0]

        # print(f'self.simulate_env_root: {self.simulate_env_root.legal_actions}')
        # print(f'mcts_search_cnt: {self.mcts_search_cnt}')
        # print('action= {}'.format(action))
        # print('action_probs= {}'.format(action_probs))
        # Return the selected action and the output probability of each action.
        return action, action_probs

    def _simulate(self, node: Node, simulate_env: Type[BaseEnv], policy_value_func: Callable) -> None:
        """
        Overview:
            Run a single playout from the root to the leaf, getting a value at the leaf and propagating it back through its parents.
            State is modified in-place, so a deepcopy must be provided.
        Arguments:
            - node (:obj:`Class Node`): Current node when performing mcts search.
            - simulate_env (:obj:`Class BaseGameEnv`): The class of simulate env.
            - policy_value_func (:obj:`Function`): The Callable to compute the action probs and state value.
        """
        while not node.is_leaf():
            # only for debug
            # print('=='*20)
            # print('node.legal_actions: ', node.legal_actions)
            # print(node.children.keys())
            # print('simulate_env.board: ', simulate_env.board)
            # print('simulate_env.legal_actions:', simulate_env.legal_actions)

            # Traverse the tree until the leaf node.
            action, node = self._select_child(node, simulate_env)
            if action is None:
                break
            # sampled related code
            simulate_env.step(action.value)

        done, winner = simulate_env.get_done_winner()
        """
        in ``self_play_mode``, the leaf_value is calculated from the perspective of player ``simulate_env.current_player``.
        in ``play_with_bot_mode``, the leaf_value is calculated from the perspective of player 1.
        """

        if not done:
            # The leaf_value here is obtained from the neural network. The perspective of this value is from the
            # player corresponding to the game state input to the neural network. For example, if the current_player
            # of the current node is player 1, the value output by the network represents the goodness of the current
            # game state from the perspective of player 1.
            leaf_value = self._expand_leaf_node(node, simulate_env, policy_value_func)

        else:
            if simulate_env.battle_mode_in_simulation_env == 'self_play_mode':
                # In a tie game, the value corresponding to a terminal node is 0.
                if winner == -1:
                    leaf_value = 0
                else:
                    # To maintain consistency with the perspective of the neural network, the value of a terminal
                    # node is also calculated from the perspective of the current_player of the terminal node,
                    # which is convenient for subsequent updates.
                    leaf_value = 1 if simulate_env.current_player == winner else -1

            if simulate_env.battle_mode_in_simulation_env == 'play_with_bot_mode':
                # in ``play_with_bot_mode``, the leaf_value should be transformed to the perspective of player 1.
                if winner == -1:
                    leaf_value = 0
                elif winner == 1:
                    leaf_value = 1
                elif winner == 2:
                    leaf_value = -1

        # Update value and visit count of nodes in this traversal.
        if simulate_env.battle_mode_in_simulation_env == 'play_with_bot_mode':
            node.update_recursive(leaf_value, simulate_env.battle_mode_in_simulation_env)
        elif simulate_env.battle_mode_in_simulation_env == 'self_play_mode':
            # NOTE: e.g.
            #       to_play: 1  ---------->  2  ---------->  1  ----------> 2
            #         state: s1 ---------->  s2 ---------->  s3 ----------> s4
            #                                     action    node
            #                                            leaf_value
            # leaf_value is calculated from the perspective of player 1, leaf_value = value_func(s3),
            # but node.value should be the value of E[q(s2, action)], i.e. calculated from the perspective of player 2.
            # thus we add the negative when call update_recursive().
            node.update_recursive(-leaf_value, simulate_env.battle_mode_in_simulation_env)

    def _select_child(self, node: Node, simulate_env: Type[BaseEnv]) -> Tuple[Union[int, float], Node]:
        """
        Overview:
            Select the child with the highest UCB score.
        Arguments:
            - node (:obj:`Class Node`): Current node.
        Returns:
            - action (:obj:`Int`): choose the action with the highest ucb score.
            - child (:obj:`Node`): the child node reached by executing the action with the highest ucb score.
        """
        action = None
        child_node = None
        best_score = -9999999
        # print(simulate_env._raw_env._go.board, simulate_env.legal_actions)
        # Iterate over each child of the current node.
        for action_tmp, child_node_tmp in node.children.items():
            # print(a, simulate_env.legal_actions)
            # print('node.legal_actions: ', node.legal_actions)
            if action_tmp.value in simulate_env.legal_actions:
                score = self._ucb_score(node, child_node_tmp)
                # Check if the score of the current child is higher than the best score so far.
                if score > best_score:
                    best_score = score
                    action = action_tmp
                    child_node = child_node_tmp
            else:
                print(f'error: {action_tmp} not in {simulate_env.legal_actions}')
        if child_node is None:
            child_node = node  # child==None, node is leaf node in play_with_bot_mode.
        if action is None:
            print('error: action is None')

        return action, child_node

    def _expand_leaf_node(self, node: Node, simulate_env: Type[BaseEnv], policy_value_func: Callable) -> float:
        """
        Overview:
            expand the node with the policy_value_func.
        Arguments:
            - node (:obj:`Class Node`): current node when performing mcts search.
            - simulate_env (:obj:`Class BaseGameEnv`): the class of simulate env.
            - policy_value_func (:obj:`Function`): the Callable to compute the action probs and state value.
        Returns:
            - leaf_value (:obj:`Bool`): the leaf node's value.
        """
        # ==============================================================
        # sampled related core code
        # ==============================================================
        if self.continuous_action_space:
            pass
        else:
            # discrete action space

            # Call the policy_value_func function to compute the action probabilities and state value, and return a
            # dictionary and the value of the leaf node.
            legal_action_probs_dict, leaf_value = policy_value_func(simulate_env)

            node.legal_actions = []

            # Extract actions and their corresponding probabilities from the dictionary
            actions = list(legal_action_probs_dict.keys())
            probabilities = list(legal_action_probs_dict.values())

            # Normalize the probabilities so they sum to 1
            probabilities = np.array(probabilities)
            probabilities /= probabilities.sum()

            # self.num_of_sampled_actions = len(actions)

            # If there are fewer legal actions than the desired number of samples,
            # adjust the number of samples to the number of legal actions
            num_samples = min(len(actions), self.num_of_sampled_actions)
            # Use numpy to randomly sample actions according to the given probabilities, without replacement
            sampled_actions = np.random.choice(actions, size=num_samples, p=probabilities, replace=False)
            sampled_actions = sampled_actions.tolist()  # Convert numpy array to list

            for action_index in range(num_samples):
                node.children[Action(sampled_actions[action_index])] = \
                    Node(
                        parent=node,
                        prior_p=legal_action_probs_dict[sampled_actions[action_index]],
                    )
                node.legal_actions.append(Action(sampled_actions[action_index]))

            # Return the value of the leaf node.
            return leaf_value

    def _ucb_score(self, parent: Node, child: Node) -> float:
        """
        Overview:
            Compute UCB score. The score for a node is based on its value, plus an exploration bonus based on the prior.
            For more details, please refer to this paper: http://gauss.ececs.uc.edu/Workshops/isaim2010/papers/rosin.pdf
            UCB = Q(s,a) + P(s,a) \cdot \frac{N(\text{parent})}{1+N(\text{child})} \cdot \left(c_1 + \log\left(\frac{N(\text{parent})+c_2+1}{c_2}\right)\right)
            - Q(s,a): value of a child node.
            - P(s,a): The prior of a child node.
            - N(parent): The number of the visiting of the parent node.
            - N(child): The number of the visiting of the child node.
            - c_1: a parameter given by self._pb_c_init to control the influence of the prior P(s,a) relative to the value Q(s,a).
            - c_2: a parameter given by self._pb_c_base to control the influence of the prior P(s,a) relative to the value Q(s,a).
        Arguments:
            - parent (:obj:`Class Node`): Current node.
            - child (:obj:`Class Node`): Current node's child.
        Returns:
            - score (:obj:`Bool`): The UCB score.
        """
        # Compute the value of parameter pb_c using the formula of the UCB algorithm.
        pb_c = math.log((parent.visit_count + self._pb_c_base + 1) / self._pb_c_base) + self._pb_c_init
        pb_c *= math.sqrt(parent.visit_count) / (child.visit_count + 1)

        # ==============================================================
        # sampled related core code
        # ==============================================================
        # TODO(pu)
        node_prior = "density"
        # node_prior = "uniform"
        if node_prior == "uniform":
            # Uniform prior for continuous action space
            prior_score = pb_c * (1 / len(parent.children))
        elif node_prior == "density":
            # TODO(pu): empirical distribution
            if self.continuous_action_space:
                # prior is log_prob
                prior_score = pb_c * (
                        torch.exp(child.prior_p) / (
                        sum([torch.exp(node.prior_p) for node in parent.children.values()]) + 1e-6)
                )
            else:
                # prior is prob
                prior_score = pb_c * (child.prior_p / (sum([node.prior_p for node in parent.children.values()]) + 1e-6))
        else:
            raise ValueError("{} is unknown prior option, choose uniform or density")

        # Compute the UCB score by combining the prior score and value score.
        value_score = child.value
        prior_score = pb_c * child.prior_p

        return prior_score + value_score

    def _add_exploration_noise(self, node: Node) -> None:
        """
        Overview:
            Add exploration noise.
        Arguments:
            - node (:obj:`Class Node`): Current node.
        """
        # Get a list of actions corresponding to the child nodes.
        actions = list(node.children.keys())
        # Create a list of alpha values for Dirichlet noise.
        alpha = [self._root_dirichlet_alpha] * len(actions)
        # Generate Dirichlet noise using the alpha values.
        noise = np.random.dirichlet(alpha)
        # Compute the weight of the exploration noise.
        frac = self._root_noise_weight
        # Update the prior probability of each child node with the exploration noise.
        for a, n in zip(actions, noise):
            node.children[a].prior_p = node.children[a].prior_p * (1 - frac) + n * frac

    def get_sampled_actions(self) -> List[int]:
        """
        Overview:
            Get the sampled actions of the root node.
        Returns:
            - output (:obj:`List`): The sampled actions of the root node.
        """
        return self.root_sampled_actions