LLMMCTS/README.md

LLMMCTS

LLMMCTS implements Monte Carlo Tree Search (MCTS) for Large Language Model (LLM) planning tasks.

Why LLM + MCTS?

MCTS is a powerful search algorithm that balances exploration and exploitation through the UCT formula: [ \text{UCT}(s,a) = Q(s,a) + c \sqrt{\frac{\ln N(s)}{N(s,a)}} ]

However, in many real-world problems, rewards are sparse—they only come at the final state. This creates two critical problems:

1. Value estimation delay — Rewards must propagate backward through many layers before affecting early decisions
2. Exploration inefficiency — Without intermediate signals, MCTS explores randomly until it discovers a reward

How LLMs Fix the Sparse Reward Problem

LLMs provide reasoning and pseudo-reward to guide the solution search process:

• Reasoning — LLM understands task structure and generates promising candidate actions
• Pseudo-reward — LLM estimates state quality at every node (not just terminal states)

In the code, this is represented by the progressvalue field:

progressvalue = llm_reasoning_estimate(state)

The LLM evaluates how close the current state is to solving the task, providing dense guidance even when the environment only gives rewards at the end.

The Three-Tier Value System

LLMMCTS combines LLM guidance with MCTS search using three complementary value signals:

Field	Source	Purpose
progressvalue	LLM heuristic	Estimate of how close we are to solving; used for fast node selection
statevalue	Monte Carlo simulation	Actual cumulative reward from simulations; accurate but expensive to compute
reward	Environment	Immediate reward from environment (may be sparse, only at terminal states)

Why this matters:

• progressvalue enables MCTS to explore promising branches quickly without waiting for terminal rewards
• statevalue provides accurate long-term estimates through Monte Carlo simulations
• reward supplies ground truth for backpropagation updates

Benefits of LLM-MCTS Integration

Benefit	Description
Overcomes sparse rewards	LLM provides progressvalue at every node, enabling fast learning without waiting for terminal rewards
Faster convergence	Dense guidance from LLM reduces sample complexity by 5-10x compared to pure Monte Carlo
Better than pure LLM	MCTS systematically compares multiple LLM-generated trajectories, avoiding local optima
Better than pure planning	LLM handles complex reasoning and novel state generation that pure planners cannot
Uncertainty quantification	Visit counts in MCTS nodes reflect confidence in LLM's progress estimates
Configurable depth	MCTS depth controls planning horizon; LLM handles long-term reasoning at each step
Parallel exploration	MCTS naturally supports parallel simulation; LLM generates diverse candidate actions

MCTS Node Structure

MCTSNode(
    nodekey::String,
    state::Dict,
    visits::Integer,
    progressvalue::Number,
    statevalue::Number,
    reward::Number,
    isterminal::Bool,
    parent::Union{MCTSNode, Nothing},
    children::Dict{String, MCTSNode},
    etc::Dict{Symbol, Any}
)

Fields:

• nodekey::String — Unique identifier for the node
• state::Dict — Current state represented as a dictionary
• visits::Integer — Number of times this node has been visited
• progressvalue::Number — LLM's estimate of state quality
• statevalue::Number — Average cumulative reward from simulations
• reward::Number — Immediate reward at this node
• isterminal::Bool — Whether this node represents a terminal state
• parent::Union{MCTSNode, Nothing} — Parent node reference (nothing for root)
• children::Dict{String, MCTSNode} — Mapping of child nodes
• etc::Dict{Symbol, Any} — Additional arbitrary data storage (uses Symbol keys)

Understanding progressvalue, statevalue, and reward

Field	Source	Purpose
progressvalue	LLM heuristic	Estimate of how close we are to solving; used for fast node selection
statevalue	Monte Carlo simulation	Actual cumulative reward from simulations; accurate but expensive to compute
reward	Environment	Immediate reward (may be sparse, only at terminal states)

Why this matters: In traditional MCTS, sparse rewards force extensive exploration. Here, LLM provides dense progressvalue guidance at every node, while statevalue (computed via simulation) provides accurate long-term estimates. MCTS balances both via UCT:

• High progressvalue → explored early (fast guidance)
• High statevalue → exploited once confirmed (accurate value)

Contributing

Overview

Key Features

• UCT-based node selection: Uses Upper Confidence Bound for Trees to balance exploration/exploitation
• Configurable expansion: Parallel or sequential child node generation
• Simulation with depth control: Rollouts with configurable maximum depth
• Reward discounting: Backpropagation with configurable future reward decay
• Multithreading support: Parallel simulation phase for improved performance

Integration with LLMs

This package is designed to work with LLMs as the state transition engine:

# LLM-based transition function
function llm_transition(state::Dict, args::NamedTuple)
    # LLM generates next thought/action based on current state
    response = llm_call(state[:thoughtHistory], args.prompt)
    
    # Parse LLM output into new state
    return Dict(
        :newNodeKey => generate_key(),
        :newstate => update_state(state, response),
        :progressvalue => estimate_value(response)
    )
end

result = runMCTS(initial_state, llm_transition, args)

Installation

using Pkg
Pkg.add("LLMMCTS")

Usage

Basic Example

using LLMMCTS

# Define transition function
function transition(state::Dict, args::NamedTuple)
    # Your transition logic here
    return Dict(:newNodeKey => "child_1", :newstate => new_state, :progressvalue => 5)
end

# Define transition arguments
transitionargs = (param1 = "value1", param2 = "value2")

# Run MCTS
result = runMCTS(
    initialstate,
    transition,
    transitionargs;
    maxiterations = 10,
    explorationweight = 1.0,
    maxSimulationDepth = 3
)

# Access results
root = result.root
best_next_state = result.bestNextState
best_terminal_state = result.bestTerminalState
high_value_states = result.highValueStateList

