Merge pull request #88 from PathOnAI/add-mcts-frontend

Add mcts frontend

Author: TataKKKL

visual-tree-search-app/components/LATSVisual.tsx

@@ -95,7 +95,7 @@ const LATSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         }
 
         // Handle tree structure updates
-        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation' || data.typ === 'tree_update_node_backpropagation') 
+        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation' || data.type === 'tree_update_node_backpropagation') 
             && Array.isArray(data.tree)) {
           // Preserve simulation flags when updating from tree
           if (updatedTreeNodes.some(node => node.isSimulated)) {

visual-tree-search-app/components/MCTSVisual.tsx

@@ -95,7 +95,7 @@ const MCTSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
         }
 
         // Handle tree structure updates
-        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation' || data.typ === 'tree_update_node_backpropagation') 
+        if ((data.type === 'tree_update_node_expansion' || data.type === 'tree_update_node_children_evaluation' || data.type === 'tree_update_node_backpropagation') 
             && Array.isArray(data.tree)) {
           // Preserve simulation flags when updating from tree
           if (updatedTreeNodes.some(node => node.isSimulated)) {
@@ -483,14 +483,14 @@ const MCTSVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
             <span className="w-3 h-3 rounded-full inline-block mr-1 bg-blue-500 dark:bg-blue-600"></span>
             <span className="text-gray-700 dark:text-gray-300">Selected</span>
           </div>
-          <div className="flex items-center">
+          {/* <div className="flex items-center">
             <span className="w-3 h-3 rounded-full inline-block mr-1 bg-green-500 dark:bg-green-600"></span>
             <span className="text-gray-700 dark:text-gray-300">Sim Start</span>
           </div>
           <div className="flex items-center">
             <span className="w-3 h-3 rounded-full inline-block mr-1 bg-orange-500 dark:bg-orange-600"></span>
             <span className="text-gray-700 dark:text-gray-300">Simulated</span>
-          </div>
+          </div> */}
         </div>
       </div>
       <div 

visual-tree-search-backend/app/api/lwats/agents_async/SearchAgents/base_agent.py

@@ -487,7 +487,8 @@ async def node_evaluation(self, node: LATSNode) -> None:
     def backpropagate(self, node: LATSNode, value: float) -> None:
         while node:
             node.visits += 1
-            node.value = (node.value * (node.visits - 1) + value) / node.visits
+            # Calculate running average: newAvg = oldAvg + (value - oldAvg) / newCount
+            node.value += (value - node.value) / node.visits
             node = node.parent
 
     # shared
@@ -558,7 +559,7 @@ async def rollout(self, node: LATSNode, websocket=None)-> tuple[float, LATSNode]
         score = confidence_score if goal_finished else 0
         await self.remove_simulated_trajectory(starting_node=node, terminal_node=terminal_node, websocket=websocket)
 
-        return score, node
+        return score, terminal_node
 
 
     # TODO: decide whether to keep the tree update

visual-tree-search-backend/app/api/lwats/agents_async/SearchAgents/lats_agent.py

@@ -85,8 +85,10 @@ async def lats_search(self, websocket=None):
             # Step 5: Backpropagation
             print(f"{GREEN}Step 5: backpropagation{RESET}")
             await self.websocket_step_start(step=5, step_name="backpropagation", websocket=websocket)
-            self.backpropagate(terminal_node, reward)
+            self.backpropagate(selected_node, reward)
             tree_data = self._get_tree_data()
+            print_entire_tree(self.root_node)
+            print(tree_data)
             if websocket:
                 await self.websocket_tree_update(type="tree_update_node_backpropagation", websocket=websocket, tree_data=tree_data)
             else:

visual-tree-search-backend/app/api/lwats/agents_async/SearchAgents/mcts_agent.py

@@ -264,11 +264,12 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
             print(f"{GREEN}Step 1: Node Selection{RESET}")
             await self.websocket_step_start(step=1, step_name="node_selection", websocket=websocket)
             selected_node = await self.node_selection(self.root_node, websocket)
-            tree_data = self._get_tree_data()
-            if websocket:
-                await self.websocket_tree_update(type="tree_update_node_selection", websocket=websocket, tree_data=tree_data)
-            else:
-                print_entire_tree(self.root_node)
+            # await self.websocket_node_selection(selected_node, websocket=websocket)
+            # tree_data = self._get_tree_data()
+            # if websocket:
+            #     await self.websocket_tree_update(type="tree_update_node_selection", websocket=websocket, tree_data=tree_data)
+            # else:
+            #     print_entire_tree(self.root_node)
 
             if selected_node is None:
                 logger.warning("All paths lead to terminal nodes. Ending search.")
@@ -338,10 +339,28 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
             for node in path:
                 old_value = node.value
                 node.visits += 1
-                node.value = (node.value * (node.visits - 1) + score) / node.visits
+                node.value += (score - node.value) / node.visits
+                # consiste with lats backpropagation
+                #node.value = (node.value * (node.visits - 1) + score) / node.visits
                 print(f"Node {node.action}:")
                 print(f"  Visits: {node.visits}")
                 print(f"  Value: {old_value:.3f} -> {node.value:.3f}")
+                # add websocket information, just use websocket here
+                # if websocket:
+                #     await websocket.send_json({
+                #         "type": "backpropagation",
+                #         "node_id": id(node),
+                #         "node_parent_id": id(node.parent),
+                #         "node_action": node.action,
+                #         "node_value": node.value,
+                #         "node_visits": node.visits,
+                #         "node_old_value": old_value,
+                #         "node_description": node.natural_language_description,
+                #     })
+
+            tree_data = self._get_tree_data()
+            print_entire_tree(self.root_node)
+            print(tree_data)
             if websocket:
                 await self.websocket_tree_update(type="tree_update_node_backpropagation", websocket=websocket, tree_data=tree_data)
             else: