update

src/LLMMCTS.jl

@@ -1,6 +1,6 @@
 module LLMMCTS
 
-  # export agent
+  export MCTSNode
 
 
   """ Order by dependencies of each file. The 1st included file must not depend on any other

src/interface.jl

@@ -9,47 +9,55 @@ using ..type, ..mcts, ..util
 # ---------------------------------------------- 100 --------------------------------------------- #
 
 
+""" Search for the best action to take for a given state and task.
 
-""" Search the best action to take for a given state and task 
+This function runs the MCTS algorithm through multiple iterations of expansion,
+simulation, and backpropagation to find optimal decisions.
+
+Does **not** mutate the input state; it creates new MCTS nodes during search.
 
 # Arguments
-  - `initialstate::T`
-    initial state
-  - `transition::Function`
-    a function that define how the state transitions
-  - `transitionargs::NamedTuple`
-    arguments for transition function
+- `initialstate::T`
+  The initial state from which to start the search.
+- `transition::Function`
+  A function that defines how the state transitions.
+- `transitionargs::NamedTuple`
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-  - `horizontalSampleExpansionPhase::Integer`
-    a number of child state MCTS sample at each node during expansion phase (default: 3)
-  - `horizontalSampleSimulationPhase::Integer`
-    a number of child state MCTS sample at each node during simulation's expansion phase (default: 3)
-  - `maxSimulationDepth::Integer`
-    a number of levels MCTS goes during simulation phase (default: 3)
-  - `maxiterations::Integer`
-    a number of iteration MCTS goes thru expansion -> simulation -> backpropagation cycle (default: 10)
-  - `explorationweight::Number`
-    exploration weight controls how much MCTS should explore new state instead of exploit 
-    a known state. 1.0 balance between exploration and exploitation like 50%-50%. 2.0 makes MCTS
-    aggressively explore new state (default: 1.0)
-  - `earlystop::Union{Function,Nothing}`
-    optional function to check early stopping condition (default: nothing)
-  - `saveSimulatedNode::Bool`
-    whether to save nodes created during simulation phase (default: false)
-  - `multithread::Bool`
-    whether to use multithreading during simulation (default: false)
+- `horizontalSampleExpansionPhase::Integer=3`
+  Number of child states sampled at each node during expansion phase.
+- `horizontalSampleSimulationPhase::Integer=3`
+  Number of child states sampled at each node during simulation's expansion phase.
+- `maxSimulationDepth::Integer=3`
+  Maximum depth MCTS goes during simulation phase.
+- `maxiterations::Integer=10`
+  Number of iterations MCTS performs through expansion → simulation → backpropagation cycles.
+- `explorationweight::Number=1.0`
+  Exploration weight controls how much MCTS explores new states versus exploiting known states.
+  A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0) 
+  encourage more aggressive exploration.
+- `earlystop::Union{Function,Nothing}=nothing`
+  Optional function to check early stopping condition. If satisfied, MCTS breaks iterations.
+- `saveSimulatedNode::Bool=false`
+  Whether to save nodes created during simulation phase.
+- `multithread::Bool=false`
+  Whether to use multithreading during simulation.
 
-# Returns
-  - `NamedTuple{(:root, :bestNextState, :bestFinalState), Tuple{MCTSNode, T, T}}`
-    - root: the complete MCTS tree with root node
-    - bestNextState: the best immediate next state
-    - bestFinalState: the best final state along the best trajectory
+# Return
+- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
+             Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}`
+  - `root`: the complete MCTS tree with root node
+  - `bestNextState`: the best immediate next state
+  - `bestTerminalState`: the best final state along the best trajectory
+  - `highValueStateList`: list of high-value terminal states (reward >= 8)
 
 # Example
-    Refers to SQLLLM package
-
-# Signature
+```jldoctest
+julia> using LLMMCTS
+julia> initialState = Dict(:reward=>0.0)
+julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5)
+```
 """
 function runMCTS(
   initialstate::T,
@@ -63,97 +71,185 @@ function runMCTS(
   explorationweight::Number=1.0,
   earlystop::Union{Function,Nothing}=nothing,
   saveSimulatedNode::Bool=false,
-  multithread=false
-  )::NamedTuple{(:root, :bestNextState, :bestFinalState),Tuple{MCTSNode,T,T}} where {T<:Any}
-
+  multithread=false,
+  )::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
+                Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any}
+  println("--> LLMMCTS runMCTS 1")
+  # 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{Symbol,Any}())
+                  Dict{String,Any}())
 
+  # 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
+    println("--> LLMMCTS runMCTS 2")
+    # Phase 1: SELECTION - Traverse tree using UCT until reaching a leaf node
+    # UCT balances exploration (new branches) vs exploitation (promising branches)
     while !isleaf(node)
+      println("--> LLMMCTS runMCTS 3")
       node = UCTselect(node, explorationweight)
     end
-
+    println("--> LLMMCTS runMCTS 4")
+    # Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
     if node.isterminal
-      # MCTS arrive at the leaf node that is also a terminal state,
-      # do nothing then go directly to backpropagation. It means the end of this iteration
+      println("--> LLMMCTS runMCTS 5")
+      # If this terminal state has high reward (>= 8), store it for later
+      if node.state[:reward] >= 8
+        println("--> LLMMCTS runMCTS 6")
+        put!(highValueState, deepcopy(node.state))
+      end
+      println("--> LLMMCTS runMCTS 7")
+      # Backpropagate the terminal node's own reward up to root
+      # This updates all ancestors with this path's outcome
       backpropagate(node, node.reward)
     else
+      println("--> LLMMCTS runMCTS 8")
+      # 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)
+      println("--> LLMMCTS runMCTS 9")
+      # Phase 4: SIMULATION + BACKPROPAGATION
+      # For each newly expanded child, run simulation and update statistics
       if multithread
+        println("--> LLMMCTS runMCTS 10")
+        # Parallel simulation: spawn threads for each child node
         @sync for (leafNodeKey, leafNode) in node.children
           @spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
                   maxSimulationDepth=maxSimulationDepth, 
                   horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
                   saveSimulatedNode=saveSimulatedNode,
-                  multithread=multithread)
+                  multithread=multithread,
+                  highValueState=highValueState,
+                  )
         end
       else
+        println("--> LLMMCTS runMCTS 11")
+        # Sequential simulation: process each child one at a time
         for (leafNodeKey, leafNode) in node.children
+          println("--> LLMMCTS runMCTS 11-1")
           simulateThenBackpropagate(leafNode, transition, transitionargs;
-                  maxSimulationDepth=maxSimulationDepth, 
-                  horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
-                  saveSimulatedNode=saveSimulatedNode,
-                  multithread=multithread)
+            maxSimulationDepth=maxSimulationDepth, 
+            horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
+            saveSimulatedNode=saveSimulatedNode,
+            multithread=multithread,
+            highValueState=highValueState)
         end
       end
     end
-
-    # stop if the early stop condition is met
+    println("--> LLMMCTS runMCTS 12")
+    # Phase 5: EARLY STOP CHECK
+    # Optional: stop search early if a condition is met
     if typeof(earlystop) <: Function && earlystop(node.state)
+      println("--> LLMMCTS runMCTS 13")
       break
     end
   end
-
-  # select the best next state and the best final state
+  println("--> LLMMCTS runMCTS 14")
+  # After all iterations, extract results from the search tree
+  # Select best immediate next state (best child of root)
   bestNextState = selectBestNextNode(root)
-  besttrajectory = selectBestTrajectoryNode(root)
+  println("--> LLMMCTS runMCTS 15")
+  # Select best terminal state along the optimal trajectory
+  bestTerminalState = selectBestTrajectoryNode(root)
 
-  return (root=root, bestNextState=bestNextState.state, bestFinalState=besttrajectory.state)
+  # Collect all high-value states from the channel into a list
+  highValueStateList = Vector{Dict{String, Any}}()
+  while !isempty(highValueState)
+    println("--> LLMMCTS runMCTS 16")
+    push!(highValueStateList, take!(highValueState))
+  end
+  println("--> LLMMCTS runMCTS 17")
+  # Return complete search results
+  result = (
+            root=root, 
+            bestNextState=bestNextState.state, 
+            bestTerminalState=bestTerminalState.state,
+            highValueStateList=highValueStateList
+            )
+
+  return result
 end
 
-""" Search the best action to take for a given state and task 
+""" Run simulation from a given node and backpropagate the reward.
+
+This function performs simulation (rollout) from the given node, collects the
+cumulative reward along the trajectory, and backpropagates it up the tree to update
+visit counts and state values.
+
+Does **not** mutate the input node's children (unless `saveSimulatedNode=true`).
 
 # Arguments
 - `node::MCTSNode`
