ton/LLMMCTS

Fork 0

Files

T

ton 2c4e5918b9 update

2026-07-04 13:00:54 +07:00

56 KiB

Raw Permalink Blame History

Workprocess Documentation for LLMMCTS

Overview

LLMMCTS implements Monte Carlo Tree Search (MCTS) for Large Language Model (LLM) planning tasks. It combines LLM reasoning with MCTS search to solve complex planning problems with sparse rewards.

Core Concept

The package addresses the sparse reward problem in MCTS by using LLMs to provide pseudo-rewards (called progressvalue) at every node, enabling faster learning without waiting for terminal rewards.

Architecture

Three-Tier Value System

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)

How reward and progressvalue Work Together

The package uses two distinct value signals that serve different purposes during MCTS search:

`reward` - Environment Feedback

Source: Environment only
When: Only at terminal states (sparse) or immediate action outcome
Purpose: Ground truth for backpropagation updates
Usage:
- Accumulated during simulation rollouts (line 287 in mcts.jl)
- Used to update statevalue via backpropagation
- Final metric for determining solution quality

`progressvalue` - LLM Heuristic

Source: LLM's reasoning about state quality
When: At every node expansion (dense guidance)
Purpose: Fast node selection without waiting for terminal rewards
Usage:
- Passed to _expand() from transition() result (line 237 in mcts.jl)
- Used as fallback in UCT when child has zero visits (util.jl:58)
- Combined with reward in selectChildNode() (line 171 in mcts.jl)

The Selection Hierarchy

1. UCT SELECTION (interface.jl → runMCTS → UCTselect)
   └── Uses: statevalue + exploration_weight * sqrt(ln(parent_visits) / child_visits)
   └── Purpose: Balance exploration vs exploitation during tree traversal

2. CHILD SELECTION (mcts.jl → selectChildNode)
   └── Uses: progressvalue + reward
   └── Purpose: Pick best child after expansion during simulation rollout

3. BEST NODE SELECTION (mcts.jl → selectBestNextNode)
   └── If statevalue sum > 0: statevalue / visits
   └── Else: progressvalue + reward
   └── Purpose: Determine final best trajectory after MCTS completes

Data Flow Example

Expansion Phase (LLM call):
┌────────────────────────────────────────────────────────────────────────┐
│ transition(state, args) → {                                            │
│   newNodeKey: "abc-123",                                               │
│   newstate: { reward: 0, isterminal: false },                          │
│   progressvalue: 7.5   ← LLM estimates this state is promising         │
│ }                                                                      │
└────────────────────────────────────────────────────────────────────────┘
                                    ↓
New Node Created:
MCTSNode(
  nodekey = "abc-123",
  state = { reward: 0, isterminal: false },
  visits = 0,
  progressvalue = 7.5,   ← LLM heuristic (fast, initial guidance)
  statevalue = 0,        ← Will be updated after simulation
  reward = 0,            ← Immediate environment reward
  ...
)

Simulation Phase:
┌────────────────────────────────────────────────────────────────────────┐
│ simulate() rolls out 3 levels:                                         │
│ Level 1: reward += 0  (non-terminal)                                   │
│ Level 2: reward += 3  (intermediate reward)                            │
│ Level 3: reward += 10 (terminal state)                                 │
│ ─────────────────────────────────────────────────────────────────────  │
│ Total simTrajectoryReward = 13                                         │
└────────────────────────────────────────────────────────────────────────┘
                                    ↓
Backpropagation:
┌────────────────────────────────────────────────────────────────────────┐
│ Update all ancestors with simTrajectoryReward = 13:                    │
│ node.visits += 1                                                       │
│ node.statevalue = (old_statevalue * (visits-1) + 13) / visits          │
│ reward *= 0.9 (discount for future rewards)                            │
└────────────────────────────────────────────────────────────────────────┘

Why Two Value Systems?

Problem	Traditional MCTS	LLMMCTS Solution
Sparse rewards	Must explore blindly until terminal	LLM provides `progressvalue` at every node
Slow learning	Need many terminal outcomes	Dense `progressvalue` enables fast guidance
Exploration cost	Random expansion is inefficient	LLM filters to promising candidates
Value uncertainty	Hand-designed heuristics may be wrong	`statevalue` (from simulation) confirms LLM estimates

Module Structure

src/
├── type.jl         # MCTSNode struct definition (MCTSNode type)
├── util.jl         # UCT selection utility function (UCTselect)
├── mcts.jl         # Core MCTS operations (select, expand, simulate, backpropagate)
├── interface.jl    # High-level interface (runMCTS, simulateThenBackpropagate)
└── LLMMCTS.jl      # Main package entry point (module exports)

Data Flow

1. Node Structure (type.jl)

MCTSNode(
    nodekey::String,           # Unique identifier
    state::Dict,               # Current state dictionary
    visits::Integer,           # Number of visits to this node
    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 is a terminal state
    parent::Union{MCTSNode, Nothing},  # Parent reference (nothing for root)
    children::Dict{String,MCTSNode},   # Child nodes mapping
    etc::Dict{String,Any}      # Additional data storage
)

2. MCTS Search Process: Rollout and Exploration

LLMMCTS explores the solution space through iterative MCTS cycles, each consisting of four phases: Selection, Expansion, Simulation, and Backpropagation.

