update docs
This commit is contained in:
+65
-47
@@ -9,46 +9,55 @@ using ..type, ..mcts, ..util
|
||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
||||
|
||||
|
||||
""" Search the best action to take for a given state and task
|
||||
""" Search for the best action to take for a given state and task.
|
||||
|
||||
This function runs the MCTS algorithm through multiple iterations of expansion,
|
||||
simulation, and backpropagation to find optimal decisions.
|
||||
|
||||
Does **not** mutate the input state; it creates new MCTS nodes during search.
|
||||
|
||||
# Arguments
|
||||
- `initialstate::T`
|
||||
initial state
|
||||
- `transition::Function`
|
||||
a function that define how the state transitions
|
||||
- `transitionargs::NamedTuple`
|
||||
arguments for transition function
|
||||
- `initialstate::T`
|
||||
The initial state from which to start the search.
|
||||
- `transition::Function`
|
||||
A function that defines how the state transitions.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments passed to the transition function.
|
||||
|
||||
# Keyword Arguments
|
||||
- `horizontalSampleExpansionPhase::Integer`
|
||||
a number of child state MCTS sample at each node during expansion phase (default: 3)
|
||||
- `horizontalSampleSimulationPhase::Integer`
|
||||
a number of child state MCTS sample at each node during simulation's expansion phase (default: 3)
|
||||
- `maxSimulationDepth::Integer`
|
||||
a number of levels MCTS goes during simulation phase (default: 3)
|
||||
- `maxiterations::Integer`
|
||||
a number of iteration MCTS goes thru expansion -> simulation -> backpropagation cycle (default: 10)
|
||||
- `explorationweight::Number`
|
||||
exploration weight controls how much MCTS should explore new state instead of exploit
|
||||
a known state. 1.0 balance between exploration and exploitation like 50%-50%. 2.0 makes MCTS
|
||||
aggressively explore new state (default: 1.0)
|
||||
- `earlystop::Union{Function,Nothing}`
|
||||
optional function to check early stopping condition if it is satisfied, MCTS will break iterations (default: nothing)
|
||||
- `saveSimulatedNode::Bool`
|
||||
whether to save nodes created during simulation phase (default: false)
|
||||
- `multithread::Bool`
|
||||
whether to use multithreading during simulation (default: false)
|
||||
- `horizontalSampleExpansionPhase::Integer=3`
|
||||
Number of child states sampled at each node during expansion phase.
|
||||
- `horizontalSampleSimulationPhase::Integer=3`
|
||||
Number of child states sampled at each node during simulation's expansion phase.
|
||||
- `maxSimulationDepth::Integer=3`
|
||||
Maximum depth MCTS goes during simulation phase.
|
||||
- `maxiterations::Integer=10`
|
||||
Number of iterations MCTS performs through expansion → simulation → backpropagation cycles.
|
||||
- `explorationweight::Number=1.0`
|
||||
Exploration weight controls how much MCTS explores new states versus exploiting known states.
|
||||
A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0)
|
||||
encourage more aggressive exploration.
|
||||
- `earlystop::Union{Function,Nothing}=nothing`
|
||||
Optional function to check early stopping condition. If satisfied, MCTS breaks iterations.
|
||||
- `saveSimulatedNode::Bool=false`
|
||||
Whether to save nodes created during simulation phase.
|
||||
- `multithread::Bool=false`
|
||||
Whether to use multithreading during simulation.
|
||||
|
||||
# Returns
|
||||
- `NamedTuple{(:root, :bestNextState, :bestFinalState), Tuple{MCTSNode, T, T}}`
|
||||
- root: the complete MCTS tree with root node
|
||||
- bestNextState: the best immediate next state
|
||||
- bestFinalState: the best final state along the best trajectory
|
||||
# Return
|
||||
- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
||||
Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}`
|
||||
- `root`: the complete MCTS tree with root node
|
||||
- `bestNextState`: the best immediate next state
|
||||
- `bestTerminalState`: the best final state along the best trajectory
|
||||
- `highValueStateList`: list of high-value terminal states (reward >= 8)
|
||||
|
||||
# Example
|
||||
Refers to SQLLLM package
|
||||
|
||||
# Signature
|
||||
```jldoctest
|
||||
julia> using LLMMCTS
|
||||
julia> initialState = Dict(:reward=>0.0)
|
||||
julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5)
|
||||
```
|
||||
"""
|
||||
function runMCTS(
|
||||
initialstate::T,
|
||||
@@ -140,28 +149,37 @@ function runMCTS(
|
||||
return result
|
||||
end
|
||||
|
||||
""" Search the best action to take for a given state and task
|
||||
""" Run simulation from a given node and backpropagate the reward.
|
||||
|
||||
This function performs simulation (rollout) from the given node, collects the
|
||||
cumulative reward along the trajectory, and backpropagates it up the tree to update
|
||||
visit counts and state values.
|
||||
|
||||
Does **not** mutate the input node's children (unless `saveSimulatedNode=true`).
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
current node to simulate from
|
||||
The current node to simulate from.
|
||||
- `transition::Function`
|
||||
a function that defines how the state transitions
|
||||
A function that defines how the state transitions.
|
||||
- `transitionargs::NamedTuple`
|
||||
arguments for transition function
|
||||
Arguments passed to the transition function.
|
||||
|
||||
# Keyword Arguments
|
||||
- `maxSimulationDepth::Integer`
|
||||
a number of levels MCTS goes during simulation phase (default: 3)
|
||||
- `horizontalSampleSimulationPhase::Integer`
|
||||
a number of child states MCTS samples at each node during simulation phase (default: 3)
|
||||
- `saveSimulatedNode::Bool`
|
||||
whether to save nodes created during simulation phase (default: false)
|
||||
- `multithread::Bool`
|
||||
whether to use multithreading during simulation (default: false)
|
||||
- `maxSimulationDepth::Integer=3`
|
||||
Maximum depth MCTS goes during simulation phase.
|
||||
- `horizontalSampleSimulationPhase::Integer=3`
|
||||
Number of child states sampled at each node during simulation phase.
|
||||
- `saveSimulatedNode::Bool=false`
|
||||
Whether to save nodes created during simulation phase. If false, children are
|
||||
cleared after simulation.
|
||||
- `multithread::Bool=false`
|
||||
Whether to use multithreading during simulation.
|
||||
|
||||
# Returns
|
||||
Nothing, but updates the node's reward and visit count through backpropagation
|
||||
# Return
|
||||
- `Nothing`
|
||||
|
||||
# Signature
|
||||
"""
|
||||
function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
||||
maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,
|
||||
|
||||
+104
-106
@@ -10,17 +10,19 @@ using ..type
|
||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
||||
|
||||
|
||||
""" Select the best next node based on the highest value metric
|
||||
""" Select the best child node based on the highest value metric.
|
||||
|
||||
The selection metric depends on the node's state values:
|
||||
- If the sum of statevalues is non-zero, uses `statevalue/visits` ratio.
|
||||
- Otherwise, uses `progressvalue + reward`.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
node of a search tree to evaluate
|
||||
- `node::MCTSNode`
|
||||
The node whose children will be evaluated.
|
||||
|
||||
# Return
|
||||
- `childNode::MCTSNode`
|
||||
the child node with highest value based on either:
|
||||
- statevalue/visits ratio if any nodes have non-zero statevalue
|
||||
- progressvalue + reward otherwise
|
||||
- `childNode::MCTSNode`
|
||||
The child node with the highest value according to the selection metric.
|
||||
"""
|
||||
function selectBestNextNode(node::MCTSNode)::MCTSNode
|
||||
highestProgressValue = -1
|
||||
@@ -56,18 +58,18 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
|
||||
end
|
||||
|
||||
|
||||
""" Select the best trajectory node based on the highest reward
|
||||
""" Select the best node along the optimal trajectory.
|
||||
|
||||
Traverses down the tree from the given node by repeatedly applying `selectBestNextNode`
|
||||
until reaching a leaf node, returning the highest-value node found along the path.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
node of a search tree to evaluate
|
||||
- `node::MCTSNode`
|
||||
The node to start trajectory selection from.
|
||||
|
||||
# Return
|
||||
- `childNode::MCTSNode`
|
||||
the highest value child node found by traversing down the tree using selectBestNextNode
|
||||
until reaching a leaf node
|
||||
|
||||
# Signature
|
||||
- `childNode::MCTSNode`
|
||||
The highest-value node found by following the optimal trajectory to a leaf.
|
||||
"""
|
||||
function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
|
||||
while !isleaf(node)
|
||||
@@ -78,22 +80,26 @@ function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
|
||||
end
|
||||
|
||||
|
||||
""" Backpropagate reward along the simulation chain
|
||||
""" Backpropagate reward along the simulation chain.
|
||||
|
||||
Updates visit counts and state values for all nodes along the path from the given
|
||||
leaf node to the root, applying reward discounting for future rewards.
|
||||
|
||||
**Modifies nodes in place.**
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
leaf node of a search tree
|
||||
- `simTrajectoryReward::T`
|
||||
total reward from trajectory simulation
|
||||
- `discountRewardCoeff::AbstractFloat`
|
||||
A discount reward coefficient to reduce future reward. The futher in the future the lower
|
||||
reward it is now.
|
||||
|
||||
# Return
|
||||
- `Nothing`
|
||||
This function modifies the nodes in place and returns nothing
|
||||
- `node::MCTSNode`
|
||||
The leaf node from which to start backpropagation.
|
||||
- `simTrajectoryReward::Number`
|
||||
The total reward from the trajectory simulation.
|
||||
|
||||
# Signature
|
||||
# Keyword Arguments
|
||||
- `discountRewardCoeff::AbstractFloat=0.9`
|
||||
Discount coefficient applied to future rewards. Larger distances from the leaf
|
||||
receive progressively lower discounted rewards.
|
||||
|
||||
# Return
|
||||
- `Nothing`
|
||||
"""
|
||||
function backpropagate(node::MCTSNode, simTrajectoryReward::T;
|
||||
discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
|
||||
@@ -106,65 +112,55 @@ function backpropagate(node::MCTSNode, simTrajectoryReward::T;
|
||||
end
|
||||
end
|
||||
|
||||
""" Determine whether a node is a leaf node of a search tree.
|
||||
""" Determine whether a node is a leaf node.
|
||||
|
||||
A leaf node has no children.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
a search tree node
|
||||
- `node::MCTSNode`
|
||||
The search tree node to check.
|
||||
|
||||
# Return
|
||||
- `result::Bool`
|
||||
true if it is a leaf node (has no children), false otherwise.
|
||||
- `result::Bool`
|
||||
`true` if the node has no children, `false` otherwise.
|
||||
|
||||
# Example
|
||||
```jldoctest
|
||||
julia> using Revise
|
||||
julia> using YiemAgent, DataStructures
|
||||
julia> initialState = Dict{String, Any}(
|
||||
"customerinfo"=> Dict{String, Any}(),
|
||||
"storeinfo"=> Dict{String, Any}(),
|
||||
|
||||
"thoughtHistory"=> OrderedDict{String, Any}(
|
||||
"question"=> "How are you?",
|
||||
)
|
||||
)
|
||||
julia> statetype = typeof(initialState)
|
||||
julia> root = YiemAgent.MCTSNode(initialState, 0, 0.0, Dict{statetype, YiemAgent.MCTSNode}())
|
||||
julia> YiemAgent.isleaf(root)
|
||||
julia> using LLMMCTS
|
||||
julia> node = MCTSNode("leaf", Dict(:reward=>1.0), 0, 0, 0, 1.0, true, nothing, Dict(), Dict())
|
||||
julia> isleaf(node)
|
||||
true
|
||||
```
|
||||
|
||||
# Signature
|
||||
"""
|
||||
isleaf(node::MCTSNode)::Bool = isempty(node.children)
|
||||
|
||||
""" Determine wheter a given node is a root node
|
||||
""" Determine whether a given node is a root node.
|
||||
|
||||
The root node is identified by having `"root"` as its `nodekey`.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
node of a search tree
|
||||
- `node::MCTSNode`
|
||||
The search tree node to check.
|
||||
|
||||
# Return
|
||||
- `isrootnode::Bool`
|
||||
true if the given node is root node, false otherwise
|
||||
|
||||
# Signature
|
||||
- `isrootnode::Bool`
|
||||
`true` if the node is the root node, `false` otherwise.
|
||||
"""
|
||||
isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false
|
||||
|
||||
|
||||
|
||||
""" Select child node based on the highest statevalue
|
||||
""" Select the child node with the highest value.
|
||||
|
||||
Uses `progressvalue + reward` as the selection metric.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
node of a search tree
|
||||
- `node::MCTSNode`
|
||||
The node whose children will be evaluated.
|
||||
|
||||
# Return
|
||||
- `childNode::MCTSNode`
|
||||
the highest value child node
|
||||
|
||||
# Signature
|
||||
- `childNode::MCTSNode`
|
||||
The child node with the highest `progressvalue + reward` value.
|
||||
"""
|
||||
function selectChildNode(node::MCTSNode)::MCTSNode
|
||||
highestProgressValue = -1
|
||||
@@ -183,26 +179,27 @@ function selectChildNode(node::MCTSNode)::MCTSNode
|
||||
end
|
||||
|
||||
|
||||
""" Expand selected node.
|
||||
""" Expand a node by generating new child nodes.
|
||||
|
||||
Creates new child nodes by applying the transition function multiple times
|
||||
(horizontally samples) from the current node.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
MCTS node to expand
|
||||
- `transition::Function`
|
||||
A function that handles state transition.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments for transition()
|
||||
- `node::MCTSNode`
|
||||
The MCTS node to expand.
|
||||
- `transition::Function`
|
||||
A function that handles state transition.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments passed to the transition function.
|
||||
|
||||
# Keyword Arguments
|
||||
- `horizontalSample::Integer`
|
||||
Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3.
|
||||
- `multithread::Bool`
|
||||
Whether to run expansion in parallel using multiple threads. Defaults to false.
|
||||
|
||||
# Return
|
||||
- None
|
||||
- `horizontalSample::Integer=3`
|
||||
Number of child nodes to generate.
|
||||
- `multithread::Bool=false`
|
||||
Whether to run expansion in parallel using multiple threads.
|
||||
|
||||
# Signature
|
||||
# Return
|
||||
- `Nothing`
|
||||
"""
|
||||
function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple;
|
||||
horizontalSample::Integer=3, multithread=false)
|
||||
@@ -219,18 +216,19 @@ end
|
||||
|
||||
""" Helper function to expand a single child node.
|
||||
|
||||
Creates one new child node from the parent node using the transition function.
|
||||
Checks for semantically equivalent states (dejavu) to avoid duplicates.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
Parent MCTS node to expand from
|
||||
- `transition::Function`
|
||||
A function that handles state transition
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments for transition()
|
||||
- `node::MCTSNode`
|
||||
The parent MCTS node to expand from.
|
||||
- `transition::Function`
|
||||
A function that handles state transition.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments passed to the transition function.
|
||||
|
||||
# Return
|
||||
- None
|
||||
|
||||
# Signature
|
||||
- `Nothing`
|
||||
"""
|
||||
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
|
||||
result = transition(node.state, transitionargs)
|
||||
@@ -252,31 +250,31 @@ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple
|
||||
end
|
||||
end
|
||||
|
||||
""" Simulate interactions between agent and environment
|
||||
""" Simulate interactions between agent and environment.
|
||||
|
||||
Performs a rollout from the given node up to the maximum simulation depth,
|
||||
sampling child nodes at each level and accumulating rewards along the way.
|
||||
|
||||
# Arguments
|
||||
- `node::MCTSNode`
|
||||
node that will be a simulation starting point.
|
||||
- `transition::Function`
|
||||
A user function that handles how state transition.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments for everything the user will use within transition().
|
||||
- `maxSimulationDepth::Integer`
|
||||
maximum depth level MCTS goes vertically during simulation.
|
||||
- `horizontalSample::Integer`
|
||||
Total number to sample from the current node (i.e. expand new node horizontally)
|
||||
- `node::MCTSNode`
|
||||
The node to start simulation from.
|
||||
- `transition::Function`
|
||||
A user function that handles state transition.
|
||||
- `transitionargs::NamedTuple`
|
||||
Arguments passed to the transition function.
|
||||
|
||||
# Keyword Arguments
|
||||
- `multithread::Bool`
|
||||
Whether to run expansion in parallel using multiple threads. Defaults to false.
|
||||
|
||||
# Return
|
||||
- `simTrajectoryReward::Number`
|
||||
Cumulative reward collected along the simulation trajectory
|
||||
- `terminalstate::Union{Dict{String, Any}, Nothing}`
|
||||
Final state if terminal state reached, nothing otherwise
|
||||
- `maxSimulationDepth::Integer=3`
|
||||
Maximum depth level MCTS goes vertically during simulation.
|
||||
- `horizontalSample::Integer=3`
|
||||
Number of child nodes sampled at each node during simulation.
|
||||
- `multithread::Bool=false`
|
||||
Whether to run expansion in parallel using multiple threads.
|
||||
|
||||
# Signature
|
||||
# Return
|
||||
- `NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}`
|
||||
- `simTrajectoryReward`: cumulative reward collected along the simulation trajectory
|
||||
- `terminalstate`: final state if a terminal state was reached, `nothing` otherwise
|
||||
"""
|
||||
function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
||||
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
||||
|
||||
+32
-14
@@ -6,26 +6,44 @@ using ..type
|
||||
|
||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
||||
|
||||
""" Select a node based on UCT score
|
||||
""" 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`
|
||||
mcts node
|
||||
- `w::T`
|
||||
exploration weight. Value is usually between 1 to 2.
|
||||
Value 1.0 makes MCTS balance between exploration and exploitation like 50%-50%.
|
||||
Value 2.0 makes MCTS aggressively search the tree.
|
||||
- `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`
|
||||
child node with highest UCT score. UCT score balances between exploitation (state value)
|
||||
and exploration (visit count) based on the exploration weight w.
|
||||
- `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).
|
||||
|
||||
# Example
|
||||
# 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).
|
||||
|
||||
# Examples
|
||||
```jldoctest
|
||||
julia>
|
||||
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(...)
|
||||
```
|
||||
|
||||
# Signature
|
||||
"""
|
||||
function UCTselect(node::MCTSNode, w::T)::MCTSNode where {T<:AbstractFloat}
|
||||
maxUCT = -Inf
|
||||
|
||||
Reference in New Issue
Block a user