-    current node to simulate from
+  The current node to simulate from.
 - `transition::Function`
-    a function that defines how the state transitions
+  A function that defines how the state transitions.
 - `transitionargs::NamedTuple`
-    arguments for transition function
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-- `maxSimulationDepth::Integer`
-    a number of levels MCTS goes during simulation phase (default: 3)
-- `horizontalSampleSimulationPhase::Integer`
-    a number of child states MCTS samples at each node during simulation phase (default: 3)
-- `saveSimulatedNode::Bool`
-    whether to save nodes created during simulation phase (default: false)
-- `multithread::Bool`
-    whether to use multithreading during simulation (default: false)
+- `maxSimulationDepth::Integer=3`
+  Maximum depth MCTS goes during simulation phase.
+- `horizontalSampleSimulationPhase::Integer=3`
+  Number of child states sampled at each node during simulation phase.
+- `saveSimulatedNode::Bool=false`
+  Whether to save nodes created during simulation phase. If false, children are
+  cleared after simulation.
+- `multithread::Bool=false`
+  Whether to use multithreading during simulation.
 
-# Returns
-  Nothing, but updates the node's reward and visit count through backpropagation
+# Return
+- `Nothing`
+
+# Signature
 """
 function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; 
                 maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,
                 saveSimulatedNode::Bool=false,
-                multithread=false)
-  simTrajectoryReward, terminalstate = simulate(node, transition, transitionargs;
-                                                maxSimulationDepth=maxSimulationDepth, 
-                                                horizontalSample=horizontalSampleSimulationPhase,
-                                                multithread=multithread)
+                multithread=false,
+                highValueState=Union{Nothing,Any}=nothing)
+  println("--> LLMMCTS simulateThenBackpropagate 1")
+  # 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)
+  println("--> LLMMCTS simulateThenBackpropagate 2")
+  # 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
+    println("--> LLMMCTS simulateThenBackpropagate 3")
+    put!(highValueState, deepcopy(terminalstate))
+  end
+  println("--> LLMMCTS simulateThenBackpropagate 4")
+  # 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
+  println("--> LLMMCTS simulateThenBackpropagate 5")
+  # 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
+    println("--> LLMMCTS simulateThenBackpropagate 6")
     node.children = Dict{String, MCTSNode}()
   end
+  println("--> LLMMCTS simulateThenBackpropagate 7")
 end
 
 

src/mcts.jl

@@ -1,7 +1,7 @@
 module mcts
   
 export selectBestNextNode, selectBestTrajectoryNode, backpropagate, isleaf, isroot, selectChildNode,
-        expand, simulate, makeNewState
+        expand, simulate
 using Base.Threads
 using GeneralUtils
 
@@ -10,29 +10,33 @@ using ..type
 # ---------------------------------------------- 100 --------------------------------------------- #
 
 
-""" Select the best next node based on the highest value metric
+""" Select the best child node based on the highest value metric.
+
+The selection metric depends on the node's state values:
+- If the sum of statevalues is non-zero, uses `statevalue/visits` ratio.
+- Otherwise, uses `progressvalue + reward`.
 
 # Arguments
-  - `node::MCTSNode`
-    node of a search tree to evaluate
+- `node::MCTSNode`
+  The node whose children will be evaluated.
 
 # Return
-  - `childNode::MCTSNode`
-    the child node with highest value based on either:
-    - statevalue/visits ratio if any nodes have non-zero statevalue
-    - progressvalue + reward otherwise
+- `childNode::MCTSNode`
+  The child node with the highest value according to the selection metric.
 """
 function selectBestNextNode(node::MCTSNode)::MCTSNode
   highestProgressValue = -1
   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
@@ -41,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
 
@@ -56,20 +61,22 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
 end
 
 
-""" Select the best trajectory node based on the highest reward
+""" Select the best node along the optimal trajectory.
+
+Traverses down the tree from the given node by repeatedly applying `selectBestNextNode`
+until reaching a leaf node, returning the highest-value node found along the path.
 
 # Arguments
-  - `node::MCTSNode`
-    node of a search tree to evaluate
+- `node::MCTSNode`
+  The node to start trajectory selection from.
 
 # Return