The Four Phases of MCTS

Phase 1: SELECTION
└── Start from root node
└── Use UCT (Upper Confidence Bound for Trees) to select child nodes
    ├── UCT formula: Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))
    ├── Exploitation term (Q): Uses progressvalue or statevalue
    └── Exploration term: Prefers less-visited nodes
└── Continue until reaching a leaf node

Phase 2: EXPANSION
└── If leaf is not terminal:
    ├── Apply transition function multiple times (horizontal sampling)
    ├── Generate new child nodes with LLM assistance
    ├── Each child represents a potential next state
    └── Dejavu detection prevents duplicate states

Phase 3: SIMULATION (Rollout)
└── From expanded leaf node:
    ├── Perform rollout up to maxSimulationDepth
    ├── At each level:
    │   ├── Accumulate reward from current node
    │   ├── Check if terminal state reached
    │   ├── Expand node to generate children
    │   └── Select best child using progressvalue + reward
    └── Return cumulative trajectory reward and terminal state

Phase 4: BACKPROPAGATION
└── Update statistics along the path to root:
    ├── Increment visit count for each node
    ├── Update statevalue (running average of rewards)
    ├── Apply discount to future rewards
    └── Propagate reward upward

Detailed Expansion Process (mcts.jl → expand/_expand)

Horizontal Sampling: At each node, the algorithm generates multiple child nodes by applying the transition function several times.

Current Node (state S)
├── transition(S, args) → Child 1 (state S1, progressvalue=P1)
├── transition(S, args) → Child 2 (state S2, progressvalue=P2)
├── transition(S, args) → Child 3 (state S3, progressvalue=P3)
└── ... (horizontalSample times)

Each child is evaluated by:
- progressvalue: LLM's estimate of state quality (fast, heuristic)
- reward: Immediate environment reward

Dejavu Detection: Before adding a child node, the algorithm checks if the newNodeKey already exists in the current node's children. If it does, the node is skipped to avoid duplicates. This handles semantically equivalent states.

Detailed Simulation Process (mcts.jl → simulate)

The simulation phase performs a rollout from a given node to explore the solution space vertically:

Input: Starting node, transition function, maxSimulationDepth

simTrajectoryReward = 0
terminalstate = nothing

for depth in 1:maxSimulationDepth:
    # Step 1: Accumulate current node's reward
    simTrajectoryReward += node.reward
    
    # Step 2: Check if terminal state
    if node.isterminal:
        terminalstate = node.state
        break
    
    # Step 3: Expand node horizontally (generate children)
    expand(node, transition, transitionargs;
           horizontalSample=horizontalSample,
           multithread=multithread)
    
    # Step 4: Select best child for next step
    node = selectChildNode(node)  # Uses progressvalue + reward

Return: (simTrajectoryReward, terminalstate)

Why this matters: The simulation phase estimates the value of a node by looking ahead maxSimulationDepth levels. Each rollout provides an estimate of the cumulative reward that can be obtained from the current state.

Detailed Backpropagation Process (mcts.jl → backpropagate)

After simulation, rewards are propagated back up the tree to update all visited nodes:

Input: Starting node (leaf from simulation), cumulative trajectory reward

while !isroot(node):
    # Update visit count
    node.visits += 1
    
    # Update statevalue (running average)
    # Formula: new_avg = (old_avg * (n-1) + new_value) / n
    node.statevalue = (node.statevalue * (node.visits-1) + simTrajectoryReward) / node.visits
    
    # Discount reward for future states
    # Future rewards are less certain, so they receive lower weight
    simTrajectoryReward *= discountRewardCoeff
    
    # Move to parent
    node = node.parent

Discounting rationale: Rewards further from the current state are discounted because:

Future rewards are uncertain
The longer the horizon, the more opportunities for suboptimal decisions
This creates a more realistic value estimate

Selection Phase: UCT (util.jl → UCTselect)

The UCT formula balances exploration vs. exploitation:

UCT(s,a) = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))

Where:
- Q(s,a) = childNode.statevalue  (exploitation: current estimate)
- c = explorationweight          (controls exploration vs exploitation)
- N(s) = parent.visits           (total visits to parent node)
- N(s,a) = childNode.visits      (visits to this specific child)

Behavior:

If childNode.visits = 0: Exploration term becomes undefined, so use progressvalue as fallback
If childNode.visits is low: Exploration term is high → encourages visiting unexplored branches
If childNode.visits is high: Exploration term approaches 0 → exploits known high-value nodes

Selection priority:

High progressvalue + low visits: Explored first (fast LLM guidance)
High statevalue + high visits: Exploited once confirmed (accurate value)
Balance: Controlled by explorationweight parameter

3. Main Workflow (interface.jl → runMCTS)

