Merge pull request #14 from fadynakhla/arthur/sentient

Arthur/sentient
fadynakhla · Jul 30, 2023 · 50e408d · 50e408d
2 parents 1e927b4 + b09fff6
commit 50e408d
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 6 deletions.
diff --git a/dr_claude/mcts/__init__.py → dr_claude/mcts_module/__init__.py b/dr_claude/mcts/__init__.py → dr_claude/mcts_module/__init__.py
diff --git a/dr_claude/mcts/action_states.py → dr_claude/mcts_module/action_states.py b/dr_claude/mcts/action_states.py → dr_claude/mcts_module/action_states.py
@@ -5,11 +5,11 @@
 import random
 from typing import Collection, Dict, Generic, List, Optional, Set, TypeVar, Union
 from typing_extensions import Self
-
+import mcts
 import numpy as np
 
 from dr_claude import datamodels
-from dr_claude.mcts import probability_calcs
+from dr_claude.mcts_module import probability_calcs
 
 
 T = TypeVar("T")
@@ -200,3 +200,26 @@ class ArgMaxDiagnosisRolloutPolicy(RollOutPolicy):
     def __call__(self, state: SimulationNextActionState) -> float:
         actions_space = state.getConditionProbabilityDict()
         return max(actions_space.values())
+
+
+"""
+The MCTS algorithm
+"""
+
+
+class MCTS(mcts.mcts):
+    def getBestChild(self, node: mcts.treeNode, explorationValue: float):
+        bestValue = float("-inf")
+        bestNodes = []
+        for child in node.children.values():
+            nodeValue = (
+                node.state.getCurrentPlayer() * child.totalReward / child.numVisits
+                + explorationValue
+                * math.sqrt(2 * math.log(node.numVisits) / child.numVisits)
+            )
+            if nodeValue > bestValue:
+                bestValue = nodeValue
+                bestNodes = [child]
+            elif nodeValue == bestValue:
+                bestNodes.append(child)
+        return random.choice(bestNodes)
diff --git a/dr_claude/mcts/probability_calcs.py → dr_claude/mcts_module/probability_calcs.py b/dr_claude/mcts/probability_calcs.py → dr_claude/mcts_module/probability_calcs.py
diff --git a/tests/test_conditional_proba.py b/tests/test_conditional_proba.py
@@ -1,5 +1,5 @@
 from dr_claude import datamodels
-from dr_claude.mcts import probability_calcs
+from dr_claude.mcts_module import probability_calcs
 
 
 def test_conditional_condition_proba():

diff --git a/tests/test_mcts.py b/tests/test_mcts.py
@@ -1,10 +1,12 @@
 import copy
+import random
 from typing import Callable, Collection, Self
 
 import loguru
 
 from dr_claude import datamodels
-from dr_claude.mcts import action_states
+from dr_claude import mcts_module
+from dr_claude.mcts_module import action_states
 
 import mcts
 
@@ -71,13 +73,14 @@ def test_convergence():
     ]
     state = ConvergenceTestState(matrix, discount_rate=1e-9)
     state.set_condition(the_condition, the_symptoms)
+    state.pertinent_pos.add(random.choice(the_symptoms))
 
     ## Rollout policy
-    rollout_policy = action_states.RandomRollOutPolicy()
+    rollout_policy = action_states.ArgMaxDiagnosisRolloutPolicy()
     rollout_policy = logtrueconditionhook(rollout_policy)
 
     ## create the initial state
-    searcher = mcts.mcts(timeLimit=3000, rolloutPolicy=rollout_policy)
+    searcher = action_states.MCTS(timeLimit=3000, rolloutPolicy=rollout_policy)
 
     action = None
     while not isinstance(action, datamodels.Condition):