-  - `childNode::MCTSNode`
-    the highest value child node found by traversing down the tree using selectBestNextNode
-    until reaching a leaf node
-
-# Signature
+- `childNode::MCTSNode`
+  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
@@ -78,99 +85,108 @@ function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
 end
 
 
-""" Backpropagate reward along the simulation chain
+""" Backpropagate reward along the simulation chain.
+
+Updates visit counts and state values for all nodes along the path from the given
+leaf node to the root, applying reward discounting for future rewards.
+
+**Modifies nodes in place.**
 
 # Arguments
-  - `node::MCTSNode`
-    leaf node of a search tree
-  - `simTrajectoryReward::T`
-    total reward from trajectory simulation
-  - `discountRewardCoeff::AbstractFloat`
-    A discount reward coefficient to reduce future reward. The futher in the future the lower
-    reward it is now.
-  
-# Return
-  - `Nothing`
-    This function modifies the nodes in place and returns nothing
+- `node::MCTSNode`
+  The leaf node from which to start backpropagation.
+- `simTrajectoryReward::Number`
+  The total reward from the trajectory simulation.
 
-# Signature
+# Keyword Arguments
+- `discountRewardCoeff::AbstractFloat=0.9`
+  Discount coefficient applied to future rewards. Larger distances from the leaf
+  receive progressively lower discounted rewards.
+
+# Return
+- `Nothing`
 """
 function backpropagate(node::MCTSNode, simTrajectoryReward::T;
-                      discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
+                       discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
+  println("--> LLMMCTS backpropagate 1")
+  # 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
-    node.visits += 1  # Increment visit count for this node
+    println("--> LLMMCTS backpropagate 2")
+    # Increment visit count - this simulation passed through this node
+    node.visits += 1
+    println("--> LLMMCTS backpropagate 3")
     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
+  println("--> LLMMCTS backpropagate 4")
 end
 
-""" Determine whether a node is a leaf node of a search tree.
+""" Determine whether a node is a leaf node.
+
+A leaf node has no children.
 
 # Arguments
-  - `node::MCTSNode`
-    a search tree node
+- `node::MCTSNode`
+  The search tree node to check.
 
 # Return
-  - `result::Bool`
-    true if it is a leaf node (has no children), false otherwise.
+- `result::Bool`
+  `true` if the node has no children, `false` otherwise.
 
 # Example
 ```jldoctest
-julia> using Revise
-julia> using YiemAgent, DataStructures
-julia> initialState = Dict{Symbol, Any}(
-          :customerinfo=> Dict{Symbol, Any}(),
-          :storeinfo=> Dict{Symbol, Any}(),  
-
-          :thoughtHistory=> OrderedDict{Symbol, Any}(
-            :question=> "How are you?",
-            )  
-          )
-julia> statetype = typeof(initialState)
-julia> root = YiemAgent.MCTSNode(initialState, 0, 0.0, Dict{statetype, YiemAgent.MCTSNode}())
-julia> YiemAgent.isleaf(root)
+julia> using LLMMCTS
+julia> node = MCTSNode("leaf", Dict(:reward=>1.0), 0, 0, 0, 1.0, true, nothing, Dict(), Dict())
+julia> isleaf(node)
 true
 ```
-
-# Signature
 """
 isleaf(node::MCTSNode)::Bool = isempty(node.children)
 
-""" Determine wheter a given node is a root node
+""" Determine whether a given node is a root node.
+
+The root node is identified by having `"root"` as its `nodekey`.
 
 # Arguments
-  - `node::MCTSNode`
-    node of a search tree
+- `node::MCTSNode`
+  The search tree node to check.
 
 # Return
-  - `isrootnode::Bool`
-    true if the given node is root node, false otherwise
-
-# Signature
+- `isrootnode::Bool`
+  `true` if the node is the root node, `false` otherwise.
 """
 isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false
 
 
 
-""" Select child node based on the highest statevalue
+""" Select the child node with the highest value.
+
+Uses `progressvalue + reward` as the selection metric.
 
 # Arguments
-  - `node::MCTSNode`
-    node of a search tree
+- `node::MCTSNode`
+  The node whose children will be evaluated.
 
 # Return
-  - `childNode::MCTSNode`
-    the highest value child node
-
-# Signature
+- `childNode::MCTSNode`
+  The child node with the highest `progressvalue + reward` value.
 """
 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