runMCTS(initialstate, transition, args)
├── Initialize root node with initialstate
│   └── root.visits = 0, root.statevalue = 0, root.children = {}
│
├── For iteration 1 to maxiterations:
│   │
│   ├── PHASE 1: SELECTION
│   │   ├── Start at root node
│   │   ├── While not leaf node:
│   │   │   ├── Apply UCTselect() with explorationweight
│   │   │   │   ├── UCT = statevalue + w * sqrt(ln(parent_visits) / child_visits)
│   │   │   │   └── Select child with highest UCT score
│   │   │   └── node = selected child
│   │   └── node is now a leaf node
│   │
│   ├── PHASE 2: TERMINAL CHECK
│   │   ├── If node.isterminal == true:
│   │   │   ├── If node.state[:reward] >= 8:
│   │   │   │   └── Store in highValueState channel
│   │   │   └── backpropagate(node, node.reward)
│   │   │       └── Update all ancestors up to root
│   │   └── Else (non-terminal leaf):
│   │       └── Continue to EXPANSION
│   │
│   ├── PHASE 3: EXPANSION
│   │   ├── expand(node, transition, transitionargs;
│   │   │         horizontalSample=horizontalSampleExpansionPhase)
│   │   │
│   │   ├── For each child generated:
│   │   │   ├── newNode = transition(current_state, args)
│   │   │   ├── newNodeKey = result[:newNodeKey]
│   │   │   ├── newstate = result[:newstate]
│   │   │   ├── progressvalue = result[:progressvalue]
│   │   │   └── Dejavu check: skip if newNodeKey already exists
│   │   │
│   │   └── Each child gets initialized with:
│   │       ├── visits = 0
│   │       ├── statevalue = 0 (no simulations yet)
│   │       ├── progressvalue = LLM estimate
│   │       ├── reward = newstate[:reward]
│   │       └── parent = current node
│   │
│   ├── PHASE 4: SIMULATION + BACKPROPAGATION
│   │   ├── If multithread == true:
│   │   │   └── Spawn parallel simulateThenBackpropagate() for each child
│   │   │
│   │   └── For each leafNode in node.children:
│   │       ├── simulateThenBackpropagate(leafNode, ...)
│   │       │   ├── simulate(leafNode, ...):
│   │       │   │   └── Rollout up to maxSimulationDepth
│   │       │   │       └── Returns (simTrajectoryReward, terminalstate)
│   │       │   │
│   │       │   ├── If terminalstate[:reward] >= 8:
│   │       │   │   └── Store in highValueState channel
│   │       │   │
│   │       │   └── backpropagate(leafNode, simTrajectoryReward):
│   │       │       └── Update visits and statevalue for all ancestors
│   │       │
│   │       └── If saveSimulatedNode == false:
│   │           └── Clear children (free memory for next iteration)
│   │
│   └── PHASE 5: EARLY STOP CHECK
│       ├── If earlystop(node.state) == true:
│       │   └── Break out of iteration loop
│       └── Continue to next iteration
│
└── After all iterations, select best result:
    ├── bestNextState = selectBestNextNode(root)
    │   └── Uses statevalue/visits or progressvalue + reward
    │
    ├── bestTerminalState = selectBestTrajectoryNode(root)
    │   └── Follows optimal trajectory to leaf
    │
    └── highValueStateList = collect from highValueState channel

7. Selection Phase Details (util.jl → UCTselect)

Purpose: Select the best child node using UCT formula that balances exploration and exploitation.

Process:

UCTselect(node, w)

Input: Parent node, exploration weight w
Output: Child node with highest UCT score

maxUCT = -Inf
selectedNode = nothing

for each childNode in node.children:
    ┌── Calculate UCT value
    │   └── if childNode.visits != 0:
    │           ├── weightedterm = w * sqrt(ln(node.visits) / childNode.visits)
    │           ├── UCTvalue = childNode.statevalue + weightedterm
    │           └── Exploration term encourages low-visited children
    │
    │       else:  # childNode.visits == 0
    │           ├── UCTvalue = childNode.progressvalue
    │           └── No exploration term (division by zero)
    │
    ├── Compare with current max
    │   └── if UCTvalue > maxUCT:
    │           ├── maxUCT = UCTvalue
    │           └── selectedNode = childNode

return selectedNode

Behavior analysis:

Scenario	UCT Value	Behavior
Child never visited (`visits=0`)	`progressvalue`	Strong exploration → will be tried
Child visited often, high reward	High `statevalue`, low exploration	Exploitation dominates
Child visited often, low reward	Low `statevalue`, low exploration	Exploitation avoids this node
Child visited rarely, moderate reward	Moderate `statevalue`, high exploration	May explore further

8. Expansion Phase Details (mcts.jl → expand/_expand)

Purpose: Generate new child nodes by applying the transition function multiple times.

Process:

expand(node, transition, transitionargs;
       horizontalSample=3, multithread=false)

Input: Node to expand, transition function, arguments
Output: node.children populated with child nodes

if multithread == true:
    └── @sync for i in 1 to horizontalSample:
        └── @spawn _expand(node, transition, transitionargs)
        └── Run expansion in parallel threads
else:
    └── for i in 1 to horizontalSample:
        └── _expand(node, transition, transitionargs)
        └── Sequential expansion

_expand(node, transition, transitionargs)

Input: Single node to expand
Output: One child node added to node.children

result = transition(node.state, transitionargs)
    └── LLM generates next state:
        └── response = llm_call(state[:thoughtHistory], args.prompt)
        └── Parse response into new state structure

newNodeKey = result[:newNodeKey]
newstate = result[:newstate]
progressvalue = result[:progressvalue]

