蒙特卡洛树搜索（Monte Carlo Tree Search）的技术演进(2)： AlphaZero

class ValueNeuralNet():
    @classmethod
    def evaluate(self, game_state):
        return np.random.random()

class PolicyNeuralNet():
    @classmethod
    def evaluate(self, game_state, action_size=4):
        return np.random.random([action_size])

class AlphaZeroNeuralNet():
    @classmethod
    def evaluate(self, game_state):
        state_value = ValueNeuralNet.evaluate(game_state)
        action_prob = PolicyNeuralNet.evaluate(game_state)
        return action_prob, state_value

def alphazero_mcts(game_state, num_reads, mdp):
    root = UCTNode(game_state)
    for _ in range(num_reads):
        leaf = root.select_leaf()
        child_priors,value_estimate = AlphaZeroNeuralNet.evaluate(leaf.game_state)
        leaf.expand(child_priors, mdp)
        leaf.backup(value_estimate)
    return max(root.children.items(),key=lambda item: item[1].number_visits)

import numpy as n

def filter_and_normalize(probs, legal_actions):
    """
    过滤合法动作并归一化概率
    
    :param probs: 策略网络输出的所有动作的概率，形状为 (num_actions,)
    :param legal_actions: 合法动作的列表，例如 [0, 1, 3]
    :return: 归一化后的合法动作概率，形状为 (num_legal_actions,)
    """
    # 过滤合法动作的概率
    legal_probs = probs[legal_actions]
    
    # 归一化概率
    legal_probs /= np.sum(legal_probs)
    
    return legal_probs

all_probs = np.array([0.1, 0.3, 0.2, 0.4])

# 合法动作
legal_actions = [0, 1, 3]

# 过滤并归一化
normalized_legal_probs = filter_and_normalize(all_probs, legal_actions)

print("归一化后的合法动作概率:", normalized_legal_probs)

import numpy as np
import math

def filter_and_normalize(probs, legal_actions):
    """
    过滤合法动作并归一化概率
    
    :param probs: 策略网络输出的所有动作的概率，形状为 (num_actions,)
    :param legal_actions: 合法动作的列表，例如 [0, 1, 3]
    :return: 归一化后的合法动作概率，形状为 (num_legal_actions,)
             legal_actions: 合法动作的列表（用于映射）
    """
    # 过滤合法动作的概率
    legal_probs = probs[legal_actions]
    
    # 归一化概率
    legal_probs /= np.sum(legal_probs)
    
    return legal_probs, legal_actions

def map_to_original_space(legal_probs, legal_actions, num_actions):
    """
    将过滤后的概率映射回原始动作空间
    
    :param legal_probs: 过滤并归一化后的合法动作概率，形状为 (num_legal_actions,)
    :param legal_actions: 合法动作的列表
    :param num_actions: 原始动作空间的大小
    :return: 映射回原始动作空间的概率，形状为 (num_actions,)
    """
    # 初始化原始动作空间的概率数组
    original_probs = np.zeros(num_actions)
    
    # 将过滤后的概率值填充到对应的原始动作位置
    for action, prob in zip(legal_actions, legal_probs):
        original_probs[action] = prob
    
    return original_probs


class UCTNode:
    def __init__(self, state, parent=None, prior=0):
        self.state = state  # Current state in the MDP
        self.is_expanded = False
        self.parent = parent  # Parent node
        self.children = {}  # Dictionary of actions to child nodes
        self.prior = prior  # Prior probability of the node
        self.total_value = 0  # Cumulative reward from simulations
        self.number_visits = 0  # Number of times the node has been visited

    def Q(self):
        """Compute the average reward (exploitation term)."""
        return self.total_value / (1 + self.number_visits)

    def U(self):
        """Compute the exploration term."""
        if self.parent is None:
            return 0  # Root node has no exploration term
        return (math.sqrt(self.parent.number_visits) * self.prior) / (1 + self.number_visits)

    def best_child(self):
          return max(self.children.values(),
               key=lambda node: node.Q() + node.U())

    def select_leaf(self):
        current = self
        while current.is_expanded:
            current = current.best_child()
        return current
    
    def expand(self, prior_probability, mdp):
        """Expand the node by generating child nodes for all legal actions."""
        if mdp.is_terminal(self.state):  # 如果是终结状态，直接返回
            self.is_expanded = False
            return
        legal_actions = mdp.get_legal_actions(self.state)
        if not legal_actions:  # 如果没有合法动作，标记为已扩展
            self.is_expanded = False
            return
        
        normalized_legal_probs, legal_actions = filter_and_normalize(prior_probability, legal_actions)
        # 映射回原始动作空间
        num_actions = len(prior_probability)
        original_space_probs = map_to_original_space(normalized_legal_probs, legal_actions, num_actions)

        
        for action in legal_actions:
            next_state, reward = mdp.step(self.state, action)
            prior = original_space_probs[action]
            self.children[action] = UCTNode(next_state, parent=self, prior=prior)
        self.is_expanded = True

    def backpropagate(self, value_estimate):
        """Backpropagate the value estimate up the tree."""
        current = self
        while current is not None:
            current.number_visits += 1
            current.total_value += value_estimate
            current = current.parent

class ValueNeuralNet():
    @classmethod
    def evaluate(self, game_state):
        return np.random.random()

class PolicyNeuralNet():
    @classmethod
    def evaluate(self, game_state, action_size=4):
        return np.random.random([action_size])

class AlphaZeroNeuralNet():
    @classmethod
    def evaluate(self, game_state):
        state_value = ValueNeuralNet.evaluate(game_state)
        action_prob = PolicyNeuralNet.evaluate(game_state)
        return action_prob, state_value

def alphazero_mcts(game_state, num_reads, mdp):
    root = UCTNode(game_state)
    for _ in range(num_reads):
        leaf = root.select_leaf()
        child_priors,value_estimate = AlphaZeroNeuralNet.evaluate(leaf.game_state)
        leaf.expand(child_priors, mdp)
        leaf.backup(value_estimate)
    return max(root.children.items(),key=lambda item: item[1].number_visits)