diff --git a/mcts/searcher/mcts.py b/mcts/searcher/mcts.py index 885d400..6905290 100755 --- a/mcts/searcher/mcts.py +++ b/mcts/searcher/mcts.py @@ -3,6 +3,10 @@ import math import random import time +import copy +from multiprocessing import Process, Queue, Manager +from multiprocessing import Pool +from multiprocessing.managers import BaseManager, NamespaceProxy from mcts.base.base import BaseState @@ -17,16 +21,58 @@ def random_policy(state: BaseState) -> float: return state.get_reward() +# -- this fails: sometimes the .join() will hang +# - this is potential due to the Queue. -> Using Process() directly without +# using Queue instead (seems to be a tiny bit slower tho..) +class RolloutProcess(Process): + def __init__(self, queue, rollout_policy): + self.queue = queue + self.rollout_policy = rollout_policy + super().__init__() + + def run(self): + while not self.queue.empty(): + inputs = self.queue.get() + input_id = inputs[0] + state = inputs[1] + return_dict = inputs[2] + reward = self.rollout_policy(state) + return_dict[input_id] = reward + + +class TreeNodeProxy(NamespaceProxy): + _exposed_ = ('__getattribute__', '__setattr__', '__delattr__', + 'all_child_have_at_least_one_visit', 'add_child') + + def all_child_have_at_least_one_visit(self): + callmethod = object.__getattribute__(self, '_callmethod') + return callmethod('all_child_have_at_least_one_visit') + #TODO: need debug + def add_child(self, action, node): + callmethod = object.__getattribute__(self, '_callmethod') + #print(callmethod) + return callmethod('add_child', args=(action, node)) + + + class TreeNode: def __init__(self, state, parent): self.state = state self.is_terminal = state.is_terminal() self.is_fully_expanded = self.is_terminal + self.all_child_have_been_explored = False self.parent = parent self.numVisits = 0 self.totalReward = 0 self.children = {} + def all_child_have_at_least_one_visit(self,) -> bool: + if not self.all_child_have_been_explored: + if len(self.children) > 0: + self.all_child_have_been_explored = all([child.numVisits>0 for child in self.children.values()]) + else: return False + return self.all_child_have_been_explored + def __str__(self): s = ["totalReward: %s" % self.totalReward, "numVisits: %d" % self.numVisits, @@ -90,6 +136,61 @@ def search(self, initialState: BaseState = None, initial_state: BaseState = None else: return action + + def search_parallel(self, initialState: BaseState = None, initial_state: BaseState = None, needDetails: bool = False, + need_details: bool = None, n_jobs: int = 1): + initial_state = initialState if initial_state is None else initial_state + need_details = needDetails if need_details is None else need_details + self.root = TreeNode(initial_state, None) + + if self.limit_type == 'time': + raise NotImplementedError("Cannot use search parallel with time limit type") + + manager = Manager() + BaseManager.register('TreeNode', TreeNode, TreeNodeProxy) + basemanager = BaseManager() + basemanager.start() + root_shared = basemanager.TreeNode(initial_state, None) + + #import pdb; pdb.set_trace() + + for _ in range(0, self.search_limit, n_jobs): + + input_nodes = [] + processes = [] + return_dict = manager.dict() + + for index in range(n_jobs): + tmp_node = self.select_node(self.root) + input_nodes.append(tmp_node) + processes.append(Process(target=self.execute_rollout_parallel,args=(index, tmp_node.state, return_dict))) + #processes.append(Process(target=self.execute_rollout_parallel,args=(index,root_shared,return_dict))) + + for p in processes: p.start() + + for p in processes: p.join() + + for index, reward in return_dict.items(): + #for index, rez in return_dict.items(): + self.backpropogate(input_nodes[index], reward) + #node, reward = rez + #self.backpropogate(node, reward) + + best_child = self.get_best_child(self.root, 0) + #best_child = self.get_best_child(root_shared, 0) + action = (action for action, node in self.root.children.items() if node is best_child).__next__() + #action = (action for action, node in root_shared.children.items() if node is best_child).__next__() + if need_details: + return action, best_child.totalReward / best_child.numVisits + else: + return action + + def execute_rollout_parallel(self, input_id, state, return_dict): + #def execute_rollout_parallel(self, input_id, root, return_dict): + #node = self.select_node(root) + reward = self.rollout_policy(state) + return_dict[input_id] = reward + def execute_round(self): """ execute a selection-expansion-simulation-backpropagation round @@ -101,7 +202,10 @@ def execute_round(self): def select_node(self, node: TreeNode) -> TreeNode: while not node.is_terminal: if node.is_fully_expanded: - node = self.get_best_child(node, self.exploration_constant) + if node.all_child_have_at_least_one_visit(): + node = self.get_best_child(node, self.exploration_constant) + else: + node = self.get_random_child(node) else: return self.expand(node) return node @@ -137,3 +241,7 @@ def get_best_child(self, node: TreeNode, explorationValue: float, exploration_va elif node_value == best_value: best_nodes.append(child) return random.choice(best_nodes) + + def get_random_child(self, node: TreeNode) -> TreeNode: + return random.choice(list(node.children.values())) +