if newNodeKey ∉ keys(node.children):
    ┌── Create new MCTSNode
    │   └── newNode = MCTSNode(
    │       ├── nodekey = newNodeKey
    │       ├── state = newstate
    │       ├── visits = 0
    │       ├── progressvalue = progressvalue
    │       ├── statevalue = 0 (no simulations yet)
    │       ├── reward = newstate[:reward]
    │       ├── isterminal = newstate[:isterminal]
    │       ├── parent = node
    │       ├── children = {}
    │       └── etc = {}
    │
    └── node.children[newNodeKey] = newNode

Dejavu detection: Before adding a child, check if newNodeKey already exists in node.children. If so, skip to avoid duplicate states.

Example: Expanding a node with horizontalSample=3

Current Node: "Math problem: Solve x^2 = 16"
    └── Attempt 1 (LLM):
        └── Thought: "Take square root of both sides"
        └── Action: "x = sqrt(16)"
        └── New state: "Solved: x = 4"
        └── newNodeKey = "abc-123"
        └── Added to children: "abc-123" → Node

    └── Attempt 2 (LLM):
        └── Thought: "Consider negative root"
        └── Action: "x = -sqrt(16)"
        └── New state: "Solved: x = -4"
        └── newNodeKey = "def-456"
        └── Added to children: "def-456" → Node

    └── Attempt 3 (LLM):
        └── Thought: "Check both solutions"
        └── Action: "Verify x=4 and x=-4"
        └── New state: "Verified: x=4, x=-4"
        └── newNodeKey = "ghi-789"
        └── Added to children: "ghi-789" → Node

Result: Node has 3 children with different solution approaches

9. Backpropagation Phase Details (mcts.jl → backpropagate)

Purpose: Update the statistics of all nodes along the simulation path with the observed reward.

Process:

backpropagate(node, simTrajectoryReward;
              discountRewardCoeff=0.9)

Input: Starting node (leaf from simulation), cumulative reward
Output: Updates visits and statevalue for all ancestors

while !isroot(node):
    ┌── Update visit count
    │   └── node.visits += 1
    │
    ├── Update statevalue (running average formula)
    │   │
    │   ├── Current average: node.statevalue
    │   ├── Current count: node.visits - 1
    │   ├── New value to add: simTrajectoryReward
    │   │
    │   └── Formula: new_avg = (old_avg * (n-1) + new_value) / n
    │       node.statevalue = (node.statevalue * (node.visits-1) + simTrajectoryReward) / node.visits
    │
    ├── Apply discount to future reward
    │   └── simTrajectoryReward *= discountRewardCoeff
    │       └── Reward gets 10% smaller for each level up
    │
    └── Move to parent
        └── node = node.parent

# After loop, root node still needs update
# (root node's parent is nothing, so loop stops but root is still updated)

Example: Backpropagating reward=20 with discount=0.9

Depth 3 (leaf): Node A2b2
    └── visits = 1, statevalue = 20/1 = 20.0
    └── reward for parent = 20 * 0.9 = 18.0

Depth 2: Node A2b
    └── visits = 1, statevalue = (0*0 + 18)/1 = 18.0
    └── reward for parent = 18 * 0.9 = 16.2

Depth 1: Node A2
    └── visits = 1, statevalue = (0*0 + 16.2)/1 = 16.2
    └── reward for parent = 16.2 * 0.9 = 14.58

Depth 0 (root): Node A
    └── visits = 1, statevalue = (0*0 + 14.58)/1 = 14.58

Why discounting matters: Future rewards are discounted because:

Uncertainty: The longer the horizon, the more uncertain the outcome
Temporal credit assignment: Immediate rewards should have higher value
Realistic evaluation: A reward 10 steps away is worth less than immediate reward

10. Solution Space Exploration Strategy

LLMMCTS explores the solution space using a combination of guided exploration and systematic exploitation:

Exploration vs. Exploitation Trade-off

Iteration 1: Initial exploration (all nodes have visits=0)
├── UCTselect() uses progressvalue for all children
├── LLM provides guidance: which states seem promising?
├── High progressvalue nodes get visited first
└── Results stored in statevalue

Iterations 2-5: Early exploitation
├── Some nodes have higher statevalue confirmed by simulations
├── UCT balances:
│   ├── Exploitation: High statevalue nodes
│   └── Exploration: Nodes with low visits (high UCT exploration term)
└── Nodes with high progressvalue + high statevalue dominate

Iterations 6+: Refinement
├── Well-visited nodes have reliable statevalue estimates
├── Poor nodes (low statevalue) are explored less
├── Still explore unvisited branches (UCT exploration term)
└── Tree converges to optimal trajectory

Horizontal vs. Vertical Exploration

Horizontal exploration (expansion):

Parameter: horizontalSample
Purpose: Generate multiple candidate next states from current state
Method: Apply transition function multiple times
LLM role: Each call generates a different thought/action

Vertical exploration (simulation):

Parameter: maxSimulationDepth
Purpose: Evaluate the long-term value of a state
Method: Rollout from current node to leaf
Reward accumulation: Sum all rewards along trajectory

Dejavu Detection and State Pruning

Dejavu detection (mcts.jl → _expand):

if newNodeKey ∉ keys(node.children)
    newNode = MCTSNode(...)
    node.children[newNodeKey] = newNode
end

Why it matters:

Prevents infinite loops in cyclic state spaces
Handles semantically equivalent states as duplicates
Reduces tree size and memory usage
Improves search efficiency

Example:

Current Node: "Problem solving step 5"
├── Attempt 1: Generates "Next step: Apply formula A"
├── Attempt 2: Generates "Next step: Apply formula A" (same as attempt 1)
│   └── Dejavu detected: skip duplicate
└── Attempt 3: Generates "Next step: Apply formula B"

11. Complete MCTS Iteration Flow

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 1                                                            │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ 1. SELECTION (UCT-based traversal):                                         │
│    Root (visits=1)                                                          │
│    └── UCTselect() on root (has no children yet, so uses progressvalue)     │
│    └── Select child with highest progressvalue                              │
│    └── Continue until reaching leaf node                                    │
│                                                                             │
│ 2. EXPANSION (horizontal sampling):                                         │
│    Leaf node                                                                │
│    └── expand(horizontalSample=3)                                           │
│    └── Generate 3 child nodes via LLM transition                            │
│    └── Each child gets: progressvalue, reward, parent=leaf                  │
│                                                                             │
│ 3. SIMULATION (vertical rollout):                                           │
│    Each child node                                                          │
│    └── simulate(maxSimulationDepth=3)                                       │
│    └── Rollout 3 levels deep, accumulating rewards                          │
│    └── Return (simTrajectoryReward, terminalstate)                          │
│                                                                             │
│ 4. BACKPROPAGATION (update statistics):                                     │
│    simTrajectoryReward                                                      │
│    └── backpropagate() up to root                                           │
│    └── Update visits and statevalue for all ancestors                       │
│    └── Apply discount to future rewards                                     │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 2                                                            │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ 1. SELECTION:                                                               │
│    Root (visits=2)                                                          │
│    └── UCTselect() now considers:                                           │
│        ├── statevalue (from iteration 1)                                    │
│        └── exploration term (low visits on unexpanded branches)             │
│    └── May select different path than iteration 1                           │
│                                                                             │
│ 2. EXPANSION:                                                               │
│    └── May expand different node or same node                               │
│    └── New children added to node.children                                  │
│                                                                             │
│ 3. SIMULATION:                                                              │
│    └── Different rollout trajectory                                         │
│    └── New reward estimate added to statistics                              │
│                                                                             │
│ 4. BACKPROPAGATION:                                                         │
│    └── Statistics updated with new information                              │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 3 to N                                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ Pattern repeats, with increasingly informed selection:                      │
│                                                                             │
│ • Nodes with high statevalue (confirmed by many simulations)                │
│   → Exploited (selected frequently)                                         │
│                                                                             │
│ • Nodes with low visits but promising progressvalue                         │
│   → Explored (UCT exploration term encourages tries)                        │
│                                                                             │
│ • Tree grows: more branches explored, more statistics accumulated           │
│                                                                             │
│ • Best trajectory emerges from accumulated statistics                       │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

simulate(node, transition, transitionargs;
         maxSimulationDepth=3, horizontalSample=3, multithread=false)

Input: Starting node at depth 0
Output: (simTrajectoryReward, terminalstate)

simTrajectoryReward = 0.0
terminalstate = nothing

for depth in 1 to maxSimulationDepth:
    ┌── Step 1: Accumulate reward
    │   └── simTrajectoryReward += node.reward
    │
    ├── Step 2: Check for terminal state
    │   └── if node.isterminal:
    │           └── terminalstate = deepcopy(node.state)
    │               break
    │
    ├── Step 3: Expand node (horizontal sampling)
    │   └── expand(node, transition, transitionargs;
    │               horizontalSample=horizontalSample)
    │       └── For i in 1 to horizontalSample:
    │           └── _expand(node, transition, transitionargs)
    │               ├── result = transition(node.state, args)
    │               ├── newNodeKey = result[:newNodeKey]
    │               ├── newstate = result[:newstate]
    │               ├── progressvalue = result[:progressvalue]
    │               └── if newNodeKey not in node.children:
    │                   └── Create new MCTSNode with:
    │                       ├── state = newstate
    │                       ├── progressvalue = progressvalue
    │                       ├── reward = newstate[:reward]
    │                       ├── parent = node
    │                       └── children = {}
    │
    └── Step 4: Select best child for next iteration
        └── node = selectChildNode(node)
            └── Uses: progressvalue + reward (no UCT here)

return (simTrajectoryReward, terminalstate)

Example: With maxSimulationDepth=3 and horizontalSample=3:

Depth 0: Node A (reward=2)
    └── Expand 3 children: A1, A2, A3
    └── Select A2 (highest progressvalue + reward)
    
Depth 1: Node A2 (reward=3)
    └── Expand 3 children: A2a, A2b, A2c
    └── Select A2b (highest progressvalue + reward)
    
Depth 2: Node A2b (reward=5)
    └── Expand 3 children: A2b1, A2b2, A2b3
    └── Select A2b2 (highest progressvalue + reward)
    
Depth 3: Node A2b2 (reward=10, isterminal=true)
    └── Terminal state reached!
    
Return: (simTrajectoryReward=20, terminalstate=A2b2.state)
        = (2 + 3 + 5 + 10, ...)

6. Backpropagation (mcts.jl → backpropagate)

Purpose: Update the statistics of all nodes along the simulation path with the observed reward.

Process:

backpropagate(node, simTrajectoryReward;
              discountRewardCoeff=0.9)

Input: Starting node (leaf from simulation), cumulative reward
Output: Updates visits and statevalue for all ancestors

while !isroot(node):
    ┌── Update visit count
    │   └── node.visits += 1
    │
    ├── Update statevalue (running average formula)
    │   │
    │   ├── Current average: node.statevalue
    │   ├── Current count: node.visits - 1
    │   ├── New value to add: simTrajectoryReward
    │   │
    │   └── Formula: new_avg = (old_avg * (n-1) + new_value) / n
    │       node.statevalue = (node.statevalue * (node.visits-1) + simTrajectoryReward) / node.visits
    │
    ├── Apply discount to future reward
    │   └── simTrajectoryReward *= discountRewardCoeff
    │       └── Reward gets 10% smaller for each level up
    │
    └── Move to parent
        └── node = node.parent

# After loop, root node still needs update
# (root node's parent is nothing, so loop stops but root is still updated)

Example: Backpropagating reward=20 with discount=0.9

Depth 3 (leaf): Node A2b2
    └── visits = 1, statevalue = 20/1 = 20.0
    └── reward for parent = 20 * 0.9 = 18.0

Depth 2: Node A2b
    └── visits = 1, statevalue = (0*0 + 18)/1 = 18.0
    └── reward for parent = 18 * 0.9 = 16.2

Depth 1: Node A2
    └── visits = 1, statevalue = (0*0 + 16.2)/1 = 16.2
    └── reward for parent = 16.2 * 0.9 = 14.58

Depth 0 (root): Node A
    └── visits = 1, statevalue = (0*0 + 14.58)/1 = 14.58

Why discounting matters: Future rewards are discounted because:

Uncertainty: The longer the horizon, the more uncertain the outcome
Temporal credit assignment: Immediate rewards should have higher value
Realistic evaluation: A reward 10 steps away is worth less than immediate reward

7. Solution Space Exploration Strategy

LLMMCTS explores the solution space using a combination of guided exploration and systematic exploitation:

Exploration vs. Exploitation Trade-off

Iteration 1: Initial exploration (all nodes have visits=0)
├── UCTselect() uses progressvalue for all children
├── LLM provides guidance: which states seem promising?
├── High progressvalue nodes get visited first
└── Results stored in statevalue

Iterations 2-5: Early exploitation
├── Some nodes have higher statevalue confirmed by simulations
├── UCT balances:
│   ├── Exploitation: High statevalue nodes
│   └── Exploration: Nodes with low visits (high UCT exploration term)
└── Nodes with high progressvalue + high statevalue dominate

Iterations 6+: Refinement
├── Well-visited nodes have reliable statevalue estimates
├── Poor nodes (low statevalue) are explored less
├── Still explore unvisited branches (UCT exploration term)
└── Tree converges to optimal trajectory

Horizontal vs. Vertical Exploration

Horizontal exploration (expansion):

Parameter: horizontalSample
Purpose: Generate multiple candidate next states from current state
Method: Apply transition function multiple times
LLM role: Each call generates a different thought/action

Vertical exploration (simulation):

Parameter: maxSimulationDepth
Purpose: Evaluate the long-term value of a state
Method: Rollout from current node to leaf
Reward accumulation: Sum all rewards along trajectory

Dejavu Detection and State Pruning

Dejavu detection (mcts.jl → _expand):

if newNodeKey ∉ keys(node.children)
    newNode = MCTSNode(...)
    node.children[newNodeKey] = newNode
end

Why it matters:

Prevents infinite loops in cyclic state spaces
Handles semantically equivalent states as duplicates
Reduces tree size and memory usage
Improves search efficiency

Example:

Current Node: "Problem solving step 5"
├── Attempt 1: Generates "Next step: Apply formula A"
├── Attempt 2: Generates "Next step: Apply formula A" (same as attempt 1)
│   └── Dejavu detected: skip duplicate
└── Attempt 3: Generates "Next step: Apply formula B"

8. Selection Phase Details (util.jl → UCTselect)

9. High-Level Interface (interface.jl → runMCTS)

Purpose: Execute the complete MCTS search algorithm with all phases.

Process:

runMCTS(initialstate, transition, transitionargs;
        horizontalSampleExpansionPhase=3,
        horizontalSampleSimulationPhase=3,
        maxSimulationDepth=3,
        maxiterations=10,
        explorationweight=1.0,
        earlystop=nothing,
        saveSimulatedNode=false,
        multithread=false)

Output: NamedTuple with:
    ├── root: Complete MCTS tree
    ├── bestNextState: Best immediate next state
    ├── bestTerminalState: Best final state along optimal trajectory
    └── highValueStateList: List of high-value terminal states (reward >= 8)

Step-by-step execution:
    └── root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, {}, {})
    └── highValueState = Channel{Any}(100)

    for iteration = 1 to maxiterations:
        ┌── node = root
        │   └── node.visits += 1  # Increment root visits
        │
        │   ┌── PHASE 1: SELECTION (until leaf)
        │   │   while !isleaf(node):
        │   │       └── node = UCTselect(node, explorationweight)
        │   │
        │   ├── PHASE 2: TERMINAL CHECK
        │   │   if node.isterminal:
        │   │       ├── if node.state[:reward] >= 8:
        │   │       │   └── put!(highValueState, deepcopy(node.state))
        │   │       └── backpropagate(node, node.reward)
        │   │           └── Update statistics up to root
        │   │       └── continue to next iteration
        │   │
        │   └── PHASE 3: EXPANSION
        │       ├── expand(node, transition, transitionargs;
        │       │         horizontalSample=horizontalSampleExpansionPhase,
        │       │         multithread=multithread)
        │       │
        │       └── For each leafNode in node.children:
        │           └── simulateThenBackpropagate(leafNode, ...)
        │               ├── simulate() → rollout and get reward
        │               ├── If terminal state with reward >= 8:
        │               │   └── Store in highValueState
        │               └── backpropagate() → update statistics
        │
        └── PHASE 4: EARLY STOP CHECK
            └── if earlystop !== nothing && earlystop(node.state):
                └── break  # Exit iteration loop

    ┌── After all iterations, select best result:
    │   ├── bestNextState = selectBestNextNode(root)
    │   │   └── Uses statevalue/visits or progressvalue + reward
    │   │
    │   ├── bestTerminalState = selectBestTrajectoryNode(root)
    │   │   └── Follows optimal trajectory to leaf
    │   │
    │   └── highValueStateList = collect from channel
    │       └── while !isempty(highValueState):
    │               push!(highValueStateList, take!(highValueState))
    │
    └── return (root=root, bestNextState=..., bestTerminalState=..., highValueStateList=...)

