update comments

src/interface.jl

@@ -75,33 +75,50 @@ function runMCTS(
   )::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
                 Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any}
 
+  # Initialize the MCTS tree with a root node representing the initial state
+  # root.visits=0: no visits yet
+  # root.statevalue=0: no simulation results yet
   root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, Dict{String,MCTSNode}(),
                   Dict{String,Any}())
 
-  # storage for holding all high reward terminal nodes
+  # Channel to collect high-value terminal states (reward >= 8)
+  # These are "good solutions" that can be returned to the user
   highValueState = Channel{Any}(100)
 
+  # Main MCTS loop: perform iterations to build the search tree
+  # Each iteration: SELECTION → EXPANSION → SIMULATION → BACKPROPAGATION
   for nth in 1:maxiterations
+    # Start from root and traverse down using UCT selection
     node = root
-    node.visits += 1
+    node.visits += 1  # Count this iteration's visit to root
 
+    # Phase 1: SELECTION - Traverse tree using UCT until reaching a leaf node
+    # UCT balances exploration (new branches) vs exploitation (promising branches)
     while !isleaf(node)
       node = UCTselect(node, explorationweight)
     end
 
+    # Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
     if node.isterminal
+      # If this terminal state has high reward (>= 8), store it for later
       if node.state[:reward] >= 8
-        put!(highrewardNode, deepcopy(node.state))
+        put!(highValueState, deepcopy(node.state))
       end
 
-      # MCTS arrive at the leaf node that is also a terminal state,
+      # Backpropagate the terminal node's own reward up to root
-      # do nothing then go directly to backpropagation. It means the end of this iteration
+      # This updates all ancestors with this path's outcome
       backpropagate(node, node.reward)
     else
+      # Phase 3: EXPANSION - Generate children for this non-terminal leaf
+      # Horizontal sampling: create multiple child nodes via LLM transition
       _ = expand(node, transition, transitionargs;
                 horizontalSample=horizontalSampleExpansionPhase,
                 multithread=multithread)
+      
+      # Phase 4: SIMULATION + BACKPROPAGATION
+      # For each newly expanded child, run simulation and update statistics
       if multithread
+        # Parallel simulation: spawn threads for each child node
         @sync for (leafNodeKey, leafNode) in node.children
           @spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
                   maxSimulationDepth=maxSimulationDepth, 
@@ -112,6 +129,7 @@ function runMCTS(
                   )
         end
       else
