Mcts

MCTSConfig

Configuration parameters for the MCTS optimizer.

Fields

• num_simulations::Int: Number of MCTS simulations to run per step
• exploration_constant::Float64: UCB1 exploration constant (usually √2 ≈ 1.41)
• degree::Int: Degree for child polydisc generation (passed to children function)
• max_children::Union{Int, Nothing}: Maximum number of children to consider per expansion (nothing = all)
• strict::Bool: If true, use single-branch descent; if false, use full children
• value_transform::Function: Transform from loss to value (default: sigmoid_transform())
• selection_mode::SelectionMode: Strategy for selecting the next step (VisitCount, BestValue, or BestLoss)
• persist_tree::Bool: If true, reuse the subtree rooted at the selected child across steps (default: true)

MCTSConfig(; kwargs...)

Create an MCTS configuration with sensible defaults.

Keyword Arguments

• num_simulations::Int=100: Number of simulations per step
• exploration_constant::Float64=1.41: UCB1 exploration constant
• degree::Int=1: Child generation degree
• max_children::Union{Int, Nothing}=nothing: Max children to consider (nothing = all)
• strict::Bool=false: Whether to use single-branch descent
• value_transform::Function=sigmoid_transform(): Loss to value transformation (see sigmoid_transform, tanh_transform, negation_transform)
• selection_mode::SelectionMode=VisitCount: Child selection strategy (VisitCount, BestValue, or BestLoss)
• persist_tree::Bool=true: If true, reuse subtree across optimization steps

MCTSNode{S,T,N}

A node in the MCTS search tree.

Fields

• polydisc::ValuationPolydisc{S,T,N}: The polydisc at this node
• parent::Union{MCTSNode{S,T,N}, Nothing}: Parent node (nothing for root)
• children::Vector{MCTSNode{S,T,N}}: Child nodes that have been expanded
• visits::Int: Number of times this node has been visited
• total_value::Float64: Sum of all values backpropagated through this node
• is_expanded::Bool: Whether this node's children have been generated
• min_loss::Float64: Minimum raw loss seen in this node's subtree
• is_terminal::Bool: True if expansion produces no children (precision limit reached)
• is_solved::Bool: True if terminal, or expanded with all children solved
• proven_value::Float64: Exact value once solved (NaN if unsolved)
• unsolved_children_count::Int: Number of children not yet marked solved

MCTSNode(polydisc::ValuationPolydisc{S,T,N}, parent=nothing) where {S,T,N}

Create a new MCTS node with the given polydisc and optional parent.

MCTSState{S,T,N}

State maintained across MCTS optimization steps.

Fields

• root::MCTSNode{S,T,N}: The current root node of the search tree
• next_branch::Int: Next branch index for strict mode
• step_count::Int: Number of optimization steps taken

SelectionMode

Enum for MCTS child selection strategy after simulations complete.

Values

• VisitCount: Select child with highest visit count (standard MCTS, robust)
• BestValue: Select child leading to best average value in tree (greedy MCTS)
• BestLoss: Select child leading to leaf with minimum raw loss (greedy, ignores visit averaging)

average_value(node::MCTSNode)

Compute the average value of a node (total_value / visits). Returns 0.0 if node has not been visited.

backpropagate!(node::MCTSNode, value::Float64)

Backpropagate a value from a leaf node up to the root.

Updates visits and total_value for all nodes on the path.

check_solved!(node::MCTSNode)

Check if a node should be marked as solved based on its children's status. A node is solved when it is expanded AND all children are solved. Sets proven_value to the max of children's proven_values (single-agent: maximize value = minimize loss).

Returns true if the node was newly marked as solved.

Note on max_children

When max_children is set in MCTSConfig, only a random subset of children is generated. Currently, a node with sampled children can still be marked solved. This is a known limitation — to fix, add an is_fully_expanded field and gate solved status on it.

evaluate_node(node::MCTSNode{S,T,N}, loss::Loss, config::MCTSConfig) where {S,T,N}

Evaluate a node using the loss function.

Returns the transformed value (by default, -loss).