10. Complete MCTS Iteration Flow

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 1                                                            │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ 1. SELECTION (UCT-based traversal):                                         │
│    Root (visits=1)                                                          │
│    └── UCTselect() on root (has no children yet, so uses progressvalue)     │
│    └── Select child with highest progressvalue                              │
│    └── Continue until reaching leaf node                                    │
│                                                                             │
│ 2. EXPANSION (horizontal sampling):                                         │
│    Leaf node                                                                │
│    └── expand(horizontalSample=3)                                           │
│    └── Generate 3 child nodes via LLM transition                            │
│    └── Each child gets: progressvalue, reward, parent=leaf                  │
│                                                                             │
│ 3. SIMULATION (vertical rollout):                                           │
│    Each child node                                                          │
│    └── simulate(maxSimulationDepth=3)                                       │
│    └── Rollout 3 levels deep, accumulating rewards                          │
│    └── Return (simTrajectoryReward, terminalstate)                          │
│                                                                             │
│ 4. BACKPROPAGATION (update statistics):                                     │
│    simTrajectoryReward                                                      │
│    └── backpropagate() up to root                                           │
│    └── Update visits and statevalue for all ancestors                       │
│    └── Apply discount to future rewards                                     │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 2                                                            │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ 1. SELECTION:                                                               │
│    Root (visits=2)                                                          │
│    └── UCTselect() now considers:                                           │
│        ├── statevalue (from iteration 1)                                    │
│        └── exploration term (low visits on unexpanded branches)             │
│    └── May select different path than iteration 1                           │
│                                                                             │
│ 2. EXPANSION:                                                               │
│    └── May expand different node or same node                               │
│    └── New children added to node.children                                  │
│                                                                             │
│ 3. SIMULATION:                                                              │
│    └── Different rollout trajectory                                         │
│    └── New reward estimate added to statistics                              │
│                                                                             │
│ 4. BACKPROPAGATION:                                                         │
│    └── Statistics updated with new information                              │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ MCTS ITERATION 3 to N                                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│ Pattern repeats, with increasingly informed selection:                      │
│                                                                             │
│ • Nodes with high statevalue (confirmed by many simulations)                │
│   → Exploited (selected frequently)                                         │
│                                                                             │
│ • Nodes with low visits but promising progressvalue                         │
│   → Explored (UCT exploration term encourages tries)                        │
│                                                                             │
│ • Tree grows: more branches explored, more statistics accumulated           │
│                                                                             │
│ • Best trajectory emerges from accumulated statistics                       │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