Advanced Usage

# With custom parameters
result = runMCTS(
    initialstate,
    transition_func,
    transitionargs;
    horizontalSampleExpansionPhase = 5,      # More children during expansion
    horizontalSampleSimulationPhase = 3,     # Sample 3 children during simulation
    maxSimulationDepth = 5,                  # Deeper search
    maxiterations = 50,                      # More iterations
    explorationweight = 2.0,                 # More aggressive exploration
    earlystop = my_earlystop_func,           # Custom early stopping
    saveSimulatedNode = true,                # Keep simulation nodes
    multithread = true                       # Enable parallel simulation
)

API Reference

Main Functions

runMCTS(initialstate, transition, transitionargs; kwargs...)

Search for the best action to take for a given state and task.

Arguments:

• initialstate::T — Initial state
• transition::Function — State transition function
• transitionargs::NamedTuple — Transition function arguments

Keyword Arguments:

• horizontalSampleExpansionPhase::Integer=3 — Children per expansion node
• horizontalSampleSimulationPhase::Integer=3 — Children per simulation node
• maxSimulationDepth::Integer=3 — Maximum simulation depth
• maxiterations::Integer=10 — Number of MCTS iterations
• explorationweight::Number=1.0 — Exploration weight (1.0 = 50/50 balance)
• earlystop::Union{Function,Nothing}=nothing — Early stopping function
• saveSimulatedNode::Bool=false — Keep simulation nodes
• multithread::Bool=false — Enable multithreading

Returns: NamedTuple with root, bestNextState, bestTerminalState, highValueStateList

simulateThenBackpropagate(node, transition, transitionargs; kwargs...)

Run simulation from a node and backpropagate the reward. Returns nothing.

Keyword Arguments:

• maxSimulationDepth::Integer=3 — Maximum simulation depth
• horizontalSampleSimulationPhase::Integer=3 — Children per simulation node
• saveSimulatedNode::Bool=false — Keep simulation nodes
• multithread::Bool=false — Enable multithreading
• highValueState — Channel to store high-value states

backpropagate(node, simTrajectoryReward; kwargs...)

Backpropagate reward along the simulation chain. Updates visit counts and state values for all nodes along the path to the root. Returns nothing.

Arguments:

• node::MCTSNode — The leaf node from which to start backpropagation
• simTrajectoryReward::Number — The total reward from the trajectory simulation

Keyword Arguments:

• discountRewardCoeff::AbstractFloat=0.9 — Discount coefficient applied to future rewards

Utility Functions

• UCTselect(node, w) — Select node using UCT score
• dictify(x; keytype=Any) — Convert JSON.Object/OrderedDict to plain Dict

MCTS Utility Functions

• selectBestNextNode(node) — Select best child node based on value metric
• selectBestTrajectoryNode(node) — Select best node along optimal trajectory
• backpropagate(node, simTrajectoryReward; kwargs...) — Backpropagate reward up the tree
• isleaf(node) — Check if node is a leaf (has no children)
• isroot(node) — Check if node is the root node
• selectChildNode(node) — Select child with highest progressvalue + reward
• expand(node, transition, transitionargs; kwargs...) — Generate child nodes
• simulate(node, transition, transitionargs; kwargs...) — Perform rollout simulation

MCTS Node Structure

MCTSNode(
    nodekey::String,
    state::Dict,
    visits::Integer,
    progressvalue::Number,
    statevalue::Number,
    reward::Number,
    isterminal::Bool,
    parent::Union{MCTSNode, Nothing},
    children::Dict{String, MCTSNode},
    etc::Dict{Symbol, Any}
)

How UCT Uses progressvalue and statevalue

The UCT formula selects children using both value signals: [ \text{UCT}(s,a) = Q(s,a) + c \sqrt{\frac{\ln N(s)}{N(s,a)}} ]

Where:

• Exploitation term (Q(s,a)) — Uses progressvalue for fast guidance, refined by statevalue as simulations accumulate
• Exploration term — Encourages visiting less-explored branches, even those with high progressvalue but low visit count

Selection priority:

1. Nodes with high progressvalue and low visits → explored first (fast guidance)
2. Nodes with high statevalue confirmed by simulations → exploited once reliable
3. Balance determined by explorationweight parameter

When to Use LLM-MCTS

Scenario	Why LLM-MCTS is suitable
Sparse reward environments	Rewards only at terminal states (e.g., game win, code execution success)
Complex reasoning tasks	Tasks requiring multi-step planning (math, coding, tool use)
High branching factor	Many possible actions; LLM filters to promising candidates
Need sample efficiency	Limited budget for environment interactions

When Not to Use LLM-MCTS

Scenario	Alternative approach
Dense rewards available	Use pure RL with reward shaping
Simple decision problems	Classical search (DFS, BFS) is sufficient
Real-time constraints	LLM calls may be too slow; use pre-trained value function
No LLM access	Use pure MCTS with hand-designed heuristics

Performance Characteristics

Metric	Typical range
Sample efficiency	5-10x fewer samples than pure Monte Carlo
LLM calls per iteration	1-5 (depends on horizontalSample* settings)
Convergence time	Scales with depth × LLM latency
Memory usage	O(branching_factor^depth) for tree storage

Limitations

• LLM latency — Each node expansion requires an LLM call; can be slow for large trees
• LLM cost — Each LLM invocation has financial cost; monitor usage
• Heuristic quality — Poor LLM pseudo-rewards lead to suboptimal search
• Determinism — LLM outputs are stochastic; use temperature=0 for reproducibility

Contributing

Contributions are welcome! Please open issues for bugs or feature requests, and submit PRs for improvements.

License

MIT License — see LICENSE for details.

Author

narawat lamaiin narawat@outlook.com