diff --git a/src/DecMCTS.py b/src/DecMCTS.py index 022348587ee124fdbeffacba9eb287310f4effb5..c734297aea36d3e5ef8ae31fca23a8a86397281d 100644 --- a/src/DecMCTS.py +++ b/src/DecMCTS.py @@ -128,20 +128,35 @@ class Tree: return list(self.graph.successors(node_id)) - def _update_distribution(): - return "TODO" + def _update_distribution(self): + """ + Get the top n Action sequences and their "probabilities" + and store them for communication + """ + + # For now, just using q = mu**2 + temp=nx.get_node_attributes(self.graph, "mu") + + top_n_nodes = sorted(temp, key=temp.get, reverse=True)[:self.comm_n] + X = [self.graph.nodes[n]["action_seq"] for n in top_n_nodes] + q = [self.graph.nodes[n]["mu"]**2 for n in top_n_nodes] + self.my_act_dist = ActionDistribution(X,q) + return True + def _get_state(self, node_id): """ - Randomly select 1 path taken by every other robot - & path taken by this robot to get to this node + Randomly select 1 path taken by every other robot & path taken by + this robot to get to this node + + Returns tuple where first element is path of current robot, + and second element is a dictionary of the other paths """ node_path = self.graph.nodes[node_id]["action_seq"] - other_paths = [] - for + other_paths = {k:self.comms[k].random_action() for k in self.comms} - return "TODO" + return (node_path, other_paths) def grow(self, nsims=10, ntime=None): @@ -159,7 +174,7 @@ class Tree: options = self.available_actions(self.data, state) def send_comms(self): - return self.my_action_distribution + return self.my_act_dist def receive_comms(self, comms_in, robot_id): """