+        # Sequential simulation: process each child one at a time
         for (leafNodeKey, leafNode) in node.children
           simulateThenBackpropagate(leafNode, transition, transitionargs;
             maxSimulationDepth=maxSimulationDepth, 
@@ -123,22 +141,27 @@ function runMCTS(
       end
     end
 
-    # stop if the early stop condition is met
+    # Phase 5: EARLY STOP CHECK
+    # Optional: stop search early if a condition is met
     if typeof(earlystop) <: Function && earlystop(node.state)
       break
     end
   end
 
-  # select the best next state and the best terminal state along the best trajectory
+  # After all iterations, extract results from the search tree
+  # Select best immediate next state (best child of root)
   bestNextState = selectBestNextNode(root)
+  
+  # Select best terminal state along the optimal trajectory
   bestTerminalState = selectBestTrajectoryNode(root)
 
-  # take all high value state from highValueState channel and put it in a list
+  # Collect all high-value states from the channel into a list
   highValueStateList = Vector{Dict{String, Any}}()
   while !isempty(highValueState)
     push!(highValueStateList, take!(highValueState))
   end
 
+  # Return complete search results
   result = (
             root=root, 
             bestNextState=bestNextState.state, 
@@ -186,12 +209,17 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit
                 saveSimulatedNode::Bool=false,
                 multithread=false,
                 highValueState=Union{Nothing,Any}=nothing)
+  # Phase 1: RUN SIMULATION (rollout)
+  # Perform a rollout from this node, accumulating rewards along the way
   simTrajectoryReward, terminalstate = 
     simulate(node, transition, transitionargs;
             maxSimulationDepth=maxSimulationDepth, 
             horizontalSample=horizontalSampleSimulationPhase,
             multithread=multithread)
-  # if a node has state value >= 8, store it in highValueState
+  
+  # Phase 2: HIGH-VALUE STATE TRACKING
+  # If we reached a terminal state with high reward (>= 8), store it
+  # This allows users to access multiple good solutions, not just the best one
   if highValueState !== nothing && 
     terminalstate !== nothing &&
     terminalstate[:reward] >= 8
@@ -199,9 +227,14 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit
     put!(highValueState, deepcopy(terminalstate))
   end
 
+  # Phase 3: BACKPROPAGATE
+  # Update statistics (visits, statevalue) for all ancestors up to root
+  # The simulation result is now incorporated into the tree
   backpropagate(node, simTrajectoryReward)
 
-  # check if the user wants to keep the simulated node
+  # Phase 4: MEMORY MANAGEMENT
+  # Clear children unless user wants to keep them for analysis
+  # This frees memory for the next iteration while preserving tree structure
   if saveSimulatedNode == false
     node.children = Dict{String, MCTSNode}()
   end

src/mcts.jl

@@ -29,12 +29,14 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
   nodekey = nothing
 
   # Calculate sum of statevalues across all child nodes
+  # This determines whether to use statevalue/visits (exploitation) or progressvalue+reward (exploration)
   stateValueSum = sum([v.statevalue for (k, v) in node.children])
 
   # If any nodes have non-zero statevalue, use statevalue/visits as selection metric
+  # This means simulations have confirmed node values - use exploitation
   if stateValueSum != 0
     for (k, childnode) in node.children
-      # Calculate average statevalue per visit
+      # Calculate average statevalue per visit (running average from simulations)
       potential = childnode.statevalue / childnode.visits
 
       if potential > highestProgressValue
@@ -43,7 +45,8 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
       end
     end
   else
-    # Otherwise use progressvalue + reward as selection metric
+    # No simulations yet - use progressvalue + reward for initial guidance
+    # This allows LLM heuristics to guide early search before simulations provide data
     for (k, childnode) in node.children
       potential = childnode.progressvalue + childnode.reward
 
@@ -72,6 +75,8 @@ until reaching a leaf node, returning the highest-value node found along the pat
   The highest-value node found by following the optimal trajectory to a leaf.
 """
 function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
+  # Follow the optimal path down the tree by repeatedly selecting the best child
+  # This gives us the highest-value trajectory from the starting node to a leaf
   while !isleaf(node)
     node = selectBestNextNode(node)
   end
@@ -102,15 +107,23 @@ leaf node to the root, applying reward discounting for future rewards.
 - `Nothing`
 """
 function backpropagate(node::MCTSNode, simTrajectoryReward::T;
-                      discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
+                       discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
+  # Propagate the simulation result back up the tree to update all ancestor nodes
+  # Each node's statistics are updated with the cumulative reward from the simulation
   while !isroot(node)
-    # Update the statistics of the current node based on the result of the playout
+    # Increment visit count - this simulation passed through this node
-    node.visits += 1  # Increment visit count for this node
+    node.visits += 1
+    
     node.statevalue += ((node.statevalue * (node.visits-1)) + simTrajectoryReward) / node.visits  # Update running average of state value
-    simTrajectoryReward *= discountRewardCoeff  # discount because future reward is uncertain
+    
-    node = node.parent  # Move up to parent node for next iteration
+    # Apply discount to future rewards - rewards further from the current state are worth less
+    # This reflects temporal uncertainty: distant future rewards are less certain
+    simTrajectoryReward *= discountRewardCoeff
+    
+    # Move up to parent node to continue propagation
+    node = node.parent
   end
-end
+end
 
 """ Determine whether a node is a leaf node.
 
@@ -166,7 +179,11 @@ function selectChildNode(node::MCTSNode)::MCTSNode
   highestProgressValue = -1
   nodekey = nothing
 
-  # loop thought node children dictionary to find the highest progress value
+  # During simulation rollout, we need to pick which child to explore next
+  # Use progressvalue + reward as the selection metric (no UCT here)
+  # - progressvalue: LLM's estimate of how promising this state is
+  # - reward: immediate environment feedback
+  # Together they guide fast exploration during simulation
   for (k, childNode) in node.children
     potential = childNode.progressvalue + childNode.reward
     if potential > highestProgressValue
@@ -203,6 +220,10 @@ Creates new child nodes by applying the transition function multiple times
 """
 function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple; 
                 horizontalSample::Integer=3, multithread=false)
+  # Generate child nodes by applying the transition function multiple times
+  # This is called "horizontal sampling" - we branch out horizontally in the tree
+  # - multithread=true: spawn parallel threads for each expansion
+  # - multithread=false: sequential expansion (default, simpler)
   if multithread
     @sync for i in 1:horizontalSample 
       @spawn _expand(node, transition, transitionargs)
@@ -231,23 +252,24 @@ Checks for semantically equivalent states (dejavu) to avoid duplicates.
 - `Nothing`
 """
 function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
-    result = transition(node.state, transitionargs)
+  # Generate one child node from the parent using the transition function
-    newNodeKey::AbstractString = result[:newNodeKey]
+  result = transition(node.state, transitionargs)
-    newstate::AbstractDict = result[:newstate]
+  newNodeKey::AbstractString = result[:newNodeKey]
-    progressvalue::Integer = result[:progressvalue]
+  newstate::AbstractDict = result[:newstate]
+  progressvalue::Integer = result[:progressvalue]
 
-    """
+  # Dejavu detection: avoid adding duplicate states
-    [] newNodeKey ∉ keys(node.children). 
+  # If newNodeKey already exists, skip - this handles semantically equivalent states
-          New state may have semantic vector close enought to 
+  if newNodeKey ∉ keys(node.children)
-        one of existing child state. Which can be assume that they are the same state 
+    # Create new MCTS node with:
-        semantically-wise i.e. De javu. This could be used to recall lessons for this 
+    # - visits=0: no simulations yet
-        similar situation to improve decisionMaker and evaluator. 
+    # - statevalue=0: will be updated after simulation
-    """
+    # - progressvalue: LLM's estimate (fast heuristic)
-    if newNodeKey ∉ keys(node.children)
+    # - reward: immediate environment feedback
-      newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward], 
+    newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward], 
-          newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{String, Any}())
+        newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{String, Any}())
-      node.children[newNodeKey] = newNode
+    node.children[newNodeKey] = newNode
-    end
+  end
 end
 
 """ Simulate interactions between agent and environment.
@@ -280,18 +302,31 @@ function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTup
   maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
   )::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}
 
