From 47d7e9f303c011a92310eee74a197d0fee21a2b3 Mon Sep 17 00:00:00 2001 From: narawat Date: Sat, 4 Jul 2026 12:48:16 +0700 Subject: [PATCH] update --- src/interface.jl | 40 ++++++++++++++++++++++++++-------------- src/mcts.jl | 30 ++++++++++++++++++++++-------- 2 files changed, 48 insertions(+), 22 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index 13a1d76..e48d81d 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,33 +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; @@ -129,8 +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, @@ -140,27 +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, @@ -209,6 +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 = @@ -216,28 +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 - + 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 8d2d0e6..5b65ee7 100644 --- a/src/mcts.jl +++ b/src/mcts.jl @@ -108,12 +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 @@ -123,7 +125,8 @@ function backpropagate(node::MCTSNode, simTrajectoryReward::T; # Move up to parent node to continue propagation node = node.parent end -end + println("--> LLMMCTS backpropagate 4") +end """ Determine whether a node is a leaf node. @@ -224,12 +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 @@ -252,23 +258,27 @@ 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 # - 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}()) + 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 @@ -301,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 @@ -312,25 +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