# LLMMCTS

[![Version](https://img.shields.io/badge/version-0.1.4-blue.svg)](https://github.com/narawat/LLMMCTS.jl)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](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{String, Any}
)
```

### 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

### Utility Functions

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

### 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{String, 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>