+  # Perform a rollout simulation from the given node:
+  # 1. Accumulate rewards along the trajectory
+  # 2. Expand nodes horizontally at each level
+  # 3. Select children to explore vertically down the tree
+  # Returns cumulative reward and whether a terminal state was reached
+  
   simTrajectoryReward = 0.0
   terminalstate = nothing
   
   for depth in 1:maxSimulationDepth
+    # Accumulate the current node's reward to the trajectory total
     simTrajectoryReward += node.reward
+    
+    # Check if we've reached a terminal state
     if node.isterminal
       terminalstate = node.state
       break
     else
+      # Expand current node to generate children (horizontal sampling)
       _ = expand(node, transition, transitionargs;
                 horizontalSample=horizontalSample,
                 multithread=multithread)
+      
+      # Select best child to continue the rollout (vertical exploration)
+      # Uses progressvalue + reward for fast selection during simulation
       node = selectChildNode(node)
     end 
   end

src/util.jl

@@ -27,12 +27,25 @@ Does **not** mutate the input node.
   The child node with the highest UCT score. Returns `nothing` if the node has no
   children (though this would indicate an error since UCTselect is called on non-leaves).
 
-# Notes
+# The UCT Formula
-- The UCT formula used is: `statevalue + w * sqrt(log(parent_visits) / child_visits)`
+
-- When a child has zero visits (`child_visits == 0`), the exploration term becomes
+```
-  undefined, so the function returns the child's `progressvalue` as a fallback.
+UCT(s,a) = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))
-- This function assumes the calling code only invokes it on non-leaf nodes (i.e.,
+
-  nodes with children).
+Where:
+  Q(s,a) = childNode.statevalue        (exploitation: accumulated reward)
+  c = w (explorationweight)            (controls exploration vs exploitation)
+  N(s) = node.visits                   (parent visits - total visits to parent)
+  N(s,a) = childNode.visits            (child visits - visits to this specific action)
+```
+
+# Behavior
+
+| Child visits | Exploration term | Behavior |
+|-------------|------------------|----------|
+| 0 (never visited) | Undefined | Uses `progressvalue` (LLM heuristic) |
+| Low (few visits) | High | Encourages exploring new branches |
+| High (many visits) | Near 0 | Exploits known good branches |
 
 # Examples
 ```jldoctest
@@ -46,26 +59,46 @@ MCTSNode(...)
 ```
 """
 function UCTselect(node::MCTSNode, w::T)::MCTSNode where {T<:AbstractFloat}
+  # UCT (Upper Confidence Bound for Trees) selects the best child using:
+  # UCT = statevalue + exploration_weight * sqrt(ln(parent_visits) / child_visits)
+  #
+  # The two terms balance:
+  # - Exploitation (statevalue): choose children that performed well in simulations
+  # - Exploration (sqrt term): encourage trying less-visited children
+  #
+  # The exploration weight `w` controls this balance:
+  # - w=1.0: equal emphasis on exploration and exploitation
+  # - w>1.0: more aggressive exploration (try new branches)
+  # - w<1.0: more exploitation (stick with known good branches)
+  
   maxUCT = -Inf
   selectedNode = nothing
 
   for (childState, childNode) in node.children
+    # Calculate UCT value for this child
     UCTvalue =
     if childNode.visits != 0  
-      weightedterm = w * sqrt(log(node.visits) / childNode.visits) # explore term
+      # Child has been visited before - use statevalue with exploration bonus
-      childNode.statevalue + weightedterm
+      # Exploration bonus = w * sqrt(ln(parent_visits) / child_visits)
-    else  # node.visits == 0 makes sqrt() in explore term error
+      # High child_visits = small bonus (exploitation dominates)
-      childNode.progressvalue  # exploit term
+      # Low child_visits = large bonus (encourages exploration)
+      weightedterm = w * sqrt(log(node.visits) / childNode.visits)
+      UCTvalue = childNode.statevalue + weightedterm
+    else  
+      # Child has never been visited - exploration term undefined
+      # Fall back to progressvalue (LLM heuristic) as exploitation term
+      # This allows LLM guidance to direct early search
+      UCTvalue = childNode.progressvalue
     end
-                
+                 
     if UCTvalue > maxUCT
       maxUCT = UCTvalue
-        selectedNode = childNode
+      selectedNode = childNode
     end
   end
 
   return selectedNode
-end
+end