@@ -183,35 +199,43 @@ function selectChildNode(node::MCTSNode)::MCTSNode
 end
 
 
-""" Expand selected node.
+""" Expand a node by generating new child nodes.
+
+Creates new child nodes by applying the transition function multiple times
+(horizontally samples) from the current node.
 
 # Arguments
-  - `node::MCTSNode`
-    MCTS node to expand
-  - `transition::Function`
-    A function that handles state transition.
-  - `transitionargs::NamedTuple`
-    Arguments for transition()
+- `node::MCTSNode`
+  The MCTS node to expand.
+- `transition::Function`
+  A function that handles state transition.
+- `transitionargs::NamedTuple`
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-  - `horizontalSample::Integer`
-    Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3.
-  - `multithread::Bool`
-    Whether to run expansion in parallel using multiple threads. Defaults to false.
-    
-# Return
-  - None
+- `horizontalSample::Integer=3`
+  Number of child nodes to generate.
+- `multithread::Bool=false`
+  Whether to run expansion in parallel using multiple threads.
 
-# Signature
+# Return
+- `Nothing`
 """
 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)
+  println("--> LLMMCTS expand 1")
   if multithread
     @sync for i in 1:horizontalSample 
       @spawn _expand(node, transition, transitionargs)
     end
   else
+    println("--> LLMMCTS expand 2")
     for i in 1:horizontalSample
+      println("--> LLMMCTS expand 3")
       _expand(node, transition, transitionargs)
     end
   end
@@ -219,166 +243,190 @@ end
 
 """ Helper function to expand a single child node.
 
+Creates one new child node from the parent node using the transition function.
+Checks for semantically equivalent states (dejavu) to avoid duplicates.
+
 # Arguments
-  - `node::MCTSNode`
-    Parent MCTS node to expand from
-  - `transition::Function`
-    A function that handles state transition
-  - `transitionargs::NamedTuple`
-    Arguments for transition()
+- `node::MCTSNode`
+  The parent MCTS node to expand from.
+- `transition::Function`
+  A function that handles state transition.
+- `transitionargs::NamedTuple`
+  Arguments passed to the transition function.
 
 # Return
-  - None
-
-# Signature
+- `Nothing`
 """
 function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
-    result = transition(node.state, transitionargs)
-    newNodeKey::AbstractString = result[:newNodeKey]
-    newstate::AbstractDict = result[:newstate]
-    progressvalue::Integer = result[:progressvalue]
-
-    """
-    [] newNodeKey ∉ keys(node.children). 
-          New state may have semantic vector close enought to 
-        one of existing child state. Which can be assume that they are the same state 
-        semantically-wise i.e. De javu. This could be used to recall lessons for this 
-        similar situation to improve decisionMaker and evaluator. 
-    """
-    if newNodeKey ∉ keys(node.children)
-      newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward], 
-          newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{Symbol, Any}())
-      node.children[newNodeKey] = newNode
-    end
+  println("--> LLMMCTS _expand 1")
+  # Generate one child node from the parent using the transition function
+  result = transition(node.state, transitionargs)
+  newNodeKey::AbstractString = result[:newNodeKey]
+  newstate::AbstractDict = result[:newstate]
+  progressvalue::Integer = result[:progressvalue]
+  println("--> LLMMCTS _expand 2")
+  # Dejavu detection: avoid adding duplicate states
+  # If newNodeKey already exists, skip - this handles semantically equivalent states
+  if newNodeKey ∉ keys(node.children)
+    println("--> LLMMCTS _expand 3")
+    # Create new MCTS node with:
+    # - visits=0: no simulations yet
+    # - statevalue=0: will be updated after simulation
+    # - progressvalue: LLM's estimate (fast heuristic)
+    # - reward: immediate environment feedback
+    newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate["reward"], 
+        newstate["isterminal"], node, Dict{String, MCTSNode}(), Dict{String, Any}())
+    println("--> LLMMCTS _expand 4")
+    node.children[newNodeKey] = newNode
+    println("--> LLMMCTS _expand 5")
+  end
 end
 
-""" Simulate interactions between agent and environment
+""" Simulate interactions between agent and environment.
+
+Performs a rollout from the given node up to the maximum simulation depth,
+sampling child nodes at each level and accumulating rewards along the way.
 
 # Arguments
-  - `node::MCTSNode`
-    node that will be a simulation starting point.
-  - `transition::Function`
-    A user function that handles how state transition.
-  - `transitionargs::NamedTuple`
-    Arguments for everything the user will use within transition().
-  - `maxSimulationDepth::Integer`
-    maximum depth level MCTS goes vertically during simulation.
-  - `horizontalSample::Integer`
-    Total number to sample from the current node (i.e. expand new node horizontally)
+- `node::MCTSNode`
+  The node to start simulation from.
+- `transition::Function`
+  A user function that handles state transition.
+- `transitionargs::NamedTuple`
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-  - `multithread::Bool`
-    Whether to run expansion in parallel using multiple threads. Defaults to false.
-    
-# Return
-  - `simTrajectoryReward::Number`
-    Cumulative reward collected along the simulation trajectory
-  - `terminalstate::Union{Dict{Symbol, Any}, Nothing}`
-    Final state if terminal state reached, nothing otherwise
+- `maxSimulationDepth::Integer=3`
+  Maximum depth level MCTS goes vertically during simulation.
+- `horizontalSample::Integer=3`
+  Number of child nodes sampled at each node during simulation.
+- `multithread::Bool=false`
+  Whether to run expansion in parallel using multiple threads.
 
-# Signature
+# Return
+- `NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}`
+  - `simTrajectoryReward`: cumulative reward collected along the simulation trajectory
+  - `terminalstate`: final state if a terminal state was reached, `nothing` otherwise
 """
 function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; 
   maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