Key Functions

Core MCTS (mcts.jl)

Function	Purpose
`selectBestNextNode(node)`	Select best child based on `statevalue/visits` or `progressvalue + reward`
`selectBestTrajectoryNode(node)`	Traverse down tree to find highest-value leaf
`selectChildNode(node)`	Select child with highest `progressvalue + reward`
`expand(node, ...)`	Generate child nodes using transition function
`_expand(node, ...)`	Helper to create single child node
`simulate(node, ...)`	Perform rollout simulation
`backpropagate(node, reward)`	Update statistics up to root
`isleaf(node)`	Check if node has no children
`isroot(node)`	Check if node is root (`nodekey == "root"`)

Interface (interface.jl)

Function	Purpose
`runMCTS(initialstate, transition, args; kwargs...)`	Main MCTS search function that executes all four phases across multiple iterations
`simulateThenBackpropagate(node, ...)`	Run simulation and backpropagate reward for a single node

Utilities (util.jl)

Function	Purpose
`UCTselect(node, w)`	Select node using UCT score with exploration weight `w`
`selectBestNextNode(node)`	Select best child based on `statevalue/visits` or `progressvalue + reward`
`selectBestTrajectoryNode(node)`	Traverse down tree to find highest-value leaf
`selectChildNode(node)`	Select child with highest `progressvalue + reward`
`expand(node, ...)`	Generate child nodes using transition function
`_expand(node, ...)`	Helper to create single child node
`simulate(node, ...)`	Perform rollout simulation
`backpropagate(node, reward)`	Update statistics up to root
`isleaf(node)`	Check if node has no children
`isroot(node)`	Check if node is root (`nodekey == "root"`)

