448 lines
12 KiB
Julia
448 lines
12 KiB
Julia
module mcts
|
|
|
|
export selectBestNextNode, selectBestTrajectoryNode, backpropagate, isleaf, isroot, selectChildNode,
|
|
expand, simulate
|
|
using Base.Threads
|
|
using GeneralUtils
|
|
|
|
using ..type
|
|
|
|
# ---------------------------------------------- 100 --------------------------------------------- #
|
|
|
|
|
|
""" Select the best next node based on the highest value metric
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
node of a search tree to evaluate
|
|
|
|
# Return
|
|
- `childNode::MCTSNode`
|
|
the child node with highest value based on either:
|
|
- statevalue/visits ratio if any nodes have non-zero statevalue
|
|
- progressvalue + reward otherwise
|
|
"""
|
|
function selectBestNextNode(node::MCTSNode)::MCTSNode
|
|
highestProgressValue = -1
|
|
nodekey = nothing
|
|
|
|
# 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
|
|
highestProgressValue = potential
|
|
nodekey = childnode.nodekey
|
|
end
|
|
end
|
|
else
|
|
# Otherwise use progressvalue + reward as selection metric
|
|
for (k, childnode) in node.children
|
|
potential = childnode.progressvalue + childnode.reward
|
|
|
|
if potential > highestProgressValue
|
|
highestProgressValue = potential
|
|
nodekey = childnode.nodekey
|
|
end
|
|
end
|
|
end
|
|
|
|
return node.children[nodekey]
|
|
end
|
|
|
|
|
|
""" Select the best trajectory node based on the highest reward
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
node of a search tree to evaluate
|
|
|
|
# Return
|
|
- `childNode::MCTSNode`
|
|
the highest value child node found by traversing down the tree using selectBestNextNode
|
|
until reaching a leaf node
|
|
|
|
# Signature
|
|
"""
|
|
function selectBestTrajectoryNode(node::MCTSNode)::MCTSNode
|
|
while !isleaf(node)
|
|
node = selectBestNextNode(node)
|
|
end
|
|
|
|
return node
|
|
end
|
|
|
|
|
|
""" Backpropagate reward along the simulation chain
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
leaf node of a search tree
|
|
- `simTrajectoryReward::T`
|
|
total reward from trajectory simulation
|
|
- `discountRewardCoeff::AbstractFloat`
|
|
A discount reward coefficient to reduce future reward. The futher in the future the lower
|
|
reward it is now.
|
|
|
|
# Return
|
|
- `Nothing`
|
|
This function modifies the nodes in place and returns nothing
|
|
|
|
# Signature
|
|
"""
|
|
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 # 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 # 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 (has no children), false otherwise.
|
|
|
|
# Example
|
|
```jldoctest
|
|
julia> using Revise
|
|
julia> using YiemAgent, DataStructures
|
|
julia> initialState = Dict{String, Any}(
|
|
"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
|
|
```
|
|
|
|
# Signature
|
|
"""
|
|
isleaf(node::MCTSNode)::Bool = isempty(node.children)
|
|
|
|
""" Determine wheter a given node is a root node
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
node of a search tree
|
|
|
|
# Return
|
|
- `isrootnode::Bool`
|
|
true if the given node is root node, false otherwise
|
|
|
|
# Signature
|
|
"""
|
|
isroot(node::MCTSNode)::Bool = node.nodekey == "root" ? true : false
|
|
|
|
|
|
|
|
""" Select child node based on the highest statevalue
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
node of a search tree
|
|
|
|
# Return
|
|
- `childNode::MCTSNode`
|
|
the highest value child node
|
|
|
|
# Signature
|
|
"""
|
|
function selectChildNode(node::MCTSNode)::MCTSNode
|
|
highestProgressValue = -1
|
|
nodekey = nothing
|
|
|
|
# loop thought node children dictionary to find the highest progress value
|
|
for (k, childNode) in node.children
|
|
potential = childNode.progressvalue + childNode.reward
|
|
if potential > highestProgressValue
|
|
highestProgressValue = potential
|
|
nodekey = childNode.nodekey
|
|
end
|
|
end
|
|
|
|
return node.children[nodekey]
|
|
end
|
|
|
|
|
|
""" Expand selected node.
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
MCTS node to expand
|
|
- `transition::Function`
|
|
A function that handles state transition.
|
|
- `transitionargs::NamedTuple`
|
|
Arguments for transition()
|
|
|
|
# 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
|
|
|
|
# Signature
|
|
"""
|
|
function expand(node::MCTSNode,transition::Function, transitionargs::NamedTuple;
|
|
horizontalSample::Integer=3, multithread=false)
|
|
if multithread
|
|
@sync for i in 1:horizontalSample
|
|
@spawn _expand(node, transition, transitionargs)
|
|
end
|
|
else
|
|
for i in 1:horizontalSample
|
|
_expand(node, transition, transitionargs)
|
|
end
|
|
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]
|
|
newstate::AbstractDict = result[:newstate]
|
|
progressvalue::Integer = result[:progressvalue]
|
|
|
|
"""
|
|
[] newNodeKey ∉ keys(node.children).
|
|
New state may have semantic vector close enought to
|
|
one of existing child state. Which can be assume that they are the same state
|
|
semantically-wise i.e. De javu. This could be used to recall lessons for this
|
|
similar situation to improve decisionMaker and evaluator.
|
|
"""
|
|
if newNodeKey ∉ keys(node.children)
|
|
newNode = MCTSNode(newNodeKey, newstate, 0, progressvalue, 0, newstate[:reward],
|
|
newstate[:isterminal], node, Dict{String, MCTSNode}(), Dict{String, Any}())
|
|
node.children[newNodeKey] = newNode
|
|
end
|
|
end
|
|
|
|
""" Simulate interactions between agent and environment
|
|
|
|
# Arguments
|
|
- `node::MCTSNode`
|
|
node that will be a simulation starting point.
|
|
- `transition::Function`
|
|
A user function that handles how state transition.
|
|
- `transitionargs::NamedTuple`
|
|
Arguments for everything the user will use within transition().
|
|
- `maxSimulationDepth::Integer`
|
|
maximum depth level MCTS goes vertically during simulation.
|
|
- `horizontalSample::Integer`
|
|
Total number to sample from the current node (i.e. expand new node horizontally)
|
|
|
|
# Keyword Arguments
|
|
- `multithread::Bool`
|
|
Whether to run expansion in parallel using multiple threads. Defaults to false.
|
|
|
|
# Return
|
|
- `simTrajectoryReward::Number`
|
|
Cumulative reward collected along the simulation trajectory
|
|
- `terminalstate::Union{Dict{String, Any}, Nothing}`
|
|
Final state if terminal state reached, nothing otherwise
|
|
|
|
# Signature
|
|
"""
|
|
function simulate(node::MCTSNode, transition::Function, transitionargs::NamedTuple;
|
|
maxSimulationDepth::Integer=3, horizontalSample::Integer=3, multithread=false
|
|
)::NamedTuple{(:simTrajectoryReward, :terminalstate), Tuple{<:Number, Union{Dict{String, Any}, Nothing}}}
|
|
|
|
simTrajectoryReward = 0.0
|
|
terminalstate = nothing
|
|
|
|
for depth in 1:maxSimulationDepth
|
|
simTrajectoryReward += node.reward
|
|
if node.isterminal
|
|
terminalstate = node.state
|
|
break
|
|
else
|
|
_ = expand(node, transition, transitionargs;
|
|
horizontalSample=horizontalSample,
|
|
multithread=multithread)
|
|
node = selectChildNode(node)
|
|
end
|
|
end
|
|
|
|
return (simTrajectoryReward=simTrajectoryReward,
|
|
terminalstate=terminalstate)
|
|
end
|
|
|
|
# """ Make new state
|
|
|
|
# # Arguments
|
|
# - `currentstate::T1`
|
|
# Current state dictionary containing thought history and metadata
|
|
# - `thoughtDict::T4`
|
|
# Dictionary containing new thought and action
|
|
# - `response::T2`
|
|
# Response string from the environment
|
|
# - `select::Union{T3, Nothing}`
|
|
# Selection value or nothing
|
|
# - `reward::T3`
|
|
# Reward value for this state
|
|
# - `isterminal::Bool`
|
|
# Whether this state is terminal
|
|
|
|
# # Return
|
|
# - `Tuple{String, Dict{String, <:Any}}`
|
|
# A tuple containing:
|
|
# - A unique node key string
|
|
# - A new state dictionary with updated thought history and metadata
|
|
|
|
# # Example
|
|
# ```jldoctest
|
|
# julia>
|
|
# ```
|
|
|
|
# # Signature
|
|
# """
|
|
# function makeNewState(currentstate::T1, thoughtDict::T4, response::T2, select::Union{T3, Nothing},
|
|
# reward::T3, isterminal::Bool
|
|
# )::Tuple{String, Dict{String, <:Any}} where {T1<:AbstractDict, T2<:AbstractString, T3<:Number, T4<:AbstractDict}
|
|
|
|
# # Find the latest thought key and index from current state's thought history
|
|
# currentstate_latestThoughtKey, currentstate_latestThoughtIndice =
|
|
# GeneralUtils.findHighestIndexKey(currentstate[:thoughtHistory], "thought")
|
|
# # Calculate next index for new thought/action
|
|
# currentstate_nextIndice =
|
|
# currentstate_latestThoughtKey == :NA ? 1 : currentstate_latestThoughtIndice + 1
|
|
# # Create new keys for thought and action based on next index
|
|
# currentstate_latestThoughtKey = Symbol("thought_$currentstate_nextIndice")
|
|
# latestActionKey = Symbol("action_$currentstate_nextIndice")
|
|
|
|
# # Find the latest thought index from input thought dictionary
|
|
# _, thoughtDict_latestThoughtIndice =
|
|
# GeneralUtils.findHighestIndexKey(thoughtDict, "thought")
|
|
|
|
# # Determine thought and action keys from thought dictionary
|
|
# thoughtDict_latestThoughtKey, thoughtDict_latestActionKey =
|
|
# if thoughtDict_latestThoughtIndice == -1
|
|
# (:thought, :action)
|
|
# else
|
|
# (
|
|
# Symbol("thought_$thoughtDict_latestThoughtIndice"),
|
|
# Symbol("action_$thoughtDict_latestThoughtIndice"),
|
|
# )
|
|
# end
|
|
|
|
# # Create new state by deep copying current state
|
|
# newstate = deepcopy(currentstate)
|
|
# # Update thought history with new thought
|
|
# newstate[:thoughtHistory][currentstate_latestThoughtKey] =
|
|
# thoughtDict[thoughtDict_latestThoughtKey]
|
|
# # Update thought history with new action
|
|
# newstate[:thoughtHistory][latestActionKey] = thoughtDict[thoughtDict_latestActionKey]
|
|
# # Create and add new observation to thought history
|
|
# newObservationKey = Symbol("observation_$(currentstate_nextIndice)")
|
|
# newstate[:thoughtHistory][newObservationKey] = response
|
|
# # Update state metadata
|
|
# newstate[:reward] = reward
|
|
# newstate[:select] = select
|
|
# newstate[:isterminal] = isterminal
|
|
|
|
# # Generate unique ID for new node
|
|
# newNodeKey = GeneralUtils.uuid4snakecase()
|
|
|
|
# return (newNodeKey, newstate)
|
|
# end
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
end # module mcts |