-
Notifications
You must be signed in to change notification settings - Fork 32
/
monte_carlo_tree_search.py
147 lines (120 loc) · 4.91 KB
/
monte_carlo_tree_search.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import torch
import math
import numpy as np
def ucb_score(parent, child):
"""
The score for an action that would transition between the parent and child.
"""
prior_score = child.prior * math.sqrt(parent.visit_count) / (child.visit_count + 1)
if child.visit_count > 0:
# The value of the child is from the perspective of the opposing player
value_score = -child.value()
else:
value_score = 0
return value_score + prior_score
class Node:
def __init__(self, prior, to_play):
self.visit_count = 0
self.to_play = to_play
self.prior = prior
self.value_sum = 0
self.children = {}
self.state = None
def expanded(self):
return len(self.children) > 0
def value(self):
if self.visit_count == 0:
return 0
return self.value_sum / self.visit_count
def select_action(self, temperature):
"""
Select action according to the visit count distribution and the temperature.
"""
visit_counts = np.array([child.visit_count for child in self.children.values()])
actions = [action for action in self.children.keys()]
if temperature == 0:
action = actions[np.argmax(visit_counts)]
elif temperature == float("inf"):
action = np.random.choice(actions)
else:
# See paper appendix Data Generation
visit_count_distribution = visit_counts ** (1 / temperature)
visit_count_distribution = visit_count_distribution / sum(visit_count_distribution)
action = np.random.choice(actions, p=visit_count_distribution)
return action
def select_child(self):
"""
Select the child with the highest UCB score.
"""
best_score = -np.inf
best_action = -1
best_child = None
for action, child in self.children.items():
score = ucb_score(self, child)
if score > best_score:
best_score = score
best_action = action
best_child = child
return best_action, best_child
def expand(self, state, to_play, action_probs):
"""
We expand a node and keep track of the prior policy probability given by neural network
"""
self.to_play = to_play
self.state = state
for a, prob in enumerate(action_probs):
if prob != 0:
self.children[a] = Node(prior=prob, to_play=self.to_play * -1)
def __repr__(self):
"""
Debugger pretty print node info
"""
prior = "{0:.2f}".format(self.prior)
return "{} Prior: {} Count: {} Value: {}".format(self.state.__str__(), prior, self.visit_count, self.value())
class MCTS:
def __init__(self, game, model, args):
self.game = game
self.model = model
self.args = args
def run(self, model, state, to_play):
root = Node(0, to_play)
# EXPAND root
action_probs, value = model.predict(state)
valid_moves = self.game.get_valid_moves(state)
action_probs = action_probs * valid_moves # mask invalid moves
action_probs /= np.sum(action_probs)
root.expand(state, to_play, action_probs)
for _ in range(self.args['num_simulations']):
node = root
search_path = [node]
# SELECT
while node.expanded():
action, node = node.select_child()
search_path.append(node)
parent = search_path[-2]
state = parent.state
# Now we're at a leaf node and we would like to expand
# Players always play from their own perspective
next_state, _ = self.game.get_next_state(state, player=1, action=action)
# Get the board from the perspective of the other player
next_state = self.game.get_canonical_board(next_state, player=-1)
# The value of the new state from the perspective of the other player
value = self.game.get_reward_for_player(next_state, player=1)
if value is None:
# If the game has not ended:
# EXPAND
action_probs, value = model.predict(next_state)
valid_moves = self.game.get_valid_moves(next_state)
action_probs = action_probs * valid_moves # mask invalid moves
action_probs /= np.sum(action_probs)
node.expand(next_state, parent.to_play * -1, action_probs)
self.backpropagate(search_path, value, parent.to_play * -1)
return root
def backpropagate(self, search_path, value, to_play):
"""
At the end of a simulation, we propagate the evaluation all the way up the tree
to the root.
"""
for node in reversed(search_path):
node.value_sum += value if node.to_play == to_play else -value
node.visit_count += 1