Merge pull request #89 from PathOnAI/improve_mcts

Improve mcts

Author: TataKKKL

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

@@ -3,11 +3,14 @@ import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
 import { Info, ChevronDown, ChevronUp } from "lucide-react";
+import { Checkbox } from "@/components/ui/checkbox";
 
 interface SearchParams {
   startingUrl: string;
   goal: string;
+  maxDepth: number;
   iterations: number;
+  set_prior_value: boolean;
 }
 
 interface ControlPanelProps {
@@ -105,6 +108,19 @@ const ControlPanelMCTS: React.FC<ControlPanelProps> = ({
                   />
                 </div>
 
+                <div className="space-y-2">
+                  <Label htmlFor="maxDepth" className="text-slate-700 dark:text-slate-300 font-medium">Max Depth</Label>
+                  <Input
+                    id="maxDepth"
+                    type="number"
+                    min={1}
+                    max={10}
+                    value={searchParams.maxDepth}
+                    onChange={(e) => handleParamChange('maxDepth', parseInt(e.target.value))}
+                    className="border-slate-300 dark:border-slate-600 focus:ring-cyan-500 focus:border-cyan-500"
+                  />
+                </div>
+
                 <div className="space-y-2">
                   <Label htmlFor="iterations" className="text-slate-700 dark:text-slate-300 font-medium">Iterations</Label>
                   <Input
@@ -118,6 +134,26 @@ const ControlPanelMCTS: React.FC<ControlPanelProps> = ({
                   />
                 </div>
               </div>
+              
+              {/* Add prior_value checkbox */}
+              <div className="mt-4">
+                <div className="flex items-center space-x-2">
+                  <Checkbox 
+                    id="set_prior_value" 
+                    checked={searchParams.set_prior_value}
+                    onCheckedChange={(checked) => handleParamChange('set_prior_value', checked === true)}
+                  />
+                  <Label 
+                    htmlFor="set_prior_value" 
+                    className="text-slate-700 dark:text-slate-300 font-medium cursor-pointer"
+                  >
+                    Use Prior Value
+                  </Label>
+                </div>
+                <p className="mt-1 ml-6 text-xs text-slate-500 dark:text-slate-400">
+                  When enabled, RMCTS will use an LLM to generate an initial value as a prior value for each newly generated node.
+                </p>
+              </div>
             </div>
           )}
         </div>

visual-tree-search-app/pages/MCTSAgent.tsx

@@ -9,8 +9,8 @@ interface SearchParams {
   startingUrl: string;
   goal: string;
   maxDepth: number;
-  num_simulations: number;
   iterations: number;
+  set_prior_value: boolean;
 }
 
 interface Message {
@@ -32,7 +32,7 @@ const MCTSAgent = () => {
     startingUrl: 'http://xwebarena.pathonai.org:7770/',
     goal: 'search running shoes, click on the first result',
     maxDepth: 3,
-    num_simulations: 1,
+    set_prior_value: false,
     iterations: 1
   });
 
@@ -96,8 +96,11 @@ const MCTSAgent = () => {
             starting_url: searchParams.startingUrl,
             goal: searchParams.goal,
             search_algorithm: "mcts",
-            iterations: searchParams.iterations
+            iterations: searchParams.iterations,
+            set_prior_value: searchParams.set_prior_value,
+            max_depth: searchParams.maxDepth
           };
+          console.log(request);
 
           wsRef.current?.send(JSON.stringify(request));
           logMessage(request, 'outgoing');
@@ -141,7 +144,9 @@ const MCTSAgent = () => {
         starting_url: searchParams.startingUrl,
         goal: searchParams.goal,
         search_algorithm: "mcts",
-        iterations: searchParams.iterations
+        iterations: searchParams.iterations,
+        set_prior_value: searchParams.set_prior_value,
+        max_depth: searchParams.maxDepth
       };
 
       wsRef.current?.send(JSON.stringify(request));

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

@@ -486,10 +486,11 @@ async def node_evaluation(self, node: LATSNode) -> None:
     ## TODO: check the logic of updating value/ reward, is the input value?
     def backpropagate(self, node: LATSNode, value: float) -> None:
         while node:
-            node.visits += 1
-            # Calculate running average: newAvg = oldAvg + (value - oldAvg) / newCount
-            node.value += (value - node.value) / node.visits
-            node = node.parent
+            if node.depth != 0:
+                node.visits += 1
+                # Calculate running average: newAvg = oldAvg + (value - oldAvg) / newCount
+                node.value += (value - node.value) / node.visits
+                node = node.parent
 
     # shared
     async def simulation(self, node: LATSNode, websocket=None) -> tuple[float, LATSNode]:

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

@@ -290,6 +290,17 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
                 print_entire_tree(self.root_node)
 
 
+            # optional: prior value
+            if self.config.set_prior_value:
+                await self.websocket_step_start(step=2, step_name="node_children_evaluation", websocket=websocket)
+                await self.node_children_evaluation(selected_node)
+                tree_data = self._get_tree_data()
+                if websocket:
+                    await self.websocket_tree_update(type="tree_update_node_children_evaluation", websocket=websocket, tree_data=tree_data)
+                else:
+                    print("after evaluation")
+                    print_entire_tree(self.root_node)
+
             # Step 3: simulation using the current node, (generate a path using the current node, and score the path)
             # TODO: implement simulation using openai
             print(f"{GREEN}Step 3: Simulation{RESET}")
@@ -337,14 +348,15 @@ async def mcts_search(self, websocket=None) -> Optional[LATSNode]:
             print(f"{GREEN}Step 5: Backpropagation{RESET}")
             await self.websocket_step_start(step=5, step_name="backpropagation", websocket=websocket)
             for node in path:
-                old_value = node.value
-                node.visits += 1
-                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}")
+                if node != self.root_node:
+                    old_value = node.value
+                    node.visits += 1
+                    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({

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

@@ -25,8 +25,12 @@ class AgentConfig:
     num_simulations: int = 1
     account_reset: bool = True
 
+    # for LATS
     simulation_score: float = 0.75
+
+    # for MCTS
     reflection_score: float = 0.75
+    set_prior_value: bool = False
 
     # Features
     features: List[str] = field(default_factory=lambda: ['axtree'])

visual-tree-search-backend/app/api/routes/tree_search_websocket.py

@@ -79,6 +79,7 @@ async def handle_search_request(websocket: WebSocket, message: Dict[str, Any]):
         storage_state = message.get("storage_state", "app/api/shopping.json")
         iterations = message.get("iterations", 1)  # Extract iterations parameter
         num_simulations=message.get("num_simulations", 1)
+        set_prior_value = message.get("set_prior_value", False)
 
         # Send status update
         await websocket.send_json({
@@ -95,8 +96,10 @@ async def handle_search_request(websocket: WebSocket, message: Dict[str, Any]):
             storage_state=storage_state,
             headless=False,
             iterations=iterations,
-            num_simulations=num_simulations
+            num_simulations=num_simulations,
+            set_prior_value=set_prior_value
         )
+        print(config)
 
         # Send status update
         await websocket.send_json({