remove trace
This commit is contained in:
+25
-25
@@ -74,7 +74,7 @@ function runMCTS(
|
|||||||
multithread=false,
|
multithread=false,
|
||||||
)::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
)::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
||||||
Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any}
|
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
|
# Initialize the MCTS tree with a root node representing the initial state
|
||||||
# root.visits=0: no visits yet
|
# root.visits=0: no visits yet
|
||||||
# root.statevalue=0: no simulation results yet
|
# root.statevalue=0: no simulation results yet
|
||||||
@@ -91,38 +91,38 @@ function runMCTS(
|
|||||||
# Start from root and traverse down using UCT selection
|
# Start from root and traverse down using UCT selection
|
||||||
node = root
|
node = root
|
||||||
node.visits += 1 # Count this iteration's visit to 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
|
# Phase 1: SELECTION - Traverse tree using UCT until reaching a leaf node
|
||||||
# UCT balances exploration (new branches) vs exploitation (promising branches)
|
# UCT balances exploration (new branches) vs exploitation (promising branches)
|
||||||
while !isleaf(node)
|
while !isleaf(node)
|
||||||
println("--> LLMMCTS runMCTS 3")
|
|
||||||
node = UCTselect(node, explorationweight)
|
node = UCTselect(node, explorationweight)
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS runMCTS 4")
|
|
||||||
# Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
|
# Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
|
||||||
if node.isterminal
|
if node.isterminal
|
||||||
println("--> LLMMCTS runMCTS 5")
|
|
||||||
# If this terminal state has high reward (>= 8), store it for later
|
# If this terminal state has high reward (>= 8), store it for later
|
||||||
if node.state[:reward] >= 8
|
if node.state[:reward] >= 8
|
||||||
println("--> LLMMCTS runMCTS 6")
|
|
||||||
put!(highValueState, deepcopy(node.state))
|
put!(highValueState, deepcopy(node.state))
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS runMCTS 7")
|
|
||||||
# Backpropagate the terminal node's own reward up to root
|
# Backpropagate the terminal node's own reward up to root
|
||||||
# This updates all ancestors with this path's outcome
|
# This updates all ancestors with this path's outcome
|
||||||
backpropagate(node, node.reward)
|
backpropagate(node, node.reward)
|
||||||
else
|
else
|
||||||
println("--> LLMMCTS runMCTS 8")
|
|
||||||
# Phase 3: EXPANSION - Generate children for this non-terminal leaf
|
# Phase 3: EXPANSION - Generate children for this non-terminal leaf
|
||||||
# Horizontal sampling: create multiple child nodes via LLM transition
|
# Horizontal sampling: create multiple child nodes via LLM transition
|
||||||
_ = expand(node, transition, transitionargs;
|
_ = expand(node, transition, transitionargs;
|
||||||
horizontalSample=horizontalSampleExpansionPhase,
|
horizontalSample=horizontalSampleExpansionPhase,
|
||||||
multithread=multithread)
|
multithread=multithread)
|
||||||
println("--> LLMMCTS runMCTS 9")
|
|
||||||
# Phase 4: SIMULATION + BACKPROPAGATION
|
# Phase 4: SIMULATION + BACKPROPAGATION
|
||||||
# For each newly expanded child, run simulation and update statistics
|
# For each newly expanded child, run simulation and update statistics
|
||||||
if multithread
|
if multithread
|
||||||
println("--> LLMMCTS runMCTS 10")
|
|
||||||
# Parallel simulation: spawn threads for each child node
|
# Parallel simulation: spawn threads for each child node
|
||||||
@sync for (leafNodeKey, leafNode) in node.children
|
@sync for (leafNodeKey, leafNode) in node.children
|
||||||
@spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
|
@spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
|
||||||
@@ -134,10 +134,10 @@ function runMCTS(
|
|||||||
)
|
)
|
||||||
end
|
end
|
||||||
else
|
else
|
||||||
println("--> LLMMCTS runMCTS 11")
|
|
||||||
# Sequential simulation: process each child one at a time
|
# Sequential simulation: process each child one at a time
|
||||||
for (leafNodeKey, leafNode) in node.children
|
for (leafNodeKey, leafNode) in node.children
|
||||||
println("--> LLMMCTS runMCTS 11-1")
|
|
||||||
simulateThenBackpropagate(leafNode, transition, transitionargs;
|
simulateThenBackpropagate(leafNode, transition, transitionargs;
|
||||||
maxSimulationDepth=maxSimulationDepth,
|
maxSimulationDepth=maxSimulationDepth,
|
||||||
horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
|
horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
|
||||||
@@ -147,29 +147,29 @@ function runMCTS(
|
|||||||
end
|
end
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS runMCTS 12")
|
|
||||||
# Phase 5: EARLY STOP CHECK
|
# Phase 5: EARLY STOP CHECK
|
||||||
# Optional: stop search early if a condition is met
|
# Optional: stop search early if a condition is met
|
||||||
if typeof(earlystop) <: Function && earlystop(node.state)
|
if typeof(earlystop) <: Function && earlystop(node.state)
|
||||||
println("--> LLMMCTS runMCTS 13")
|
|
||||||
break
|
break
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS runMCTS 14")
|
|
||||||
# After all iterations, extract results from the search tree
|
# After all iterations, extract results from the search tree
|
||||||
# Select best immediate next state (best child of root)
|
# Select best immediate next state (best child of root)
|
||||||
bestNextState = selectBestNextNode(root)
|
bestNextState = selectBestNextNode(root)
|
||||||
println("--> LLMMCTS runMCTS 15")
|
|
||||||
# Select best terminal state along the optimal trajectory
|
# Select best terminal state along the optimal trajectory
|
||||||
bestTerminalState = selectBestTrajectoryNode(root)
|
bestTerminalState = selectBestTrajectoryNode(root)
|
||||||
|
|
||||||
# Collect all high-value states from the channel into a list
|
# Collect all high-value states from the channel into a list
|
||||||
highValueStateList = Vector{Dict{String, Any}}()
|
highValueStateList = Vector{Dict{String, Any}}()
|
||||||
while !isempty(highValueState)
|
while !isempty(highValueState)
|
||||||
println("--> LLMMCTS runMCTS 16")
|
|
||||||
push!(highValueStateList, take!(highValueState))
|
push!(highValueStateList, take!(highValueState))
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS runMCTS 17")
|
|
||||||
# Return complete search results
|
# Return complete search results
|
||||||
result = (
|
result = (
|
||||||
root=root,
|
root=root,
|
||||||
@@ -218,7 +218,7 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit
|
|||||||
saveSimulatedNode::Bool=false,
|
saveSimulatedNode::Bool=false,
|
||||||
multithread=false,
|
multithread=false,
|
||||||
highValueState=Union{Nothing,Any}=nothing)
|
highValueState=Union{Nothing,Any}=nothing)
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 1")
|
|
||||||
# Phase 1: RUN SIMULATION (rollout)
|
# Phase 1: RUN SIMULATION (rollout)
|
||||||
# Perform a rollout from this node, accumulating rewards along the way
|
# Perform a rollout from this node, accumulating rewards along the way
|
||||||
simTrajectoryReward, terminalstate =
|
simTrajectoryReward, terminalstate =
|
||||||
@@ -226,30 +226,30 @@ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transit
|
|||||||
maxSimulationDepth=maxSimulationDepth,
|
maxSimulationDepth=maxSimulationDepth,
|
||||||
horizontalSample=horizontalSampleSimulationPhase,
|
horizontalSample=horizontalSampleSimulationPhase,
|
||||||
multithread=multithread)
|
multithread=multithread)
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 2")
|
|
||||||
# Phase 2: HIGH-VALUE STATE TRACKING
|
# Phase 2: HIGH-VALUE STATE TRACKING
|
||||||
# If we reached a terminal state with high reward (>= 8), store it
|
# 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
|
# This allows users to access multiple good solutions, not just the best one
|
||||||
if highValueState !== nothing &&
|
if highValueState !== nothing &&
|
||||||
terminalstate !== nothing &&
|
terminalstate !== nothing &&
|
||||||
terminalstate["reward"] >= 8
|
terminalstate["reward"] >= 8
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 3")
|
|
||||||
put!(highValueState, deepcopy(terminalstate))
|
put!(highValueState, deepcopy(terminalstate))
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 4")
|
|
||||||
# Phase 3: BACKPROPAGATE
|
# Phase 3: BACKPROPAGATE
|
||||||
# Update statistics (visits, statevalue) for all ancestors up to root
|
# Update statistics (visits, statevalue) for all ancestors up to root
|
||||||
# The simulation result is now incorporated into the tree
|
# The simulation result is now incorporated into the tree
|
||||||
backpropagate(node, simTrajectoryReward)
|
backpropagate(node, simTrajectoryReward)
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 5")
|
|
||||||
# Phase 4: MEMORY MANAGEMENT
|
# Phase 4: MEMORY MANAGEMENT
|
||||||
# Clear children unless user wants to keep them for analysis
|
# Clear children unless user wants to keep them for analysis
|
||||||
# This frees memory for the next iteration while preserving tree structure
|
# This frees memory for the next iteration while preserving tree structure
|
||||||
if saveSimulatedNode == false
|
if saveSimulatedNode == false
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 6")
|
|
||||||
node.children = Dict{String, MCTSNode}()
|
node.children = Dict{String, MCTSNode}()
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS simulateThenBackpropagate 7")
|
|
||||||
end
|
end
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+19
-19
@@ -108,14 +108,14 @@ leaf node to the root, applying reward discounting for future rewards.
|
|||||||
"""
|
"""
|
||||||
function backpropagate(node::MCTSNode, simTrajectoryReward::T;
|
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
|
# 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
|
# Each node's statistics are updated with the cumulative reward from the simulation
|
||||||
while !isroot(node)
|
while !isroot(node)
|
||||||
println("--> LLMMCTS backpropagate 2")
|
|
||||||
# Increment visit count - this simulation passed through this node
|
# Increment visit count - this simulation passed through this node
|
||||||
node.visits += 1
|
node.visits += 1
|
||||||
println("--> LLMMCTS backpropagate 3")
|
|
||||||
node.statevalue += ((node.statevalue * (node.visits-1)) + simTrajectoryReward) / node.visits # Update running average of state value
|
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
|
# 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
|
# Move up to parent node to continue propagation
|
||||||
node = node.parent
|
node = node.parent
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS backpropagate 4")
|
|
||||||
end
|
end
|
||||||
|
|
||||||
""" Determine whether a node is a leaf node.
|
""" 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
|
# This is called "horizontal sampling" - we branch out horizontally in the tree
|
||||||
# - multithread=true: spawn parallel threads for each expansion
|
# - multithread=true: spawn parallel threads for each expansion
|
||||||
# - multithread=false: sequential expansion (default, simpler)
|
# - multithread=false: sequential expansion (default, simpler)
|
||||||
println("--> LLMMCTS expand 1")
|
|
||||||
if multithread
|
if multithread
|
||||||
@sync for i in 1:horizontalSample
|
@sync for i in 1:horizontalSample
|
||||||
@spawn _expand(node, transition, transitionargs)
|
@spawn _expand(node, transition, transitionargs)
|
||||||
end
|
end
|
||||||
else
|
else
|
||||||
println("--> LLMMCTS expand 2")
|
|
||||||
for i in 1:horizontalSample
|
for i in 1:horizontalSample
|
||||||
println("--> LLMMCTS expand 3")
|
|
||||||
_expand(node, transition, transitionargs)
|
_expand(node, transition, transitionargs)
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
@@ -258,17 +258,17 @@ Checks for semantically equivalent states (dejavu) to avoid duplicates.
|
|||||||
- `Nothing`
|
- `Nothing`
|
||||||
"""
|
"""
|
||||||
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
|
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
|
||||||
println("--> LLMMCTS _expand 1")
|
|
||||||
# Generate one child node from the parent using the transition function
|
# Generate one child node from the parent using the transition function
|
||||||
result = transition(node.state, transitionargs)
|
result = transition(node.state, transitionargs)
|
||||||
newNodeKey::AbstractString = result[:newNodeKey]
|
newNodeKey::AbstractString = result[:newNodeKey]
|
||||||
newstate::AbstractDict = result[:newstate]
|
newstate::AbstractDict = result[:newstate]
|
||||||
progressvalue::Integer = result[:progressvalue]
|
progressvalue::Integer = result[:progressvalue]
|
||||||
println("--> LLMMCTS _expand 2")
|
|
||||||
# Dejavu detection: avoid adding duplicate states
|
# Dejavu detection: avoid adding duplicate states
|
||||||
# If newNodeKey already exists, skip - this handles semantically equivalent states
|
# If newNodeKey already exists, skip - this handles semantically equivalent states
|
||||||
if newNodeKey ∉ keys(node.children)
|
if newNodeKey ∉ keys(node.children)
|
||||||
println("--> LLMMCTS _expand 3")
|
|
||||||
# Create new MCTS node with:
|
# Create new MCTS node with:
|
||||||
# - visits=0: no simulations yet
|
# - visits=0: no simulations yet
|
||||||
# - statevalue=0: will be updated after simulation
|
# - statevalue=0: will be updated after simulation
|
||||||
@@ -276,9 +276,9 @@ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple
|
|||||||
# - reward: immediate environment feedback
|
# - reward: immediate environment feedback
|
||||||
newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate["reward"],
|
newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate["reward"],
|
||||||
newstate["isterminal"], node, Dict{String, MCTSNode}(), Dict{String, Any}())
|
newstate["isterminal"], node, Dict{String, MCTSNode}(), Dict{String, Any}())
|
||||||
println("--> LLMMCTS _expand 4")
|
|
||||||
node.children[newNodeKey] = newNode
|
node.children[newNodeKey] = newNode
|
||||||
println("--> LLMMCTS _expand 5")
|
|
||||||
end
|
end
|
||||||
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;
|
function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
||||||
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
||||||
)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}
|
)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}
|
||||||
println("--> LLMMCTS simulate 1")
|
|
||||||
# Perform a rollout simulation from the given node:
|
# Perform a rollout simulation from the given node:
|
||||||
# 1. Accumulate rewards along the trajectory
|
# 1. Accumulate rewards along the trajectory
|
||||||
# 2. Expand nodes horizontally at each level
|
# 2. Expand nodes horizontally at each level
|
||||||
@@ -322,29 +322,29 @@ function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTup
|
|||||||
terminalstate = nothing
|
terminalstate = nothing
|
||||||
|
|
||||||
for depth in 1:maxSimulationDepth
|
for depth in 1:maxSimulationDepth
|
||||||
println("--> LLMMCTS simulate 2")
|
|
||||||
# Accumulate the current node's reward to the trajectory total
|
# Accumulate the current node's reward to the trajectory total
|
||||||
simTrajectoryReward += node.reward
|
simTrajectoryReward += node.reward
|
||||||
|
|
||||||
# Check if we've reached a terminal state
|
# Check if we've reached a terminal state
|
||||||
if node.isterminal
|
if node.isterminal
|
||||||
println("--> LLMMCTS simulate 3")
|
|
||||||
terminalstate = node.state
|
terminalstate = node.state
|
||||||
break
|
break
|
||||||
else
|
else
|
||||||
println("--> LLMMCTS simulate 4")
|
|
||||||
# Expand current node to generate children (horizontal sampling)
|
# Expand current node to generate children (horizontal sampling)
|
||||||
_ = expand(node, transition, transitionargs;
|
_ = expand(node, transition, transitionargs;
|
||||||
horizontalSample=horizontalSample,
|
horizontalSample=horizontalSample,
|
||||||
multithread=multithread)
|
multithread=multithread)
|
||||||
println("--> LLMMCTS simulate 5")
|
|
||||||
# Select best child to continue the rollout (vertical exploration)
|
# Select best child to continue the rollout (vertical exploration)
|
||||||
# Uses progressvalue + reward for fast selection during simulation
|
# Uses progressvalue + reward for fast selection during simulation
|
||||||
node = selectChildNode(node)
|
node = selectChildNode(node)
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS simulate 6")
|
|
||||||
end
|
end
|
||||||
println("--> LLMMCTS simulate 7")
|
|
||||||
return (simTrajectoryReward=simTrajectoryReward,
|
return (simTrajectoryReward=simTrajectoryReward,
|
||||||
terminalstate=terminalstate)
|
terminalstate=terminalstate)
|
||||||
end
|
end
|
||||||
|
|||||||
Reference in New Issue
Block a user