1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
use core::f64;
use std::{
rc::{Rc, Weak},
time::Instant,
};
use crate::{
action::Action, mdp::MDP, node::Node, policy::RolloutPolicy, rand::genrand, strategy::Strategy,
ucb1::UCB1,
};
pub struct MCTS<M, S, A, P>
where
M: MDP<S, A>,
A: Action,
S: Clone,
P: RolloutPolicy<M, S, A>,
{
mdp: M,
root: Rc<Node<S, A>>,
bandit: UCB1,
policy: P,
}
impl<M, S, A, P> MCTS<M, S, A, P>
where
M: MDP<S, A>,
A: Action,
S: Clone + Eq + PartialEq,
P: RolloutPolicy<M, S, A>,
{
pub fn new(mdp: M, policy: P) -> Self {
let state = mdp.get_initial_state();
Self {
root: Rc::new(Node::new(state, None, None, Weak::new())),
mdp,
bandit: UCB1::default(),
policy,
}
}
/// Execute the MCTS algorithm from the initial state given, with timeout in seconds
/// After how many milliseconds, the mcts should timeout
/// TODO: Move this to be more dynamic, and support max-depth timeout
pub fn mcts(&mut self, timeout: u128) {
let start_time = Instant::now();
while start_time.elapsed().as_millis() < timeout {
// Find a state node to expand
let selected_node = self.root.select(&self.mdp, &self.bandit);
// let xx = !self.mdp.is_terminal(&selected_node.state);
if !self.mdp.is_terminal(&selected_node.state) {
let child = selected_node.expand(&self.mdp, &self.policy);
let reward = self.simulate(&child, start_time, timeout);
child.back_propagate(reward, &mut self.bandit);
}
}
}
/// TODO: This would eventually be moved to a trait that must be implemented on state!, this MCTS or whatever!
pub(crate) fn heuristic_eval(&self, _state: &S) -> f64 {
0.0
}
/// Simulate until a terminal state
pub(crate) fn simulate(
&self,
node: &Rc<Node<S, A>>,
start_time: Instant,
timeout: u128,
) -> f64 {
let mut state = node.state.clone();
let mut cumulative_reward = 0.0;
// let mut depth = 0;
while !self.mdp.is_terminal(&state) && start_time.elapsed().as_millis() < timeout {
let actions = self.mdp.get_actions(&state);
// Choose an action to execute
let action = self.policy.pick(&state, &actions);
// Execute the action
let (next_state, reward, ..) = self.mdp.execute(&state, &action);
// Discount the reward
// cumulative_reward += f64::powi(self.mdp.get_discount_factor(), depth) * reward;
cumulative_reward += reward;
// depth += 1;
state = next_state;
}
if !self.mdp.is_terminal(&state) {
// todo! this needs to be a trait
cumulative_reward += self.heuristic_eval(&state);
}
return cumulative_reward;
}
pub fn best_action(&self, strategy: Strategy) -> Option<A> {
let root = &self.root;
let children = root.children.borrow();
if children.is_empty() {
return None;
}
match strategy {
Strategy::MostVisited => children
.iter()
.max_by_key(|c| *c.visits.borrow())
.and_then(|c| c.action),
Strategy::HighestQValue => children
.iter()
.max_by(|a, b| {
a.q_value()
.partial_cmp(&b.q_value())
.unwrap_or(std::cmp::Ordering::Equal)
})
.and_then(|c| c.action),
Strategy::Probabilistic => {
// Softmax over Q-values
let qvalues = children.iter().map(|c| c.q_value()).collect::<Vec<_>>();
let maxq = qvalues.iter().copied().fold(f64::NEG_INFINITY, f64::max);
// subtract maxq for numerical stability
let expq: Vec<f64> = qvalues.iter().map(|q| (q - maxq).exp()).collect();
let sum = expq.iter().sum::<f64>().max(f64::MIN_POSITIVE);
let probs = expq.iter().map(|x| x / sum).collect::<Vec<_>>();
// sample based on probabilities
let mut r = genrand(0, 10_000) as f64 / 10_000.0;
for (i, p) in probs.iter().enumerate() {
r -= p;
if r <= 0.0 {
return children[i].action;
}
}
// fallback
children[0].action
}
Strategy::HeuristicWin => {
// prioritize terminal winning moves
let mut winning_mvs = vec![];
let mut bestq = f64::NEG_INFINITY;
let mut best_mvs = vec![];
for child in children.iter() {
let q = child.q_value();
// if child is terminal with positive reward (win)
// if let Some(reward) = child.score.borrow() {}
if *child.score.borrow() > 0.0 {
winning_mvs.push(child);
continue;
}
if q > bestq {
bestq = q;
best_mvs = vec![child];
} else if (q - bestq).abs() < 1e-9 {
best_mvs.push(child);
}
}
let chosen = if !winning_mvs.is_empty() {
&winning_mvs[genrand(0, winning_mvs.len())]
} else {
&best_mvs[genrand(0, best_mvs.len())]
};
chosen.action
}
}
}
}