-)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{Symbol, Any}, Nothing}}}
-
+  )::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}
+  println("--> LLMMCTS simulate 1")
+  # 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
+    println("--> LLMMCTS simulate 2")
+    # Accumulate the current node's reward to the trajectory total
     simTrajectoryReward += node.reward
+    
+    # Check if we've reached a terminal state
     if node.isterminal
+      println("--> LLMMCTS simulate 3")
       terminalstate = node.state
       break
     else
+      println("--> LLMMCTS simulate 4")
+      # Expand current node to generate children (horizontal sampling)
       _ = expand(node, transition, transitionargs;
                 horizontalSample=horizontalSample,
                 multithread=multithread)
+      println("--> LLMMCTS simulate 5")
+      # Select best child to continue the rollout (vertical exploration)
+      # Uses progressvalue + reward for fast selection during simulation
       node = selectChildNode(node)
     end 
+    println("--> LLMMCTS simulate 6")
   end
-
-  return (simTrajectoryReward=simTrajectoryReward, terminalstate=terminalstate)
+  println("--> LLMMCTS simulate 7")
+  return (simTrajectoryReward=simTrajectoryReward, 
+          terminalstate=terminalstate)
 end
 
-""" Make new state
+# """ Make new state
 
-# Arguments
-  - `currentstate::T1`
-    Current state dictionary containing thought history and metadata
-  - `thoughtDict::T4`
-    Dictionary containing new thought and action
-  - `response::T2`
-    Response string from the environment
-  - `select::Union{T3, Nothing}`
-    Selection value or nothing
-  - `reward::T3`
-    Reward value for this state
-  - `isterminal::Bool`
-    Whether this state is terminal
+# # Arguments
+#   - `currentstate::T1`
+#     Current state dictionary containing thought history and metadata
+#   - `thoughtDict::T4`
+#     Dictionary containing new thought and action
+#   - `response::T2`
+#     Response string from the environment
+#   - `select::Union{T3, Nothing}`
+#     Selection value or nothing
+#   - `reward::T3`
+#     Reward value for this state
+#   - `isterminal::Bool`
+#     Whether this state is terminal
 
-# Return
-  - `Tuple{String, Dict{Symbol, <:Any}}`
-    A tuple containing:
-    - A unique node key string
-    - A new state dictionary with updated thought history and metadata
+# # Return
+#   - `Tuple{String, Dict{String, <:Any}}`
+#     A tuple containing:
+#     - A unique node key string
+#     - A new state dictionary with updated thought history and metadata
 
-# Example
-```jldoctest
-julia> 
-```
+# # Example
+# ```jldoctest
+# julia> 
+# ```
 