Configuration Parameters

Keyword Arguments for runMCTS()

Parameter	Default	Description
`horizontalSampleExpansionPhase`	3	Number of child nodes generated at each expansion (horizontal sampling)
`horizontalSampleSimulationPhase`	3	Number of child nodes generated during simulation rollout
`maxSimulationDepth`	3	Maximum depth of simulation rollout (vertical exploration)
`maxiterations`	10	Number of MCTS iterations (each iteration: selection → expansion → simulation → backpropagation)
`explorationweight`	1.0	UCT exploration weight (controls exploration vs exploitation balance)
`earlystop`	nothing	Optional function to check early stopping condition (takes node.state, returns bool)
`saveSimulatedNode`	false	Whether to keep nodes created during simulation (true = keep for analysis, false = free memory)
`multithread`	false	Enable parallel simulation across child nodes (true = use Julia threads)

How Parameters Affect Search

Parameter	Low Value	High Value	Trade-off
`horizontalSample`	Fast, less exploration	Slow, more exploration	More children = better coverage but more LLM calls
`maxSimulationDepth`	Quick estimates, short horizon	Slow, long horizon	Deeper rollouts = more accurate but expensive
`maxiterations`	Quick search	Thorough search	More iterations = better solution but slower
`explorationweight`	Greedy (exploitation)	Aggressive (exploration)	Higher = explores more novel paths

