simple_mcts/engine/
engine.rs

1use crate::tree::*;
2
3/// Represents an index corresponding to a specific legal move or action.
4#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
5pub struct Action(usize);
6
7impl Action{
8    /// Constructs a new `Action`.
9    ///
10    /// # Arguments
11    /// * `action` - The integer index representing the move.
12    ///
13    /// # Returns
14    /// A new `Action` instance.
15    pub fn new(action: usize) -> Self{
16        Action(action)
17    }
18
19    /// Returns the underlying integer index of the action.
20    ///
21    /// # Returns
22    /// The `usize` value stored inside the `Action`.
23    pub fn action(&self) -> usize{
24        self.0
25    }
26}
27
28/// Represents the internal tactical state of a node in the MCTS tree.
29#[derive(Copy, Clone, Debug, PartialEq)]
30enum MctsNodeState{
31    /// The node is active and can be expanded or traversed further.
32    Active,
33    /// The node represents a terminal state with a fixed final game score.
34    Terminal(f32)
35}
36
37/// Internal container representing a single state within the Monte Carlo Tree Search.
38///
39/// It utilizes a Structure of Arrays (SoA) layout for its children's statistics
40/// (`children_scores`, `children_visits`, `children_policies`) to guarantee
41/// cache-friendly memory lookups during the intensive selection phase.
42#[derive(Debug, Copy, Clone, PartialEq)]
43struct MctsNode<const N: usize>{
44    children_scores: [f32; N],
45    children_visits: [i32; N],
46    children_policies: [f32; N],
47
48    visits: i32,
49    state: MctsNodeState,
50    action: Action,
51}
52
53
54/// Evaluates a child node during the tree selection phase.
55///
56/// This function determines whether the current node should be prioritized.
57/// It typically implements an exploration/exploitation
58/// balance formula such as UCB1.
59///
60/// # Arguments
61///
62/// * `score` - The score value of the child node (w).
63/// * `node_visits` - The number of visits to the currently evaluated child node (n).
64/// * `parent_visits` - The total number of visits to the parent node (N).
65/// * `policy` - The exploration policy value.
66/// * `c` - The exploration constant used to adjust the ratio (often sqrt(2)).
67///
68/// # Returns
69///
70/// Returns a floating-point score representing the node's priority
71pub trait SelectionFunction: Fn(f32, i32, i32, f32, f32) -> f32 {}
72impl<T: Fn(f32, i32, i32, f32, f32) -> f32> SelectionFunction for T {}
73
74/// Transforms a score during the backpropagation phase.
75///
76/// This is typically used in alternating-turn games (like Chess or Tic-Tac-Toe)
77/// to invert the score (e.g., `|score| -score`) from one depth to another.
78pub trait ScoreTransformer: FnMut(f32) -> f32 {}
79impl<T: FnMut(f32) -> f32> ScoreTransformer for T {}
80
81/// A buffer used to store the sequence of actions traversed during the selection phase.
82pub trait ResettableBuffer {
83    /// Pushes an action to the end of the buffer.
84    fn push(&mut self, value: Action);
85    /// Clears the buffer, removing all elements.
86    fn clear(&mut self);
87}
88
89/// Implements `ResettableBuffer` for standard collections.
90#[macro_export]
91macro_rules! impl_resettable_buffer {
92    ($type:ty) => {
93        impl ResettableBuffer for $type {
94            fn push(&mut self, value: Action) {
95                self.push(value);
96            }
97            fn clear(&mut self) {
98                self.clear();
99            }
100        }
101    };
102}
103
104impl_resettable_buffer!(Vec<Action>);
105
106impl<const N: usize> MctsNode<N> {
107    /// Constructs a new, unvisited `MctsNode` with prior policies, an incoming action, and a state.
108    fn new(policies: [f32; N], action: Action, state: MctsNodeState) -> Self {
109        MctsNode{
110            children_scores: [0.; N],
111            children_visits:  [0; N],
112            children_policies: policies,
113            visits: 0,
114            action,
115            state
116        }
117    }
118
119    /// Selects the best child action index based on the provided selection heuristic.
120    ///
121    /// It systematically filters out invalid actions (where policy <= 0.0) and uses
122    /// `total_cmp` to safely sort floating-point numbers.
123    ///
124    /// # Panics
125    ///
126    /// Panics if no legal actions remain (i.e., all policies are less than or equal to 0.0),
127    /// indicating a logical mismatch with the game wrapper's state.
128    fn best_child(&self, score_f: &impl SelectionFunction, c: f32) -> usize {
129        self.children_scores.iter()
130            .zip(self.children_visits.iter())
131            .zip(self.children_policies.iter())
132            .enumerate()
133            .filter_map(
134            |(i, ((c_scores, c_visits), policy))| {
135                if *policy <= 0. { None } else { Some((i, score_f(*c_scores, *c_visits, self.visits, *policy, c))) }
136            }
137        ).max_by(
138            |(_a, a), (_b, b)| { a.total_cmp(b) }
139        ).expect("Error during the selection of the best action.").0
140    }
141}
142
143/// A strictly typed identifier for an instantiated node within the MCTS tree.
144#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
145pub struct MctsNodeId(NodeId);
146
147/// The core Engine handling the Monte Carlo Tree Search algorithm.
148///
149/// It manages the internal tree structure and provides methods to traverse,
150/// expand, and backpropagate statistics.
151pub struct Engine<const N: usize> {
152    tree: Tree<MctsNode<N>, N>
153}
154
155/// Represents the outcome of a tree selection phase.
156#[derive(Copy, Clone, Debug, PartialEq)]
157pub enum SelectionResult {
158    /// The tree is entirely empty.
159    Empty,
160    /// A leaf node was reached. Provides the parent ID and the selected action.
161    Active(MctsNodeId, Action),
162    /// A terminal node was reached during traversal. Provides the node ID and final score.
163    Terminal(MctsNodeId, f32)
164}
165
166/// Represents the evaluation of a game state, usually provided by a heuristic function
167/// or a neural network (e.g., in AlphaZero/MuZero architectures).
168///
169/// This enum strictly binds the game's termination status with its corresponding data,
170/// preventing the representation of invalid states (like a terminal state having future policies).
171///
172/// # Note on Score Perspective
173/// The `score` value provided here (both in `Active` and `Terminal` states) must
174/// always be expressed from the perspective of the player who **just made the move**
175/// that led to this state. This ensures correct backpropagation in alternating-turn games
176/// when using the engine's `score_updater` logic.
177#[derive(Copy, Clone, Debug, PartialEq)]
178pub enum StateEvaluation<const N: usize> {
179    /// The game is ongoing. Contains the current board evaluation score and the prior
180    /// probabilities (policies) for the next `N` possible actions.
181    Active(f32, [f32; N]),
182
183    /// The game has ended. Contains only the final fixed score (e.g., 1.0 for win, -1.0 for loss).
184    Terminal(f32)
185}
186
187impl<const N: usize> StateEvaluation<N>{
188    /// Extracts the score from the evaluation, regardless of whether the state is active or terminal.
189    ///
190    /// # Returns
191    ///
192    /// A floating-point value (`f32`) representing the underlying score of this state.
193    pub fn score(&self) -> f32{
194        match self {
195            StateEvaluation::Active(score, _) => *score,
196            StateEvaluation::Terminal(score) => *score
197        }
198    }
199}
200
201/// Errors that can occur during the execution of the MCTS engine.
202#[derive(Debug, Clone, Copy, PartialEq)]
203pub enum MctsEngineError{
204    UnknownError,
205    InvalidMctsNode,
206    ChildAlreadyExists,
207    InvalidAction,
208    SelectionIsTerminal
209}
210
211impl<const N: usize> Engine<N> {
212    /// Creates a new, empty MCTS engine.
213    ///
214    /// # Examples
215    ///
216    /// ```
217    /// # use crate::simple_mcts::Engine;
218    /// let mut engine = Engine::<7>::new();
219    /// ```
220    pub fn new() -> Self {
221        Engine::<N> {
222            tree: Tree::<MctsNode<N>, N>::new()
223        }
224    }
225
226    /// Creates a new MCTS engine with a pre-allocated capacity for the underlying tree.
227    ///
228    /// This is recommended to avoid reallocations when running millions of simulations.
229    ///
230    /// # Examples
231    ///
232    /// ```
233    /// # use crate::simple_mcts::Engine;
234    /// let mut engine = Engine::<7>::with_capacity(1_000_000);
235    /// ```
236    pub fn with_capacity(capacity: usize) -> Self {
237        Engine::<N> {
238            tree: Tree::<MctsNode<N>, N>::with_capacity(capacity)
239        }
240    }
241
242    /// Traverses the tree from the root to a leaf node according to the selection function.
243    ///
244    /// The traversal path (sequence of actions) is recorded in `path_out`.
245    ///
246    /// # Arguments
247    ///
248    /// * `path_out` - A mutable reference to a buffer that will be cleared and filled with the chosen actions.
249    /// * `score_f` - The selection heuristic (e.g., PUCT formula).
250    /// * `c` - The exploration hyperparameter.
251    ///
252    /// # Returns
253    ///
254    /// A `SelectionResult` indicating whether a new leaf was found, a terminal state was hit, or the tree is empty.
255    ///
256    /// # Panics
257    ///
258    /// Panics if the internal tree structure is corrupted (i.e., a parent node ID points to a non-existent element).
259    /// Also panics if an active node has no legal moves (all policies are <= 0.0).
260    ///
261    /// # Examples
262    ///
263    /// ```
264    /// # use crate::simple_mcts::Engine;
265    /// # use crate::simple_mcts::puct;
266    /// let mut path = Vec::new();
267    /// let mut engine = Engine::<7>::new();
268    /// let result = engine.select(&mut path, &puct, 1.414);
269    /// ```
270    pub fn select(&self, path_out: &mut impl ResettableBuffer, score_f: &impl SelectionFunction, c: f32) -> SelectionResult{
271        path_out.clear();
272
273        if let Some(mut current_id) = self.tree.root(){
274            loop {
275                let current = self.tree.get(current_id).unwrap();
276
277                match current.data().state {
278                    MctsNodeState::Terminal(score) => {
279                        return SelectionResult::Terminal(MctsNodeId(current_id), score);
280                    }
281                    MctsNodeState::Active => {
282                        let action = current.data().best_child(score_f, c);
283
284                        path_out.push(Action(action));
285                        if let Some(child_id) = current.child(action){
286                            current_id = child_id;
287                        }
288                        else{
289                            return SelectionResult::Active(MctsNodeId(current_id), Action(action));
290                        }
291                    }
292                }
293            }
294        }
295        else{ SelectionResult::Empty }
296    }
297
298    /// Expands the tree by attaching a new child node to the result of a previous selection.
299    ///
300    /// # Arguments
301    ///
302    /// * `evaluation` - The evaluated data of the new state (score and policies if active, or just the final score if terminal).
303    /// * `selection` - The `SelectionResult` obtained from a previous call to `select`.
304    ///
305    /// # Returns
306    ///
307    /// Returns the newly minted `MctsNodeId` on success, or an `MctsEngineError`.
308    ///
309    /// # Panics
310    ///
311    /// Panics if the `SelectionResult` contains an out-of-bounds action index (`>= N`).
312    /// Because valid selections are exclusively generated by the engine's `select` method,
313    /// this panic indicates that the provided `SelectionResult` was either manually forged,
314    /// corrupted, or used entirely out of context.
315    ///
316    /// # Examples
317    ///
318    /// ```
319    /// # use crate::simple_mcts::{Engine, StateEvaluation};
320    /// # use crate::simple_mcts::puct;
321    /// let mut path = Vec::new();
322    /// let mut engine = Engine::<2>::new();
323    /// let result = engine.select(&mut path, &puct, 1.414);
324    ///
325    /// // Expand an ongoing game state
326    /// let evaluation = StateEvaluation::Active(-0.2, [0.3, 0.7]);
327    /// engine.expand(evaluation, result).unwrap();
328    /// ```
329    pub fn expand(&mut self, evaluation: StateEvaluation<N>, selection: SelectionResult) -> Result<MctsNodeId, MctsEngineError>{
330        let (state, policies) = match evaluation {
331            StateEvaluation::Active(_, policies) => (MctsNodeState::Active, policies),
332            StateEvaluation::Terminal(score) => (MctsNodeState::Terminal(score), [0.; N])
333        };
334
335        match selection {
336            SelectionResult::Empty => {
337                match self.tree.set_root(MctsNode::new(policies, Action(0), state)) {
338                    Err(TreeError::RootAlreadyExists) => Err(MctsEngineError::ChildAlreadyExists),
339                    Err(_) => Err(MctsEngineError::UnknownError),
340                    Ok(node) => Ok(MctsNodeId(node)),
341                }
342            },
343            SelectionResult::Active(child_id, action) => {
344                match self.tree.add(child_id.0, action.0, MctsNode::new(policies, action, state)) {
345                    Err(TreeError::ParentDoesntExist) => Err(MctsEngineError::InvalidMctsNode),
346                    Err(TreeError::ChildAlreadyExists) => Err(MctsEngineError::ChildAlreadyExists),
347                    Err(_) => Err(MctsEngineError::UnknownError),
348                    Ok(node) => Ok(MctsNodeId(node))
349                }
350            },
351            SelectionResult::Terminal(_, _) => Err(MctsEngineError::SelectionIsTerminal)
352        }
353    }
354
355    /// Backpropagates a simulation score up to the root of the tree.
356    ///
357    /// Every node encountered on the way up will have its visits incremented and
358    /// its children statistics updated. The score is transformed at every step using `score_updater`.
359    ///
360    /// # Arguments
361    ///
362    /// * `node` - The `MctsNodeId` from which to start the backpropagation (usually the newly expanded leaf).
363    /// * `score` - The initial score to propagate.
364    /// * `score_updater` - A function or closure applied to the score at each step (e.g., negating it).
365    ///
366    /// # Returns
367    ///
368    /// Returns `Ok(())` on success, or an error if an invalid node ID is encountered.
369    ///
370    /// # Examples
371    ///
372    /// ```
373    /// # use simple_mcts::StateEvaluation;
374    /// use crate::simple_mcts::Engine;
375    /// # use crate::simple_mcts::puct;
376    /// let mut path = Vec::new();
377    /// let mut engine = Engine::<2>::new();
378    /// let result = engine.select(&mut path, &puct, 1.414);
379    /// let evaluation = StateEvaluation::Active(-0.2, [0.3, 0.7]);
380    ///
381    /// // Invert the score at each step for an alternating-turn game
382    /// let node = engine.expand(evaluation, result).unwrap();
383    /// engine.backpropagate(node, 1.0, |s| -s).unwrap();
384    /// ```
385    pub fn backpropagate(&mut self, node: MctsNodeId, score: f32, mut score_updater: impl ScoreTransformer) -> Result<(), MctsEngineError> {
386        let mut current = Some(node.0);
387        let mut score = score;
388
389        while let Some(node) = current {
390            let node = self.tree.get_mut(node).map_err(|_| MctsEngineError::InvalidMctsNode)?;
391            node.data_mut().visits += 1;
392
393            let action = node.data().action.0;
394            current = node.parent();
395
396            if let Some(parent) = current {
397                let parent = self.tree.get_mut(parent).map_err(|_| MctsEngineError::InvalidMctsNode)?;
398                parent.data_mut().children_visits[action] += 1;
399                parent.data_mut().children_scores[action] += score;
400            }
401
402            score = score_updater(score);
403        }
404
405        Ok(())
406    }
407
408    /// A convenience method that sequentially performs `expand` and `backpropagate`.
409    ///
410    /// If the selection provided was already terminal, expansion is skipped and
411    /// backpropagation starts immediately from the terminal node.
412    ///
413    /// # Arguments
414    ///
415    /// * `evaluation` - The evaluated data of the leaf node (score and policies).
416    /// * `selection` - The result returned by a prior `select` call.
417    /// * `score_updater` - The transformation applied to the score during backpropagation (e.g., inverting it).
418    ///
419    /// # Examples
420    ///
421    /// ```
422    /// # use crate::simple_mcts::{Engine, StateEvaluation};
423    /// # use crate::simple_mcts::puct;
424    /// let mut path = Vec::new();
425    /// let mut engine = Engine::<2>::new();
426    /// let result = engine.select(&mut path, &puct, 1.414);
427    ///
428    /// // Expand and backpropagate in one step, inverting the score at each depth
429    /// let evaluation = StateEvaluation::Active(-0.2, [0.3, 0.7]);
430    /// engine.update(evaluation, result, |s| -s).unwrap();
431    /// ```
432    pub fn update(&mut self, evaluation: StateEvaluation<N>, selection: SelectionResult, score_updater: impl ScoreTransformer) -> Result<(), MctsEngineError> {
433        let score = evaluation.score();
434
435        let node = match selection {
436            SelectionResult::Empty => self.expand(evaluation, selection)?,
437            SelectionResult::Active(_child_id, _action) => self.expand(evaluation, selection)?,
438            SelectionResult::Terminal(child_id, _score) => child_id
439        };
440
441        self.backpropagate(node, score, score_updater)?;
442
443        Ok(())
444    }
445
446    /// Retrieves the visit counts of all possible actions from the root node.
447    ///
448    /// This array is typically used at the end of the MCTS cycle to decide the actual
449    /// move to play in the game.
450    ///
451    /// # Returns
452    ///
453    /// An array of integers representing the number of visits for each child branch.
454    /// Returns an array of zeros if the tree is currently empty.
455    ///
456    /// # Panics
457    ///
458    /// Panics if the internal tree has a root ID but the root node cannot be retrieved.
459    ///
460    /// # Examples
461    ///
462    /// ```
463    ///  # use crate::simple_mcts::Engine;
464    /// let mut engine = Engine::<3>::new();
465    /// let scores = engine.scores();
466    /// assert_eq!(scores, [0, 0, 0]);
467    /// ```
468    pub fn scores(&self) -> [i32; N]{
469        if let Some(root) = self.tree.root() {
470            self.tree.get(root).unwrap().data().children_visits
471        }
472        else {
473            [0; N]
474        }
475    }
476
477    /// Promotes the child node corresponding to the given action to the new root of the tree.
478    ///
479    /// This method updates the MCTS tree to reflect a move played on the actual game board.
480    /// It performs an amortized memory cleanup (compacting) if the number of unreachable
481    /// nodes exceeds a heuristic threshold (twice the number of visits of the new root).
482    ///
483    /// If the action leads to a path that has not been explored by the MCTS, the current
484    /// tree is cleared, as it no longer contains valid information for the new state.
485    ///
486    /// # Arguments
487    ///
488    /// * `action` - The action that was performed on the game board.
489    ///
490    /// # Panics
491    ///
492    /// Panics if:
493    /// - The action index is out of bounds (i.e., `action.0 >= N`).
494    /// - The internal tree structure is corrupted (e.g., attempting to move the root to a non-existent node).
495    ///
496    /// # Note
497    ///
498    /// This is an amortized operation. `self.tree.compact()` is only called when
499    /// memory overhead becomes significant, ensuring high performance during game play.
500    pub fn commit_action(&mut self, action: Action) {
501        if let Some(root) = self.tree.root() {
502            let child = self.tree.child(root, action.0).unwrap();
503
504            if let Some(new_root_id) = child {
505                self.tree.move_root_to(new_root_id).unwrap();
506
507                let new_root_data = self.tree.data(new_root_id).unwrap();
508                if self.tree.allocated_nodes() > 2*new_root_data.visits as usize {
509                    self.tree.compact();
510                }
511            }
512            else{
513                self.tree.clear();
514            }
515        }
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use crate::{puct, negate_score};
522    use super::*;
523
524    #[test]
525    fn test_action_creation() {
526        let action = Action::new(42);
527        assert_eq!(action.action(), 42);
528    }
529
530    #[test]
531    fn test_action_equality() {
532        let a1 = Action::new(10);
533        let a2 = Action::new(10);
534        let a3 = Action::new(11);
535
536        assert_eq!(a1, a2, "Two actions with the same index should be equal");
537        assert_ne!(a1, a3, "Actions with different indices should not be equal");
538    }
539
540    #[test]
541    fn test_action_hashability() {
542        use std::collections::HashSet;
543
544        let mut set = HashSet::new();
545        let a1 = Action::new(5);
546        let a2 = Action::new(5);
547
548        set.insert(a1);
549
550        assert!(set.contains(&a2), "Action should be hashable and work in a HashSet");
551        assert_eq!(set.len(), 1, "HashSet should handle identical Actions correctly");
552    }
553
554    fn dummy_select(_score: f32, _node_visits: i32, _parent_visits: i32, policy: f32, _c: f32) -> f32 {
555        policy
556    }
557
558    fn identity_score(s: f32) -> f32 {
559        s
560    }
561
562    const N: usize = 2;
563
564    #[test]
565    fn test_engine_initialization() {
566        let engine = Engine::<N>::new();
567        assert!(engine.tree.root().is_none());
568    }
569
570    #[test]
571    fn test_select_on_empty_tree_returns_empty() {
572        let engine = Engine::<N>::new();
573        let mut path = Vec::new();
574        let selection = engine.select(&mut path, &dummy_select, 1.0);
575
576        assert_eq!(selection, SelectionResult::Empty);
577        assert!(path.is_empty());
578    }
579
580    #[test]
581    fn test_update_empty_creates_root() {
582        let mut engine = Engine::<N>::new();
583        let mut path = Vec::new();
584
585        let selection = engine.select(&mut path, &dummy_select, 1.0);
586
587        let evaluation = StateEvaluation::Active(1.0, [0.7, 0.3]);
588        let result = engine.update(evaluation, selection, negate_score);
589        assert!(result.is_ok());
590
591        let root_id = engine.tree.root().unwrap();
592        let root = engine.tree.get(root_id).unwrap();
593        assert_eq!(root.data().visits, 1);
594        assert_eq!(root.data().state, MctsNodeState::Active);
595    }
596
597    #[test]
598    fn test_select_active_node_and_expand_child() {
599        let mut engine = Engine::<N>::new();
600        let mut path = Vec::new();
601
602        let selection = engine.select(&mut path, &dummy_select, 1.0);
603        let evaluation = StateEvaluation::Active(0.0, [0.7, 0.3]);
604        engine.update(evaluation, selection, negate_score).unwrap();
605
606        let selection2 = engine.select(&mut path, &dummy_select, 1.0);
607
608        match selection2 {
609            SelectionResult::Active(id, action) => {
610                assert_eq!(id.0, engine.tree.root().unwrap());
611                assert_eq!(action, Action(0));
612            },
613            _ => panic!("Expected Active selection"),
614        }
615        assert_eq!(path, vec![Action(0)]);
616
617        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
618        let result = engine.update(evaluation, selection2, negate_score);
619        assert!(result.is_ok());
620    }
621
622    #[test]
623    fn test_backpropagate_values_correctly() {
624        let mut engine = Engine::<N>::new();
625        let mut path = Vec::new();
626
627        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
628        let evaluation = StateEvaluation::Active(0.0, [0.7, 0.3]);
629        engine.update(evaluation, sel_empty, negate_score).unwrap();
630
631        let sel_active = engine.select(&mut path, &dummy_select, 1.0);
632        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
633        engine.update(evaluation, sel_active, negate_score).unwrap();
634
635        let root_id = engine.tree.root().unwrap();
636        let root = engine.tree.get(root_id).unwrap();
637
638        assert_eq!(root.data().visits, 2);
639        assert_eq!(root.data().children_visits[0], 1);
640        assert_eq!(root.data().children_scores[0], 1.0);
641    }
642
643    #[test]
644    fn test_terminal_node_selection_and_state() {
645        let mut engine = Engine::<N>::new();
646        let mut path = Vec::new();
647
648        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
649        let evaluation = StateEvaluation::Terminal(42.0);
650        engine.update(evaluation, sel_empty, identity_score).unwrap();
651
652        let root_id = engine.tree.root().unwrap();
653        let root = engine.tree.get(root_id).unwrap();
654        assert_eq!(root.data().state, MctsNodeState::Terminal(42.0));
655
656        let sel_terminal = engine.select(&mut path, &dummy_select, 1.0);
657
658        match sel_terminal {
659            SelectionResult::Terminal(id, score) => {
660                assert_eq!(id.0, root_id);
661                assert_eq!(score, 42.0);
662            },
663            _ => panic!("Expected Terminal selection"),
664        }
665    }
666
667    #[test]
668    #[should_panic(expected = "Error during the selection of the best action.")]
669    fn test_best_child_panics_if_no_legal_moves() {
670        let mut engine = Engine::<N>::new();
671        let mut path = Vec::new();
672
673        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
674        let evaluation = StateEvaluation::Active(0.0, [0.0, 0.0]);
675        engine.update(evaluation, sel_empty, negate_score).unwrap();
676
677        engine.select(&mut path, &dummy_select, 1.0);
678    }
679
680    #[test]
681    fn test_expand_on_terminal_returns_error() {
682        let mut engine = Engine::<N>::new();
683        let mut path = Vec::new();
684
685        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
686        let evaluation = StateEvaluation::Terminal(1.0);
687        engine.update(evaluation, sel_empty, identity_score).unwrap();
688
689        let sel_terminal = engine.select(&mut path, &dummy_select, 1.0);
690        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
691        let result = engine.expand(evaluation, sel_terminal);
692
693        assert!(matches!(result, Err(MctsEngineError::SelectionIsTerminal)));
694    }
695
696    #[test]
697    fn test_adding_existing_child_returns_error() {
698        let mut engine = Engine::<N>::new();
699        let mut path = Vec::new();
700
701        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
702        let evaluation = StateEvaluation::Active(0.0, [0.7, 0.3]);
703        engine.update(evaluation, sel_empty, negate_score).unwrap();
704
705        let sel_active = engine.select(&mut path, &dummy_select, 1.0);
706        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
707        engine.update(evaluation, sel_active, negate_score).unwrap();
708
709        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
710        let result = engine.update(evaluation, sel_active, negate_score);
711        assert!(matches!(result, Err(MctsEngineError::ChildAlreadyExists)));
712    }
713
714    #[test]
715    fn test_scores_empty_engine() {
716        let engine = Engine::<N>::new();
717        // Vérifie qu'un arbre vide renvoie bien un tableau de zéros
718        assert_eq!(engine.scores(), [0; N]);
719    }
720
721    #[test]
722    fn test_scores_after_root_expansion() {
723        let mut engine = Engine::<N>::new();
724        let mut path = Vec::new();
725
726        let sel = engine.select(&mut path, &dummy_select, 1.0);
727        let evaluation = StateEvaluation::Active(0.0, [0.5, 0.5]);
728        engine.update(evaluation, sel, identity_score).unwrap();
729
730        assert_eq!(engine.scores(), [0; N]);
731    }
732
733    #[test]
734    fn test_scores_after_multiple_updates() {
735        let mut engine = Engine::<N>::new();
736        let mut path = Vec::new();
737
738        let sel_empty = engine.select(&mut path, &dummy_select, 1.0);
739        let evaluation = StateEvaluation::Active(0.0, [0.7, 0.3]);
740        engine.update(evaluation, sel_empty, identity_score).unwrap();
741
742        for _ in 0..3 {
743            let sel = engine.select(&mut path, &dummy_select, 1.0);
744            let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
745            engine.update(evaluation, sel, identity_score).unwrap();
746        }
747
748        let sel = engine.select(&mut path, &dummy_select, 1.0);
749        let evaluation = StateEvaluation::Active(1.0, [0.5, 0.5]);
750        engine.update(evaluation, sel, identity_score).unwrap();
751
752        let expected = [4, 0];
753        assert_eq!(engine.scores(), expected);
754    }
755
756    #[test]
757    #[allow(unused_assignments)]
758    fn test_scores_are_independent_copies() {
759        let mut engine = Engine::<N>::new();
760        let mut path = Vec::new();
761
762        let sel = engine.select(&mut path, &dummy_select, 1.0);
763        let evaluation = StateEvaluation::Active(0.0, [0.5, 0.5]);
764        engine.update(evaluation, sel, identity_score).unwrap();
765
766        let mut scores = engine.scores();
767
768
769        scores[0] = 999;
770
771        assert_ne!(engine.scores()[0], 999);
772        assert_eq!(engine.scores()[0], 0);
773    }
774
775    #[test]
776    fn test_scnerario_1(){
777        let c = std::f32::consts::SQRT_2;
778        let mut path = Vec::new();
779        let mut engine = Engine::<3>::new();
780
781        let selection = engine.select(&mut path, &puct, c);
782        let evaluation = StateEvaluation::Active(0.1, [0.5, 0.3, 0.2]);
783        let node  = engine.expand(evaluation, selection).unwrap();
784        engine.backpropagate(node, 0.1, negate_score).unwrap();
785
786        assert_eq!(engine.scores(), [0, 0, 0]);
787        assert_eq!(path.len(), 0);
788
789        let selection = engine.select(&mut path, &puct, c);
790        let evaluation = StateEvaluation::Active(0.5, [1., 0., 0.]);
791        let node  = engine.expand(evaluation, selection).unwrap();
792        engine.backpropagate(node, 0.5, negate_score).unwrap();
793
794        assert_eq!(engine.scores(), [1, 0, 0]);
795        assert_eq!(path.as_slice(), &[Action(0)]);
796
797        let selection = engine.select(&mut path, &puct, c);
798        let evaluation = StateEvaluation::Active(1., [1., 0., 0.]);
799        let node  = engine.expand(evaluation, selection).unwrap();
800        engine.backpropagate(node, 1.0, negate_score).unwrap();
801
802        assert_eq!(engine.scores(), [2, 0, 0]);
803        assert_eq!(path.as_slice(), &[Action(0), Action(0)]);
804
805        let selection = engine.select(&mut path, &puct, c);
806        let evaluation = StateEvaluation::Terminal(-0.9);
807        let node  = engine.expand(evaluation, selection).unwrap();
808        engine.backpropagate(node, -0.9, negate_score).unwrap();
809
810        assert_eq!(engine.scores(), [2, 1, 0]);
811        assert_eq!(path.as_slice(), &[Action(1)]);
812
813        let selection = engine.select(&mut path, &puct, c);
814        let evaluation = StateEvaluation::Terminal(-0.9);
815        let node  = engine.expand(evaluation, selection).unwrap();
816        engine.backpropagate(node, -0.9, negate_score).unwrap();
817
818        assert_eq!(engine.scores(), [2, 1, 1]);
819        assert_eq!(path.as_slice(), &[Action(2)]);
820    }
821
822    fn setup_engine() -> Engine<3> {
823        let mut engine = Engine::<3>::new();
824        let mut path = Vec::new();
825
826        let sel = engine.select(&mut path, &|_,_,_,_,_| 0.0, 1.0);
827        let evaluation = StateEvaluation::Active(0., [1., 0., 0.]);
828        let node = engine.expand(evaluation, sel).unwrap();
829        engine.backpropagate(node, 0.0, negate_score).unwrap();
830
831
832        let sel = engine.select(&mut path, &|_,_,_,_,_| 0.0, 1.0);
833        let evaluation = StateEvaluation::Active(1., [1., 0., 0.]);
834        let node = engine.expand(evaluation, sel).unwrap();
835        engine.backpropagate(node, 1.0, negate_score).unwrap();
836
837        engine
838    }
839
840    #[test]
841    fn test_commit_action_success() {
842        let mut engine = setup_engine();
843
844        engine.tree.root().unwrap();
845        engine.commit_action(Action(0));
846
847        let new_root = engine.tree.root().unwrap();
848
849        assert_eq!(engine.tree.data(new_root).unwrap().visits, 1);
850        assert_eq!(engine.tree.get(new_root).unwrap().data().action, Action(0));
851    }
852
853    #[test]
854    fn test_commit_action_unexplored_leads_to_clear() {
855        let mut engine = setup_engine();
856
857        engine.commit_action(Action(2));
858        assert!(engine.tree.root().is_none());
859    }
860
861    #[test]
862    fn test_commit_action_no_root_does_nothing() {
863        let mut engine = Engine::<3>::new();
864
865        engine.commit_action(Action(0));
866
867        assert!(engine.tree.root().is_none());
868    }
869
870    #[test]
871    #[should_panic(expected = "index out of bounds")]
872    fn test_commit_action_out_of_bounds_panics() {
873        let mut engine = setup_engine();
874
875        engine.commit_action(Action(5));
876    }
877
878    #[test]
879    fn test_compact_trigger() {
880        let mut engine = setup_engine();
881
882        let root = engine.tree.root().unwrap();
883        engine.tree.get_mut(root).unwrap().data_mut().visits = 100;
884
885        let initial_nodes = engine.tree.allocated_nodes();
886
887        engine.commit_action(Action(0));
888        assert!(engine.tree.allocated_nodes() <= initial_nodes);
889    }
890}
simple_mcts/engine/engine.rs

simple_mcts/engine/
engine.rs