expand_node!(node::MCTSNode{S,T,N}, config::MCTSConfig) where {S,T,N}

Expand a node by generating its child polydiscs.

Uses the same children generation as greedy/gradient descent.

find_best_node_in_tree(node::MCTSNode)

Recursively find the node with the best average value in the entire tree.

Returns the node with highest average_value, considering only visited nodes.

find_min_loss_node_in_tree(node::MCTSNode)

Recursively find the node with the minimum raw loss in the entire tree.

Returns the node with lowest min_loss, considering only visited nodes.

get_tree_size(node::MCTSNode)

Count the total number of nodes in the MCTS tree.

mcts_best_child_by_value(root::MCTSNode)

Return the child with the best average value (alternative to visit-count selection).

mcts_descent(loss::Loss, param::ValuationPolydisc{S,T,N}, state::MCTSState{S,T,N}, config::MCTSConfig) where {S,T,N}

Perform one step of MCTS optimization.

This function follows the same interface as greedy_descent and gradient_descent, making it compatible with OptimSetup.

Arguments

• loss::Loss: The loss function structure
• param::ValuationPolydisc{S,T,N}: Current parameter values
• state::MCTSState{S,T,N}: MCTS state (includes the search tree)
• config::MCTSConfig: Configuration parameters

Returns

Tuple{ValuationPolydisc{S,T,N}, MCTSState{S,T,N}, Bool}: New parameters, updated state, and convergence status

mcts_descent_init(param::ValuationPolydisc{S,T,N}, loss::Loss, config::MCTSConfig=MCTSConfig()) where {S,T,N}

Initialize an optimization setup for MCTS.

Arguments

• param::ValuationPolydisc{S,T,N}: Initial parameter values
• loss::Loss: The loss function structure
• config::MCTSConfig: MCTS configuration (uses defaults if not provided)

Returns

OptimSetup: Configured optimization setup for MCTS

mcts_search(root::MCTSNode{S,T,N}, loss::Loss, config::MCTSConfig) where {S,T,N}

Run MCTS from a root node and return the best child.

Performs config.num_simulations iterations of:

1. Selection: Follow UCB1 to a leaf
2. Expansion: Expand the leaf if not terminal
3. Evaluation: Compute value using -loss
4. Backpropagation: Update all nodes on the path

Returns the selected child and whether the root has converged.

print_tree_stats(node::MCTSNode, depth::Int=0, max_depth::Int=3)

Print statistics about the MCTS tree for debugging.

propagate_solved_up!(node::MCTSNode)

Starting from a solved node, walk up the parent chain decrementing each ancestor's unsolved_children_count and checking if the ancestor becomes solved. Stops as soon as a parent is not newly solved.

select_best_child(root::MCTSNode, config::MCTSConfig)

Select the best child of root according to the configured selection mode.

Selection Modes

• VisitCount: Returns child with highest visit count (standard MCTS, robust)
• BestValue: Finds node with best average value in tree, returns root's child leading to it (greedy)
• BestLoss: Finds leaf with minimum raw loss, returns root's child leading to it

Arguments

• root::MCTSNode: The root node with expanded children
• config::MCTSConfig: Configuration specifying selection mode

Returns

The selected child node.

select_child(node::MCTSNode, exploration_constant::Float64)

Select the child with the highest UCB1 score.

select_path(root::MCTSNode, exploration_constant::Float64)

Select a path from root to a leaf node using UCB1.

Returns the leaf node reached by following UCB1 selections.

trace_to_root_child(node::MCTSNode, root::MCTSNode)

Trace back from a node to find which direct child of root lies on the path.

Arguments

• node::MCTSNode: The node to trace back from
• root::MCTSNode: The root node

Returns

The direct child of root that is an ancestor of node, or node itself if it's a direct child. Returns nothing if node is the root or not in the tree.

ucb1_score(node::MCTSNode, parent_visits::Int, exploration_constant::Float64)

Compute the UCB1 score for a node.

UCB1(node) = averagevalue(node) + c * √(ln(parentvisits) / node_visits)

Higher scores indicate nodes that should be explored (either high value or underexplored).