From 6a18591a0bcd7219ebe0007972836a6db7695709 Mon Sep 17 00:00:00 2001 From: narawat Date: Tue, 30 Jun 2026 15:59:00 +0700 Subject: [PATCH] update comments --- src/interface.jl | 53 ++++++++++++++++++++++++------ src/mcts.jl | 85 ++++++++++++++++++++++++++++++++++-------------- src/util.jl | 59 +++++++++++++++++++++++++-------- 3 files changed, 149 insertions(+), 48 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index c0c64f9..13a1d76 100644 --- a/src/interface.jl +++ b/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, - # do nothing then go directly to backpropagation. It means the end of this iteration + # Backpropagate the terminal node's own reward up to root + # 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 diff --git a/src/mcts.jl b/src/mcts.jl index 1ac464f..8d2d0e6 100644 --- a/src/mcts.jl +++ b/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 - node.visits += 1 # Increment visit count for this node + # Increment visit count - this simulation passed through 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) - newNodeKey::AbstractString = result[:newNodeKey] - newstate::AbstractDict = result[:newstate] - progressvalue::Integer = result[:progressvalue] + # 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] - """ - [] 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{String, Any}()) - node.children[newNodeKey] = newNode - end + # Dejavu detection: avoid adding duplicate states + # If newNodeKey already exists, skip - this handles semantically equivalent states + if newNodeKey ∉ keys(node.children) + # 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}()) + node.children[newNodeKey] = newNode + 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 diff --git a/src/util.jl b/src/util.jl index 5a9491c..db0cafa 100644 --- a/src/util.jl +++ b/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 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. -- This function assumes the calling code only invokes it on non-leaf nodes (i.e., - nodes with children). +# 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) +``` + +# 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 - 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