simple_mcts/wrapper/mcts.rs
1use crate::{Action, Engine, MctsEngineError, ScoreTransformer, SelectionFunction, SelectionResult, StateEvaluation};
2
3/// Configuration for the MCTS batch manager.
4///
5/// Contains the hyperparameters and functions required to run the selection
6/// and backpropagation phases.
7#[derive(Debug)]
8pub struct MctsConfig<F: SelectionFunction, S: ScoreTransformer>{
9 c: f32,
10 selection_function: F,
11 score_transformer: S,
12}
13
14impl<F: SelectionFunction, S: ScoreTransformer> MctsConfig<F, S> {
15 /// Creates a new configuration for the MCTS engine.
16 ///
17 /// # Arguments
18 ///
19 /// * `c` - The exploration hyperparameter (often sqrt(2)).
20 /// * `selection_function` - The heuristic used to select nodes (e.g., UCB1 or PUCT).
21 /// * `score_transformer` - The function applied to scores during backpropagation.
22 ///
23 /// # Returns
24 ///
25 /// A new `MctsConfig` instance.
26 pub fn new(c: f32, selection_function: F, score_transformer: S) -> Self {
27 Self { c, selection_function, score_transformer }
28 }
29}
30
31/// A unique, strongly-typed identifier representing a specific `Engine` instance
32/// within the `MctsBatch` manager.
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
34pub struct MctsId(usize);
35
36/// Represents the output of a selection phase for a specific engine.
37pub struct MctsSelectionResult {
38 pub id: MctsId,
39 pub selection: SelectionResult,
40 pub path: Vec<Action>,
41}
42
43/// A structure bundling an MCTS evaluation with its targeted engine ID.
44pub struct MctsStateEvaluation<const N: usize> {
45 pub id: MctsId,
46 pub selection: SelectionResult,
47 pub evaluation: StateEvaluation<N>
48}
49
50/// A structure bundling the final score array with its targeted engine ID.
51pub struct MctsScore<const N: usize> {
52 pub id: MctsId,
53 pub scores: [i32; N]
54}
55
56/// A structure bundling a game action with its targeted engine ID.
57pub struct MctsAction {
58 pub id: MctsId,
59 pub action: Action
60}
61
62/// Errors that can occur when interacting with the `MctsBatch` manager.
63#[derive(Copy, Clone, Debug, PartialEq)]
64pub enum MctsError{
65 /// An error occurred deep within the individual MCTS tree engine.
66 EngineError(MctsEngineError, MctsId),
67 /// The provided `MctsId` does not correspond to any active engine.
68 InvalidMctsId(MctsId),
69}
70
71type EngineCollection<const N: usize> = Vec<Option<Engine<N>>>;
72pub type IdCollection = Vec<MctsId>;
73pub type ResultCollection = Vec<MctsSelectionResult>;
74pub type StateEvaluationCollection<const N: usize> = Vec<MctsStateEvaluation<N>>;
75pub type ScoreCollection<const N: usize> = Vec<MctsScore<N>>;
76pub type ActionCollection = Vec<MctsAction>;
77
78/// The Batch Manager for MCTS Engines.
79///
80/// It acts as a memory-efficient Arena Allocator (Slot Map) that can manage
81/// thousands of independent MCTS trees simultaneously without reallocating memory.
82pub struct MctsBatch<F: SelectionFunction, S: ScoreTransformer, const N: usize> {
83 config: MctsConfig<F, S>,
84 engines: EngineCollection<N>,
85 free_list: Vec<MctsId>,
86
87 sessions: usize
88}
89
90impl<F: SelectionFunction, S: ScoreTransformer, const N: usize> MctsBatch<F, S, N> {
91 /// Initializes a new, empty batch manager from a given configuration.
92 ///
93 /// # Arguments
94 ///
95 /// * `config` - The `MctsConfig` specifying exploration rate and heuristics.
96 ///
97 /// # Returns
98 ///
99 /// A new, empty `MctsBatch` instance.
100 ///
101 /// # Examples
102 ///
103 /// ```
104 /// # use crate::simple_mcts::{MctsBatch, MctsConfig};
105 /// let config = MctsConfig::new(1.414, |w, n, p_n, p, c| 0.0, |s| -s);
106 /// let mut batch = MctsBatch::<_, _, 9>::from_config(config);
107 /// ```
108 pub fn from_config(config: MctsConfig<F, S>) -> Self {
109 Self {
110 config,
111 engines: Vec::new(),
112 sessions: 0,
113 free_list: Vec::new(),
114 }
115 }
116
117 /// Allocates a new MCTS engine in the batch.
118 ///
119 /// Memory slots are recycled using a free-list to guarantee O(1) amortized performance.
120 ///
121 /// # Returns
122 ///
123 /// Returns the uniquely generated `MctsId` for the new engine.
124 pub fn add(&mut self) -> MctsId{
125 self.sessions += 1;
126
127 if let Some(id) = self.free_list.pop(){
128 self.engines[id.0] = Some(Engine::new());
129 id
130 }
131 else{
132 let id = self.engines.len();
133 self.engines.push(Some(Engine::new()));
134 MctsId(id)
135 }
136 }
137
138 /// Pre-allocates and populates the batch with a specified number of new engines.
139 ///
140 /// # Arguments
141 ///
142 /// * `n` - The number of engines to instantiate.
143 ///
144 /// # Returns
145 ///
146 /// A collection (`IdCollection`) containing all the newly minted `MctsId`s.
147 pub fn populate(&mut self, n: usize) -> IdCollection {
148 let mut ids = Vec::with_capacity(n);
149
150 for _ in 0..n {
151 ids.push(self.add());
152 }
153
154 ids
155 }
156
157 /// Safely deallocates a specific engine and marks its memory slot for reuse.
158 ///
159 /// Silently ignores the operation if the provided ID is already deleted or out of bounds.
160 ///
161 /// # Arguments
162 ///
163 /// * `id` - The `MctsId` of the engine to remove.
164 pub fn remove_one(&mut self, id: MctsId){
165 if let Some(slot) = self.engines.get_mut(id.0){
166 if slot.take().is_some() {
167 self.sessions -= 1;
168 self.free_list.push(id);
169 };
170 }
171 }
172
173 /// Deallocates a collection of engines.
174 ///
175 /// # Arguments
176 ///
177 /// * `ids` - The collection of `MctsId`s to remove from the batch.
178 pub fn remove(&mut self, ids: IdCollection) {
179 for id in ids {
180 self.remove_one(id);
181 }
182 }
183
184 /// Deallocates all active engines and resets the free-list.
185 ///
186 /// # Note
187 ///
188 /// This completely invalidates any `MctsId` currently held by the user.
189 pub fn clear(&mut self) {
190 self.engines.clear();
191 self.free_list.clear();
192 self.sessions = 0;
193 }
194
195 /// Returns the number of currently active engines in the batch.
196 ///
197 /// # Returns
198 ///
199 /// An `usize` representing the active session count.
200 pub fn sessions(&self) -> usize {
201 self.sessions
202 }
203
204 /// Triggers the selection phase across all active engines.
205 ///
206 /// # Returns
207 ///
208 /// A `ResultCollection` containing the traversal paths and selection outcomes
209 /// for every active engine.
210 ///
211 /// # Examples
212 ///
213 /// ```
214 /// # use crate::simple_mcts::{MctsBatch, MctsConfig};
215 /// # let config = MctsConfig::new(1.414, |w, n, p_n, p, c| 0.0, |s| -s);
216 /// # let mut batch = MctsBatch::<_, _, 9>::from_config(config);
217 /// batch.populate(3);
218 /// let selections = batch.selection();
219 /// assert_eq!(selections.len(), 3);
220 /// ```
221 pub fn selection(&self) -> ResultCollection {
222 self.engines.iter().enumerate()
223 .filter_map(|(id, opt)| {
224 let engine = opt.as_ref()?;
225
226 let mut path = Vec::new();
227 let selection = engine.select(&mut path, &self.config.selection_function, self.config.c);
228
229 Some(MctsSelectionResult {
230 id: MctsId(id),
231 selection,
232 path
233 })
234 }).collect()
235 }
236
237 /// Updates a single engine with an evaluation result using a custom, localized score transformer.
238 ///
239 /// This method acts as an "escape hatch" for complex or asymmetric games (e.g., 3-player games,
240 /// hidden information) where the backpropagation logic depends on the specific context of the
241 /// ongoing game, rather than the universal rule defined in `MctsConfig`.
242 ///
243 /// # Arguments
244 ///
245 /// * `evaluation` - The evaluated state to be integrated into the specific tree.
246 /// * `score_transformer` - A mutable reference to a custom closure or struct implementing `ScoreTransformer`.
247 ///
248 /// # Errors
249 ///
250 /// Returns an `MctsError::InvalidMctsId` if the ID doesn't point to an active engine.
251 /// Returns an `MctsError::EngineError` if the internal tree expansion or backpropagation fails.
252 ///
253 /// # Examples
254 ///
255 /// ```
256 /// # use crate::simple_mcts::{MctsBatch, MctsConfig, MctsStateEvaluation, StateEvaluation};
257 /// # let config = MctsConfig::new(1.414, |w, n, p_n, p, c| 0.0, |s| -s);
258 /// # let mut batch = MctsBatch::<_, _, 2>::from_config(config);
259 /// # let id = batch.add();
260 /// # let mut selections = batch.selection();
261 /// # let sel = selections.pop().unwrap();
262 /// let eval = MctsStateEvaluation {
263 /// id,
264 /// selection: sel.selection,
265 /// evaluation: StateEvaluation::Terminal(1.0)
266 /// };
267 ///
268 /// // Using a custom transformer capturing a local state/context
269 /// let mut local_scores = [1.0, -0.5, -0.5]; // e.g., 3-player score array
270 ///
271 /// batch.update_one_with_transformer(eval, &mut |score| {
272 /// // Custom asymmetric logic using the local context
273 /// score * local_scores[0]
274 /// }).unwrap();
275 /// ```
276 pub fn update_one_with_transformer(&mut self, evaluation: MctsStateEvaluation<N>, score_transformer: &mut impl ScoreTransformer) -> Result<(), MctsError> {
277 let opt = self.engines.get_mut(evaluation.id.0);
278
279 if let Some(Some(engine)) = opt {
280 engine
281 .update(evaluation.evaluation, evaluation.selection, score_transformer)
282 .map_err(|err| MctsError::EngineError(err, evaluation.id))?
283 }
284 else{
285 return Err(MctsError::InvalidMctsId(evaluation.id));
286 }
287
288 Ok(())
289 }
290
291 /// Updates a single engine with an evaluation result and backpropagates the score.
292 ///
293 /// # Arguments
294 ///
295 /// * `evaluation` - The evaluated state to be integrated into the specific tree.
296 ///
297 /// # Errors
298 ///
299 /// Returns an `MctsError::InvalidMctsId` if the ID doesn't point to an active engine.
300 /// Returns an `MctsError::EngineError` if the internal tree expansion or backpropagation fails.
301 pub fn update_one(&mut self, evaluation: MctsStateEvaluation<N>) -> Result<(), MctsError> {
302 let opt = self.engines.get_mut(evaluation.id.0);
303
304 if let Some(Some(engine)) = opt {
305 engine
306 .update(evaluation.evaluation, evaluation.selection, &mut self.config.score_transformer)
307 .map_err(|err| MctsError::EngineError(err, evaluation.id))?
308 }
309 else{
310 return Err(MctsError::InvalidMctsId(evaluation.id));
311 }
312
313 Ok(())
314 }
315
316 /// Transactionally updates multiple engines with their respective evaluations.
317 ///
318 /// This function consumes the provided vector. If an error occurs, the function
319 /// aborts and pushes the problematic evaluation back into the vector to prevent data loss.
320 ///
321 /// # Arguments
322 ///
323 /// * `evaluations` - A mutable reference to a collection of evaluations to process.
324 ///
325 /// # Errors
326 ///
327 /// Returns an `MctsError` if any single engine update fails or if an ID is invalid.
328 ///
329 /// # Examples
330 ///
331 /// ```
332 /// # use crate::simple_mcts::{MctsBatch, MctsConfig, MctsStateEvaluation, StateEvaluation};
333 /// # let config = MctsConfig::new(1.414, |w, n, p_n, p, c| 0.0, |s| -s);
334 /// # let mut batch = MctsBatch::<_, _, 2>::from_config(config);
335 /// let id = batch.add();
336 /// let mut selections = batch.selection();
337 /// let sel = selections.pop().unwrap();
338 ///
339 /// let mut evals = vec![MctsStateEvaluation {
340 /// id,
341 /// selection: sel.selection,
342 /// evaluation: StateEvaluation::Terminal(1.0)
343 /// }];
344 ///
345 /// batch.update(&mut evals).unwrap();
346 /// ```
347 pub fn update(&mut self, evaluations: &mut StateEvaluationCollection<N>) -> Result<(), MctsError> {
348 while let Some(evaluation) = evaluations.pop() {
349 let opt = self.engines.get_mut(evaluation.id.0);
350
351 if let Some(Some(engine)) = opt {
352 engine
353 .update(evaluation.evaluation, evaluation.selection, &mut self.config.score_transformer)
354 .map_err(|err| {
355 let id = evaluation.id;
356 evaluations.push(evaluation);
357 MctsError::EngineError(err, id)
358 })?;
359 }
360 else{
361 let id = evaluation.id;
362 evaluations.push(evaluation);
363 return Err(MctsError::InvalidMctsId(id));
364 }
365 }
366
367 Ok(())
368 }
369
370 /// Retrieves the visit counts of all possible actions for a specific engine.
371 ///
372 /// # Arguments
373 ///
374 /// * `id` - The unique identifier of the target engine.
375 ///
376 /// # Returns
377 ///
378 /// Returns an array `[i32; N]` representing the visit distribution of the root node.
379 ///
380 /// # Errors
381 ///
382 /// Returns `MctsError::InvalidMctsId` if the specified engine is inactive or deleted.
383 pub fn scores_one(&self, id: MctsId) -> Result<[i32; N], MctsError> {
384 let opt = self.engines.get(id.0);
385
386 if let Some(Some(engine)) = opt {
387 Ok(engine.scores())
388 }
389 else{
390 Err(MctsError::InvalidMctsId(id))
391 }
392 }
393
394 /// Retrieves the visit counts of all possible actions for all active engines.
395 ///
396 /// # Returns
397 ///
398 /// A `ScoreCollection` containing the score arrays paired with their respective engine IDs.
399 pub fn scores(&self) -> ScoreCollection<N> {
400 self.engines.iter().enumerate()
401 .filter_map(|(id, opt)| {
402 let engine = opt.as_ref()?;
403 Some(MctsScore {
404 id: MctsId(id),
405 scores: engine.scores()
406 })
407 }).collect()
408 }
409
410 /// Transactionally commits a list of played actions, updating the internal roots
411 /// of the corresponding engines.
412 ///
413 /// # Arguments
414 ///
415 /// * `actions` - A mutable reference to the collection of actions to commit.
416 ///
417 /// # Errors
418 ///
419 /// Returns `MctsError::InvalidMctsId` if an action targets a non-existent engine.
420 ///
421 /// # Panics
422 ///
423 /// Panics if an action index is out of bounds (i.e., `>= N`) or if the internal
424 /// tree structure of an engine is corrupted.
425 pub fn commit_actions(&mut self, actions: &mut ActionCollection) -> Result<(), MctsError> {
426 while let Some(action) = actions.pop() {
427 if let Some(Some(engine)) = self.engines.get_mut(action.id.0) {
428 engine.commit_action(action.action);
429 }
430 else {
431 let id = action.id;
432 actions.push(action);
433 return Err(MctsError::InvalidMctsId(id));
434 }
435 }
436
437 Ok(())
438 }
439
440 /// Commits a single played action, moving the root of the specified engine.
441 ///
442 /// # Arguments
443 ///
444 /// * `action` - The `MctsAction` containing the target ID and the action index.
445 ///
446 /// # Errors
447 ///
448 /// Returns `MctsError::InvalidMctsId` if the target engine does not exist.
449 ///
450 /// # Panics
451 ///
452 /// Panics if the action index is out of bounds (i.e., `>= N`) or if the internal
453 /// tree structure is corrupted.
454 pub fn commit_action(&mut self, action: MctsAction) -> Result<(), MctsError> {
455 if let Some(Some(engine)) = self.engines.get_mut(action.id.0) {
456 engine.commit_action(action.action);
457 Ok(())
458 }
459 else { Err(MctsError::InvalidMctsId(action.id)) }
460 }
461}
462
463#[cfg(test)]
464mod tests {
465 use super::*;
466
467 fn dummy_selection(_: f32, _: i32, _: i32, _: f32, _: f32) -> f32 { 0.0 }
468 fn dummy_score_transformer(score: f32) -> f32 { -score }
469
470 fn create_test_batch() -> MctsBatch<impl SelectionFunction, impl ScoreTransformer, 9> {
471 let config = MctsConfig::new(1.0, dummy_selection, dummy_score_transformer);
472 MctsBatch::from_config(config)
473 }
474
475 #[test]
476 fn test_lifecycle_add_and_populate() {
477 let mut batch = create_test_batch();
478
479 let id1 = batch.add();
480 assert_eq!(id1, MctsId(0));
481 assert_eq!(batch.sessions(), 1);
482
483 let ids = batch.populate(3);
484 assert_eq!(ids, vec![MctsId(1), MctsId(2), MctsId(3)]);
485 assert_eq!(batch.sessions(), 4);
486 }
487
488 #[test]
489 fn test_freelist_reuse() {
490 let mut batch = create_test_batch();
491 batch.populate(3);
492
493 batch.remove_one(MctsId(1));
494 assert_eq!(batch.sessions(), 2);
495
496 let new_id = batch.add();
497 assert_eq!(new_id, MctsId(1));
498 assert_eq!(batch.sessions(), 3);
499 }
500
501 #[test]
502 fn test_remove_idempotence() {
503 let mut batch = create_test_batch();
504 batch.populate(2);
505
506 batch.remove_one(MctsId(0));
507 assert_eq!(batch.sessions(), 1);
508 batch.remove_one(MctsId(0));
509 assert_eq!(batch.sessions(), 1);
510 assert_eq!(batch.free_list.len(), 1);
511 }
512
513 #[test]
514 fn test_clear() {
515 let mut batch = create_test_batch();
516 batch.populate(5);
517 batch.remove_one(MctsId(2));
518
519 batch.clear();
520 assert_eq!(batch.sessions(), 0);
521 assert_eq!(batch.engines.len(), 0);
522 assert_eq!(batch.free_list.len(), 0);
523 }
524
525 #[test]
526 fn test_selection_success() {
527 let mut batch = create_test_batch();
528 batch.populate(2);
529
530 let selections = batch.selection();
531 assert_eq!(selections.len(), 2);
532 assert_eq!(selections[0].id, MctsId(0));
533 assert_eq!(selections[1].id, MctsId(1));
534 }
535
536 #[test]
537 fn test_update_and_scores_success() {
538 let mut batch = create_test_batch();
539 let id = batch.add();
540
541 let mut selections = batch.selection();
542 let sel = selections.pop().unwrap();
543
544 let eval = MctsStateEvaluation {
545 id,
546 selection: sel.selection,
547 evaluation: StateEvaluation::Terminal(1.0),
548 };
549
550 let mut evaluations = vec![eval];
551 let result = batch.update(&mut evaluations);
552
553 assert!(result.is_ok());
554 assert!(evaluations.is_empty(), "The evaluation should have been completed.");
555
556 let scores = batch.scores_one(id);
557 assert!(scores.is_ok());
558 }
559
560 #[test]
561 fn test_update_one_invalid_id() {
562 let mut batch = create_test_batch();
563
564 let eval = MctsStateEvaluation {
565 id: MctsId(99),
566 selection: SelectionResult::Empty,
567 evaluation: StateEvaluation::Terminal(1.0),
568 };
569
570 let result = batch.update_one(eval);
571 assert!(matches!(result, Err(MctsError::InvalidMctsId(MctsId(99)))));
572 }
573
574 #[test]
575 fn test_update_batch_transactional_rollback_on_invalid_id() {
576 let mut batch = create_test_batch();
577 let id_valid = batch.add();
578
579 let eval_valid = MctsStateEvaluation {
580 id: id_valid,
581 selection: SelectionResult::Empty,
582 evaluation: StateEvaluation::Terminal(1.0),
583 };
584 let eval_invalid = MctsStateEvaluation {
585 id: MctsId(99),
586 selection: SelectionResult::Empty,
587 evaluation: StateEvaluation::Terminal(-1.0),
588 };
589
590 let mut evaluations = vec![eval_valid, eval_invalid];
591
592 let result = batch.update(&mut evaluations);
593
594 assert!(matches!(result, Err(MctsError::InvalidMctsId(MctsId(99)))));
595 assert_eq!(evaluations.len(), 2, "The transaction failed; no data must be lost.");
596 assert_eq!(evaluations.last().unwrap().id, MctsId(99));
597 }
598
599 #[test]
600 fn test_commit_actions_transactional_rollback() {
601 let mut batch = create_test_batch();
602 let id_valid = batch.add();
603
604 let action_valid = MctsAction { id: id_valid, action: Action::new(0) };
605 let action_invalid = MctsAction { id: MctsId(99), action: Action::new(0) };
606
607 let mut actions = vec![action_valid, action_invalid];
608
609 let result = batch.commit_actions(&mut actions);
610
611 assert!(matches!(result, Err(MctsError::InvalidMctsId(MctsId(99)))));
612 assert_eq!(actions.len(), 2);
613 }
614
615 #[test]
616 fn test_scores_invalid_id() {
617 let batch = create_test_batch();
618 let result = batch.scores_one(MctsId(42));
619 assert!(matches!(result, Err(MctsError::InvalidMctsId(MctsId(42)))));
620 }
621
622 #[test]
623 fn test_update_one_with_transformer_success() {
624 let mut batch = create_test_batch();
625 let id = batch.add();
626
627 let mut selections = batch.selection();
628 let sel = selections.pop().unwrap();
629
630 let eval = MctsStateEvaluation {
631 id,
632 selection: sel.selection,
633 evaluation: StateEvaluation::Terminal(1.0),
634 };
635
636 let mut calls_count = 0;
637 let my_local_multiplier = 5.0;
638
639 let mut custom_transformer = |score: f32| {
640 calls_count += 1;
641 score * my_local_multiplier
642 };
643
644 let result = batch.update_one_with_transformer(eval, &mut custom_transformer);
645
646 assert!(result.is_ok(), "The update should have succeeded.");
647 assert_eq!(calls_count, 1, "The custom transformer was not called the correct number of times.");
648 }
649
650 #[test]
651 fn test_update_one_with_transformer_invalid_id() {
652 let mut batch = create_test_batch();
653
654 let eval = MctsStateEvaluation {
655 id: MctsId(99),
656 selection: SelectionResult::Empty,
657 evaluation: StateEvaluation::Terminal(1.0),
658 };
659
660 let mut custom_transformer = |score: f32| score;
661
662 let result = batch.update_one_with_transformer(eval, &mut custom_transformer);
663
664 assert!(matches!(result, Err(MctsError::InvalidMctsId(MctsId(99)))));
665 }
666}