update docs

src/interface.jl

@@ -9,46 +9,55 @@ using ..type, ..mcts, ..util
 # ---------------------------------------------- 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
-  - `initialstate::T`
-    initial state
-  - `transition::Function`
-    a function that define how the state transitions
-  - `transitionargs::NamedTuple`
-    arguments for transition function
+- `initialstate::T`
+  The initial state from which to start the search.
+- `transition::Function`
+  A function that defines how the state transitions.
+- `transitionargs::NamedTuple`
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-  - `horizontalSampleExpansionPhase::Integer`
-    a number of child state MCTS sample at each node during expansion phase (default: 3)
-  - `horizontalSampleSimulationPhase::Integer`
-    a number of child state MCTS sample at each node during simulation's expansion phase (default: 3)
-  - `maxSimulationDepth::Integer`
-    a number of levels MCTS goes during simulation phase (default: 3)
-  - `maxiterations::Integer`
-    a number of iteration MCTS goes thru expansion -> simulation -> backpropagation cycle (default: 10)
-  - `explorationweight::Number`
-    exploration weight controls how much MCTS should explore new state instead of exploit 
-    a known state. 1.0 balance between exploration and exploitation like 50%-50%. 2.0 makes MCTS
-    aggressively explore new state (default: 1.0)
-  - `earlystop::Union{Function,Nothing}`
-    optional function to check early stopping condition if it is satisfied, MCTS will break iterations (default: nothing)
-  - `saveSimulatedNode::Bool`
-    whether to save nodes created during simulation phase (default: false)
-  - `multithread::Bool`
-    whether to use multithreading during simulation (default: false)
+- `horizontalSampleExpansionPhase::Integer=3`
+  Number of child states sampled at each node during expansion phase.
+- `horizontalSampleSimulationPhase::Integer=3`
+  Number of child states sampled at each node during simulation's expansion phase.
+- `maxSimulationDepth::Integer=3`
+  Maximum depth MCTS goes during simulation phase.
+- `maxiterations::Integer=10`
+  Number of iterations MCTS performs through expansion → simulation → backpropagation cycles.
+- `explorationweight::Number=1.0`
+  Exploration weight controls how much MCTS explores new states versus exploiting known states.
+  A value of 1.0 balances exploration and exploitation equally. Higher values (e.g., 2.0) 
+  encourage more aggressive exploration.
+- `earlystop::Union{Function,Nothing}=nothing`
+  Optional function to check early stopping condition. If satisfied, MCTS breaks iterations.
+- `saveSimulatedNode::Bool=false`
+  Whether to save nodes created during simulation phase.
+- `multithread::Bool=false`
+  Whether to use multithreading during simulation.
 
-# Returns
-  - `NamedTuple{(:root, :bestNextState, :bestFinalState), Tuple{MCTSNode, T, T}}`
-    - root: the complete MCTS tree with root node
-    - bestNextState: the best immediate next state
-    - bestFinalState: the best final state along the best trajectory
+# Return
+- `NamedTuple{(:root, :bestNextState, :bestTerminalState, :highValueStateList),
+             Tuple{MCTSNode,T,T,Vector{Dict{String,Any}}}}`
+  - `root`: the complete MCTS tree with root node
+  - `bestNextState`: the best immediate next state
+  - `bestTerminalState`: the best final state along the best trajectory
+  - `highValueStateList`: list of high-value terminal states (reward >= 8)
 
 # Example
-    Refers to SQLLLM package
-
-# Signature
+```jldoctest
+julia> using LLMMCTS
+julia> initialState = Dict(:reward=>0.0)
+julia> result = runMCTS(initialState, transition_func, transition_args; maxiterations=5)
+```
 """
 function runMCTS(
   initialstate::T,
@@ -140,28 +149,37 @@ function runMCTS(
   return result
 end
 
-""" Search the best action to take for a given state and task 
+""" Run simulation from a given node and backpropagate the reward.
+
+This function performs simulation (rollout) from the given node, collects the
+cumulative reward along the trajectory, and backpropagates it up the tree to update
+visit counts and state values.
+
+Does **not** mutate the input node's children (unless `saveSimulatedNode=true`).
 
 # Arguments
 - `node::MCTSNode`
-    current node to simulate from
+  The current node to simulate from.
 - `transition::Function`
-    a function that defines how the state transitions
+  A function that defines how the state transitions.
 - `transitionargs::NamedTuple`
-    arguments for transition function
+  Arguments passed to the transition function.
 
 # Keyword Arguments
-- `maxSimulationDepth::Integer`
-    a number of levels MCTS goes during simulation phase (default: 3)
-- `horizontalSampleSimulationPhase::Integer`
-    a number of child states MCTS samples at each node during simulation phase (default: 3)
-- `saveSimulatedNode::Bool`
-    whether to save nodes created during simulation phase (default: false)
-- `multithread::Bool`
-    whether to use multithreading during simulation (default: false)
+- `maxSimulationDepth::Integer=3`
+  Maximum depth MCTS goes during simulation phase.
+- `horizontalSampleSimulationPhase::Integer=3`
+  Number of child states sampled at each node during simulation phase.
+- `saveSimulatedNode::Bool=false`
+  Whether to save nodes created during simulation phase. If false, children are
+  cleared after simulation.
+- `multithread::Bool=false`
+  Whether to use multithreading during simulation.
 
-# Returns
-  Nothing, but updates the node's reward and visit count through backpropagation
+# Return
+- `Nothing`
+
+# Signature
 """
 function simulateThenBackpropagate(node::MCTSNode, transition::Function, transitionargs::NamedTuple; 
                 maxSimulationDepth::Integer=3, horizontalSampleSimulationPhase::Integer=3,