update distribution and get state

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):
         """