def add_binary_tree_6(depth, parent_nodes, num_leafs):
parent_node = parent_nodes[0]
# process while_loop_1
while_2_parent_nodes = []
for i in range(depth):
#print("creating depth {} of loop1 of mcts tree".format(i))
parent_node, while_node_0 = add_binary_tree_block_6(parent_node)
while_2_parent_nodes.append(while_node_0)
# add terminating while_loop_1 nodes
terminate_while_node_0 = MctsTreeNode(parent=parent_node,
children=[],
val=0)
parent_node.children = [terminate_while_node_0]
while_2_parent_nodes.append(terminate_while_node_0)
# process while_loop_2
for parent_node in while_2_parent_nodes:
for i in range(depth):
#print("creating depth {} of loop2 of mcts tree".format(i))
parent_node, _ = add_binary_tree_block_6(parent_node)
# add terminating while_loop_2 nodes
terminate_while_node_0 = MctsTreeNode(parent=parent_node,
children=[],
val=0)
parent_node.children = [terminate_while_node_0]
# add terminating while_loop_2 nodes as leaf nodes
num_leafs[0] += 1
def add_binary_tree_7(depth, parent_nodes, num_leafs):
while_nodes_0 = []
for i in range(depth):
#print("creating depth {} of mcts tree".format(i))
next_parent_nodes = []
for parent_node in parent_nodes:
next_parent_nodes_new, while_node_0_new = add_binary_tree_block_7(
parent_node)
next_parent_nodes += next_parent_nodes_new
while_nodes_0 += while_node_0_new
parent_nodes = next_parent_nodes
# add terminating while nodes
terminate_while_nodes_0 = []
for parent_node in parent_nodes:
terminate_while_node_0 = MctsTreeNode(parent=parent_node,
children=[],
val=0)
parent_node.children = [terminate_while_node_0]
terminate_while_nodes_0.append(terminate_while_node_0)
# append last if node to while_nodes_0 and terminate_while_node_0
for node in (while_nodes_0 + terminate_while_nodes_0):
if_node_1 = MctsTreeNode(parent=node, children=[], val=1)
if_node_0 = MctsTreeNode(parent=node, children=[], val=0)
node.children = [if_node_1, if_node_0]
# add if_node_1 and if_node_0 as leaf nodes
num_leafs[0] += 2
def add_binary_tree_block_6(parent_node):
# add while nodes
while_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
while_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [while_node_1, while_node_0]
next_parent_node = while_node_1
return next_parent_node, while_node_0
def add_binary_tree_block_5(parent_node):
# add if nodes
if_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
if_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [if_node_1, if_node_0]
next_parent_nodes = [if_node_1, if_node_0]
return next_parent_nodes
def add_binary_tree_block_2(parent_node, num_leafs):
# add while nodes
while_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
while_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [while_node_1, while_node_0]
# add while_node_0 as leaf node
num_leafs[0] += 1
next_parent_node = while_node_1
return next_parent_node
def grow_type5_loc(node, max_depth):
# add if nodes
if_node_1 = MctsTreeNode(node_type="if1",
parent=node,
children=[],
val=1,
depth=(node.depth + 1))
if_node_0 = MctsTreeNode(node_type="if0",
parent=node,
children=[],
val=0,
depth=(node.depth + 1))
node.children = [if_node_1, if_node_0]
def grow_type3_loc(node, max_depth):
# add while nodes
while_node_1 = MctsTreeNode(node_type="while1",
parent=node,
children=[],
val=1,
depth=(node.depth + 1))
while_node_0 = MctsTreeNode(node_type="while0",
parent=node,
children=[],
val=0,
depth=(node.depth + 1))
node.children = [while_node_1, while_node_0]
def generate_bin_tree(tree_type, depth, use_subtree_flag):
num_leafs = [0]
# create root node
root = MctsTreeNode(parent=None, children=[], val="root")
# add 20 location children nodes or 1 loc node is agentloc and agentdir is fixed
loc_nodes = get_loc_nodes(root, tree_type, use_subtree_flag)
root.children = loc_nodes
# add binary tree for each loc node
for loc_node in root.children:
get_tree_gen_func(tree_type)(depth=depth,
parent_nodes=[loc_node],
num_leafs=num_leafs)
print("created tree for {}".format(loc_node.val))
return root, num_leafs[0]
def add_binary_tree_block_4(parent_node, num_leafs):
# add while nodes
while_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
while_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [while_node_1, while_node_0]
# add while_node_0 as leaf node
num_leafs[0] += 1
# add a nodes
a_node_1 = MctsTreeNode(parent=while_node_1, children=[], val=1)
a_node_0 = MctsTreeNode(parent=while_node_1, children=[], val=0)
while_node_1.children = [a_node_1, a_node_0]
# add bc nodes
bc_node_1 = MctsTreeNode(parent=a_node_0, children=[], val=1)
bc_node_0 = MctsTreeNode(parent=a_node_0, children=[], val=0)
a_node_0.children = [bc_node_1, bc_node_0]
next_parent_nodes = [a_node_1, bc_node_1, bc_node_0]
return next_parent_nodes
def get_loc_nodes(parent_node, tree_type, use_subtree_flag):
if (use_subtree_flag):
direction_choices = ["fixed_loc_dir"]
else:
direction_choices = get_direction_choices(tree_type)
loc_nodes = []
for dir in direction_choices:
loc_node = MctsTreeNode(parent=parent_node, children=[], val=dir)
loc_nodes.append(loc_node)
return loc_nodes
def grow_type2_while1(node, max_depth):
if (node.depth < max_depth):
# add while nodes
while_node_1 = MctsTreeNode(node_type="while1",
parent=node,
children=[],
val=1,
depth=(node.depth + 1))
while_node_0 = MctsTreeNode(node_type="while0",
parent=node,
children=[],
val=0,
depth=(node.depth + 1))
node.children = [while_node_1, while_node_0]
else:
# add terminating while0 node
while_node_0 = MctsTreeNode(node_type="while0",
parent=node,
children=[],
val=0,
depth=(node.depth + 1))
node.children = [while_node_0]
def add_binary_tree_2(depth, parent_nodes, num_leafs):
parent_node = parent_nodes[0]
for i in range(depth):
#print("creating depth {} of mcts tree".format(i))
parent_node = add_binary_tree_block_2(parent_node, num_leafs)
# add terminating while nodes
terminate_while_node_0 = MctsTreeNode(parent=parent_node,
children=[],
val=0)
parent_node.children = [terminate_while_node_0]
# add terminate_while_node_0 as leaf node
num_leafs[0] += 1
def get_loc_nodes_onthefly(parent_node, use_subtree_flag):
if (use_subtree_flag):
direction_choices = ["fixed_loc_dir"]
else:
direction_choices = get_direction_choices()
loc_nodes = []
for dir in direction_choices:
loc_node = MctsTreeNode(node_type="loc",
parent=parent_node,
children=[],
val=dir,
depth=(parent_node.depth + 1))
loc_nodes.append(loc_node)
return loc_nodes
def add_binary_tree_3(depth, parent_nodes, num_leafs):
for i in range(depth):
#print("creating depth {} of mcts tree".format(i))
next_parent_nodes = []
for parent_node in parent_nodes:
next_parent_nodes += add_binary_tree_block_3(
parent_node, num_leafs)
parent_nodes = next_parent_nodes
# add terminating while nodes
for parent_node in parent_nodes:
terminate_while_node_0 = MctsTreeNode(parent=parent_node,
children=[],
val=0)
parent_node.children = [terminate_while_node_0]
# add terminate_while_node_0 as leaf node
num_leafs[0] += 1
def add_binary_tree_block_7(parent_node):
# add while nodes
while_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
while_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [while_node_1, while_node_0]
# add if nodes
if_node_1 = MctsTreeNode(parent=while_node_1, children=[], val=1)
if_node_0 = MctsTreeNode(parent=while_node_1, children=[], val=0)
while_node_1.children = [if_node_1, if_node_0]
next_parent_nodes = [if_node_1, if_node_0]
return next_parent_nodes, [while_node_0]
def add_binary_tree_block_3(parent_node, num_leafs):
# add while nodes
while_node_1 = MctsTreeNode(parent=parent_node, children=[], val=1)
while_node_0 = MctsTreeNode(parent=parent_node, children=[], val=0)
parent_node.children = [while_node_1, while_node_0]
# add while_node_0 as leaf node
num_leafs[0] += 1
# add if nodes
if_node_1 = MctsTreeNode(parent=while_node_1, children=[], val=1)
if_node_0 = MctsTreeNode(parent=while_node_1, children=[], val=0)
while_node_1.children = [if_node_1, if_node_0]
next_parent_nodes = [if_node_1, if_node_0]
return next_parent_nodes
def train_mcts_tree(tree_type, num_train_iterations, input_code,
input_task_data, env_type, env_type_initial, Z, expl_const,
program, depth):
# create root node of mcts tree
root = MctsTreeNode(node_type="root",
parent=None,
children=[],
val="root",
depth=0)
# snapshot_t stores at which t to take a snapshot of the mcts tree
#snapshot_t = [16, 32, 64, 128, 1000, 2500, 5000, 7500, 10000]
#snapshot_roots = {}
flag_no_valid_children = False
# no of times scores were reused from cached score table
score_table_hits = 0
unique_traces = []
all_traces = []
cnt_invalid_traces = 0
cnt_node_expansions = [0]
cnt_node_grown = [0]
t = 1
# store a list of expanded nodes to be reset for diversity
#expanded_nodes = []
#grown_nodes = []
# expand root node once for all t training iterations
#expand_node(root, cnt_node_expansions, expanded_nodes)
grow_node(root, tree_type, cnt_node_grown, depth)
# initialize tqdm progress bar
pbar = tqdm(total=num_train_iterations + 1)
# for each training iteration do
while (t <= num_train_iterations):
# increment num_iterations of root
# needed since parent's node num_iterations is used in node selection
root.num_iterations += 1
# initialize
curr_node = root
mcts_trace = []
# index in trace to start backpropogation update from (inclusive)
# this index is the node from which simulation begins in mcts terminology
backprop_idx = None
# selection - compute a mcts trace from root to a leaf
while (curr_node.children != []):
# before selection check each child node is in expanded state
#for child in curr_node.children:
# expand_node(child, cnt_node_expansions, expanded_nodes)
# simulation stage
if (backprop_idx != None):
# select node at random
selected_node = select_node(curr_node.children,
mcts_trace,
input_code,
tree_type,
env_type,
env_type_initial,
expl_const,
program,
random=True)
if (selected_node == None):
curr_node.valid = False
flag_no_valid_children = True
break
else:
# selection stage
selected_node = select_node(curr_node.children, mcts_trace,
input_code, tree_type, env_type,
env_type_initial, expl_const,
program)
if (selected_node == None):
curr_node.valid = False
flag_no_valid_children = True
break
# expansion to include the first num_iterations=0 node
if ((selected_node.num_iterations == 0)
and (backprop_idx == None)):
# since backprop_idx is the len it will include the new node to be updated added to mcts_trace later
backprop_idx = len(mcts_trace)
mcts_trace.append(selected_node)
curr_node = selected_node
# grow curr_node for next selection / simulation step
grow_node(curr_node, tree_type, cnt_node_grown, depth)
if (flag_no_valid_children):
flag_no_valid_children = False
cnt_invalid_traces += 1
# continue without increment mcts iteration counter t
continue
else:
# compute score of mcts trace
# reuse cached leaf scores if previously visited
if (mcts_trace[-1].num_picked == 0):
score_dict = score_mcts_trace(mcts_trace, tree_type,
input_code, input_task_data,
env_type, env_type_initial, Z,
program)
mcts_trace[-1].score_dict = score_dict
unique_traces.append(mcts_trace)
score = score_dict["score_total"]
#all_traces.append(mcts_trace)
else:
score_table_hits += 1
score = mcts_trace[-1].score_dict["score_total"]
#all_traces.append(mcts_trace)
mcts_trace[-1].num_picked += 1
# backpropagation
# update average scores of selected nodes
update_node_scores(mcts_trace, score, backprop_idx)
# update iteration counts of selected nodes
update_node_num_iterations(mcts_trace, backprop_idx)
# take snapshots of evolving tree at different t
#if( t+1 in snapshot_t ):
# snapshot_roots[t] = copy.deepcopy(root)
# increment mcts iteration counter t
t += 1
# update tqdm progress bar
pbar.update(1)
# close tqdm progress bar
pbar.close()
# remove individual variables above, access keys in dictionary directly
mcts_info = {
"mcts_run_totaliters": num_train_iterations,
"mcts_run_num_unique_traces": len(unique_traces),
"mcts_run_num_repeat_traces": score_table_hits,
"mcts_run_num_invalid_traces": cnt_invalid_traces
}
#return root, score_table_hits, unique_traces, mcts_info, cnt_node_expansions, expanded_nodes, all_traces
return score_table_hits, unique_traces, mcts_info, cnt_node_grown, all_traces