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

Integrating LLMs with MCTS creates a powerful planning system that combines:
- **LLM reasoning**: Understand complex tasks, generate creative solutions, and reason about states
- **MCTS optimization**: Systematically explore solution spaces and find optimal trajectories

### The Sparse Reward Problem in Traditional MCTS

In traditional reinforcement learning and MCTS, **sparse rewards** are a fundamental challenge that severely limits performance:

| Challenge | Traditional MCTS | LLM-MCTS Integration |
|-----------|------------------|----------------------|
| **Reward signal** | Only at terminal states (e.g., win/loss) | **Pseudo-rewards at every state** via LLM |
| **Value estimation** | Relies on Monte Carlo sampling (high variance, slow convergence) | LLM provides **progress value** as heuristic |
| **Sample efficiency** | Low—requires many samples to discover reward | High—LLM guides search toward promising regions |
| **Exploration strategy** | Blind exploration until reward discovered | LLM suggests promising actions to try |

#### Why Sparse Rewards Are Problematic

1. **High variance in value estimates** — With sparse rewards, Monte Carlo estimates have high variance because few samples contribute to each node's value

2. **Slow learning** — Rewards must propagate backward through many layers before affecting early decisions, requiring many iterations

3. **Poor exploration** — Without intermediate signals, MCTS explores randomly until it偶然 discovers a reward, wasting computation

4. **Local optima** — Without guidance, MCTS may get stuck in suboptimal regions of the search space

#### How LLM-MCTS Solves This

**Progress Value vs. Reward:**
- `progressvalue` — LLM's estimate of how close we are to solving the task (pseudo-reward)
- `statevalue` — Actual cumulative reward from Monte Carlo simulations
- `reward` — Immediate reward from environment (may be sparse, only at terminal states)

**The three-tier value system:**

```julia
# LLM provides progress value at every node (dense, fast guidance)
progressvalue = llm_estimate(state)  # Heuristic, available immediately

# MCTS computes statevalue via simulation (sparse but accurate)
statevalue = monte_carlo_average(simulations)  # Accurate but expensive

# Environment provides immediate reward (may be sparse)
reward = environment_reward(state)  # Only at terminal states
```

**How MCTS uses these values together:**

| Phase | Which value used | Why |
|-------|------------------|-----|
| Node selection (UCT) | `progressvalue` + `statevalue` | Dense guidance for fast exploration |
| Simulation | `statevalue` + `reward` | Accurate long-term estimates |
| Backpropagation | `reward` | Ground truth updates |

**Why this matters:** LLMs provide dense `progressvalue` guidance at every node, allowing MCTS to focus computation on promising trajectories. The simulation phase confirms these estimates with accurate `statevalue` computed via Monte Carlo. This combination solves the sparse reward problem while maintaining accuracy.

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

## 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
transition_args = (param1 = "value1", param2 = "value2")

# Run MCTS
result = runMCTS(
    initialstate,
    transition,
    transition_args;
    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,
    transition_args;
    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.

### Utility Functions

- `UCTselect(node, w)` — Select node using UCT score
- `dictify(x; keytype=Any, stringkey=false)` — 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}
)
```

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

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>