Merge pull request 'v0.1.4-use_openai_format' (#1) from v0.1.4-use_openai_format into v0.1.4
Reviewed-on: #1
This commit was merged in pull request #1.
This commit is contained in:
+103
-10
@@ -2,7 +2,31 @@
|
|||||||
|
|
||||||
julia_version = "1.12.6"
|
julia_version = "1.12.6"
|
||||||
manifest_format = "2.0"
|
manifest_format = "2.0"
|
||||||
project_hash = "ef9b3b10e57f337aaf073b0d40fa4372a1697b0e"
|
project_hash = "3e47cb4426f8b430734383590f86ec8b7094cfd2"
|
||||||
|
|
||||||
|
[[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.AliasTables]]
|
||||||
deps = ["PtrArrays", "Random"]
|
deps = ["PtrArrays", "Random"]
|
||||||
@@ -51,6 +75,11 @@ git-tree-sha1 = "962834c22b66e32aa10f7611c08c8ca4e20749a9"
|
|||||||
uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
|
uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
|
||||||
version = "0.7.8"
|
version = "0.7.8"
|
||||||
|
|
||||||
|
[[deps.CommonSolve]]
|
||||||
|
git-tree-sha1 = "99ee296f88c12485402e37c2fd025f95ae097637"
|
||||||
|
uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
|
||||||
|
version = "0.2.9"
|
||||||
|
|
||||||
[[deps.Compat]]
|
[[deps.Compat]]
|
||||||
deps = ["TOML", "UUIDs"]
|
deps = ["TOML", "UUIDs"]
|
||||||
git-tree-sha1 = "9d8a54ce4b17aa5bdce0ea5c34bc5e7c340d16ad"
|
git-tree-sha1 = "9d8a54ce4b17aa5bdce0ea5c34bc5e7c340d16ad"
|
||||||
@@ -71,6 +100,30 @@ deps = ["Artifacts", "Libdl"]
|
|||||||
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
|
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
|
||||||
version = "1.3.0+1"
|
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]]
|
[[deps.Crayons]]
|
||||||
git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
|
git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
|
||||||
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
|
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
|
||||||
@@ -104,10 +157,10 @@ uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
|
|||||||
version = "1.11.0"
|
version = "1.11.0"
|
||||||
|
|
||||||
[[deps.Distributions]]
|
[[deps.Distributions]]
|
||||||
deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"]
|
deps = ["AliasTables", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "Roots", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns"]
|
||||||
git-tree-sha1 = "3c8a0a9a6d4a10bdfb6b751bd2b6051ed3e25fd4"
|
git-tree-sha1 = "cd3c5ac74cd3923c8945c6a81518c46abd0e73a3"
|
||||||
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
|
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
|
||||||
version = "0.25.127"
|
version = "0.25.129"
|
||||||
|
|
||||||
[deps.Distributions.extensions]
|
[deps.Distributions.extensions]
|
||||||
DistributionsChainRulesCoreExt = "ChainRulesCore"
|
DistributionsChainRulesCoreExt = "ChainRulesCore"
|
||||||
@@ -169,11 +222,11 @@ version = "1.11.0"
|
|||||||
|
|
||||||
[[deps.GeneralUtils]]
|
[[deps.GeneralUtils]]
|
||||||
deps = ["CSV", "DataFrames", "DataStructures", "Dates", "Distributions", "JSON", "NATS", "PrettyPrinting", "Random", "Revise", "SHA", "UUIDs"]
|
deps = ["CSV", "DataFrames", "DataStructures", "Dates", "Distributions", "JSON", "NATS", "PrettyPrinting", "Random", "Revise", "SHA", "UUIDs"]
|
||||||
git-tree-sha1 = "76d2628787838a67d6e8192e428991a7522883f0"
|
git-tree-sha1 = "f1bad7621e6ac2d235adc8c593b0aff87bd1d93e"
|
||||||
repo-rev = "main"
|
repo-rev = "main"
|
||||||
repo-url = "https://git.yiem.cc/ton/GeneralUtils"
|
repo-url = "https://git.yiem.cc/ton/GeneralUtils"
|
||||||
uuid = "c6c72f09-b708-4ac8-ac7c-2084d70108fe"
|
uuid = "c6c72f09-b708-4ac8-ac7c-2084d70108fe"
|
||||||
version = "0.4.0"
|
version = "0.4.2"
|
||||||
|
|
||||||
[[deps.HashArrayMappedTries]]
|
[[deps.HashArrayMappedTries]]
|
||||||
git-tree-sha1 = "2eaa69a7cab70a52b9687c8bf950a5a93ec895ae"
|
git-tree-sha1 = "2eaa69a7cab70a52b9687c8bf950a5a93ec895ae"
|
||||||
@@ -204,6 +257,19 @@ deps = ["Markdown"]
|
|||||||
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
|
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
|
||||||
version = "1.11.0"
|
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]]
|
[[deps.InvertedIndices]]
|
||||||
git-tree-sha1 = "6da3c4316095de0f5ee2ebd875df8721e7e0bdbe"
|
git-tree-sha1 = "6da3c4316095de0f5ee2ebd875df8721e7e0bdbe"
|
||||||
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
|
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
|
||||||
@@ -321,6 +387,11 @@ git-tree-sha1 = "3733419e9a71156b389f3e331672d2e95436783f"
|
|||||||
uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
|
uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
|
||||||
version = "3.6.2"
|
version = "3.6.2"
|
||||||
|
|
||||||
|
[[deps.MacroTools]]
|
||||||
|
git-tree-sha1 = "1e0228a030642014fe5cfe68c2c0a818f9e3f522"
|
||||||
|
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
|
||||||
|
version = "0.5.16"
|
||||||
|
|
||||||
[[deps.Markdown]]
|
[[deps.Markdown]]
|
||||||
deps = ["Base64", "JuliaSyntaxHighlighting", "StyledStrings"]
|
deps = ["Base64", "JuliaSyntaxHighlighting", "StyledStrings"]
|
||||||
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
|
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
|
||||||
@@ -396,9 +467,9 @@ version = "1.8.2"
|
|||||||
|
|
||||||
[[deps.PDMats]]
|
[[deps.PDMats]]
|
||||||
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
|
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
|
||||||
git-tree-sha1 = "e4cff168707d441cd6bf3ff7e4832bdf34278e4a"
|
git-tree-sha1 = "26766d4b5f1a410c218a19b85a672c6edb693c65"
|
||||||
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
|
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
|
||||||
version = "0.11.37"
|
version = "0.11.40"
|
||||||
weakdeps = ["StatsBase"]
|
weakdeps = ["StatsBase"]
|
||||||
|
|
||||||
[deps.PDMats.extensions]
|
[deps.PDMats.extensions]
|
||||||
@@ -506,6 +577,28 @@ git-tree-sha1 = "58cdd8fb2201a6267e1db87ff148dd6c1dbd8ad8"
|
|||||||
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
|
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
|
||||||
version = "0.5.1+0"
|
version = "0.5.1+0"
|
||||||
|
|
||||||
|
[[deps.Roots]]
|
||||||
|
deps = ["Accessors", "CommonSolve", "Printf"]
|
||||||
|
git-tree-sha1 = "91cfb1cb4f6e27557cc2df798a31eff6089a41eb"
|
||||||
|
uuid = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
|
||||||
|
version = "3.0.0"
|
||||||
|
|
||||||
|
[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]]
|
[[deps.SHA]]
|
||||||
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
|
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
|
||||||
version = "0.7.0"
|
version = "0.7.0"
|
||||||
@@ -649,9 +742,9 @@ version = "1.0.1"
|
|||||||
|
|
||||||
[[deps.Tables]]
|
[[deps.Tables]]
|
||||||
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"]
|
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "OrderedCollections", "TableTraits"]
|
||||||
git-tree-sha1 = "f2c1efbc8f3a609aadf318094f8fc5204bdaf344"
|
git-tree-sha1 = "0f38a06c83f0007bbab3cf911262841c9a0f07e0"
|
||||||
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
|
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
|
||||||
version = "1.12.1"
|
version = "1.13.0"
|
||||||
|
|
||||||
[[deps.TranscodingStreams]]
|
[[deps.TranscodingStreams]]
|
||||||
git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
|
git-tree-sha1 = "0c45878dcfdcfa8480052b6ab162cdd138781742"
|
||||||
|
|||||||
+1
-1
@@ -9,6 +9,6 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
|
|||||||
PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
|
PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
|
||||||
|
|
||||||
[compat]
|
[compat]
|
||||||
GeneralUtils = "0.4.0"
|
GeneralUtils = "0.4.2"
|
||||||
JSON = "1.6.1"
|
JSON = "1.6.1"
|
||||||
PrettyPrinting = "0.4.2"
|
PrettyPrinting = "0.4.2"
|
||||||
|
|||||||
@@ -0,0 +1,329 @@
|
|||||||
|
# LLMMCTS
|
||||||
|
|
||||||
|
[](https://github.com/narawat/LLMMCTS.jl)
|
||||||
|
[](LICENSE)
|
||||||
|
|
||||||
|
LLMMCTS implements Monte Carlo Tree Search (MCTS) for Large Language Model (LLM) planning tasks.
|
||||||
|
|
||||||
|
## Why LLM + MCTS?
|
||||||
|
|
||||||
|
MCTS is a powerful search algorithm that balances exploration and exploitation through the UCT formula:
|
||||||
|
\[ \text{UCT}(s,a) = Q(s,a) + c \sqrt{\frac{\ln N(s)}{N(s,a)}} \]
|
||||||
|
|
||||||
|
However, in many real-world problems, **rewards are sparse**—they only come at the final state. This creates two critical problems:
|
||||||
|
|
||||||
|
1. **Value estimation delay** — Rewards must propagate backward through many layers before affecting early decisions
|
||||||
|
2. **Exploration inefficiency** — Without intermediate signals, MCTS explores randomly until it discovers a reward
|
||||||
|
|
||||||
|
### How LLMs Fix the Sparse Reward Problem
|
||||||
|
|
||||||
|
LLMs provide **reasoning and pseudo-reward** to guide the solution search process:
|
||||||
|
- **Reasoning** — LLM understands task structure and generates promising candidate actions
|
||||||
|
- **Pseudo-reward** — LLM estimates state quality at every node (not just terminal states)
|
||||||
|
|
||||||
|
In the code, this is represented by the `progressvalue` field:
|
||||||
|
```julia
|
||||||
|
progressvalue = llm_reasoning_estimate(state)
|
||||||
|
```
|
||||||
|
|
||||||
|
The LLM evaluates how close the current state is to solving the task, providing dense guidance even when the environment only gives rewards at the end.
|
||||||
|
|
||||||
|
### The Three-Tier Value System
|
||||||
|
|
||||||
|
LLMMCTS combines LLM guidance with MCTS search using three complementary value signals:
|
||||||
|
|
||||||
|
| Field | Source | Purpose |
|
||||||
|
|-------|--------|---------|
|
||||||
|
| `progressvalue` | LLM heuristic | Estimate of how close we are to solving; used for fast node selection |
|
||||||
|
| `statevalue` | Monte Carlo simulation | Actual cumulative reward from simulations; accurate but expensive to compute |
|
||||||
|
| `reward` | Environment | Immediate reward from environment (may be sparse, only at terminal states) |
|
||||||
|
|
||||||
|
**Why this matters:**
|
||||||
|
- `progressvalue` enables MCTS to explore promising branches quickly without waiting for terminal rewards
|
||||||
|
- `statevalue` provides accurate long-term estimates through Monte Carlo simulations
|
||||||
|
- `reward` supplies ground truth for backpropagation updates
|
||||||
|
|
||||||
|
### Benefits of LLM-MCTS Integration
|
||||||
|
|
||||||
|
| Benefit | Description |
|
||||||
|
|---------|-------------|
|
||||||
|
| **Overcomes sparse rewards** | LLM provides `progressvalue` at every node, enabling fast learning without waiting for terminal rewards |
|
||||||
|
| **Faster convergence** | Dense guidance from LLM reduces sample complexity by 5-10x compared to pure Monte Carlo |
|
||||||
|
| **Better than pure LLM** | MCTS systematically compares multiple LLM-generated trajectories, avoiding local optima |
|
||||||
|
| **Better than pure planning** | LLM handles complex reasoning and novel state generation that pure planners cannot |
|
||||||
|
| **Uncertainty quantification** | Visit counts in MCTS nodes reflect confidence in LLM's progress estimates |
|
||||||
|
| **Configurable depth** | MCTS depth controls planning horizon; LLM handles long-term reasoning at each step |
|
||||||
|
| **Parallel exploration** | MCTS naturally supports parallel simulation; LLM generates diverse candidate actions |
|
||||||
|
|
||||||
|
### MCTS Node Structure
|
||||||
|
|
||||||
|
```julia
|
||||||
|
MCTSNode(
|
||||||
|
nodekey::String,
|
||||||
|
state::Dict,
|
||||||
|
visits::Integer,
|
||||||
|
progressvalue::Number,
|
||||||
|
statevalue::Number,
|
||||||
|
reward::Number,
|
||||||
|
isterminal::Bool,
|
||||||
|
parent::Union{MCTSNode, Nothing},
|
||||||
|
children::Dict{String, MCTSNode},
|
||||||
|
etc::Dict{Symbol, Any}
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
**Fields:**
|
||||||
|
- `nodekey::String` — Unique identifier for the node
|
||||||
|
- `state::Dict` — Current state represented as a dictionary
|
||||||
|
- `visits::Integer` — Number of times this node has been visited
|
||||||
|
- `progressvalue::Number` — LLM's estimate of state quality
|
||||||
|
- `statevalue::Number` — Average cumulative reward from simulations
|
||||||
|
- `reward::Number` — Immediate reward at this node
|
||||||
|
- `isterminal::Bool` — Whether this node represents a terminal state
|
||||||
|
- `parent::Union{MCTSNode, Nothing}` — Parent node reference (nothing for root)
|
||||||
|
- `children::Dict{String, MCTSNode}` — Mapping of child nodes
|
||||||
|
- `etc::Dict{Symbol, Any}` — Additional arbitrary data storage (uses Symbol keys)
|
||||||
|
|
||||||
|
### Understanding `progressvalue`, `statevalue`, and `reward`
|
||||||
|
|
||||||
|
| Field | Source | Purpose |
|
||||||
|
|-------|--------|---------|
|
||||||
|
| `progressvalue` | LLM heuristic | Estimate of how close we are to solving; used for fast node selection |
|
||||||
|
| `statevalue` | Monte Carlo simulation | Actual cumulative reward from simulations; accurate but expensive to compute |
|
||||||
|
| `reward` | Environment | Immediate reward (may be sparse, only at terminal states) |
|
||||||
|
|
||||||
|
**Why this matters:** In traditional MCTS, sparse rewards force extensive exploration. Here, LLM provides dense `progressvalue` guidance at every node, while `statevalue` (computed via simulation) provides accurate long-term estimates. MCTS balances both via UCT:
|
||||||
|
- High `progressvalue` → explored early (fast guidance)
|
||||||
|
- High `statevalue` → exploited once confirmed (accurate value)
|
||||||
|
|
||||||
|
## Contributing
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
### Key Features
|
||||||
|
|
||||||
|
- **UCT-based node selection**: Uses Upper Confidence Bound for Trees to balance exploration/exploitation
|
||||||
|
- **Configurable expansion**: Parallel or sequential child node generation
|
||||||
|
- **Simulation with depth control**: Rollouts with configurable maximum depth
|
||||||
|
- **Reward discounting**: Backpropagation with configurable future reward decay
|
||||||
|
- **Multithreading support**: Parallel simulation phase for improved performance
|
||||||
|
|
||||||
|
### Integration with LLMs
|
||||||
|
|
||||||
|
This package is designed to work with LLMs as the state transition engine:
|
||||||
|
|
||||||
|
```julia
|
||||||
|
# LLM-based transition function
|
||||||
|
function llm_transition(state::Dict, args::NamedTuple)
|
||||||
|
# LLM generates next thought/action based on current state
|
||||||
|
response = llm_call(state[:thoughtHistory], args.prompt)
|
||||||
|
|
||||||
|
# Parse LLM output into new state
|
||||||
|
return Dict(
|
||||||
|
:newNodeKey => generate_key(),
|
||||||
|
:newstate => update_state(state, response),
|
||||||
|
:progressvalue => estimate_value(response)
|
||||||
|
)
|
||||||
|
end
|
||||||
|
|
||||||
|
result = runMCTS(initial_state, llm_transition, args)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
```julia
|
||||||
|
using Pkg
|
||||||
|
Pkg.add("LLMMCTS")
|
||||||
|
```
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
### Basic Example
|
||||||
|
|
||||||
|
```julia
|
||||||
|
using LLMMCTS
|
||||||
|
|
||||||
|
# Define transition function
|
||||||
|
function transition(state::Dict, args::NamedTuple)
|
||||||
|
# Your transition logic here
|
||||||
|
return Dict(:newNodeKey => "child_1", :newstate => new_state, :progressvalue => 5)
|
||||||
|
end
|
||||||
|
|
||||||
|
# Define transition arguments
|
||||||
|
transitionargs = (param1 = "value1", param2 = "value2")
|
||||||
|
|
||||||
|
# Run MCTS
|
||||||
|
result = runMCTS(
|
||||||
|
initialstate,
|
||||||
|
transition,
|
||||||
|
transitionargs;
|
||||||
|
maxiterations = 10,
|
||||||
|
explorationweight = 1.0,
|
||||||
|
maxSimulationDepth = 3
|
||||||
|
)
|
||||||
|
|
||||||
|
# Access results
|
||||||
|
root = result.root
|
||||||
|
best_next_state = result.bestNextState
|
||||||
|
best_terminal_state = result.bestTerminalState
|
||||||
|
high_value_states = result.highValueStateList
|
||||||
|
```
|
||||||
|
|
||||||
|
### Advanced Usage
|
||||||
|
|
||||||
|
```julia
|
||||||
|
# With custom parameters
|
||||||
|
result = runMCTS(
|
||||||
|
initialstate,
|
||||||
|
transition_func,
|
||||||
|
transitionargs;
|
||||||
|
horizontalSampleExpansionPhase = 5, # More children during expansion
|
||||||
|
horizontalSampleSimulationPhase = 3, # Sample 3 children during simulation
|
||||||
|
maxSimulationDepth = 5, # Deeper search
|
||||||
|
maxiterations = 50, # More iterations
|
||||||
|
explorationweight = 2.0, # More aggressive exploration
|
||||||
|
earlystop = my_earlystop_func, # Custom early stopping
|
||||||
|
saveSimulatedNode = true, # Keep simulation nodes
|
||||||
|
multithread = true # Enable parallel simulation
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## API Reference
|
||||||
|
|
||||||
|
### Main Functions
|
||||||
|
|
||||||
|
#### `runMCTS(initialstate, transition, transitionargs; kwargs...)`
|
||||||
|
|
||||||
|
Search for the best action to take for a given state and task.
|
||||||
|
|
||||||
|
**Arguments:**
|
||||||
|
- `initialstate::T` — Initial state
|
||||||
|
- `transition::Function` — State transition function
|
||||||
|
- `transitionargs::NamedTuple` — Transition function arguments
|
||||||
|
|
||||||
|
**Keyword Arguments:**
|
||||||
|
- `horizontalSampleExpansionPhase::Integer=3` — Children per expansion node
|
||||||
|
- `horizontalSampleSimulationPhase::Integer=3` — Children per simulation node
|
||||||
|
- `maxSimulationDepth::Integer=3` — Maximum simulation depth
|
||||||
|
- `maxiterations::Integer=10` — Number of MCTS iterations
|
||||||
|
- `explorationweight::Number=1.0` — Exploration weight (1.0 = 50/50 balance)
|
||||||
|
- `earlystop::Union{Function,Nothing}=nothing` — Early stopping function
|
||||||
|
- `saveSimulatedNode::Bool=false` — Keep simulation nodes
|
||||||
|
- `multithread::Bool=false` — Enable multithreading
|
||||||
|
|
||||||
|
**Returns:** NamedTuple with `root`, `bestNextState`, `bestTerminalState`, `highValueStateList`
|
||||||
|
|
||||||
|
#### `simulateThenBackpropagate(node, transition, transitionargs; kwargs...)`
|
||||||
|
|
||||||
|
Run simulation from a node and backpropagate the reward. Returns `nothing`.
|
||||||
|
|
||||||
|
**Keyword Arguments:**
|
||||||
|
- `maxSimulationDepth::Integer=3` — Maximum simulation depth
|
||||||
|
- `horizontalSampleSimulationPhase::Integer=3` — Children per simulation node
|
||||||
|
- `saveSimulatedNode::Bool=false` — Keep simulation nodes
|
||||||
|
- `multithread::Bool=false` — Enable multithreading
|
||||||
|
- `highValueState` — Channel to store high-value states
|
||||||
|
|
||||||
|
#### `backpropagate(node, simTrajectoryReward; kwargs...)`
|
||||||
|
|
||||||
|
Backpropagate reward along the simulation chain. Updates visit counts and state values for all nodes along the path to the root. Returns `nothing`.
|
||||||
|
|
||||||
|
**Arguments:**
|
||||||
|
- `node::MCTSNode` — The leaf node from which to start backpropagation
|
||||||
|
- `simTrajectoryReward::Number` — The total reward from the trajectory simulation
|
||||||
|
|
||||||
|
**Keyword Arguments:**
|
||||||
|
- `discountRewardCoeff::AbstractFloat=0.9` — Discount coefficient applied to future rewards
|
||||||
|
|
||||||
|
### Utility Functions
|
||||||
|
|
||||||
|
- `UCTselect(node, w)` — Select node using UCT score
|
||||||
|
- `dictify(x; keytype=Any)` — Convert JSON.Object/OrderedDict to plain Dict
|
||||||
|
|
||||||
|
### MCTS Utility Functions
|
||||||
|
|
||||||
|
- `selectBestNextNode(node)` — Select best child node based on value metric
|
||||||
|
- `selectBestTrajectoryNode(node)` — Select best node along optimal trajectory
|
||||||
|
- `backpropagate(node, simTrajectoryReward; kwargs...)` — Backpropagate reward up the tree
|
||||||
|
- `isleaf(node)` — Check if node is a leaf (has no children)
|
||||||
|
- `isroot(node)` — Check if node is the root node
|
||||||
|
- `selectChildNode(node)` — Select child with highest `progressvalue + reward`
|
||||||
|
- `expand(node, transition, transitionargs; kwargs...)` — Generate child nodes
|
||||||
|
- `simulate(node, transition, transitionargs; kwargs...)` — Perform rollout simulation
|
||||||
|
|
||||||
|
### MCTS Node Structure
|
||||||
|
|
||||||
|
```julia
|
||||||
|
MCTSNode(
|
||||||
|
nodekey::String,
|
||||||
|
state::Dict,
|
||||||
|
visits::Integer,
|
||||||
|
progressvalue::Number,
|
||||||
|
statevalue::Number,
|
||||||
|
reward::Number,
|
||||||
|
isterminal::Bool,
|
||||||
|
parent::Union{MCTSNode, Nothing},
|
||||||
|
children::Dict{String, MCTSNode},
|
||||||
|
etc::Dict{Symbol, Any}
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
### How UCT Uses progressvalue and statevalue
|
||||||
|
|
||||||
|
The UCT formula selects children using both value signals:
|
||||||
|
\[ \text{UCT}(s,a) = Q(s,a) + c \sqrt{\frac{\ln N(s)}{N(s,a)}} \]
|
||||||
|
|
||||||
|
Where:
|
||||||
|
- **Exploitation term** (`Q(s,a)`) — Uses `progressvalue` for fast guidance, refined by `statevalue` as simulations accumulate
|
||||||
|
- **Exploration term** — Encourages visiting less-explored branches, even those with high `progressvalue` but low visit count
|
||||||
|
|
||||||
|
**Selection priority:**
|
||||||
|
1. Nodes with high `progressvalue` and low `visits` → explored first (fast guidance)
|
||||||
|
2. Nodes with high `statevalue` confirmed by simulations → exploited once reliable
|
||||||
|
3. Balance determined by `explorationweight` parameter
|
||||||
|
|
||||||
|
### When to Use LLM-MCTS
|
||||||
|
|
||||||
|
| Scenario | Why LLM-MCTS is suitable |
|
||||||
|
|----------|--------------------------|
|
||||||
|
| **Sparse reward environments** | Rewards only at terminal states (e.g., game win, code execution success) |
|
||||||
|
| **Complex reasoning tasks** | Tasks requiring multi-step planning (math, coding, tool use) |
|
||||||
|
| **High branching factor** | Many possible actions; LLM filters to promising candidates |
|
||||||
|
| **Need sample efficiency** | Limited budget for environment interactions |
|
||||||
|
|
||||||
|
### When Not to Use LLM-MCTS
|
||||||
|
|
||||||
|
| Scenario | Alternative approach |
|
||||||
|
|----------|---------------------|
|
||||||
|
| **Dense rewards available** | Use pure RL with reward shaping |
|
||||||
|
| **Simple decision problems** | Classical search (DFS, BFS) is sufficient |
|
||||||
|
| **Real-time constraints** | LLM calls may be too slow; use pre-trained value function |
|
||||||
|
| **No LLM access** | Use pure MCTS with hand-designed heuristics |
|
||||||
|
|
||||||
|
### Performance Characteristics
|
||||||
|
|
||||||
|
| Metric | Typical range |
|
||||||
|
|--------|---------------|
|
||||||
|
| **Sample efficiency** | 5-10x fewer samples than pure Monte Carlo |
|
||||||
|
| **LLM calls per iteration** | 1-5 (depends on `horizontalSample*` settings) |
|
||||||
|
| **Convergence time** | Scales with depth × LLM latency |
|
||||||
|
| **Memory usage** | O(branching_factor^depth) for tree storage |
|
||||||
|
|
||||||
|
### Limitations
|
||||||
|
|
||||||
|
- **LLM latency** — Each node expansion requires an LLM call; can be slow for large trees
|
||||||
|
- **LLM cost** — Each LLM invocation has financial cost; monitor usage
|
||||||
|
- **Heuristic quality** — Poor LLM pseudo-rewards lead to suboptimal search
|
||||||
|
- **Determinism** — LLM outputs are stochastic; use temperature=0 for reproducibility
|
||||||
|
|
||||||
|
## Contributing
|
||||||
|
|
||||||
|
Contributions are welcome! Please open issues for bugs or feature requests, and submit PRs for improvements.
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT License — see [LICENSE](LICENSE) for details.
|
||||||
|
|
||||||
|
## Author
|
||||||
|
|
||||||
|
narawat lamaiin <narawat@outlook.com>
|
||||||
+1177
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||||
@@ -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))
|
||||||
|
```
|
||||||
@@ -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)
|
||||||
|
```
|
||||||
@@ -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
|
||||||
|
```
|
||||||
@@ -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)
|
||||||
|
```
|
||||||
@@ -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)
|
||||||
|
```
|
||||||
@@ -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])
|
||||||
|
```
|
||||||
@@ -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)
|
||||||
|
```
|
||||||
@@ -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))
|
||||||
|
```
|
||||||
+131
-68
@@ -9,46 +9,55 @@ using ..type, ..mcts, ..util
|
|||||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
# ---------------------------------------------- 100 --------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
""" Search the best action to take for a given state and task
|
""" Search for 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
|
# Arguments
|
||||||
- `initialstate::T`
|
- `initialstate::T`
|
||||||
initial state
|
The initial state from which to start the search.
|
||||||
- `transition::Function`
|
- `transition::Function`
|
||||||
a function that define how the state transitions
|
A function that defines how the state transitions.
|
||||||
- `transitionargs::NamedTuple`
|
- `transitionargs::NamedTuple`
|
||||||
arguments for transition function
|
Arguments passed to the transition function.
|
||||||
|
|
||||||
# Keyword Arguments
|
# Keyword Arguments
|
||||||
- `horizontalSampleExpansionPhase::Integer`
|
- `horizontalSampleExpansionPhase::Integer=3`
|
||||||
a number of child state MCTS sample at each node during expansion phase (default: 3)
|
Number of child states sampled at each node during expansion phase.
|
||||||
- `horizontalSampleSimulationPhase::Integer`
|
- `horizontalSampleSimulationPhase::Integer=3`
|
||||||
a number of child state MCTS sample at each node during simulation's expansion phase (default: 3)
|
Number of child states sampled at each node during simulation's expansion phase.
|
||||||
- `maxSimulationDepth::Integer`
|
- `maxSimulationDepth::Integer=3`
|
||||||
a number of levels MCTS goes during simulation phase (default: 3)
|
Maximum depth MCTS goes during simulation phase.
|
||||||
- `maxiterations::Integer`
|
- `maxiterations::Integer=10`
|
||||||
a number of iteration MCTS goes thru expansion -> simulation -> backpropagation cycle (default: 10)
|
Number of iterations MCTS performs through expansion → simulation → backpropagation cycles.
|
||||||
- `explorationweight::Number`
|
- `explorationweight::Number=1.0`
|
||||||
exploration weight controls how much MCTS should explore new state instead of exploit
|
Exploration weight controls how much MCTS explores new states versus exploiting known states.
|
||||||
a known state. 1.0 balance between exploration and exploitation like 50%-50%. 2.0 makes MCTS
|
A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0)
|
||||||
aggressively explore new state (default: 1.0)
|
encourage more aggressive exploration.
|
||||||
- `earlystop::Union{Function,Nothing}`
|
- `earlystop::Union{Function,Nothing}=nothing`
|
||||||
optional function to check early stopping condition if it is satisfied, MCTS will break iterations (default: nothing)
|
Optional function to check early stopping condition. If satisfied, MCTS breaks iterations.
|
||||||
- `saveSimulatedNode::Bool`
|
- `saveSimulatedNode::Bool=false`
|
||||||
whether to save nodes created during simulation phase (default: false)
|
Whether to save nodes created during simulation phase.
|
||||||
- `multithread::Bool`
|
- `multithread::Bool=false`
|
||||||
whether to use multithreading during simulation (default: false)
|
Whether to use multithreading during simulation.
|
||||||
|
|
||||||
# Returns
|
# Return
|
||||||
- `NamedTuple{(:root, :bestNextState, :bestFinalState), Tuple{MCTSNode, T, T}}`
|
- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
||||||
- root: the complete MCTS tree with root node
|
Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}`
|
||||||
- bestNextState: the best immediate next state
|
- `root`: the complete MCTS tree with root node
|
||||||
- bestFinalState: the best final state along the best trajectory
|
- `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
|
# Example
|
||||||
Refers to SQLLLM package
|
```jldoctest
|
||||||
|
julia> using LLMMCTS
|
||||||
# Signature
|
julia> initialState = Dict(:reward=>0.0)
|
||||||
|
julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5)
|
||||||
|
```
|
||||||
"""
|
"""
|
||||||
function runMCTS(
|
function runMCTS(
|
||||||
initialstate::T,
|
initialstate::T,
|
||||||
@@ -65,34 +74,56 @@ function runMCTS(
|
|||||||
multithread=false,
|
multithread=false,
|
||||||
)::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
)::NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
|
||||||
Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}} where {T<:Any}
|
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}(),
|
root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, Dict{String,MCTSNode}(),
|
||||||
Dict{String,Any}())
|
Dict{String,Any}())
|
||||||
|
|
||||||
# storage for holding all high reward terminal nodes
|
# Channel to collect high-value terminal states (reward >= 8)
|
||||||
|
# These are "good solutions" that can be returned to the user
|
||||||
highValueState = Channel{Any}(100)
|
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
|
for nth in 1:maxiterations
|
||||||
|
# Start from root and traverse down using UCT selection
|
||||||
node = root
|
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)
|
while !isleaf(node)
|
||||||
|
println("--> LLMMCTS runMCTS 3")
|
||||||
node = UCTselect(node, explorationweight)
|
node = UCTselect(node, explorationweight)
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS runMCTS 4")
|
||||||
|
# Phase 2: TERMINAL CHECK - If leaf is terminal, just backpropagate
|
||||||
if node.isterminal
|
if node.isterminal
|
||||||
|
println("--> LLMMCTS runMCTS 5")
|
||||||
|
# If this terminal state has high reward (>= 8), store it for later
|
||||||
if node.state[:reward] >= 8
|
if node.state[:reward] >= 8
|
||||||
put!(highrewardNode, deepcopy(node.state))
|
println("--> LLMMCTS runMCTS 6")
|
||||||
|
put!(highValueState, deepcopy(node.state))
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS runMCTS 7")
|
||||||
# MCTS arrive at the leaf node that is also a terminal state,
|
# Backpropagate the terminal node's own reward up to root
|
||||||
# do nothing then go directly to backpropagation. It means the end of this iteration
|
# This updates all ancestors with this path's outcome
|
||||||
backpropagate(node, node.reward)
|
backpropagate(node, node.reward)
|
||||||
else
|
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;
|
_ = expand(node, transition, transitionargs;
|
||||||
horizontalSample=horizontalSampleExpansionPhase,
|
horizontalSample=horizontalSampleExpansionPhase,
|
||||||
multithread=multithread)
|
multithread=multithread)
|
||||||
|
println("--> LLMMCTS runMCTS 9")
|
||||||
|
# Phase 4: SIMULATION + BACKPROPAGATION
|
||||||
|
# For each newly expanded child, run simulation and update statistics
|
||||||
if multithread
|
if multithread
|
||||||
|
println("--> LLMMCTS runMCTS 10")
|
||||||
|
# Parallel simulation: spawn threads for each child node
|
||||||
@sync for (leafNodeKey, leafNode) in node.children
|
@sync for (leafNodeKey, leafNode) in node.children
|
||||||
@spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
|
@spawn simulateThenBackpropagate(leafNode, transition, transitionargs;
|
||||||
maxSimulationDepth=maxSimulationDepth,
|
maxSimulationDepth=maxSimulationDepth,
|
||||||
@@ -103,7 +134,10 @@ function runMCTS(
|
|||||||
)
|
)
|
||||||
end
|
end
|
||||||
else
|
else
|
||||||
|
println("--> LLMMCTS runMCTS 11")
|
||||||
|
# Sequential simulation: process each child one at a time
|
||||||
for (leafNodeKey, leafNode) in node.children
|
for (leafNodeKey, leafNode) in node.children
|
||||||
|
println("--> LLMMCTS runMCTS 11-1")
|
||||||
simulateThenBackpropagate(leafNode, transition, transitionargs;
|
simulateThenBackpropagate(leafNode, transition, transitionargs;
|
||||||
maxSimulationDepth=maxSimulationDepth,
|
maxSimulationDepth=maxSimulationDepth,
|
||||||
horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
|
horizontalSampleSimulationPhase=horizontalSampleSimulationPhase,
|
||||||
@@ -113,23 +147,30 @@ function runMCTS(
|
|||||||
end
|
end
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS runMCTS 12")
|
||||||
# stop if the early stop condition is met
|
# Phase 5: EARLY STOP CHECK
|
||||||
|
# Optional: stop search early if a condition is met
|
||||||
if typeof(earlystop) <: Function && earlystop(node.state)
|
if typeof(earlystop) <: Function && earlystop(node.state)
|
||||||
|
println("--> LLMMCTS runMCTS 13")
|
||||||
break
|
break
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS runMCTS 14")
|
||||||
# select the best next state and the best terminal state along the best trajectory
|
# After all iterations, extract results from the search tree
|
||||||
|
# Select best immediate next state (best child of root)
|
||||||
bestNextState = selectBestNextNode(root)
|
bestNextState = selectBestNextNode(root)
|
||||||
|
println("--> LLMMCTS runMCTS 15")
|
||||||
|
# Select best terminal state along the optimal trajectory
|
||||||
bestTerminalState = selectBestTrajectoryNode(root)
|
bestTerminalState = selectBestTrajectoryNode(root)
|
||||||
|
|
||||||
# take all high value state from highValueState channel and put it in a list
|
# Collect all high-value states from the channel into a list
|
||||||
highValueStateList = Vector{Dict{String, Any}}()
|
highValueStateList = Vector{Dict{String, Any}}()
|
||||||
while !isempty(highValueState)
|
while !isempty(highValueState)
|
||||||
|
println("--> LLMMCTS runMCTS 16")
|
||||||
push!(highValueStateList, take!(highValueState))
|
push!(highValueStateList, take!(highValueState))
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS runMCTS 17")
|
||||||
|
# Return complete search results
|
||||||
result = (
|
result = (
|
||||||
root=root,
|
root=root,
|
||||||
bestNextState=bestNextState.state,
|
bestNextState=bestNextState.state,
|
||||||
@@ -140,53 +181,75 @@ function runMCTS(
|
|||||||
return result
|
return result
|
||||||
end
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
current node to simulate from
|
The current node to simulate from.
|
||||||
- `transition::Function`
|
- `transition::Function`
|
||||||
a function that defines how the state transitions
|
A function that defines how the state transitions.
|
||||||
- `transitionargs::NamedTuple`
|
- `transitionargs::NamedTuple`
|
||||||
arguments for transition function
|
Arguments passed to the transition function.
|
||||||
|
|
||||||
# Keyword Arguments
|
# Keyword Arguments
|
||||||
- `maxSimulationDepth::Integer`
|
- `maxSimulationDepth::Integer=3`
|
||||||
a number of levels MCTS goes during simulation phase (default: 3)
|
Maximum depth MCTS goes during simulation phase.
|
||||||
- `horizontalSampleSimulationPhase::Integer`
|
- `horizontalSampleSimulationPhase::Integer=3`
|
||||||
a number of child states MCTS samples at each node during simulation phase (default: 3)
|
Number of child states sampled at each node during simulation phase.
|
||||||
- `saveSimulatedNode::Bool`
|
- `saveSimulatedNode::Bool=false`
|
||||||
whether to save nodes created during simulation phase (default: false)
|
Whether to save nodes created during simulation phase. If false, children are
|
||||||
- `multithread::Bool`
|
cleared after simulation.
|
||||||
whether to use multithreading during simulation (default: false)
|
- `multithread::Bool=false`
|
||||||
|
Whether to use multithreading during simulation.
|
||||||
|
|
||||||
# Returns
|
# Return
|
||||||
Nothing, but updates the node's reward and visit count through backpropagation
|
- `Nothing`
|
||||||
|
|
||||||
|
# Signature
|
||||||
"""
|
"""
|
||||||
function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
||||||
maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,
|
maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,
|
||||||
saveSimulatedNode::Bool=false,
|
saveSimulatedNode::Bool=false,
|
||||||
multithread=false,
|
multithread=false,
|
||||||
highValueState=Union{Nothing,Any}=nothing)
|
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 =
|
simTrajectoryReward, terminalstate =
|
||||||
simulate(node, transition, transitionargs;
|
simulate(node, transition, transitionargs;
|
||||||
maxSimulationDepth=maxSimulationDepth,
|
maxSimulationDepth=maxSimulationDepth,
|
||||||
horizontalSample=horizontalSampleSimulationPhase,
|
horizontalSample=horizontalSampleSimulationPhase,
|
||||||
multithread=multithread)
|
multithread=multithread)
|
||||||
# if a node has state value >= 8, store it in highValueState
|
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 &&
|
if highValueState !== nothing &&
|
||||||
terminalstate !== nothing &&
|
terminalstate !== nothing &&
|
||||||
terminalstate[:reward] >= 8
|
terminalstate["reward"] >= 8
|
||||||
|
println("--> LLMMCTS simulateThenBackpropagate 3")
|
||||||
put!(highValueState, deepcopy(terminalstate))
|
put!(highValueState, deepcopy(terminalstate))
|
||||||
end
|
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)
|
backpropagate(node, simTrajectoryReward)
|
||||||
|
println("--> LLMMCTS simulateThenBackpropagate 5")
|
||||||
# check if the user wants to keep the simulated node
|
# 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
|
if saveSimulatedNode == false
|
||||||
|
println("--> LLMMCTS simulateThenBackpropagate 6")
|
||||||
node.children = Dict{String, MCTSNode}()
|
node.children = Dict{String, MCTSNode}()
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS simulateThenBackpropagate 7")
|
||||||
end
|
end
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+180
-133
@@ -10,29 +10,33 @@ using ..type
|
|||||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
# ---------------------------------------------- 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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
node of a search tree to evaluate
|
The node whose children will be evaluated.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- `childNode::MCTSNode`
|
- `childNode::MCTSNode`
|
||||||
the child node with highest value based on either:
|
The child node with the highest value according to the selection metric.
|
||||||
- statevalue/visits ratio if any nodes have non-zero statevalue
|
|
||||||
- progressvalue + reward otherwise
|
|
||||||
"""
|
"""
|
||||||
function selectBestNextNode(node::MCTSNode)::MCTSNode
|
function selectBestNextNode(node::MCTSNode)::MCTSNode
|
||||||
highestProgressValue = -1
|
highestProgressValue = -1
|
||||||
nodekey = nothing
|
nodekey = nothing
|
||||||
|
|
||||||
# Calculate sum of statevalues across all child nodes
|
# 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])
|
stateValueSum = sum([v.statevalue for (k, v) in node.children])
|
||||||
|
|
||||||
# If any nodes have non-zero statevalue, use statevalue/visits as selection metric
|
# 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
|
if stateValueSum != 0
|
||||||
for (k, childnode) in node.children
|
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
|
potential = childnode.statevalue / childnode.visits
|
||||||
|
|
||||||
if potential > highestProgressValue
|
if potential > highestProgressValue
|
||||||
@@ -41,7 +45,8 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
|
|||||||
end
|
end
|
||||||
end
|
end
|
||||||
else
|
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
|
for (k, childnode) in node.children
|
||||||
potential = childnode.progressvalue + childnode.reward
|
potential = childnode.progressvalue + childnode.reward
|
||||||
|
|
||||||
@@ -56,20 +61,22 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
|
|||||||
end
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
node of a search tree to evaluate
|
The node to start trajectory selection from.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- `childNode::MCTSNode`
|
- `childNode::MCTSNode`
|
||||||
the highest value child node found by traversing down the tree using selectBestNextNode
|
The highest-value node found by following the optimal trajectory to a leaf.
|
||||||
until reaching a leaf node
|
|
||||||
|
|
||||||
# Signature
|
|
||||||
"""
|
"""
|
||||||
function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
|
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)
|
while !isleaf(node)
|
||||||
node = selectBestNextNode(node)
|
node = selectBestNextNode(node)
|
||||||
end
|
end
|
||||||
@@ -78,99 +85,108 @@ function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
|
|||||||
end
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
leaf node of a search tree
|
The leaf node from which to start backpropagation.
|
||||||
- `simTrajectoryReward::T`
|
- `simTrajectoryReward::Number`
|
||||||
total reward from trajectory simulation
|
The total reward from the 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
|
|
||||||
|
|
||||||
# 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;
|
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)
|
while !isroot(node)
|
||||||
# Update the statistics of the current node based on the result of the playout
|
println("--> LLMMCTS backpropagate 2")
|
||||||
node.visits += 1 # Increment visit count for this node
|
# 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
|
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
|
end
|
||||||
|
println("--> LLMMCTS backpropagate 4")
|
||||||
end
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
a search tree node
|
The search tree node to check.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- `result::Bool`
|
- `result::Bool`
|
||||||
true if it is a leaf node (has no children), false otherwise.
|
`true` if the node has no children, `false` otherwise.
|
||||||
|
|
||||||
# Example
|
# Example
|
||||||
```jldoctest
|
```jldoctest
|
||||||
julia> using Revise
|
julia> using LLMMCTS
|
||||||
julia> using YiemAgent, DataStructures
|
julia> node = MCTSNode("leaf", Dict(:reward=>1.0), 0, 0, 0, 1.0, true, nothing, Dict(), Dict())
|
||||||
julia> initialState = Dict{String, Any}(
|
julia> isleaf(node)
|
||||||
"customerinfo"=> Dict{String, Any}(),
|
|
||||||
"storeinfo"=> Dict{String, Any}(),
|
|
||||||
|
|
||||||
"thoughtHistory"=> OrderedDict{String, 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)
|
|
||||||
true
|
true
|
||||||
```
|
```
|
||||||
|
|
||||||
# Signature
|
|
||||||
"""
|
"""
|
||||||
isleaf(node::MCTSNode)::Bool = isempty(node.children)
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
node of a search tree
|
The search tree node to check.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- `isrootnode::Bool`
|
- `isrootnode::Bool`
|
||||||
true if the given node is root node, false otherwise
|
`true` if the node is the root node, `false` otherwise.
|
||||||
|
|
||||||
# Signature
|
|
||||||
"""
|
"""
|
||||||
isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
node of a search tree
|
The node whose children will be evaluated.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- `childNode::MCTSNode`
|
- `childNode::MCTSNode`
|
||||||
the highest value child node
|
The child node with the highest `progressvalue + reward` value.
|
||||||
|
|
||||||
# Signature
|
|
||||||
"""
|
"""
|
||||||
function selectChildNode(node::MCTSNode)::MCTSNode
|
function selectChildNode(node::MCTSNode)::MCTSNode
|
||||||
highestProgressValue = -1
|
highestProgressValue = -1
|
||||||
nodekey = nothing
|
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
|
for (k, childNode) in node.children
|
||||||
potential = childNode.progressvalue + childNode.reward
|
potential = childNode.progressvalue + childNode.reward
|
||||||
if potential > highestProgressValue
|
if potential > highestProgressValue
|
||||||
@@ -183,35 +199,43 @@ function selectChildNode(node::MCTSNode)::MCTSNode
|
|||||||
end
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
MCTS node to expand
|
The MCTS node to expand.
|
||||||
- `transition::Function`
|
- `transition::Function`
|
||||||
A function that handles state transition.
|
A function that handles state transition.
|
||||||
- `transitionargs::NamedTuple`
|
- `transitionargs::NamedTuple`
|
||||||
Arguments for transition()
|
Arguments passed to the transition function.
|
||||||
|
|
||||||
# Keyword Arguments
|
# Keyword Arguments
|
||||||
- `horizontalSample::Integer`
|
- `horizontalSample::Integer=3`
|
||||||
Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3.
|
Number of child nodes to generate.
|
||||||
- `multithread::Bool`
|
- `multithread::Bool=false`
|
||||||
Whether to run expansion in parallel using multiple threads. Defaults to false.
|
Whether to run expansion in parallel using multiple threads.
|
||||||
|
|
||||||
# Return
|
|
||||||
- None
|
|
||||||
|
|
||||||
# Signature
|
# Return
|
||||||
|
- `Nothing`
|
||||||
"""
|
"""
|
||||||
function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple;
|
function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple;
|
||||||
horizontalSample::Integer=3, multithread=false)
|
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
|
if multithread
|
||||||
@sync for i in 1:horizontalSample
|
@sync for i in 1:horizontalSample
|
||||||
@spawn _expand(node, transition, transitionargs)
|
@spawn _expand(node, transition, transitionargs)
|
||||||
end
|
end
|
||||||
else
|
else
|
||||||
|
println("--> LLMMCTS expand 2")
|
||||||
for i in 1:horizontalSample
|
for i in 1:horizontalSample
|
||||||
|
println("--> LLMMCTS expand 3")
|
||||||
_expand(node, transition, transitionargs)
|
_expand(node, transition, transitionargs)
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
@@ -219,85 +243,108 @@ end
|
|||||||
|
|
||||||
""" Helper function to expand a single child node.
|
""" 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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
Parent MCTS node to expand from
|
The parent MCTS node to expand from.
|
||||||
- `transition::Function`
|
- `transition::Function`
|
||||||
A function that handles state transition
|
A function that handles state transition.
|
||||||
- `transitionargs::NamedTuple`
|
- `transitionargs::NamedTuple`
|
||||||
Arguments for transition()
|
Arguments passed to the transition function.
|
||||||
|
|
||||||
# Return
|
# Return
|
||||||
- None
|
- `Nothing`
|
||||||
|
|
||||||
# Signature
|
|
||||||
"""
|
"""
|
||||||
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
|
function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
|
||||||
result = transition(node.state, transitionargs)
|
println("--> LLMMCTS _expand 1")
|
||||||
newNodeKey::AbstractString = result[:newNodeKey]
|
# Generate one child node from the parent using the transition function
|
||||||
newstate::AbstractDict = result[:newstate]
|
result = transition(node.state, transitionargs)
|
||||||
progressvalue::Integer = result[:progressvalue]
|
newNodeKey::AbstractString = result[:newNodeKey]
|
||||||
|
newstate::AbstractDict = result[:newstate]
|
||||||
"""
|
progressvalue::Integer = result[:progressvalue]
|
||||||
[] newNodeKey ∉ keys(node.children).
|
println("--> LLMMCTS _expand 2")
|
||||||
New state may have semantic vector close enought to
|
# Dejavu detection: avoid adding duplicate states
|
||||||
one of existing child state. Which can be assume that they are the same state
|
# If newNodeKey already exists, skip - this handles semantically equivalent states
|
||||||
semantically-wise i.e. De javu. This could be used to recall lessons for this
|
if newNodeKey ∉ keys(node.children)
|
||||||
similar situation to improve decisionMaker and evaluator.
|
println("--> LLMMCTS _expand 3")
|
||||||
"""
|
# Create new MCTS node with:
|
||||||
if newNodeKey ∉ keys(node.children)
|
# - visits=0: no simulations yet
|
||||||
newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward],
|
# - statevalue=0: will be updated after simulation
|
||||||
newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{String, Any}())
|
# - progressvalue: LLM's estimate (fast heuristic)
|
||||||
node.children[newNodeKey] = newNode
|
# - reward: immediate environment feedback
|
||||||
end
|
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
|
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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
node that will be a simulation starting point.
|
The node to start simulation from.
|
||||||
- `transition::Function`
|
- `transition::Function`
|
||||||
A user function that handles how state transition.
|
A user function that handles state transition.
|
||||||
- `transitionargs::NamedTuple`
|
- `transitionargs::NamedTuple`
|
||||||
Arguments for everything the user will use within transition().
|
Arguments passed to the transition function.
|
||||||
- `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)
|
|
||||||
|
|
||||||
# Keyword Arguments
|
# Keyword Arguments
|
||||||
- `multithread::Bool`
|
- `maxSimulationDepth::Integer=3`
|
||||||
Whether to run expansion in parallel using multiple threads. Defaults to false.
|
Maximum depth level MCTS goes vertically during simulation.
|
||||||
|
- `horizontalSample::Integer=3`
|
||||||
# Return
|
Number of child nodes sampled at each node during simulation.
|
||||||
- `simTrajectoryReward::Number`
|
- `multithread::Bool=false`
|
||||||
Cumulative reward collected along the simulation trajectory
|
Whether to run expansion in parallel using multiple threads.
|
||||||
- `terminalstate::Union{Dict{String, Any}, Nothing}`
|
|
||||||
Final state if terminal state reached, nothing otherwise
|
|
||||||
|
|
||||||
# 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;
|
function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
||||||
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
||||||
)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, 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
|
simTrajectoryReward = 0.0
|
||||||
terminalstate = nothing
|
terminalstate = nothing
|
||||||
|
|
||||||
for depth in 1:maxSimulationDepth
|
for depth in 1:maxSimulationDepth
|
||||||
|
println("--> LLMMCTS simulate 2")
|
||||||
|
# Accumulate the current node's reward to the trajectory total
|
||||||
simTrajectoryReward += node.reward
|
simTrajectoryReward += node.reward
|
||||||
|
|
||||||
|
# Check if we've reached a terminal state
|
||||||
if node.isterminal
|
if node.isterminal
|
||||||
|
println("--> LLMMCTS simulate 3")
|
||||||
terminalstate = node.state
|
terminalstate = node.state
|
||||||
break
|
break
|
||||||
else
|
else
|
||||||
|
println("--> LLMMCTS simulate 4")
|
||||||
|
# Expand current node to generate children (horizontal sampling)
|
||||||
_ = expand(node, transition, transitionargs;
|
_ = expand(node, transition, transitionargs;
|
||||||
horizontalSample=horizontalSample,
|
horizontalSample=horizontalSample,
|
||||||
multithread=multithread)
|
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)
|
node = selectChildNode(node)
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS simulate 6")
|
||||||
end
|
end
|
||||||
|
println("--> LLMMCTS simulate 7")
|
||||||
return (simTrajectoryReward=simTrajectoryReward,
|
return (simTrajectoryReward=simTrajectoryReward,
|
||||||
terminalstate=terminalstate)
|
terminalstate=terminalstate)
|
||||||
end
|
end
|
||||||
|
|||||||
+73
-22
@@ -6,48 +6,99 @@ using ..type
|
|||||||
|
|
||||||
# ---------------------------------------------- 100 --------------------------------------------- #
|
# ---------------------------------------------- 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
|
# Arguments
|
||||||
- `node::MCTSNode`
|
- `node::MCTSNode`
|
||||||
mcts node
|
The MCTS node whose children will be evaluated.
|
||||||
- `w::T`
|
- `w::AbstractFloat`
|
||||||
exploration weight. Value is usually between 1 to 2.
|
Exploration weight. Typical values range from 1.0 to 2.0. A value of 1.0 balances
|
||||||
Value 1.0 makes MCTS balance between exploration and exploitation like 50%-50%.
|
exploration and exploitation equally. Higher values (e.g., 2.0) encourage more
|
||||||
Value 2.0 makes MCTS aggressively search the tree.
|
exploration of less-visited nodes.
|
||||||
# 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.
|
|
||||||
|
|
||||||
# Example
|
# Return
|
||||||
```jldoctest
|
- `selectedNode::MCTSNode`
|
||||||
julia>
|
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}
|
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
|
maxUCT = -Inf
|
||||||
selectedNode = nothing
|
selectedNode = nothing
|
||||||
|
|
||||||
for (childState, childNode) in node.children
|
for (childState, childNode) in node.children
|
||||||
|
# Calculate UCT value for this child
|
||||||
UCTvalue =
|
UCTvalue =
|
||||||
if childNode.visits != 0
|
if childNode.visits != 0
|
||||||
weightedterm = w * sqrt(log(node.visits) / childNode.visits) # explore term
|
# Child has been visited before - use statevalue with exploration bonus
|
||||||
childNode.statevalue + weightedterm
|
# Exploration bonus = w * sqrt(ln(parent_visits) / child_visits)
|
||||||
else # node.visits == 0 makes sqrt() in explore term error
|
# High child_visits = small bonus (exploitation dominates)
|
||||||
childNode.progressvalue # exploit term
|
# 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
|
end
|
||||||
|
|
||||||
if UCTvalue > maxUCT
|
if UCTvalue > maxUCT
|
||||||
maxUCT = UCTvalue
|
maxUCT = UCTvalue
|
||||||
selectedNode = childNode
|
selectedNode = childNode
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
return selectedNode
|
return selectedNode
|
||||||
end
|
end
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user