diff --git a/src/interface.jl b/src/interface.jl index e48d81d..c28a3ec 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -74,7 +74,7 @@ function runMCTS( 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 @@ -91,38 +91,38 @@ function runMCTS( # Start from root and traverse down using UCT selection node = root 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 - 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; @@ -134,10 +134,10 @@ function runMCTS( ) 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, @@ -147,29 +147,29 @@ function runMCTS( end end end - 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 - 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) - println("--> LLMMCTS runMCTS 15") + # Select best terminal state along the optimal trajectory bestTerminalState = selectBestTrajectoryNode(root) # 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, @@ -218,7 +218,7 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit saveSimulatedNode::Bool=false, 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 = @@ -226,30 +226,30 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit 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) - 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 diff --git a/src/mcts.jl b/src/mcts.jl index 5b65ee7..ef2be7a 100644 --- a/src/mcts.jl +++ b/src/mcts.jl @@ -108,14 +108,14 @@ leaf node to the root, applying reward discounting for future rewards. """ function backpropagate(node::MCTSNode, simTrajectoryReward::T; 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) - 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 # Apply discount to future rewards - rewards further from the current state are worth less @@ -125,7 +125,7 @@ function backpropagate(node::MCTSNode, simTrajectoryReward::T; # 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. @@ -227,15 +227,15 @@ function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple; # 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 @@ -258,17 +258,17 @@ Checks for semantically equivalent states (dejavu) to avoid duplicates. - `Nothing` """ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple) - 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 @@ -276,9 +276,9 @@ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple # - 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 @@ -311,7 +311,7 @@ sampling child nodes at each level and accumulating rewards along the way. function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false )::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 @@ -322,29 +322,29 @@ function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTup 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 - println("--> LLMMCTS simulate 7") + return (simTrajectoryReward=simTrajectoryReward, terminalstate=terminalstate) end