update
This commit is contained in:
@@ -0,0 +1,329 @@
|
||||
# LLMMCTS
|
||||
|
||||
[](https://github.com/narawat/LLMMCTS.jl)
|
||||
[](LICENSE)
|
||||
|
||||
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:
|
||||
```julia
|
||||
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
|
||||
|
||||
```julia
|
||||
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:
|
||||
|
||||
```julia
|
||||
# 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
|
||||
|
||||
```julia
|
||||
using Pkg
|
||||
Pkg.add("LLMMCTS")
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Example
|
||||
|
||||
```julia
|
||||
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
|
||||
|
||||
```julia
|
||||
# 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
|
||||
|
||||
```julia
|
||||
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](LICENSE) for details.
|
||||
|
||||
## Author
|
||||
|
||||
narawat lamaiin <narawat@outlook.com>
|
||||
Reference in New Issue
Block a user