add parameter to run the algorithm for a specified amount of time (#12)

mctspy/tree/search.py

@@ -1,3 +1,4 @@
+import time
 class MonteCarloTreeSearch(object):
 
     def __init__(self, node):
@@ -9,22 +10,34 @@ def __init__(self, node):
         """
         self.root = node
 
-    def best_action(self, simulations_number):
+    def best_action(self, simulations_number=None, total_simulation_seconds=None):
         """
 
         Parameters
         ----------
         simulations_number : int
             number of simulations performed to get the best action
 
+        total_simulation_seconds : float
+            Amount of time the algorithm has to run. Specified in seconds
+
         Returns
         -------
 
         """
-        for _ in range(0, simulations_number):            
-            v = self._tree_policy()
-            reward = v.rollout()
-            v.backpropagate(reward)
+
+        if simulations_number is None :
+            assert(total_simulation_seconds is not None)
+            end_time = time.time() + total_simulation_seconds
+            while time.time() < end_time:
+                v = self._tree_policy()
+                reward = v.rollout()
+                v.backpropagate(reward)
+        else :
+            for _ in range(0, simulations_number):            
+                v = self._tree_policy()
+                reward = v.rollout()
+                v.backpropagate(reward)
         # to select best child go for exploitation only
         return self.root.best_child(c_param=0.)