diff --git a/Manifest.toml b/Manifest.toml index 852f9d9..f7c7bc8 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,31 @@ julia_version = "1.12.6" manifest_format = "2.0" -project_hash = "d65e3fb9ee71e7ba2bf20b111f9f950fd4d5183a" +project_hash = "eefceed38d1e3b7216d3f052806b5d83d4fd80b7" + +[[deps.Accessors]] +deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "MacroTools"] +git-tree-sha1 = "7063ad1083578215c7c4bf410368150abe8d5524" +uuid = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +version = "0.1.45" + + [deps.Accessors.extensions] + AxisKeysExt = "AxisKeys" + IntervalSetsExt = "IntervalSets" + LinearAlgebraExt = "LinearAlgebra" + StaticArraysExt = "StaticArrays" + StructArraysExt = "StructArrays" + TestExt = "Test" + UnitfulExt = "Unitful" + + [deps.Accessors.weakdeps] + AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5" + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" + Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [[deps.AliasTables]] deps = ["PtrArrays", "Random"] @@ -51,6 +75,11 @@ git-tree-sha1 = "962834c22b66e32aa10f7611c08c8ca4e20749a9" uuid = "944b1d66-785c-5afd-91f1-9de20f533193" version = "0.7.8" +[[deps.CommonSolve]] +git-tree-sha1 = "99ee296f88c12485402e37c2fd025f95ae097637" +uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" +version = "0.2.9" + [[deps.Compat]] deps = ["TOML", "UUIDs"] git-tree-sha1 = "9d8a54ce4b17aa5bdce0ea5c34bc5e7c340d16ad" @@ -71,6 +100,30 @@ deps = ["Artifacts", "Libdl"] uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" version = "1.3.0+1" +[[deps.CompositionsBase]] +git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad" +uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b" +version = "0.1.2" +weakdeps = ["InverseFunctions"] + + [deps.CompositionsBase.extensions] + CompositionsBaseInverseFunctionsExt = "InverseFunctions" + +[[deps.ConstructionBase]] +git-tree-sha1 = "b4b092499347b18a015186eae3042f72267106cb" +uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" +version = "1.6.0" + + [deps.ConstructionBase.extensions] + ConstructionBaseIntervalSetsExt = "IntervalSets" + ConstructionBaseLinearAlgebraExt = "LinearAlgebra" + ConstructionBaseStaticArraysExt = "StaticArrays" + + [deps.ConstructionBase.weakdeps] + IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" + LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + [[deps.Crayons]] git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15" uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f" @@ -104,10 +157,10 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" version = "1.11.0" [[deps.Distributions]] -deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] -git-tree-sha1 = "3c8a0a9a6d4a10bdfb6b751bd2b6051ed3e25fd4" +deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "Roots", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"] +git-tree-sha1 = "cd3c5ac74cd3923c8945c6a81518c46abd0e73a3" uuid = "31c24e10-a181-5473-b8eb-7969acd0382f" -version = "0.25.127" +version = "0.25.129" [deps.Distributions.extensions] DistributionsChainRulesCoreExt = "ChainRulesCore" @@ -126,6 +179,11 @@ git-tree-sha1 = "7442a5dfe1ebb773c29cc2962a8980f47221d76c" uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" version = "0.9.5" +[[deps.EnumX]] +git-tree-sha1 = "c49898e8438c828577f04b92fc9368c388ac783c" +uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56" +version = "1.0.7" + [[deps.FilePathsBase]] deps = ["Compat", "Dates"] git-tree-sha1 = "3bab2c5aa25e7840a4b065805c0cdfc01f3068d2" @@ -167,13 +225,24 @@ deps = ["Random"] uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" version = "1.11.0" +[[deps.Gamma]] +git-tree-sha1 = "86f86b6168a016ed88e4ae4e64577b98c3b59e8e" +uuid = "a0844989-3bd2-4988-8bea-c9407ab0941b" +version = "1.1.0" + [[deps.GeneralUtils]] -deps = ["CSV", "DataFrames", "DataStructures", "Dates", "Distributions", "JSON", "NATS", "PrettyPrinting", "Random", "Revise", "SHA", "UUIDs"] -git-tree-sha1 = "8720a31344bc85ad610ae12f7e1247de22070765" +deps = ["CSV", "DataFrames", "DataStructures", "Dates", "Distributions", "HTTP", "JSON", "NATS", "PrettyPrinting", "Random", "Revise", "SHA", "UUIDs"] +git-tree-sha1 = "7c0600c166a5deb2c607018a491c04eb25969c2e" repo-rev = "main" repo-url = "https://git.yiem.cc/ton/GeneralUtils" uuid = "c6c72f09-b708-4ac8-ac7c-2084d70108fe" -version = "0.3.2" +version = "0.4.9" + +[[deps.HTTP]] +deps = ["Base64", "CodecZlib", "Dates", "EnumX", "PrecompileTools", "Random", "Reseau", "SHA", "URIs", "UUIDs", "Zlib_jll"] +git-tree-sha1 = "eda1d37cb55d90a17d0957c75841138c88b361a1" +uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3" +version = "2.5.4" [[deps.HashArrayMappedTries]] git-tree-sha1 = "2eaa69a7cab70a52b9687c8bf950a5a93ec895ae" @@ -181,10 +250,10 @@ uuid = "076d061b-32b6-4027-95e0-9a2c6f6d7e74" version = "0.2.0" [[deps.HypergeometricFunctions]] -deps = ["LinearAlgebra", "OpenLibm_jll", "SpecialFunctions"] -git-tree-sha1 = "68c173f4f449de5b438ee67ed0c9c748dc31a2ec" +deps = ["Gamma", "LinearAlgebra"] +git-tree-sha1 = "18d7deab5fb0440dc6a7b6993c5c27b25420de10" uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a" -version = "0.3.28" +version = "0.3.29" [[deps.InlineStrings]] git-tree-sha1 = "8f3d257792a522b4601c24a577954b0a8cd7334d" @@ -204,6 +273,19 @@ deps = ["Markdown"] uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" version = "1.11.0" +[[deps.InverseFunctions]] +git-tree-sha1 = "a779299d77cd080bf77b97535acecd73e1c5e5cb" +uuid = "3587e190-3f89-42d0-90ee-14403ec27112" +version = "0.1.17" + + [deps.InverseFunctions.extensions] + InverseFunctionsDatesExt = "Dates" + InverseFunctionsTestExt = "Test" + + [deps.InverseFunctions.weakdeps] + Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" + Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + [[deps.InvertedIndices]] git-tree-sha1 = "6da3c4316095de0f5ee2ebd875df8721e7e0bdbe" uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" @@ -261,10 +343,10 @@ uuid = "ac6e5ff7-fb65-4e79-a425-ec3bc9c03011" version = "1.12.0" [[deps.LLMMCTS]] -deps = ["GeneralUtils", "JSON", "PrettyPrinting"] +deps = ["JSON", "PrettyPrinting"] path = "." uuid = "d76c5a4d-449e-4835-8cc4-dd86ec44f241" -version = "0.1.3" +version = "0.1.4" [[deps.LaTeXStrings]] git-tree-sha1 = "dda21b8cbd6a6c40d9d02a73230f9d70fed6918c" @@ -317,9 +399,14 @@ version = "1.11.0" [[deps.LoweredCodeUtils]] deps = ["CodeTracking", "Compiler", "JuliaInterpreter"] -git-tree-sha1 = "0aad96d7b987a5600e260eec50147b254d5ff7e6" +git-tree-sha1 = "3733419e9a71156b389f3e331672d2e95436783f" uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b" -version = "3.6.0" +version = "3.6.2" + +[[deps.MacroTools]] +git-tree-sha1 = "1e0228a030642014fe5cfe68c2c0a818f9e3f522" +uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" +version = "0.5.16" [[deps.Markdown]] deps = ["Base64", "JuliaSyntaxHighlighting", "StyledStrings"] @@ -396,9 +483,9 @@ version = "1.8.2" [[deps.PDMats]] deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "e4cff168707d441cd6bf3ff7e4832bdf34278e4a" +git-tree-sha1 = "26766d4b5f1a410c218a19b85a672c6edb693c65" uuid = "90014a1f-27ba-587c-ab20-58faa44d9150" -version = "0.11.37" +version = "0.11.40" weakdeps = ["StatsBase"] [deps.PDMats.extensions] @@ -482,6 +569,12 @@ git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" uuid = "189a3867-3050-52da-a836-e630ba90ab69" version = "1.2.2" +[[deps.Reseau]] +deps = ["NetworkOptions", "OpenSSL_jll", "PrecompileTools", "Random", "SHA"] +git-tree-sha1 = "0eab6d95ed40c2ef3992255c1c71e4f9748932b5" +uuid = "802f3686-a58f-41ce-bb0c-3c43c75bba36" +version = "1.3.1" + [[deps.Revise]] deps = ["CRC32c", "CodeTracking", "FileWatching", "InteractiveUtils", "JuliaInterpreter", "LibGit2", "LoweredCodeUtils", "OrderedCollections", "Preferences", "REPL", "UUIDs"] git-tree-sha1 = "27e3ee13fc8739a59b380d6163d6a82f52c03bd7" @@ -506,6 +599,28 @@ git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8" uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f" version = "0.5.1+0" +[[deps.Roots]] +deps = ["Accessors", "CommonSolve", "Printf"] +git-tree-sha1 = "ed45bcc7cf3c8887595b973f2b1efbe91dcc50ec" +uuid = "f2b01f46-fcfa-551c-844a-d8ac1e96c665" +version = "3.0.1" + + [deps.Roots.extensions] + RootsChainRulesCoreExt = "ChainRulesCore" + RootsForwardDiffExt = "ForwardDiff" + RootsIntervalRootFindingExt = "IntervalRootFinding" + RootsSymPyExt = "SymPy" + RootsSymPyPythonCallExt = "SymPyPythonCall" + RootsUnitfulExt = "Unitful" + + [deps.Roots.weakdeps] + ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" + IntervalRootFinding = "d2bf35a9-74e0-55ec-b149-d360ff49b807" + SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6" + SymPyPythonCall = "bc8888f7-b21e-4b7c-a06a-5d9c9496438c" + Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" + [[deps.SHA]] uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" version = "0.7.0" @@ -649,9 +764,9 @@ version = "1.0.1" [[deps.Tables]] deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"] -git-tree-sha1 = "f2c1efbc8f3a609aadf318094f8fc5204bdaf344" +git-tree-sha1 = "0f38a06c83f0007bbab3cf911262841c9a0f07e0" uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -version = "1.12.1" +version = "1.13.0" [[deps.TranscodingStreams]] git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742" diff --git a/Project.toml b/Project.toml index 80f97b3..34cc0e1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "LLMMCTS" uuid = "d76c5a4d-449e-4835-8cc4-dd86ec44f241" -version = "0.1.3" +version = "0.1.4" authors = ["narawat lamaiin "] [deps] @@ -9,6 +9,6 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337" [compat] -GeneralUtils = "0.3.2" +GeneralUtils = "0.4.9" JSON = "1.6.1" PrettyPrinting = "0.4.2" diff --git a/README.md b/README.md new file mode 100644 index 0000000..ea7a77d --- /dev/null +++ b/README.md @@ -0,0 +1,329 @@ +# 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{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 diff --git a/docs/workprocess.md b/docs/workprocess.md new file mode 100644 index 0000000..15770a6 --- /dev/null +++ b/docs/workprocess.md @@ -0,0 +1,1177 @@ +# 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) + +```julia +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: +1. Future rewards are uncertain +2. The longer the horizon, the more opportunities for suboptimal decisions +3. 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**: +1. **High `progressvalue` + low visits**: Explored first (fast LLM guidance) +2. **High `statevalue` + high visits**: Exploited once confirmed (accurate value) +3. **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: +1. **Uncertainty**: The longer the horizon, the more uncertain the outcome +2. **Temporal credit assignment**: Immediate rewards should have higher value +3. **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): + +```julia +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: +1. **Uncertainty**: The longer the horizon, the more uncertain the outcome +2. **Temporal credit assignment**: Immediate rewards should have higher value +3. **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): + +```julia +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 + +```julia +# 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)**: +```julia +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. + +## Related Files + +- `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 diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..a0a5b8e --- /dev/null +++ b/examples/README.md @@ -0,0 +1,57 @@ +# LLMMCTS Examples + +This directory contains example scripts demonstrating how to use LLMMCTS for various problem types. + +## Examples + +1. **simple_example.jl** - Basic MCTS usage with a simple state transition function +2. **pathfinding.jl** - Grid-based pathfinding problem +3. **math_problem.jl** - Solving math problems using MCTS-guided reasoning +4. **tool_use.jl** - Coordinating with external tools (APIs, databases) +5. **chess_game.jl** - Game playing scenario (simplified chess-like) +6. **code_generation.jl** - Guiding LLM code generation +7. **reasoning.jl** - Multi-step reasoning with chain-of-thought +8. **configuration_examples.jl** - Demonstrating different MCTS configuration options + +## Running Examples + +```bash +julia examples/simple_example.jl +julia examples/pathfinding.jl +julia examples/configuration_examples.jl +``` + +## Key Concepts + +### State +The state is represented as a `Dict{String, Any}` that contains all information needed for the problem. + +### Transition Function +The transition function takes the current state and returns: +```julia +Dict( + :newNodeKey => unique_id, + :newstate => new_state_dict, + :progressvalue => llm_estimate +) +``` + +### Progress Value +`progressvalue` is provided by LLM reasoning and guides the search without waiting for terminal rewards. + +### State Value +`statevalue` is computed through Monte Carlo simulations and provides accurate long-term estimates. + +## Configuration Parameters + +- `maxiterations` - Number of MCTS iterations (default: 10) +- `explorationweight` - UCT exploration weight (default: 1.0) +- `maxSimulationDepth` - Maximum simulation rollout depth (default: 3) +- `horizontalSampleExpansionPhase` - Children per expansion (default: 3) +- `multithread` - Enable parallel simulation (default: false) +- `saveSimulatedNode` - Keep simulation nodes (default: false) + +## See Also + +- [README.md](../README.md) - Complete package documentation +- [workprocess.md](../workprocess.md) - Detailed technical documentation diff --git a/examples/chess_game.jl b/examples/chess_game.jl new file mode 100644 index 0000000..0ca80d6 --- /dev/null +++ b/examples/chess_game.jl @@ -0,0 +1,191 @@ +# Chess-like Game Example - MCTS for Game Playing + +This example demonstrates MCTS for a simplified chess-like game where the goal is to capture the opponent's pieces. + +```julia +using LLMMCTS + +# Simple game state +# board: Dict mapping positions to pieces +# turn: :white or :black +struct GameState + board::Dict{String, String} # position => piece + turn::Symbol + piece_count::Int +end + +# Initialize a simple board +function init_board() + board = Dict{String, String}() + + # Place some pieces + board["e1"] = "K" # White King + board["e8"] = "k" # Black King + + # Random pieces + board["d4"] = "P" # White Pawn + board["d5"] = "p" # Black Pawn + + return board +end + +# Check if position is on board +function on_board(pos::String) + cols = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'] + rows = ['1', '2', '3', '4', '5', '6', '7', '8'] + length(pos) == 2 && + pos[1] in cols && + pos[2] in rows +end + +# Game transition function +function chess_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + board = state[:board] + turn = state[:turn] + + if current_step >= args.max_moves + # Max moves reached, end game + newstate = Dict( + :step => current_step + 1, + :board => board, + :turn => turn, + :reward => 0.0, + :isterminal => true + ) + return Dict( + :newNodeKey => "max_moves", + :newstate => newstate, + :progressvalue => 5.0 + ) + end + + # Generate possible moves + possible_moves = String[] + + # Find all pieces of current turn's color + turn_prefix = turn == :white ? "upper" : "lower" + + # Simple move generation: try moving each piece + for (pos, piece) in board + if !isempty(piece) + # Try moving to adjacent positions + for dx in [-1, 0, 1] + for dy in [-1, 0, 1] + if dx == 0 && dy == 0 + continue + end + + # Simple coordinate conversion + col = pos[1] + row = parse(Int, pos[2]) + + new_col = col + dx + new_row = row + dy + + if new_col >= 'a' && new_col <= 'h' && + new_row >= 1 && new_row <= 8 + new_pos = string(new_col, new_row) + if on_board(new_pos) + push!(possible_moves, pos * new_pos) + end + end + end + end + end + end + + if isempty(possible_moves) + # No moves available, game over + newstate = Dict( + :step => current_step + 1, + :board => board, + :turn => turn, + :reward => turn == :white ? 10.0 : -10.0, + :isterminal => true + ) + return Dict( + :newNodeKey => "game_over", + :newstate => newstate, + :progressvalue => turn == :white ? 10.0 : 0.0 + ) + end + + # LLM would select the best move + # For this example, pick a random valid move + move_idx = (current_step - 1) % length(possible_moves) + 1 + move = possible_moves[move_idx] + + # Simulate the move (simplified) + from_pos = move[1:2] + to_pos = move[3:4] + + new_board = copy(board) + piece = get(new_board, from_pos, "") + new_board[to_pos] = piece + delete!(new_board, from_pos) + + # Calculate reward based on capture + reward = 0.0 + if !isempty(get(new_board, to_pos, "")) + reward = 5.0 # Capture! + end + + # Progress value: estimate of game state quality + progressvalue = 5.0 + reward # Capturing is good + + # Switch turns + new_turn = turn == :white ? :black : :white + + newstate = Dict( + :step => current_step + 1, + :board => new_board, + :turn => new_turn, + :reward => reward, + :isterminal => false + ) + + return Dict( + :newNodeKey => "move_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :board => init_board(), + :turn => :white, + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = ( + max_moves = 10, +) + +# Run MCTS +result = runMCTS( + initialstate, + chess_transition, + transitionargs; + maxiterations = 30, + explorationweight = 2.0, # More exploration for game playing + maxSimulationDepth = 4, + horizontalSampleExpansionPhase = 5 +) + +# Display results +println("Chess-like Game MCTS") +println("====================") +println() +println("Best move sequence:") +println(" Initial board state") +println(" → ", result.bestTerminalState[:step], " moves") +println() +println("Final board has ", length(result.bestTerminalState[:board]), " pieces") +println("Root node visits: ", result.root.visits) +println("High value states: ", length(result.highValueStateList)) +``` diff --git a/examples/code_generation.jl b/examples/code_generation.jl new file mode 100644 index 0000000..e1cc2c3 --- /dev/null +++ b/examples/code_generation.jl @@ -0,0 +1,115 @@ +# Code Generation - MCTS for Programming Tasks + +This example shows how MCTS can guide LLM code generation by exploring different implementation strategies. + +```julia +using LLMMCTS + +# State represents the current state of code generation +# It includes the code written so far and the problem being solved + +function code_generation_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + problem = state[:problem] + code_so_far = get(state, :code, "") + + if current_step == 0 + # First step: Plan the approach + new_code = """ + # Function to solve: $(problem) + function solve_problem(input) + """ + newstate = Dict( + :step => 1, + :code => new_code, + :thought => "Plan the approach for: $(problem)", + :reward => 2.0, + :isterminal => false + ) + progressvalue = 5.0 + elseif current_step == 1 + # Second step: Implement main logic + new_code = code_so_far * """ + # Main logic implementation + result = input * 2 # Placeholder implementation + return result + end + """ + newstate = Dict( + :step => 2, + :code => new_code, + :thought => "Implement main function logic", + :reward => 3.0, + :isterminal => false + ) + progressvalue = 7.0 + elseif current_step == 2 + # Third step: Add tests + new_code = code_so_far * """ + + # Test the function + @assert solve_problem(5) == 10 + @assert solve_problem(0) == 0 + println("All tests passed!") + """ + newstate = Dict( + :step => 3, + :code => new_code, + :thought => "Add unit tests to verify implementation", + :reward => 5.0, + :isterminal => true # Code generation complete + ) + progressvalue = 10.0 + else + newstate = Dict( + :step => current_step, + :code => code_so_far, + :thought => "Code generation complete", + :reward => 10.0, + :isterminal => true + ) + progressvalue = 10.0 + end + + return Dict( + :newNodeKey => "code_step_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :problem => "Create a function that doubles its input", + :code => "", + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = (max_steps = 3,) + +# Run MCTS +result = runMCTS( + initialstate, + code_generation_transition, + transitionargs; + maxiterations = 20, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3 +) + +# Display results +println("Code Generation Example") +println("=======================") +println() +println("Problem: ", initialstate[:problem]) +println() +println("Generated code:") +println(result.bestTerminalState[:code]) +println() +println("Code generation complete! ✓") +println("Root node visits: ", result.root.visits) +``` diff --git a/examples/configuration_examples.jl b/examples/configuration_examples.jl new file mode 100644 index 0000000..87fcee6 --- /dev/null +++ b/examples/configuration_examples.jl @@ -0,0 +1,239 @@ +# MCTS Configuration Examples + +This file demonstrates different MCTS configuration options and their effects on search behavior. + +```julia +using LLMMCTS + +# Simple transition function for demonstration +function simple_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + newstate = Dict( + :step => current_step + 1, + :reward => (current_step + 1) * 2, + :isterminal => current_step >= args.max_steps - 1 + ) + progressvalue = (current_step / args.max_steps) * 10 + return Dict( + :newNodeKey => "step_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +initialstate = Dict( + :step => 0, + :reward => 0, + :isterminal => false +) + +transitionargs = (max_steps = 5,) + +# ============================================================================ +# Example 1: Balanced Search (Default) +# ============================================================================ +println("Example 1: Balanced Search (Default)") +println("=" ^ 50) + +result1 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 10, + explorationweight = 1.0, # Balanced exploration/exploitation + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3 +) + +println("Exploration weight: 1.0 (balanced)") +println("Root visits: ", result1.root.visits) +println("Best terminal step: ", result1.bestTerminalState[:step]) +println() + +# ============================================================================ +# Example 2: Aggressive Exploration +# ============================================================================ +println("Example 2: Aggressive Exploration") +println("=" * 50) + +result2 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 10, + explorationweight = 2.0, # More exploration + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 5 # More children per node +) + +println("Exploration weight: 2.0 (aggressive exploration)") +println("Root visits: ", result2.root.visits) +println("Children explored: ", length(result2.root.children)) +println() + +# ============================================================================ +# Example 3: Deep Search (Long Horizon) +# ============================================================================ +println("Example 3: Deep Search (Long Horizon)") +println("=" * 50) + +result3 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 20, + explorationweight = 1.0, + maxSimulationDepth = 5, # Deeper search + horizontalSampleExpansionPhase = 3 +) + +println("Max simulation depth: 5 (deep search)") +println("Root visits: ", result3.root.visits) +println("Search explores further into the future") +println() + +# ============================================================================ +# Example 4: Fast Search (Shallow, Many Iterations) +# ============================================================================ +println("Example 4: Fast Search (Shallow, Many Iterations)") +println("=" * 50) + +result4 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 50, # Many iterations + explorationweight = 1.0, + maxSimulationDepth = 2, # Shallow search + horizontalSampleExpansionPhase = 3 +) + +println("Many iterations (50), shallow depth (2)") +println("Root visits: ", result4.root.visits) +println("Faster but less thorough search") +println() + +# ============================================================================ +# Example 5: Parallel Simulation (Multithreading) +# ============================================================================ +println("Example 5: Parallel Simulation (Multithreading)") +println("=" * 50) + +result5 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 10, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3, + multithread = true # Enable parallel simulation +) + +println("Multithreading enabled") +println("Root visits: ", result5.root.visits) +println("Parallel simulation across child nodes") +println() + +# ============================================================================ +# Example 6: Early Stopping +# ============================================================================ +println("Example 6: Early Stopping") +println("=" * 50) + +# Define early stopping function +function early_stop(state::Dict) + # Stop when we reach a good enough solution + return get(state, :step, 0) >= 3 +end + +result6 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 20, # Would run more if not for early stop + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3, + earlystop = early_stop +) + +println("Early stopping enabled (stops at step >= 3)") +println("Actual iterations: ", result6.root.visits) +println("Early stopping saved unnecessary computation") +println() + +# ============================================================================ +# Example 7: Save Simulation Nodes (for Analysis) +# ============================================================================ +println("Example 7: Save Simulation Nodes") +println("=" * 50) + +result7 = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 5, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3, + saveSimulatedNode = true # Keep simulation nodes +) + +println("saveSimulatedNode = true") +println("Simulation nodes are preserved") +println("Root children: ", length(result7.root.children)) +println("Useful for debugging or further analysis") +println() + +# ============================================================================ +# Example 8: High-Value State Tracking +# ============================================================================ +println("Example 8: High-Value State Tracking") +println("=" * 50) + +# Transition that can produce high-value states +function high_value_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + reward = current_step * 3 + + # Occasionally produce high-value states + if current_step == 2 || current_step == 4 + reward = 9.0 # High value + end + + newstate = Dict( + :step => current_step + 1, + :reward => reward, + :isterminal => current_step >= args.max_steps - 1 + ) + progressvalue = (current_step / args.max_steps) * 10 + return Dict( + :newNodeKey => "step_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +high_value_initial = Dict( + :step => 0, + :reward => 0, + :isterminal => false +) + +result8 = runMCTS( + high_value_initial, + high_value_transition, + transitionargs; + maxiterations = 15, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3 +) + +println("High-value states found: ", length(result8.highValueStateList)) +println("States with reward >= 8 were tracked") +for (i, state) in enumerate(result8.highValueStateList) + println(" High-value state $i: step = ", state[:step]) +end +``` diff --git a/examples/math_problem.jl b/examples/math_problem.jl new file mode 100644 index 0000000..506aa11 --- /dev/null +++ b/examples/math_problem.jl @@ -0,0 +1,97 @@ +# Math Problem Solving - MCTS Example + +This example demonstrates using MCTS to solve a math problem by exploring different solution strategies. + +```julia +using LLMMCTS + +# State represents the current state of problem solving +# It contains the problem statement and the steps taken so far + +function math_problem_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + problem = state[:problem] + + # Example problem: Solve x^2 = 16 + if current_step == 0 + # First step: analyze the problem + newstate = Dict( + :step => 1, + :thought => "This is a quadratic equation x^2 = 16", + :action => "Take square root of both sides", + :reward => 2.0, + :isterminal => false + ) + progressvalue = 5.0 + elseif current_step == 1 + # Second step: solve + newstate = Dict( + :step => 2, + :thought => "Taking square root gives x = ±4", + :action => "x = sqrt(16) or x = -sqrt(16)", + :reward => 3.0, + :isterminal => false + ) + progressvalue = 7.0 + elseif current_step == 2 + # Third step: verify + newstate = Dict( + :step => 3, + :thought => "Verify both solutions work", + :action => "Check x=4: 4^2=16 ✓, Check x=-4: (-4)^2=16 ✓", + :reward => 5.0, + :isterminal => true # Problem solved! + ) + progressvalue = 10.0 + else + # Terminal state + newstate = Dict( + :step => current_step, + :thought => "Problem solved", + :action => "Solution complete", + :reward => 10.0, + :isterminal => true + ) + progressvalue = 10.0 + end + + return Dict( + :newNodeKey => "step_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :problem => "Solve x^2 = 16", + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = () + +# Run MCTS +result = runMCTS( + initialstate, + math_problem_transition, + transitionargs; + maxiterations = 15, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3 +) + +# Display results +println("Problem: ", initialstate[:problem]) +println() +println("Best solution trajectory:") +println(" Step ", result.bestTerminalState[:step]) +println(" Thought: ", result.bestTerminalState[:thought]) +println(" Action: ", result.bestTerminalState[:action]) +println() +println("Solution complete! ✓") +println("Root node visits: ", result.root.visits) +``` diff --git a/examples/pathfinding.jl b/examples/pathfinding.jl new file mode 100644 index 0000000..7e5d2e5 --- /dev/null +++ b/examples/pathfinding.jl @@ -0,0 +1,98 @@ +# Pathfinding Problem - MCTS Example + +This example shows how to use MCTS for a pathfinding problem where the goal is to reach a target location. + +```julia +using LLMMCTS + +# Grid-based pathfinding state +struct Position + x::Int + y::Int +end + +# State transition function for pathfinding +function pathfinding_transition(state::Dict, args::NamedTuple) + current_pos = Position(state[:pos_x], state[:pos_y]) + target_pos = Position(args.target_x, args.target_y) + + # Generate possible moves (up, down, left, right) + moves = [ + (0, 1), # up + (0, -1), # down + (1, 0), # right + (-1, 0) # left + ] + + # In a real scenario, LLM would select which move to try + # For this example, we'll try all moves + move_idx = state[:move_idx] % length(moves) + 1 + dx, dy = moves[move_idx] + + new_x = current_pos.x + dx + new_y = current_pos.y + dy + + # Calculate distance to target + distance = abs(new_x - target_pos.x) + abs(new_y - target_pos.y) + + # Reward: negative of distance (closer is better) + reward = -distance + + # Progress value: LLM estimate (here we use inverse distance as heuristic) + progressvalue = 10 - distance + + newstate = Dict( + :pos_x => new_x, + :pos_y => new_y, + :move_idx => state[:move_idx] + 1, + :reward => reward, + :isterminal => (new_x == target_pos.x && new_y == target_pos.y) || + (state[:move_idx] >= args.max_moves) + ) + + return Dict( + :newNodeKey => "pos_$(new_x)_$(new_y)", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :pos_x => 0, + :pos_y => 0, + :move_idx => 0, + :reward => 0, + :isterminal => false +) + +# Target position +target_x, target_y = 3, 2 + +# Transition arguments +transitionargs = ( + target_x = target_x, + target_y = target_y, + max_moves = 10 +) + +# Run MCTS +result = runMCTS( + initialstate, + pathfinding_transition, + transitionargs; + maxiterations = 20, + explorationweight = 1.5, + maxSimulationDepth = 5, + horizontalSampleExpansionPhase = 4 +) + +# Display results +println("Target: ($target_x, $target_y)") +println("Best final position: (", + result.bestTerminalState[:pos_x], ", ", + result.bestTerminalState[:pos_y], ")") +println("Final distance: ", abs(result.bestTerminalState[:pos_x] - target_x) + + abs(result.bestTerminalState[:pos_y] - target_y)) +println("Root node visits: ", result.root.visits) +``` diff --git a/examples/reasoning.jl b/examples/reasoning.jl new file mode 100644 index 0000000..14e9edc --- /dev/null +++ b/examples/reasoning.jl @@ -0,0 +1,134 @@ +# Multi-step Reasoning - MCTS with Chain of Thought + +This example demonstrates MCTS for multi-step reasoning problems, where the LLM generates chain-of-thought reasoning at each step. + +```julia +using LLMMCTS + +# State tracks the reasoning process +# thought_history: Dict mapping thought/action keys to their content + +function reasoning_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + thought_history = get(state, :thought_history, Dict{String, String}()) + problem = state[:problem] + + if current_step == 0 + # Step 1: Understand the problem + thought = "First, I need to understand what the problem is asking. The problem requires me to analyze the given information and determine the solution approach." + action = "Identify the key components of the problem" + + new_thought_history = copy(thought_history) + new_thought_history["thought_1"] = thought + new_thought_history["action_1"] = action + + newstate = Dict( + :step => 1, + :thought_history => new_thought_history, + :reward => 1.0, + :isterminal => false + ) + progressvalue = 3.0 + elseif current_step == 1 + # Step 2: Break down the problem + thought = "Next, I should break this down into smaller sub-problems. This will make it easier to solve step by step." + action = "Divide the problem into manageable parts" + + new_thought_history = copy(thought_history) + new_thought_history["thought_2"] = thought + new_thought_history["action_2"] = action + + newstate = Dict( + :step => 2, + :thought_history => new_thought_history, + :reward => 2.0, + :isterminal => false + ) + progressvalue = 5.0 + elseif current_step == 2 + # Step 3: Solve each sub-problem + thought = "Now I'll solve each sub-problem individually, using appropriate methods for each." + action = "Apply solution methods to each sub-problem" + + new_thought_history = copy(thought_history) + new_thought_history["thought_3"] = thought + new_thought_history["action_3"] = action + + newstate = Dict( + :step => 3, + :thought_history => new_thought_history, + :reward => 3.0, + :isterminal => false + ) + progressvalue = 7.0 + elseif current_step == 3 + # Step 4: Combine solutions + thought = "Finally, I'll combine all the solutions to form the complete answer to the original problem." + action = "Integrate solutions and verify the answer" + + new_thought_history = copy(thought_history) + new_thought_history["thought_4"] = thought + new_thought_history["action_4"] = action + + newstate = Dict( + :step => 4, + :thought_history => new_thought_history, + :reward => 4.0, + :isterminal => true # Reasoning complete + ) + progressvalue = 10.0 + else + newstate = Dict( + :step => current_step, + :thought_history => thought_history, + :reward => 10.0, + :isterminal => true + ) + progressvalue = 10.0 + end + + return Dict( + :newNodeKey => "reasoning_step_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :problem => "Explain how photosynthesis works", + :thought_history => Dict{String, String}(), + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = (max_steps = 4,) + +# Run MCTS +result = runMCTS( + initialstate, + reasoning_transition, + transitionargs; + maxiterations = 25, + explorationweight = 1.0, + maxSimulationDepth = 4, + horizontalSampleExpansionPhase = 3 +) + +# Display results +println("Multi-step Reasoning Example") +println("=============================") +println() +println("Problem: ", initialstate[:problem]) +println() +println("Reasoning steps:") +for (key, value) in result.bestTerminalState[:thought_history] + println(" $key: $value") +end +println() +println("Reasoning complete! ✓") +println("Root node visits: ", result.root.visits) +println("Total steps in reasoning chain: ", result.bestTerminalState[:step]) +``` diff --git a/examples/simple_example.jl b/examples/simple_example.jl new file mode 100644 index 0000000..eeea7c4 --- /dev/null +++ b/examples/simple_example.jl @@ -0,0 +1,59 @@ +# Simple MCTS Example + +This example demonstrates basic MCTS usage with a simple state transition function. + +```julia +using LLMMCTS + +# Define a simple state transition function +function simple_transition(state::Dict, args::NamedTuple) + # In a real scenario, this would call an LLM + # For this example, we'll just generate deterministic next states + + current_step = get(state, :step, 0) + new_step = current_step + 1 + + # Create new state + newstate = Dict( + :step => new_step, + :reward => new_step * 2, # Simple reward function + :isterminal => new_step >= args.max_steps + ) + + # LLM would provide progressvalue estimate + progressvalue = (new_step / args.max_steps) * 10 + + return Dict( + :newNodeKey => "step_$(new_step)", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = (max_steps = 5,) + +# Run MCTS +result = runMCTS( + initialstate, + simple_transition, + transitionargs; + maxiterations = 10, + explorationweight = 1.0, + maxSimulationDepth = 3, + horizontalSampleExpansionPhase = 3 +) + +# Access results +println("Root node visits: ", result.root.visits) +println("Best next state: ", result.bestNextState) +println("Best terminal state: ", result.bestTerminalState) +println("High value states: ", result.highValueStateList) +``` diff --git a/examples/tool_use.jl b/examples/tool_use.jl new file mode 100644 index 0000000..da4e62e --- /dev/null +++ b/examples/tool_use.jl @@ -0,0 +1,109 @@ +# Tool Use Example - MCTS with External Tools + +This example shows how MCTS can coordinate with external tools (like APIs, databases, or other services). + +```julia +using LLMMCTS + +# Simulated tool interface +struct Tool + name::String + description::String +end + +const AVAILABLE_TOOLS = [ + Tool("calculator", "Perform mathematical calculations"), + Tool("web_search", "Search the web for information"), + Tool("database_query", "Query a database") +] + +# State tracks which tools have been used and their results +function tool_use_transition(state::Dict, args::NamedTuple) + current_step = get(state, :step, 0) + tools_used = get(state, :tools_used, String[]) + + # LLM would decide which tool to use + # For this example, we try tools in order + tool_idx = (current_step - 1) % length(AVAILABLE_TOOLS) + 1 + + if tool_idx > length(AVAILABLE_TOOLS) + # All tools tried, return terminal state + newstate = Dict( + :step => current_step + 1, + :tools_used => tools_used, + :reward => 8.0, + :isterminal => true + ) + return Dict( + :newNodeKey => "all_tools_tried", + :newstate => newstate, + :progressvalue => 8.0 + ) + end + + tool = AVAILABLE_TOOLS[tool_idx] + + # Simulate tool execution + tool_result = "Tool '$(tool.name)' executed successfully" + + # Calculate reward based on progress + progress = length(tools_used) / length(AVAILABLE_TOOLS) + reward = progress * 5 + + # Progress value: LLM estimates how close we are to solving + progressvalue = progress * 10 + + new_tools_used = vcat(tools_used, tool.name) + + newstate = Dict( + :step => current_step + 1, + :tools_used => new_tools_used, + :current_tool => tool.name, + :tool_result => tool_result, + :reward => reward, + :isterminal => false + ) + + return Dict( + :newNodeKey => "tool_$(tool.name)_$current_step", + :newstate => newstate, + :progressvalue => progressvalue + ) +end + +# Initial state +initialstate = Dict( + :step => 0, + :tools_used => String[], + :reward => 0, + :isterminal => false +) + +# Transition arguments +transitionargs = (max_tools = 3,) + +# Run MCTS +result = runMCTS( + initialstate, + tool_use_transition, + transitionargs; + maxiterations = 20, + explorationweight = 1.2, + maxSimulationDepth = 4, + horizontalSampleExpansionPhase = 3 +) + +# Display results +println("Available tools:") +for tool in AVAILABLE_TOOLS + println(" - $(tool.name): $(tool.description)") +end +println() +println("Best tool usage sequence:") +for tool in result.bestTerminalState[:tools_used] + println(" → Used: $tool") +end +println() +println("Root node visits: ", result.root.visits) +println("High value states found: ", length(result.highValueStateList)) +``` diff --git a/src/LLMMCTS.jl b/src/LLMMCTS.jl index c96bbd4..c3b98de 100644 --- a/src/LLMMCTS.jl +++ b/src/LLMMCTS.jl @@ -1,6 +1,6 @@ module LLMMCTS - # export agent + export MCTSNode """ Order by dependencies of each file. The 1st included file must not depend on any other diff --git a/src/interface.jl b/src/interface.jl index 32a8b3f..e48d81d 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -9,47 +9,55 @@ using ..type, ..mcts, ..util # ---------------------------------------------- 100 --------------------------------------------- # +""" Search for the best action to take for a given state and task. -""" Search the best action to take for a given state and task +This function runs the MCTS algorithm through multiple iterations of expansion, +simulation, and backpropagation to find optimal decisions. + +Does **not** mutate the input state; it creates new MCTS nodes during search. # Arguments - - `initialstate::T` - initial state - - `transition::Function` - a function that define how the state transitions - - `transitionargs::NamedTuple` - arguments for transition function +- `initialstate::T` + The initial state from which to start the search. +- `transition::Function` + A function that defines how the state transitions. +- `transitionargs::NamedTuple` + Arguments passed to the transition function. # Keyword Arguments - - `horizontalSampleExpansionPhase::Integer` - a number of child state MCTS sample at each node during expansion phase (default: 3) - - `horizontalSampleSimulationPhase::Integer` - a number of child state MCTS sample at each node during simulation's expansion phase (default: 3) - - `maxSimulationDepth::Integer` - a number of levels MCTS goes during simulation phase (default: 3) - - `maxiterations::Integer` - a number of iteration MCTS goes thru expansion -> simulation -> backpropagation cycle (default: 10) - - `explorationweight::Number` - exploration weight controls how much MCTS should explore new state instead of exploit - a known state. 1.0 balance between exploration and exploitation like 50%-50%. 2.0 makes MCTS - aggressively explore new state (default: 1.0) - - `earlystop::Union{Function,Nothing}` - optional function to check early stopping condition (default: nothing) - - `saveSimulatedNode::Bool` - whether to save nodes created during simulation phase (default: false) - - `multithread::Bool` - whether to use multithreading during simulation (default: false) +- `horizontalSampleExpansionPhase::Integer=3` + Number of child states sampled at each node during expansion phase. +- `horizontalSampleSimulationPhase::Integer=3` + Number of child states sampled at each node during simulation's expansion phase. +- `maxSimulationDepth::Integer=3` + Maximum depth MCTS goes during simulation phase. +- `maxiterations::Integer=10` + Number of iterations MCTS performs through expansion → simulation → backpropagation cycles. +- `explorationweight::Number=1.0` + Exploration weight controls how much MCTS explores new states versus exploiting known states. + A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0) + encourage more aggressive exploration. +- `earlystop::Union{Function,Nothing}=nothing` + Optional function to check early stopping condition. If satisfied, MCTS breaks iterations. +- `saveSimulatedNode::Bool=false` + Whether to save nodes created during simulation phase. +- `multithread::Bool=false` + Whether to use multithreading during simulation. -# Returns - - `NamedTuple{(:root, :bestNextState, :bestFinalState), Tuple{MCTSNode, T, T}}` - - root: the complete MCTS tree with root node - - bestNextState: the best immediate next state - - bestFinalState: the best final state along the best trajectory +# Return +- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList), + Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}` + - `root`: the complete MCTS tree with root node + - `bestNextState`: the best immediate next state + - `bestTerminalState`: the best final state along the best trajectory + - `highValueStateList`: list of high-value terminal states (reward >= 8) # Example - Refers to SQLLLM package - -# Signature +```jldoctest +julia> using LLMMCTS +julia> initialState = Dict(:reward=>0.0) +julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5) +``` """ function runMCTS( initialstate::T, @@ -63,97 +71,185 @@ function runMCTS( explorationweight::Number=1.0, earlystop::Union{Function,Nothing}=nothing, saveSimulatedNode::Bool=false, - multithread=false - )::NamedTuple{(:root, :bestNextState, :bestFinalState),Tuple{MCTSNode,T,T}} where {T<:Any} - + multithread=false, + )::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList), + Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any} + println("--> LLMMCTS runMCTS 1") + # Initialize the MCTS tree with a root node representing the initial state + # root.visits=0: no visits yet + # root.statevalue=0: no simulation results yet root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, Dict{String,MCTSNode}(), - Dict{Symbol,Any}()) + Dict{String,Any}()) + # Channel to collect high-value terminal states (reward >= 8) + # These are "good solutions" that can be returned to the user + highValueState = Channel{Any}(100) + + # Main MCTS loop: perform iterations to build the search tree + # Each iteration: SELECTION → EXPANSION → SIMULATION → BACKPROPAGATION for nth in 1:maxiterations + # Start from root and traverse down using UCT selection node = root - node.visits += 1 - + node.visits += 1 # Count this iteration's visit to root + println("--> LLMMCTS runMCTS 2") + # Phase 1: SELECTION - Traverse tree using UCT until reaching a leaf node + # UCT balances exploration (new branches) vs exploitation (promising branches) while !isleaf(node) + println("--> LLMMCTS runMCTS 3") node = UCTselect(node, explorationweight) end - + println("--> LLMMCTS runMCTS 4") + # Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate if node.isterminal - # MCTS arrive at the leaf node that is also a terminal state, - # do nothing then go directly to backpropagation. It means the end of this iteration + println("--> LLMMCTS runMCTS 5") + # If this terminal state has high reward (>= 8), store it for later + if node.state[:reward] >= 8 + println("--> LLMMCTS runMCTS 6") + put!(highValueState, deepcopy(node.state)) + end + println("--> LLMMCTS runMCTS 7") + # Backpropagate the terminal node's own reward up to root + # This updates all ancestors with this path's outcome backpropagate(node, node.reward) else + println("--> LLMMCTS runMCTS 8") + # Phase 3: EXPANSION - Generate children for this non-terminal leaf + # Horizontal sampling: create multiple child nodes via LLM transition _ = expand(node, transition, transitionargs; horizontalSample=horizontalSampleExpansionPhase, multithread=multithread) + println("--> LLMMCTS runMCTS 9") + # Phase 4: SIMULATION + BACKPROPAGATION + # For each newly expanded child, run simulation and update statistics if multithread + println("--> LLMMCTS runMCTS 10") + # Parallel simulation: spawn threads for each child node @sync for (leafNodeKey, leafNode) in node.children @spawn simulateThenBackpropagate(leafNode, transition, transitionargs; maxSimulationDepth=maxSimulationDepth, horizontalSampleSimulationPhase=horizontalSampleSimulationPhase, saveSimulatedNode=saveSimulatedNode, - multithread=multithread) + multithread=multithread, + highValueState=highValueState, + ) end else + println("--> LLMMCTS runMCTS 11") + # Sequential simulation: process each child one at a time for (leafNodeKey, leafNode) in node.children + println("--> LLMMCTS runMCTS 11-1") simulateThenBackpropagate(leafNode, transition, transitionargs; - maxSimulationDepth=maxSimulationDepth, - horizontalSampleSimulationPhase=horizontalSampleSimulationPhase, - saveSimulatedNode=saveSimulatedNode, - multithread=multithread) + maxSimulationDepth=maxSimulationDepth, + horizontalSampleSimulationPhase=horizontalSampleSimulationPhase, + saveSimulatedNode=saveSimulatedNode, + multithread=multithread, + highValueState=highValueState) end end end - - # stop if the early stop condition is met + println("--> LLMMCTS runMCTS 12") + # Phase 5: EARLY STOP CHECK + # Optional: stop search early if a condition is met if typeof(earlystop) <: Function && earlystop(node.state) + println("--> LLMMCTS runMCTS 13") break end end - - # select the best next state and the best final state + println("--> LLMMCTS runMCTS 14") + # After all iterations, extract results from the search tree + # Select best immediate next state (best child of root) bestNextState = selectBestNextNode(root) - besttrajectory = selectBestTrajectoryNode(root) + println("--> LLMMCTS runMCTS 15") + # Select best terminal state along the optimal trajectory + bestTerminalState = selectBestTrajectoryNode(root) - return (root=root, bestNextState=bestNextState.state, bestFinalState=besttrajectory.state) + # Collect all high-value states from the channel into a list + highValueStateList = Vector{Dict{String, Any}}() + while !isempty(highValueState) + println("--> LLMMCTS runMCTS 16") + push!(highValueStateList, take!(highValueState)) + end + println("--> LLMMCTS runMCTS 17") + # Return complete search results + result = ( + root=root, + bestNextState=bestNextState.state, + bestTerminalState=bestTerminalState.state, + highValueStateList=highValueStateList + ) + + return result end -""" Search the best action to take for a given state and task +""" Run simulation from a given node and backpropagate the reward. + +This function performs simulation (rollout) from the given node, collects the +cumulative reward along the trajectory, and backpropagates it up the tree to update +visit counts and state values. + +Does **not** mutate the input node's children (unless `saveSimulatedNode=true`). # Arguments - `node::MCTSNode` - current node to simulate from + The current node to simulate from. - `transition::Function` - a function that defines how the state transitions + A function that defines how the state transitions. - `transitionargs::NamedTuple` - arguments for transition function + Arguments passed to the transition function. # Keyword Arguments -- `maxSimulationDepth::Integer` - a number of levels MCTS goes during simulation phase (default: 3) -- `horizontalSampleSimulationPhase::Integer` - a number of child states MCTS samples at each node during simulation phase (default: 3) -- `saveSimulatedNode::Bool` - whether to save nodes created during simulation phase (default: false) -- `multithread::Bool` - whether to use multithreading during simulation (default: false) +- `maxSimulationDepth::Integer=3` + Maximum depth MCTS goes during simulation phase. +- `horizontalSampleSimulationPhase::Integer=3` + Number of child states sampled at each node during simulation phase. +- `saveSimulatedNode::Bool=false` + Whether to save nodes created during simulation phase. If false, children are + cleared after simulation. +- `multithread::Bool=false` + Whether to use multithreading during simulation. -# Returns - Nothing, but updates the node's reward and visit count through backpropagation +# Return +- `Nothing` + +# Signature """ function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3, saveSimulatedNode::Bool=false, - multithread=false) - simTrajectoryReward, terminalstate = simulate(node, transition, transitionargs; - maxSimulationDepth=maxSimulationDepth, - horizontalSample=horizontalSampleSimulationPhase, - multithread=multithread) + multithread=false, + highValueState=Union{Nothing,Any}=nothing) + println("--> LLMMCTS simulateThenBackpropagate 1") + # Phase 1: RUN SIMULATION (rollout) + # Perform a rollout from this node, accumulating rewards along the way + simTrajectoryReward, terminalstate = + simulate(node, transition, transitionargs; + maxSimulationDepth=maxSimulationDepth, + horizontalSample=horizontalSampleSimulationPhase, + multithread=multithread) + println("--> LLMMCTS simulateThenBackpropagate 2") + # Phase 2: HIGH-VALUE STATE TRACKING + # If we reached a terminal state with high reward (>= 8), store it + # This allows users to access multiple good solutions, not just the best one + if highValueState !== nothing && + terminalstate !== nothing && + terminalstate["reward"] >= 8 + println("--> LLMMCTS simulateThenBackpropagate 3") + put!(highValueState, deepcopy(terminalstate)) + end + println("--> LLMMCTS simulateThenBackpropagate 4") + # Phase 3: BACKPROPAGATE + # Update statistics (visits, statevalue) for all ancestors up to root + # The simulation result is now incorporated into the tree backpropagate(node, simTrajectoryReward) - - # check if the user wants to keep the simulated node + println("--> LLMMCTS simulateThenBackpropagate 5") + # Phase 4: MEMORY MANAGEMENT + # Clear children unless user wants to keep them for analysis + # This frees memory for the next iteration while preserving tree structure if saveSimulatedNode == false + println("--> LLMMCTS simulateThenBackpropagate 6") node.children = Dict{String, MCTSNode}() end + println("--> LLMMCTS simulateThenBackpropagate 7") end diff --git a/src/mcts.jl b/src/mcts.jl index 2715baa..5b65ee7 100644 --- a/src/mcts.jl +++ b/src/mcts.jl @@ -1,7 +1,7 @@ module mcts export selectBestNextNode, selectBestTrajectoryNode, backpropagate, isleaf, isroot, selectChildNode, - expand, simulate, makeNewState + expand, simulate using Base.Threads using GeneralUtils @@ -10,29 +10,33 @@ using ..type # ---------------------------------------------- 100 --------------------------------------------- # -""" Select the best next node based on the highest value metric +""" Select the best child node based on the highest value metric. + +The selection metric depends on the node's state values: +- If the sum of statevalues is non-zero, uses `statevalue/visits` ratio. +- Otherwise, uses `progressvalue + reward`. # Arguments - - `node::MCTSNode` - node of a search tree to evaluate +- `node::MCTSNode` + The node whose children will be evaluated. # Return - - `childNode::MCTSNode` - the child node with highest value based on either: - - statevalue/visits ratio if any nodes have non-zero statevalue - - progressvalue + reward otherwise +- `childNode::MCTSNode` + The child node with the highest value according to the selection metric. """ function selectBestNextNode(node::MCTSNode)::MCTSNode highestProgressValue = -1 nodekey = nothing # Calculate sum of statevalues across all child nodes + # This determines whether to use statevalue/visits (exploitation) or progressvalue+reward (exploration) stateValueSum = sum([v.statevalue for (k, v) in node.children]) # If any nodes have non-zero statevalue, use statevalue/visits as selection metric + # This means simulations have confirmed node values - use exploitation if stateValueSum != 0 for (k, childnode) in node.children - # Calculate average statevalue per visit + # Calculate average statevalue per visit (running average from simulations) potential = childnode.statevalue / childnode.visits if potential > highestProgressValue @@ -41,7 +45,8 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode end end else - # Otherwise use progressvalue + reward as selection metric + # No simulations yet - use progressvalue + reward for initial guidance + # This allows LLM heuristics to guide early search before simulations provide data for (k, childnode) in node.children potential = childnode.progressvalue + childnode.reward @@ -56,20 +61,22 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode end -""" Select the best trajectory node based on the highest reward +""" Select the best node along the optimal trajectory. + +Traverses down the tree from the given node by repeatedly applying `selectBestNextNode` +until reaching a leaf node, returning the highest-value node found along the path. # Arguments - - `node::MCTSNode` - node of a search tree to evaluate +- `node::MCTSNode` + The node to start trajectory selection from. # Return - - `childNode::MCTSNode` - the highest value child node found by traversing down the tree using selectBestNextNode - until reaching a leaf node - -# Signature +- `childNode::MCTSNode` + The highest-value node found by following the optimal trajectory to a leaf. """ function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode + # Follow the optimal path down the tree by repeatedly selecting the best child + # This gives us the highest-value trajectory from the starting node to a leaf while !isleaf(node) node = selectBestNextNode(node) end @@ -78,99 +85,108 @@ function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode end -""" Backpropagate reward along the simulation chain +""" Backpropagate reward along the simulation chain. + +Updates visit counts and state values for all nodes along the path from the given +leaf node to the root, applying reward discounting for future rewards. + +**Modifies nodes in place.** # Arguments - - `node::MCTSNode` - leaf node of a search tree - - `simTrajectoryReward::T` - total reward from trajectory simulation - - `discountRewardCoeff::AbstractFloat` - A discount reward coefficient to reduce future reward. The futher in the future the lower - reward it is now. - -# Return - - `Nothing` - This function modifies the nodes in place and returns nothing +- `node::MCTSNode` + The leaf node from which to start backpropagation. +- `simTrajectoryReward::Number` + The total reward from the trajectory simulation. -# Signature +# Keyword Arguments +- `discountRewardCoeff::AbstractFloat=0.9` + Discount coefficient applied to future rewards. Larger distances from the leaf + receive progressively lower discounted rewards. + +# Return +- `Nothing` """ function backpropagate(node::MCTSNode, simTrajectoryReward::T; - discountRewardCoeff::AbstractFloat=0.9) where {T<:Number} + discountRewardCoeff::AbstractFloat=0.9) where {T<:Number} + println("--> LLMMCTS backpropagate 1") + # Propagate the simulation result back up the tree to update all ancestor nodes + # Each node's statistics are updated with the cumulative reward from the simulation while !isroot(node) - # Update the statistics of the current node based on the result of the playout - node.visits += 1 # Increment visit count for this node + println("--> LLMMCTS backpropagate 2") + # Increment visit count - this simulation passed through this node + node.visits += 1 + println("--> LLMMCTS backpropagate 3") node.statevalue += ((node.statevalue * (node.visits-1)) + simTrajectoryReward) / node.visits # Update running average of state value - simTrajectoryReward *= discountRewardCoeff # discount because future reward is uncertain - node = node.parent # Move up to parent node for next iteration + + # Apply discount to future rewards - rewards further from the current state are worth less + # This reflects temporal uncertainty: distant future rewards are less certain + simTrajectoryReward *= discountRewardCoeff + + # Move up to parent node to continue propagation + node = node.parent end + println("--> LLMMCTS backpropagate 4") end -""" Determine whether a node is a leaf node of a search tree. +""" Determine whether a node is a leaf node. + +A leaf node has no children. # Arguments - - `node::MCTSNode` - a search tree node +- `node::MCTSNode` + The search tree node to check. # Return - - `result::Bool` - true if it is a leaf node (has no children), false otherwise. +- `result::Bool` + `true` if the node has no children, `false` otherwise. # Example ```jldoctest -julia> using Revise -julia> using YiemAgent, DataStructures -julia> initialState = Dict{Symbol, Any}( - :customerinfo=> Dict{Symbol, Any}(), - :storeinfo=> Dict{Symbol, Any}(), - - :thoughtHistory=> OrderedDict{Symbol, Any}( - :question=> "How are you?", - ) - ) -julia> statetype = typeof(initialState) -julia> root = YiemAgent.MCTSNode(initialState, 0, 0.0, Dict{statetype, YiemAgent.MCTSNode}()) -julia> YiemAgent.isleaf(root) +julia> using LLMMCTS +julia> node = MCTSNode("leaf", Dict(:reward=>1.0), 0, 0, 0, 1.0, true, nothing, Dict(), Dict()) +julia> isleaf(node) true ``` - -# Signature """ isleaf(node::MCTSNode)::Bool = isempty(node.children) -""" Determine wheter a given node is a root node +""" Determine whether a given node is a root node. + +The root node is identified by having `"root"` as its `nodekey`. # Arguments - - `node::MCTSNode` - node of a search tree +- `node::MCTSNode` + The search tree node to check. # Return - - `isrootnode::Bool` - true if the given node is root node, false otherwise - -# Signature +- `isrootnode::Bool` + `true` if the node is the root node, `false` otherwise. """ isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false -""" Select child node based on the highest statevalue +""" Select the child node with the highest value. + +Uses `progressvalue + reward` as the selection metric. # Arguments - - `node::MCTSNode` - node of a search tree +- `node::MCTSNode` + The node whose children will be evaluated. # Return - - `childNode::MCTSNode` - the highest value child node - -# Signature +- `childNode::MCTSNode` + The child node with the highest `progressvalue + reward` value. """ function selectChildNode(node::MCTSNode)::MCTSNode highestProgressValue = -1 nodekey = nothing - # loop thought node children dictionary to find the highest progress value + # During simulation rollout, we need to pick which child to explore next + # Use progressvalue + reward as the selection metric (no UCT here) + # - progressvalue: LLM's estimate of how promising this state is + # - reward: immediate environment feedback + # Together they guide fast exploration during simulation for (k, childNode) in node.children potential = childNode.progressvalue + childNode.reward if potential > highestProgressValue @@ -183,35 +199,43 @@ function selectChildNode(node::MCTSNode)::MCTSNode end -""" Expand selected node. +""" Expand a node by generating new child nodes. + +Creates new child nodes by applying the transition function multiple times +(horizontally samples) from the current node. # Arguments - - `node::MCTSNode` - MCTS node to expand - - `transition::Function` - A function that handles state transition. - - `transitionargs::NamedTuple` - Arguments for transition() +- `node::MCTSNode` + The MCTS node to expand. +- `transition::Function` + A function that handles state transition. +- `transitionargs::NamedTuple` + Arguments passed to the transition function. # Keyword Arguments - - `horizontalSample::Integer` - Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3. - - `multithread::Bool` - Whether to run expansion in parallel using multiple threads. Defaults to false. - -# Return - - None +- `horizontalSample::Integer=3` + Number of child nodes to generate. +- `multithread::Bool=false` + Whether to run expansion in parallel using multiple threads. -# Signature +# Return +- `Nothing` """ function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple; horizontalSample::Integer=3, multithread=false) + # Generate child nodes by applying the transition function multiple times + # This is called "horizontal sampling" - we branch out horizontally in the tree + # - multithread=true: spawn parallel threads for each expansion + # - multithread=false: sequential expansion (default, simpler) + println("--> LLMMCTS expand 1") if multithread @sync for i in 1:horizontalSample @spawn _expand(node, transition, transitionargs) end else + println("--> LLMMCTS expand 2") for i in 1:horizontalSample + println("--> LLMMCTS expand 3") _expand(node, transition, transitionargs) end end @@ -219,166 +243,190 @@ end """ Helper function to expand a single child node. +Creates one new child node from the parent node using the transition function. +Checks for semantically equivalent states (dejavu) to avoid duplicates. + # Arguments - - `node::MCTSNode` - Parent MCTS node to expand from - - `transition::Function` - A function that handles state transition - - `transitionargs::NamedTuple` - Arguments for transition() +- `node::MCTSNode` + The parent MCTS node to expand from. +- `transition::Function` + A function that handles state transition. +- `transitionargs::NamedTuple` + Arguments passed to the transition function. # Return - - None - -# Signature +- `Nothing` """ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple) - result = transition(node.state, transitionargs) - newNodeKey::AbstractString = result[:newNodeKey] - newstate::AbstractDict = result[:newstate] - progressvalue::Integer = result[:progressvalue] - - """ - [] newNodeKey ∉ keys(node.children). - New state may have semantic vector close enought to - one of existing child state. Which can be assume that they are the same state - semantically-wise i.e. De javu. This could be used to recall lessons for this - similar situation to improve decisionMaker and evaluator. - """ - if newNodeKey ∉ keys(node.children) - newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward], - newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{Symbol, Any}()) - node.children[newNodeKey] = newNode - end + println("--> LLMMCTS _expand 1") + # Generate one child node from the parent using the transition function + result = transition(node.state, transitionargs) + newNodeKey::AbstractString = result[:newNodeKey] + newstate::AbstractDict = result[:newstate] + progressvalue::Integer = result[:progressvalue] + println("--> LLMMCTS _expand 2") + # Dejavu detection: avoid adding duplicate states + # If newNodeKey already exists, skip - this handles semantically equivalent states + if newNodeKey ∉ keys(node.children) + println("--> LLMMCTS _expand 3") + # Create new MCTS node with: + # - visits=0: no simulations yet + # - statevalue=0: will be updated after simulation + # - progressvalue: LLM's estimate (fast heuristic) + # - reward: immediate environment feedback + newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate["reward"], + newstate["isterminal"], node, Dict{String, MCTSNode}(), Dict{String, Any}()) + println("--> LLMMCTS _expand 4") + node.children[newNodeKey] = newNode + println("--> LLMMCTS _expand 5") + end end -""" Simulate interactions between agent and environment +""" Simulate interactions between agent and environment. + +Performs a rollout from the given node up to the maximum simulation depth, +sampling child nodes at each level and accumulating rewards along the way. # Arguments - - `node::MCTSNode` - node that will be a simulation starting point. - - `transition::Function` - A user function that handles how state transition. - - `transitionargs::NamedTuple` - Arguments for everything the user will use within transition(). - - `maxSimulationDepth::Integer` - maximum depth level MCTS goes vertically during simulation. - - `horizontalSample::Integer` - Total number to sample from the current node (i.e. expand new node horizontally) +- `node::MCTSNode` + The node to start simulation from. +- `transition::Function` + A user function that handles state transition. +- `transitionargs::NamedTuple` + Arguments passed to the transition function. # Keyword Arguments - - `multithread::Bool` - Whether to run expansion in parallel using multiple threads. Defaults to false. - -# Return - - `simTrajectoryReward::Number` - Cumulative reward collected along the simulation trajectory - - `terminalstate::Union{Dict{Symbol, Any}, Nothing}` - Final state if terminal state reached, nothing otherwise +- `maxSimulationDepth::Integer=3` + Maximum depth level MCTS goes vertically during simulation. +- `horizontalSample::Integer=3` + Number of child nodes sampled at each node during simulation. +- `multithread::Bool=false` + Whether to run expansion in parallel using multiple threads. -# Signature +# Return +- `NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}` + - `simTrajectoryReward`: cumulative reward collected along the simulation trajectory + - `terminalstate`: final state if a terminal state was reached, `nothing` otherwise """ function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false -)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{Symbol, Any}, Nothing}}} - + )::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}} + println("--> LLMMCTS simulate 1") + # Perform a rollout simulation from the given node: + # 1. Accumulate rewards along the trajectory + # 2. Expand nodes horizontally at each level + # 3. Select children to explore vertically down the tree + # Returns cumulative reward and whether a terminal state was reached + simTrajectoryReward = 0.0 terminalstate = nothing for depth in 1:maxSimulationDepth + println("--> LLMMCTS simulate 2") + # Accumulate the current node's reward to the trajectory total simTrajectoryReward += node.reward + + # Check if we've reached a terminal state if node.isterminal + println("--> LLMMCTS simulate 3") terminalstate = node.state break else + println("--> LLMMCTS simulate 4") + # Expand current node to generate children (horizontal sampling) _ = expand(node, transition, transitionargs; horizontalSample=horizontalSample, multithread=multithread) + println("--> LLMMCTS simulate 5") + # Select best child to continue the rollout (vertical exploration) + # Uses progressvalue + reward for fast selection during simulation node = selectChildNode(node) end + println("--> LLMMCTS simulate 6") end - - return (simTrajectoryReward=simTrajectoryReward, terminalstate=terminalstate) + println("--> LLMMCTS simulate 7") + return (simTrajectoryReward=simTrajectoryReward, + terminalstate=terminalstate) end -""" Make new state +# """ Make new state -# Arguments - - `currentstate::T1` - Current state dictionary containing thought history and metadata - - `thoughtDict::T4` - Dictionary containing new thought and action - - `response::T2` - Response string from the environment - - `select::Union{T3, Nothing}` - Selection value or nothing - - `reward::T3` - Reward value for this state - - `isterminal::Bool` - Whether this state is terminal +# # Arguments +# - `currentstate::T1` +# Current state dictionary containing thought history and metadata +# - `thoughtDict::T4` +# Dictionary containing new thought and action +# - `response::T2` +# Response string from the environment +# - `select::Union{T3, Nothing}` +# Selection value or nothing +# - `reward::T3` +# Reward value for this state +# - `isterminal::Bool` +# Whether this state is terminal -# Return - - `Tuple{String, Dict{Symbol, <:Any}}` - A tuple containing: - - A unique node key string - - A new state dictionary with updated thought history and metadata +# # Return +# - `Tuple{String, Dict{String, <:Any}}` +# A tuple containing: +# - A unique node key string +# - A new state dictionary with updated thought history and metadata -# Example -```jldoctest -julia> -``` +# # Example +# ```jldoctest +# julia> +# ``` -# Signature -""" -function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing}, - reward::T3, isterminal::Bool - )::Tuple{String, Dict{Symbol, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict} +# # Signature +# """ +# function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing}, +# reward::T3, isterminal::Bool +# )::Tuple{String, Dict{String, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict} - # Find the latest thought key and index from current state's thought history - currentstate_latestThoughtKey, currentstate_latestThoughtIndice = - GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought") - # Calculate next index for new thought/action - currentstate_nextIndice = - currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1 - # Create new keys for thought and action based on next index - currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice") - latestActionKey = Symbol("action_$currentstate_nextIndice") +# # Find the latest thought key and index from current state's thought history +# currentstate_latestThoughtKey, currentstate_latestThoughtIndice = +# GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought") +# # Calculate next index for new thought/action +# currentstate_nextIndice = +# currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1 +# # Create new keys for thought and action based on next index +# currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice") +# latestActionKey = Symbol("action_$currentstate_nextIndice") - # Find the latest thought index from input thought dictionary - _, thoughtDict_latestThoughtIndice = - GeneralUtils.findHighestIndexKey(thoughtDict, "thought") +# # Find the latest thought index from input thought dictionary +# _, thoughtDict_latestThoughtIndice = +# GeneralUtils.findHighestIndexKey(thoughtDict, "thought") - # Determine thought and action keys from thought dictionary - thoughtDict_latestThoughtKey, thoughtDict_latestActionKey = - if thoughtDict_latestThoughtIndice == -1 - (:thought, :action) - else - ( - Symbol("thought_$thoughtDict_latestThoughtIndice"), - Symbol("action_$thoughtDict_latestThoughtIndice"), - ) - end +# # Determine thought and action keys from thought dictionary +# thoughtDict_latestThoughtKey, thoughtDict_latestActionKey = +# if thoughtDict_latestThoughtIndice == -1 +# (:thought, :action) +# else +# ( +# Symbol("thought_$thoughtDict_latestThoughtIndice"), +# Symbol("action_$thoughtDict_latestThoughtIndice"), +# ) +# end - # Create new state by deep copying current state - newstate = deepcopy(currentstate) - # Update thought history with new thought - newstate[:thoughtHistory][currentstate_latestThoughtKey] = - thoughtDict[thoughtDict_latestThoughtKey] - # Update thought history with new action - newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey] - # Create and add new observation to thought history - newObservationKey = Symbol("observation_$(currentstate_nextIndice)") - newstate[:thoughtHistory][newObservationKey] = response - # Update state metadata - newstate[:reward] = reward - newstate[:select] = select - newstate[:isterminal] = isterminal +# # Create new state by deep copying current state +# newstate = deepcopy(currentstate) +# # Update thought history with new thought +# newstate[:thoughtHistory][currentstate_latestThoughtKey] = +# thoughtDict[thoughtDict_latestThoughtKey] +# # Update thought history with new action +# newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey] +# # Create and add new observation to thought history +# newObservationKey = Symbol("observation_$(currentstate_nextIndice)") +# newstate[:thoughtHistory][newObservationKey] = response +# # Update state metadata +# newstate[:reward] = reward +# newstate[:select] = select +# newstate[:isterminal] = isterminal - # Generate unique ID for new node - newNodeKey = GeneralUtils.uuid4snakecase() +# # Generate unique ID for new node +# newNodeKey = GeneralUtils.uuid4snakecase() - return (newNodeKey, newstate) -end +# return (newNodeKey, newstate) +# end diff --git a/src/type.jl b/src/type.jl index 0b8f46e..d3e7af5 100644 --- a/src/type.jl +++ b/src/type.jl @@ -58,7 +58,7 @@ mutable struct MCTSNode{T1<:AbstractDict, T2<:AbstractString} isterminal::Bool parent::Union{MCTSNode, Nothing} children::Dict{String, MCTSNode} - etc::Dict{Symbol, Any} # store anything + etc::Dict{String, Any} # store anything end diff --git a/src/util.jl b/src/util.jl index 512556d..db0cafa 100644 --- a/src/util.jl +++ b/src/util.jl @@ -6,48 +6,99 @@ using ..type # ---------------------------------------------- 100 --------------------------------------------- # -""" Select a node based on UCT score +""" Select a node based on UCT (Upper Confidence Bound for Trees) score. + +The function computes UCT values for all child nodes and returns the child with the +highest UCT score. The UCT formula balances exploitation (child state value) and +exploration (visit count and parent visit count) using the exploration weight `w`. + +Does **not** mutate the input node. # Arguments - - `node::MCTSNode` - mcts node - - `w::T` - exploration weight. Value is usually between 1 to 2. - Value 1.0 makes MCTS balance between exploration and exploitation like 50%-50%. - Value 2.0 makes MCTS aggressively search the tree. -# Return - - `selectedNode::MCTSNode` - child node with highest UCT score. UCT score balances between exploitation (state value) - and exploration (visit count) based on the exploration weight w. +- `node::MCTSNode` + The MCTS node whose children will be evaluated. +- `w::AbstractFloat` + Exploration weight. Typical values range from 1.0 to 2.0. A value of 1.0 balances + exploration and exploitation equally. Higher values (e.g., 2.0) encourage more + exploration of less-visited nodes. -# Example -```jldoctest -julia> +# Return +- `selectedNode::MCTSNode` + The child node with the highest UCT score. Returns `nothing` if the node has no + children (though this would indicate an error since UCTselect is called on non-leaves). + +# The UCT Formula + +``` +UCT(s,a) = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a)) + +Where: + Q(s,a) = childNode.statevalue (exploitation: accumulated reward) + c = w (explorationweight) (controls exploration vs exploitation) + N(s) = node.visits (parent visits - total visits to parent) + N(s,a) = childNode.visits (child visits - visits to this specific action) ``` -# Signature +# Behavior + +| Child visits | Exploration term | Behavior | +|-------------|------------------|----------| +| 0 (never visited) | Undefined | Uses `progressvalue` (LLM heuristic) | +| Low (few visits) | High | Encourages exploring new branches | +| High (many visits) | Near 0 | Exploits known good branches | + +# Examples +```jldoctest +julia> using LLMMCTS +julia> child1 = MCTSNode("a", Dict(:reward=>5.0), 0, 10, 50, 0, false, nothing, Dict(), Dict()) +julia> child2 = MCTSNode("b", Dict(:reward=>6.0), 0, 5, 30, 0, false, nothing, Dict(), Dict()) +julia> parent = MCTSNode("root", Dict(:reward=>0.0), 0, 15, 100, 0, false, nothing, + Dict("a"=>child1, "b"=>child2), Dict()) +julia> selected = UCTselect(parent, 1.0) +MCTSNode(...) +``` """ function UCTselect(node::MCTSNode, w::T)::MCTSNode where {T<:AbstractFloat} + # UCT (Upper Confidence Bound for Trees) selects the best child using: + # UCT = statevalue + exploration_weight * sqrt(ln(parent_visits) / child_visits) + # + # The two terms balance: + # - Exploitation (statevalue): choose children that performed well in simulations + # - Exploration (sqrt term): encourage trying less-visited children + # + # The exploration weight `w` controls this balance: + # - w=1.0: equal emphasis on exploration and exploitation + # - w>1.0: more aggressive exploration (try new branches) + # - w<1.0: more exploitation (stick with known good branches) + maxUCT = -Inf selectedNode = nothing for (childState, childNode) in node.children + # Calculate UCT value for this child UCTvalue = if childNode.visits != 0 - weightedterm = w * sqrt(log(node.visits) / childNode.visits) # explore term - childNode.statevalue + weightedterm - else # node.visits == 0 makes sqrt() in explore term error - childNode.progressvalue # exploit term + # Child has been visited before - use statevalue with exploration bonus + # Exploration bonus = w * sqrt(ln(parent_visits) / child_visits) + # High child_visits = small bonus (exploitation dominates) + # Low child_visits = large bonus (encourages exploration) + weightedterm = w * sqrt(log(node.visits) / childNode.visits) + UCTvalue = childNode.statevalue + weightedterm + else + # Child has never been visited - exploration term undefined + # Fall back to progressvalue (LLM heuristic) as exploitation term + # This allows LLM guidance to direct early search + UCTvalue = childNode.progressvalue end - + if UCTvalue > maxUCT maxUCT = UCTvalue - selectedNode = childNode + selectedNode = childNode end end return selectedNode -end +end diff --git a/test/runtest.jl b/test/runtests.jl similarity index 100% rename from test/runtest.jl rename to test/runtests.jl