update

src/mcts.jl

@@ -10,27 +10,29 @@ using ..type
 # ---------------------------------------------- 100 --------------------------------------------- #
 
 
-""" 
+""" Select the best next node based on the highest value metric
 
 # Arguments
   - `node::MCTSNode`
-    node of a search tree
+    node of a search tree to evaluate
 
 # Return
   - `childNode::MCTSNode`
-    the highest value child node
-
-# Signature
+    the child node with highest value based on either:
+    - statevalue/visits ratio if any nodes have non-zero statevalue
+    - progressvalue + reward otherwise
 """
 function selectBestNextNode(node::MCTSNode)::MCTSNode
   highestProgressValue = -1
   nodekey = nothing
 
-  # if all childnode has statevalue == 0, use progressvalue + reward to select the best node
+  # Calculate sum of statevalues across all child nodes
   stateValueSum = sum([v.statevalue for (k, v) in node.children])
 
+  # If any nodes have non-zero statevalue, use statevalue/visits as selection metric
   if stateValueSum != 0
     for (k, childnode) in node.children
+      # Calculate average statevalue per visit
       potential = childnode.statevalue / childnode.visits
 
       if potential > highestProgressValue
@@ -39,6 +41,7 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
       end
     end
   else
+    # Otherwise use progressvalue + reward as selection metric
     for (k, childnode) in node.children
       potential = childnode.progressvalue + childnode.reward
 
@@ -53,15 +56,16 @@ function selectBestNextNode(node::MCTSNode)::MCTSNode
 end
 
 
-""" 
+""" Select the best trajectory node based on the highest reward
 
 # Arguments
   - `node::MCTSNode`
-    node of a search tree
+    node of a search tree to evaluate
 
 # Return
   - `childNode::MCTSNode`
-    the highest value child node
+    the highest value child node found by traversing down the tree using selectBestNextNode
+    until reaching a leaf node
 
 # Signature
 """
@@ -86,7 +90,8 @@ end
     reward it is now.
   
 # Return
-  - `None`
+  - `Nothing`
+    This function modifies the nodes in place and returns nothing
 
 # Signature
 """
@@ -94,22 +99,23 @@ function backpropagate(node::MCTSNode, simTrajectoryReward::T;
                       discountRewardCoeff::AbstractFloat=0.9) where {T<:Number}
   while !isroot(node)
     # Update the statistics of the current node based on the result of the playout
-    node.visits += 1
-    node.statevalue += ((node.statevalue * (node.visits-1)) + simTrajectoryReward) / node.visits
+    node.visits += 1  # Increment visit count for this node
+    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
+    node = node.parent  # Move up to parent node for next iteration
   end
 end
 
-
 """ Determine whether a node is a leaf node of a search tree.
 
 # Arguments
   - `node::MCTSNode`
     a search tree node
+
 # Return
   - `result::Bool`
-    true if it is a leaf node, false otherwise.
+    true if it is a leaf node (has no children), false otherwise.
+
 # Example
 ```jldoctest
 julia> using Revise
@@ -128,14 +134,10 @@ julia> YiemAgent.isleaf(root)
 true
 ```
 
-# TODO
-  [] update docs
-
 # Signature
 """
 isleaf(node::MCTSNode)::Bool = isempty(node.children)
 
-
 """ Determine wheter a given node is a root node
 
 # Arguments
@@ -185,14 +187,18 @@ end
 
 # Arguments
   - `node::MCTSNode`
-    MCTS node
+    MCTS node to expand
   - `transition::Function`
     A function that handles state transition.
   - `transitionargs::NamedTuple`
     Arguments for transition()
-  - `totalsample::Integer`
-    Total number to sample from the current node (i.e. expand new node horizontally)
-  
+
+# Keyword Arguments
+  - `horizontalSample::Integer`
+    Total number to sample from the current node (i.e. expand new node horizontally). Defaults to 3.
+  - `multithread::Bool`
+    Whether to run expansion in parallel using multiple threads. Defaults to false.
+    
 # Return
   - None
 
@@ -211,6 +217,21 @@ function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple;
   end
 end
 
+""" Helper function to expand a single child node.
+
+# Arguments
+  - `node::MCTSNode`
+    Parent MCTS node to expand from
+  - `transition::Function`
+    A function that handles state transition
+  - `transitionargs::NamedTuple`
+    Arguments for transition()
+
+# Return
+  - None
+
+# Signature
+"""
 function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple)
     result = transition(node.state, transitionargs)
     newNodeKey::AbstractString = result[:newNodeKey]
@@ -231,7 +252,6 @@ function _expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple
     end
 end
 
-
 """ Simulate interactions between agent and environment
 
 # Arguments