191 lines
4.0 KiB
Julia
191 lines
4.0 KiB
Julia
module util
|
|
|
|
export UCTselect
|
|
|
|
using ..type
|
|
|
|
# ---------------------------------------------- 100 --------------------------------------------- #
|
|
|
|
""" 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`
|
|
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.
|
|
|
|
# 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)
|
|
```
|
|
|
|
# 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
|
|
# 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
|
|
end
|
|
end
|
|
|
|
return selectedNode
|
|
end
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
end # module util |