Usage Pattern

# 1. Define transition function
function 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

# 2. Define transition arguments
transitionargs = (prompt="Solve this math problem", other_param="value")

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

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

Key Differences from Traditional MCTS

Aspect	Traditional MCTS	LLMMCTS
Value estimation	Hand-designed heuristics	LLM-provided `progressvalue`
Reward signal	Environment only (sparse)	Environment + LLM pseudo-rewards (dense)
Expansion	Random or heuristics	LLM-guided generation
Learning speed	Slow (needs terminal rewards)	Fast (dense intermediate signals)

Performance Characteristics

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

Known Issues

Bug: Variable name mismatch in interface.jl:94

The code references highrewardNode but the channel is named highValueState. This causes a UndefVarError when a terminal state with reward >= 8 is encountered.

Current code (buggy):

highValueState = Channel{Any}(100)
# ...
if node.state[:reward] >= 8
  put!(highrewardNode, deepcopy(node.state))  # BUG: should be highValueState
end

Expected behavior: Store high-value terminal states in the highValueState channel.

src/type.jl: MCTSNode struct definition
src/util.jl: UCT selection utility
src/mcts.jl: Core MCTS operations
src/interface.jl: High-level interface
src/LLMMCTS.jl: Main package entry point
README.md: User-facing documentation

56 KiB Raw Permalink Blame History Unescape Escape

Workprocess Documentation for LLMMCTS

Overview

Core Concept

Architecture

Three-Tier Value System

How reward and progressvalue Work Together

reward - Environment Feedback

progressvalue - LLM Heuristic

The Selection Hierarchy

Data Flow Example

Why Two Value Systems?

Module Structure

Data Flow

1. Node Structure (type.jl)

2. MCTS Search Process: Rollout and Exploration

The Four Phases of MCTS

Detailed Expansion Process (mcts.jl → expand/_expand)

Detailed Simulation Process (mcts.jl → simulate)

Detailed Backpropagation Process (mcts.jl → backpropagate)

Selection Phase: UCT (util.jl → UCTselect)

3. Main Workflow (interface.jl → runMCTS)

7. Selection Phase Details (util.jl → UCTselect)

8. Expansion Phase Details (mcts.jl → expand/_expand)

9. Backpropagation Phase Details (mcts.jl → backpropagate)

10. Solution Space Exploration Strategy

Exploration vs. Exploitation Trade-off

Horizontal vs. Vertical Exploration

Dejavu Detection and State Pruning

11. Complete MCTS Iteration Flow

6. Backpropagation (mcts.jl → backpropagate)

7. Solution Space Exploration Strategy

Exploration vs. Exploitation Trade-off

Horizontal vs. Vertical Exploration

Dejavu Detection and State Pruning

8. Selection Phase Details (util.jl → UCTselect)

9. High-Level Interface (interface.jl → runMCTS)

10. Complete MCTS Iteration Flow

Key Functions

Core MCTS (mcts.jl)

Interface (interface.jl)

Utilities (util.jl)

Configuration Parameters

Keyword Arguments for runMCTS()

How Parameters Affect Search

Usage Pattern

Key Differences from Traditional MCTS

Performance Characteristics

Limitations

Known Issues

Bug: Variable name mismatch in interface.jl:94

Related Files

56 KiB

Raw Permalink Blame History

`reward` - Environment Feedback

`progressvalue` - LLM Heuristic