-# Signature
-"""
-function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing}, 
-  reward::T3, isterminal::Bool
-  )::Tuple{String, Dict{Symbol, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict}
+# # Signature
+# """
+# function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing}, 
+#   reward::T3, isterminal::Bool
+#   )::Tuple{String, Dict{String, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict}
 
-  # Find the latest thought key and index from current state's thought history
-  currentstate_latestThoughtKey, currentstate_latestThoughtIndice = 
-      GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought")
-  # Calculate next index for new thought/action
-  currentstate_nextIndice = 
-              currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1
-  # Create new keys for thought and action based on next index
-  currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice")
-  latestActionKey = Symbol("action_$currentstate_nextIndice")
+#   # Find the latest thought key and index from current state's thought history
+#   currentstate_latestThoughtKey, currentstate_latestThoughtIndice = 
+#       GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought")
+#   # Calculate next index for new thought/action
+#   currentstate_nextIndice = 
+#               currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1
+#   # Create new keys for thought and action based on next index
+#   currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice")
+#   latestActionKey = Symbol("action_$currentstate_nextIndice")
 
-  # Find the latest thought index from input thought dictionary
-  _, thoughtDict_latestThoughtIndice = 
-      GeneralUtils.findHighestIndexKey(thoughtDict, "thought")
+#   # Find the latest thought index from input thought dictionary
+#   _, thoughtDict_latestThoughtIndice = 
+#       GeneralUtils.findHighestIndexKey(thoughtDict, "thought")
   
-  # Determine thought and action keys from thought dictionary
-  thoughtDict_latestThoughtKey, thoughtDict_latestActionKey =
-  if thoughtDict_latestThoughtIndice == -1
-    (:thought, :action)
-  else
-    (
-      Symbol("thought_$thoughtDict_latestThoughtIndice"), 
-      Symbol("action_$thoughtDict_latestThoughtIndice"),
-    )
-  end
+#   # Determine thought and action keys from thought dictionary
+#   thoughtDict_latestThoughtKey, thoughtDict_latestActionKey =
+#   if thoughtDict_latestThoughtIndice == -1
+#     (:thought, :action)
+#   else
+#     (
+#       Symbol("thought_$thoughtDict_latestThoughtIndice"), 
+#       Symbol("action_$thoughtDict_latestThoughtIndice"),
+#     )
+#   end
 
-  # Create new state by deep copying current state
-  newstate = deepcopy(currentstate)
-  # Update thought history with new thought
-  newstate[:thoughtHistory][currentstate_latestThoughtKey] = 
-      thoughtDict[thoughtDict_latestThoughtKey]
-  # Update thought history with new action
-  newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey]
-  # Create and add new observation to thought history
-  newObservationKey = Symbol("observation_$(currentstate_nextIndice)")
-  newstate[:thoughtHistory][newObservationKey] = response
-  # Update state metadata
-  newstate[:reward] = reward
-  newstate[:select] = select
-  newstate[:isterminal] = isterminal
+#   # Create new state by deep copying current state
+#   newstate = deepcopy(currentstate)
+#   # Update thought history with new thought
+#   newstate[:thoughtHistory][currentstate_latestThoughtKey] = 
+#       thoughtDict[thoughtDict_latestThoughtKey]
+#   # Update thought history with new action
+#   newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey]
+#   # Create and add new observation to thought history
+#   newObservationKey = Symbol("observation_$(currentstate_nextIndice)")
+#   newstate[:thoughtHistory][newObservationKey] = response
+#   # Update state metadata
+#   newstate[:reward] = reward
+#   newstate[:select] = select
+#   newstate[:isterminal] = isterminal
 
-  # Generate unique ID for new node
-  newNodeKey = GeneralUtils.uuid4snakecase()
+#   # Generate unique ID for new node
+#   newNodeKey = GeneralUtils.uuid4snakecase()
 
-  return (newNodeKey, newstate)
-end
+#   return (newNodeKey, newstate)
+# end
 
 
 

src/type.jl

@@ -58,7 +58,7 @@ mutable struct MCTSNode{T1<:AbstractDict, T2<:AbstractString}
   isterminal::Bool
   parent::Union{MCTSNode, Nothing} 
   children::Dict{String, MCTSNode}
-  etc::Dict{Symbol, Any}  # store anything
+  etc::Dict{String, Any}  # store anything
 end
 
 

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