Remove the reward attribute of the node and add proper close session logic to all agents

Author: TataKKKL

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

@@ -102,7 +102,7 @@ def _get_tree_data(self):
                 "value": node.value,
                 "visits": node.visits,
                 "feedback": node.feedback,
-                "reward": node.reward
+                # "reward": node.reward
             }
             tree_data.append(node_data)
 
@@ -129,7 +129,7 @@ async def remove_simulated_trajectory(self, starting_node, terminal_node: LATSNo
                 "description": node.natural_language_description,
                 "visits": node.visits,
                 "value": float(f"{node.value:.3f}") if hasattr(node, 'value') else None,
-                "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
+                # "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
                 "is_terminal": node.is_terminal,
                 "feedback": node.feedback if hasattr(node, 'feedback') else None,
                 "is_root": not hasattr(node, 'parent') or node.parent is None,
@@ -159,7 +159,7 @@ def _get_trajectory_data(self, terminal_node: LATSNode):
                 "description": node.natural_language_description,
                 "visits": node.visits,
                 "value": float(f"{node.value:.3f}") if hasattr(node, 'value') else None,
-                "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
+                # "reward": float(f"{node.reward:.3f}") if hasattr(node, 'reward') else None,
                 "is_terminal": node.is_terminal,
                 "feedback": node.feedback if hasattr(node, 'feedback') else None,
                 "is_root": not hasattr(node, 'parent') or node.parent is None,
@@ -432,7 +432,7 @@ async def node_children_evaluation(self, node: LATSNode) -> None:
 
         for child, score in zip(node.children, scores):
             child.value = score
-            child.reward = score
+            # child.reward = score
 
     async def node_evaluation(self, node: LATSNode) -> None:
         """Evaluate the current node and assign its score."""
@@ -469,7 +469,7 @@ async def node_evaluation(self, node: LATSNode) -> None:
 
             # Assign the score to the node
             node.value = score
-            node.reward = score
+            # node.reward = score
 
 
         except Exception as e:

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

@@ -37,7 +37,7 @@ async def lats_search(self, websocket=None):
             await self.websocket_node_selection(node, websocket=websocket)
 
             if node is None:
-                print("All paths lead to terminal nodes with reward 0. Ending search.")
+                print("All paths lead to terminal nodes with value 0. Ending search.")
                 break
 
             # Step 2: Node Expansion
@@ -76,8 +76,10 @@ async def lats_search(self, websocket=None):
             terminal_nodes.append(terminal_node)
             await self.websocket_simulation_result(reward, terminal_node, websocket=websocket)
 
-            if reward == 1:
+            # simulation score threshold
+            if reward >= self.config.simulation_score:
                 await self.websocket_search_complete("success", reward, terminal_node.get_trajectory(), websocket=websocket)
+                await self.playwright_manager.close()
                 return terminal_node
 
             # Step 5: Backpropagation
@@ -95,8 +97,8 @@ async def lats_search(self, websocket=None):
         all_nodes_list = collect_all_nodes(self.root_node)
         all_nodes_list.extend(terminal_nodes)
 
-        ## temp change: if reward is the same, choose the deeper node
-        best_child = max(all_nodes_list, key=lambda x: (x.reward, x.depth))
+        ## temp change: if value is the same, choose the deeper node
+        best_child = max(all_nodes_list, key=lambda x: (x.value, x.depth))
 
         if best_child.value >= 0.75:
             print("Successful trajectory found")

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

@@ -88,7 +88,7 @@ def __init__(
         self.value = 0.0
         self.depth = 0 if parent is None else parent.depth + 1
         self.is_terminal = False
-        self.reward = 0.0
+        # self.reward = 0.0
         self.exhausted = False  # If all children are terminal
         self.em = 0.0  # Exact match, evaluation metric
         self.observation: Optional[Observation] = None
@@ -177,7 +177,7 @@ def to_dict(self) -> dict:
             'value': self.value,
             'depth': self.depth,
             'is_terminal': self.is_terminal,
-            'reward': self.reward,
+            # 'reward': self.reward,
             'em': self.em,
         }
 

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

@@ -302,6 +302,7 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
                 # Convert path to serializable trajectory
                 # trajectory = [node.action for node in path if node.action is not None]
                 await self.websocket_search_complete("success", score, selected_node.get_trajectory(), websocket=websocket)
+                await self.playwright_manager.close()
                 return selected_node
 
             print(f"path: {path}")
@@ -328,4 +329,5 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
              # Convert node to serializable trajectory
             # trajectory = [n.action for n in self.get_path_to_root(best_node) if n.action is not None]
             await self.websocket_search_complete("partial_success", best_node.value, best_node.get_trajectory(), websocket=websocket)
+        await self.playwright_manager.close()
         return best_node

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

@@ -110,6 +110,7 @@ async def bfs(self, websocket=None):
 
                     # Send completion update if websocket is provided
                     await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket) 
+                    await self.playwright_manager.close()
 
                     return current_node
 
@@ -120,6 +121,7 @@ async def bfs(self, websocket=None):
 
             # Send completion update if websocket is provided
             await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=websocket)
+            await self.playwright_manager.close()
 
             return best_node
 
@@ -128,6 +130,7 @@ async def bfs(self, websocket=None):
 
         # Send failure update if websocket is provided
         await self.websocket_search_complete("failure", 0, None, websocket=websocket)
+        await self.playwright_manager.close()
 
         return None
 
@@ -209,7 +212,8 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
                 print(f"Found satisfactory solution with score {score}")
 
                 # Send completion update if websocket is provided
-                await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket)                
+                await self.websocket_search_complete("success", score, current_node.get_trajectory(), websocket=websocket) 
+                await self.playwright_manager.close()               
                 return current_node
 
             # Add non-terminal children to stack in reverse order
@@ -234,6 +238,7 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
 
             # Send completion update if websocket is provided
             await self.websocket_search_complete("partial_success", best_score, best_node.get_trajectory(), websocket=websocket)
+            await self.playwright_manager.close()
 
             return best_node
 
@@ -242,6 +247,7 @@ async def dfs(self, websocket=None) -> List[Dict[str, Any]]:
 
         # Send failure update if websocket is provided
         await self.websocket_search_complete("failure", 0, None, websocket=websocket)
+        await self.playwright_manager.close()
 
         return None
 

visual-tree-search-backend/app/api/lwats/core_async/config.py

@@ -25,6 +25,7 @@ class AgentConfig:
     num_simulations: int = 1
     account_reset: bool = True
 
+    simulation_score: float = 0.75
     reflection_score: float = 0.75
 
     # Features