update

src/util.jl

@@ -6,48 +6,99 @@ using ..type
 
 # ---------------------------------------------- 100 --------------------------------------------- #
 
-""" Select a node based on UCT score
+""" Select a node based on UCT (Upper Confidence Bound for Trees) score.
+
+The function computes UCT values for all child nodes and returns the child with the
+highest UCT score. The UCT formula balances exploitation (child state value) and
+exploration (visit count and parent visit count) using the exploration weight `w`.
+
+Does **not** mutate the input node.
 
 # Arguments
-  - `node::MCTSNode`
-    mcts node
-  - `w::T`
-    exploration weight. Value is usually between 1 to 2. 
-    Value 1.0 makes MCTS balance between exploration and exploitation like 50%-50%.
-    Value 2.0 makes MCTS aggressively search the tree.
-# Return
-  - `selectedNode::MCTSNode`
-    child node with highest UCT score. UCT score balances between exploitation (state value) 
-    and exploration (visit count) based on the exploration weight w.
+- `node::MCTSNode`
+  The MCTS node whose children will be evaluated.
+- `w::AbstractFloat`
+  Exploration weight. Typical values range from 1.0 to 2.0. A value of 1.0 balances
+  exploration and exploitation equally. Higher values (e.g., 2.0) encourage more
+  exploration of less-visited nodes.
 
-# Example
-```jldoctest
-julia> 
+# Return
+- `selectedNode::MCTSNode`
+  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).
+
+# The UCT Formula
+
+```
+UCT(s,a) = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))
+
+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)
 ```
 
-# Signature
+# 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
+julia> using LLMMCTS
+julia> child1 = MCTSNode("a", Dict(:reward=>5.0), 0, 10, 50, 0, false, nothing, Dict(), Dict())
+julia> child2 = MCTSNode("b", Dict(:reward=>6.0), 0, 5, 30, 0, false, nothing, Dict(), Dict())
+julia> parent = MCTSNode("root", Dict(:reward=>0.0), 0, 15, 100, 0, false, nothing,
+                           Dict("a"=>child1, "b"=>child2), Dict())
+julia> selected = UCTselect(parent, 1.0)
+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
-      childNode.statevalue + weightedterm
-    else  # node.visits == 0 makes sqrt() in explore term error
-      childNode.progressvalue  # exploit term
+      # Child has been visited before - use statevalue with exploration bonus
+      # Exploration bonus = w * sqrt(ln(parent_visits) / child_visits)
+      # High child_visits = small bonus (exploitation dominates)
+      # 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