simple_mcts/mcts_batch.rs
1//! Batch processing implementation for Monte Carlo Tree Search (MCTS) algorithm
2//!
3//! Provides parallel simulation capabilities for multiple MCTS instances
4
5use std::{sync::atomic::{AtomicU8, Ordering}, time::{SystemTime, UNIX_EPOCH}};
6
7use rand::{rngs::StdRng, SeedableRng};
8
9use crate::{utils, Game, GameEvaluator, Mcts, MctsConfig, MctsError, MctsState, SelectionFunction};
10
11/// Configuration parameters for an `MctsBatch` manager.
12///
13/// This struct extends `MctsConfig` with batch-specific settings like the random
14/// number generator seed for reproducible batch simulations.
15///
16/// # Type Parameters
17/// - `N`: The number of possible actions in the game.
18pub struct MctsBatchConfig<const N: usize>{
19 /// The exploration coefficient for MCTS instances within the batch.
20 pub exploration_coef: f64,
21 /// The selection function for MCTS instances within the batch.
22 pub selection_function: SelectionFunction<N>,
23 // An optional seed for the random number generator used in the batch.
24 ///
25 /// Providing a `Some(value)` will initialize the RNG with a fixed seed,
26 /// ensuring reproducible simulation results across runs. If `None`,
27 /// a new seed based on the current time will be used, leading to
28 /// non-reproducible runs (but more "random" behavior).
29 pub seed: Option<u64>,
30}
31
32impl<const N: usize> MctsBatchConfig<N>{
33 /// The default configuration for an `MctsBatch`.
34 ///
35 /// It inherits the default exploration coefficient and selection function from `MctsConfig::DEFAULT`,
36 /// and uses `None` for the seed, meaning simulations will not be reproducible by default unless a seed is provided.
37 pub const DEFAULT: MctsBatchConfig<N> = MctsBatchConfig::<N>{
38 exploration_coef: MctsConfig::<N>::DEFAULT.exploration_coef,
39 selection_function: MctsConfig::<N>::DEFAULT.selection_function,
40 seed: None
41 };
42}
43
44/// Represents the **history of a single MCTS instance**, capturing key data at each step:
45/// - The **game state** at that point.
46/// - The **estimated value** of the state.
47/// - The **action probabilities** derived from the MCTS search.
48#[allow(type_alias_bounds)]
49type History<T: Game<N>, const N: usize> = Vec<(T, f64, [f64; N])>;
50
51/// Manages a collection of **MCTS simulations** that can be processed in parallel.
52///
53/// `MctsBatch` tracks the complete history of each simulation, which is valuable for
54/// training machine learning models or analyzing game outcomes.
55///
56/// # Type Parameters
57/// - `T`: The type of game being simulated.
58/// - `N`: The number of possible actions in the game.
59pub struct MctsBatch<T: Game<N>, const N: usize>{
60 /// Collection of MCTS instances with their histories
61 instances: Vec<Option<(Mcts<T, N>, History<T, N>)>>,
62 /// Number of active instances
63 count: usize,
64 /// Shared random number generator
65 rand: StdRng,
66 /// Current state of the batch processor
67 state: MctsState,
68 /// Config used for Mcts
69 config: MctsConfig<N>
70}
71
72impl<T: Game<N>, const N: usize> MctsBatch<T, N>{
73 /// Creates a new empty MCTS batch processor with the default configuration.
74 ///
75 /// The random number generator will be seeded based on the current system time.
76 #[inline]
77 pub fn new() -> Self{
78 MctsBatch::from_config(&MctsBatchConfig::<N>::DEFAULT)
79 }
80
81 /// Creates a new empty MCTS batch processor from a specified configuration.
82 ///
83 /// This allows customizing the MCTS parameters for all instances in the batch,
84 /// including the random number generator's seed for reproducibility.
85 ///
86 /// # Parameters
87 /// - `config`: The `MctsBatchConfig` to use for this batch manager.
88 #[inline]
89 pub fn from_config(config: &MctsBatchConfig<N>) -> Self{
90 MctsBatch {
91 instances: Vec::new(),
92 count: 0,
93 rand: SeedableRng::seed_from_u64(
94 if let Some(seed) = &config.seed {
95 *seed
96 } else {
97 (SystemTime::now().duration_since(UNIX_EPOCH).expect("").as_nanos()%u64::MAX as u128) as u64
98 }
99 ),
100 state: MctsState(AtomicU8::new(MctsState::USABLE)),
101 config: MctsConfig { exploration_coef: config.exploration_coef, selection_function: config.selection_function }
102 }
103 }
104
105 /// Returns the number of active MCTS instances currently managed by the batch processor.
106 ///
107 /// This count represents the simulations that are still ongoing and have not yet
108 /// reached a terminal state or been otherwise removed from the batch.
109 ///
110 /// # Returns
111 /// The `usize` representing the total number of active MCTS instances.
112 ///
113 /// # Examples
114 /// ```rust
115 /// use simple_mcts::{MctsBatch, test_utils::GameTest, MctsError, Game};
116 ///
117 /// fn main() -> Result<(), MctsError> {
118 /// let mut manager = MctsBatch::<GameTest, 4>::new();
119 /// manager.populate_from_game(vec![GameTest::new(), GameTest::new()])?;
120 ///
121 /// // Initially, the count should reflect the number of populated instances.
122 /// assert_eq!(manager.get_count(), 2);
123 ///
124 /// // After some iterations, if games finish or are processed, the count might change.
125 /// // (Assuming `next` or other operations might reduce the count)
126 /// // manager.iterate(&evaluator)?; // If an iteration causes a game to finish
127 /// // manager.next()?; // If retrieving results removes finished games
128 ///
129 /// // Example: If one game finishes and is removed
130 /// // assert_eq!(manager.get_count(), 1);
131 /// Ok(())
132 /// }
133 /// ```
134 #[inline]
135 pub fn get_count(&self) -> usize{
136 self.count
137 }
138
139 /// Returns the current operational state of the `MctsBatch` processor.
140 ///
141 /// This indicates whether the batch is `Usable` (ready for operations),
142 /// `AwaitingSimulation` (waiting for external evaluation results), or `Locked`
143 /// (temporarily busy with internal processing).
144 ///
145 /// # Returns
146 /// A clone of the current `MctsState` of the batch processor.
147 ///
148 /// # Examples
149 /// ```rust
150 /// use simple_mcts::{MctsBatch, test_utils::GameTest, MctsError, MctsState, Game};
151 ///
152 /// fn main() -> Result<(), MctsError> {
153 /// let manager = MctsBatch::<GameTest, 4>::new();
154 /// // A newly created MctsBatch is typically in the Usable state.
155 /// assert_eq!(manager.get_state(), MctsState::USABLE);
156 ///
157 /// // The state would change, for example, after calling `start_iteration`
158 /// // (if MctsBatch had such a method that changed its state to AwaitingSimulation)
159 /// // or during internal processing.
160 /// Ok(())
161 /// }
162 /// ```
163 #[inline]
164 pub fn get_state(&self) -> u8{
165 self.state.0.load(Ordering::SeqCst)
166 }
167
168 /// Populates the batch with new MCTS instances.
169 ///
170 /// This method requires the batch processor to be in a `Usable` state.
171 ///
172 /// # Parameters
173 /// - `n`: The number of new instances to add.
174 ///
175 /// # Returns
176 /// `Ok(())` if the batch is successfully populated.
177 /// `Err(MctsError::InvalidState(_))` if the batch processor is not in the `Usable` state.
178 pub fn populate(&mut self, mut n: usize) -> Result<(), MctsError>{
179 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
180 Ok(_) => {}
181 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
182 }
183
184 let mut empty_place: usize = self.instances.len() - self.get_count();
185 self.count += n;
186
187 if empty_place > 0 && n > 0{
188 for opt in &mut self.instances{
189 if opt.is_none(){
190 *opt = Some((Mcts::<T, N>::new(), Vec::new()));
191 empty_place -= 1;
192 n -= 1;
193 }
194
195 if n == 0 { return Ok(()); }
196 if empty_place == 0{ break; }
197 }
198 }
199
200 for _ in 0..n{
201 self.instances.push(Some((Mcts::<T, N>::from_config(&self.config), Vec::new())));
202 }
203
204 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
205 Ok(())
206 }
207
208 /// **Populates the batch with new MCTS instances, each initialized with an existing game state.**
209 ///
210 /// This method requires the batch processor to be in a `Usable` state.
211 ///
212 /// # Parameters
213 /// - `games`: A vector of game states to initialize the new instances.
214 ///
215 /// # Returns
216 /// - `Ok(())` if the batch is successfully populated.
217 /// - `Err(MctsError::InvalidState(_))` if the batch is not in the `Usable` state.
218 pub fn populate_from_game(&mut self, mut games: Vec<T>) -> Result<(), MctsError>{
219 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
220 Ok(_) => {}
221 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
222 }
223
224 let mut n: usize = games.len();
225 let mut empty_place: usize = self.instances.len() - self.get_count();
226
227 self.count += n;
228
229 if empty_place > 0 && n > 0{
230 for opt in &mut self.instances{
231 if opt.is_none(){
232 *opt = Some((Mcts::<T, N>::from_game_with_config(games.pop().unwrap(), &self.config), Vec::new()));
233 empty_place -= 1;
234 n -= 1;
235 }
236
237 if n == 0 { return Ok(()); }
238 if empty_place == 0{ break; }
239 }
240 }
241
242 for _ in 0..n{
243 self.instances.push(Some((Mcts::<T, N>::from_game(games.pop().unwrap()), Vec::new())));
244 }
245
246 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
247 Ok(())
248 }
249
250 /// **Clears all instances from the batch**, resetting it to an empty state.
251 ///
252 /// This method requires the batch processor to be in a `Usable` state.
253 ///
254 /// # Returns
255 /// - `Ok(())` if the batch is successfully cleared.
256 /// - `Err(MctsError::InvalidState(_))` if the batch is not in the `Usable` state.
257 pub fn clear(&mut self) -> Result<(), MctsError>{
258 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
259 Ok(_) => {}
260 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
261 }
262
263 self.instances.clear();
264 self.count = 0;
265
266 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
267 Ok(())
268 }
269
270 /// Performs **one full MCTS iteration** on all active instances in the batch.
271 ///
272 /// This method requires the batch processor to be in a `Usable` state.
273 ///
274 /// # Parameters
275 /// - `evaluator`: The policy/value evaluator used for simulations.
276 ///
277 /// # Returns
278 /// - `Ok(())` if all active instances complete an iteration successfully.
279 /// - `Err(MctsError::InvalidState(_))`: If the batch is not in the `Usable` state.
280 /// - Propagated errors from MCTS instances, such as:
281 /// - `MctsError::SearchAlreadyOver`
282 /// - `MctsError::InvalidEvaluationCount`
283 /// - `MctsError::ActionOutOfRange`
284 /// - `MctsError::InvalidAction`
285 /// - `MctsError::UnexploredAction`
286 pub fn iterate(&mut self, evaluator: &dyn GameEvaluator<T, N>) -> Result<(), MctsError>{
287 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
288 Ok(_) => {}
289 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
290 }
291
292 for opt in &mut self.instances{
293 if let Some((mcts, _history)) = opt{
294 if !mcts.get_game().is_finish() {
295 mcts.iterate(evaluator)?;
296 }
297 }
298 }
299
300 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
301 Ok(())
302 }
303
304 /// Initiates the selection and expansion phases for all active MCTS instances in the batch.
305 ///
306 /// This method requires the batch processor to be in a `Usable` state
307 /// and transitions it to `AwaitingSimulation`.
308 ///
309 /// # Returns
310 /// `Ok(Vec<T::State>)` containing the game states from each instance that require external evaluation.
311 /// `Err(MctsError::InvalidState(_))` if the batch processor is not in the `Usable` state.
312 /// `Err(MctsError::SearchAlreadyOver)` if an MCTS instance's root node indicates the game is already finished.
313 pub fn start_iteration(&mut self) -> Result<Vec<T::State>, MctsError>{
314 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
315 Ok(_) => {}
316 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
317 }
318
319 let mut game_states: Vec<T::State> = Vec::with_capacity(self.get_count());
320
321 for opt in &mut self.instances{
322 if let Some((mcts, _history)) = opt{
323 if !mcts.get_game().is_finish() {
324 game_states.push(mcts.start_iteration()?);
325 }
326 }
327 }
328
329 self.state.0.store(MctsState::AWAITING_SIMULATION, Ordering::Relaxed);
330 Ok(game_states)
331 }
332
333 /// Applies the results of external simulations and performs backpropagation for all active MCTS instances.
334 ///
335 /// This method requires the batch processor to be in the `AwaitingSimulation` state
336 /// and transitions it back to `Usable`.
337 ///
338 /// # Parameters
339 /// - `evaluations`: A vector of tuples, where each tuple contains the estimated value (f64)
340 /// and action probabilities ([f64; N]) for an MCTS instance.
341 /// The order of evaluations should correspond to the order of game states
342 /// returned by `start_iteration`.
343 ///
344 /// # Returns
345 /// `Ok(())` if all simulations are successfully applied and backpropagation completes.
346 /// `Err(MctsError::InvalidState(_))` if the batch processor is not in the `AwaitingSimulation` state.
347 /// `Err(MctsError::InvalidEvaluationCount(expected, received))` if the number of provided
348 /// evaluations does not match the number of active instances
349 /// that were awaiting simulation.
350 pub fn apply_simulation(&mut self, mut evaluations : Vec<(f64, [f64; N])>) -> Result<(), MctsError>{
351 match self.state.0.compare_exchange(MctsState::AWAITING_SIMULATION, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
352 Ok(_) => {}
353 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
354 }
355
356 if evaluations.len() != self.get_count() {
357 self.state.0.store(MctsState::AWAITING_SIMULATION, Ordering::Relaxed);
358 return Err(MctsError::InvalidEvaluationCount(self.get_count(), evaluations.len()));
359 }
360
361 for opt in &mut self.instances.iter_mut().rev(){
362 if let Some((mcts, history)) = opt{
363 if mcts.get_game().is_finish() { continue; }
364 let game = mcts.get_game().clone();
365
366 mcts.apply_simulation(evaluations.pop().unwrap())?;
367 let (value, policy) = mcts.get_result();
368
369 let action = utils::sample(&policy, &mut self.rand);
370 mcts.play(action)?;
371
372 history.push((game, value, policy));
373 }
374 }
375
376 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
377 Ok(())
378 }
379
380 /// Retrieves the final results (score and policy) for any MCTS instances that have completed
381 /// their search (game is finished at the root).
382 ///
383 /// Completed instances are removed from the batch. This method can only be called when
384 /// the batch is in a `Usable` state.
385 ///
386 /// # Returns
387 /// `Ok(Vec<(f64, [f64; N])>)` containing tuples of final score and action probabilities
388 /// for completed games. Returns an empty `Vec` if no instances have finished.
389 /// `Err(MctsError::InvalidState(_))` if the batch processor is not in the `Usable` state.
390 pub fn next(&mut self) -> Result<Vec<History<T, N>>, MctsError>{
391 match self.state.0.compare_exchange(MctsState::USABLE, MctsState::LOCKED, Ordering::SeqCst, Ordering::SeqCst){
392 Ok(_) => {}
393 Err(current_state) => { return Err(MctsError::InvalidState(current_state)); }
394 }
395
396 let mut result = Vec::new();
397
398 for opt in &mut self.instances{
399 if let Some((mcts, _history)) = &opt{
400 if mcts.get_game().is_finish(){
401 {
402 let (mcts, history) = opt.as_mut().unwrap();
403
404 let mut score: f64 = -mcts.get_game().get_result().unwrap();
405 for (_game, value, _policy) in history.iter_mut().rev(){
406 *value=score;
407 score = -score;
408 }
409 }
410 result.push(opt.take().unwrap().1);
411 self.count -= 1;
412 }
413 else if mcts.count_visit() > 0{
414 let (mcts, history) = opt.as_mut().unwrap();
415
416 let (value, policy) = mcts.get_result();
417 history.push((mcts.get_game().clone(), value, policy));
418
419 let action: usize = utils::sample(&policy, &mut self.rand);
420
421 mcts.play(action)?;
422 }
423 }
424 }
425
426 self.state.0.store(MctsState::USABLE, Ordering::Relaxed);
427 Ok(result)
428 }
429}
430
431#[cfg(test)]
432mod tests {
433 use crate::{test_utils::{GameEvaluatorTest2, GameTest}, Game, MctsBatch, MctsError};
434
435 #[test]
436 fn test_batch_iterate_1() -> Result<(), MctsError>{
437 let mut manager = MctsBatch::<GameTest, 4>::new();
438 let evaluator: GameEvaluatorTest2 = GameEvaluatorTest2::new();
439
440 let a: GameTest = GameTest::new();
441 let b: GameTest = GameTest::new();
442
443 manager.populate_from_game(vec![a, b])?;
444
445 let mut i: usize = 0;
446 let mut result = Vec::new();
447 while result.is_empty() {
448 manager.iterate(&evaluator)?;
449 result = manager.next()?;
450
451 i += 1;
452 }
453
454 assert_eq!(result.len(), 2);
455 assert_eq!(i, 4+1);
456
457 Ok(())
458 }
459
460 #[test]
461 fn test_batch_iterate_2() -> Result<(), MctsError>{
462 let mut manager = MctsBatch::<GameTest, 4>::new();
463 let evaluator: GameEvaluatorTest2 = GameEvaluatorTest2::new();
464
465 let a: GameTest = GameTest::new();
466 let mut b: GameTest = GameTest::new();
467 b.play(4);
468
469 manager.populate_from_game(vec![a, b])?;
470 assert_eq!(manager.get_count(), 2);
471
472 let mut i: usize = 0;
473 let mut result = Vec::new();
474 while result.is_empty() {
475 manager.iterate(&evaluator)?;
476 result = manager.next()?;
477
478 i += 1;
479 }
480
481 assert_eq!(result.len(), 1);
482 assert_eq!(i, 3+1);
483 assert_eq!(manager.get_count(), 1);
484
485 manager.iterate(&evaluator)?;
486 result = manager.next()?;
487
488 assert_eq!(result.len(), 1);
489 assert_eq!(manager.get_count(), 0);
490 Ok(())
491 }
492
493 #[test]
494 fn test_batch_iterate_3() -> Result<(), MctsError>{
495 let mut manager = MctsBatch::<GameTest, 4>::new();
496 let evaluator: GameEvaluatorTest2 = GameEvaluatorTest2::new();
497
498 manager.populate_from_game(vec![GameTest::new()])?;
499
500 let mut i: usize = 0;
501 let mut result = Vec::new();
502 while result.is_empty() {
503 for _ in 0..18{
504 manager.iterate(&evaluator)?;
505 }
506 result = manager.next()?;
507
508 i += 1;
509 }
510
511 assert_eq!(result.len(), 1);
512 assert_eq!(i, 4+1);
513
514 assert_eq!((result[0][0].0).get_actions(), [true, true, true, true]);
515 assert_eq!((result[0][0].1), 1.0);
516 Ok(())
517 }
518}