Scale exploration with Gini impurity of policy values

src/mcts.rs

@@ -364,7 +364,8 @@ impl<'a> Searcher<'a> {
 
         let cpuct = SearchHelpers::get_cpuct(self.params, node_stats, is_root);
         let fpu = SearchHelpers::get_fpu(node_stats);
-        let expl_scale = SearchHelpers::get_explore_scaling(self.params, node_stats);
+        let expl_scale =
+            SearchHelpers::get_explore_scaling(self.params, node_stats, &self.tree[ptr]);
 
         let expl = cpuct * expl_scale;
 

src/mcts/helpers.rs

@@ -2,7 +2,7 @@ use std::time::Instant;
 
 use crate::{
     mcts::{MctsParams, Searcher},
-    tree::{ActionStats, Edge},
+    tree::{ActionStats, Edge, Node},
 };
 
 pub struct SearchHelpers;
@@ -35,10 +35,25 @@ impl SearchHelpers {
     /// Exploration Scaling
     ///
     /// Larger value implies more exploration.
-    pub fn get_explore_scaling(params: &MctsParams, node_stats: &ActionStats) -> f32 {
+    fn base_explore_scaling(params: &MctsParams, node_stats: &ActionStats) -> f32 {
         (params.expl_tau() * (node_stats.visits().max(1) as f32).ln()).exp()
     }
 
+    #[allow(unused_variables)]
+    pub fn get_explore_scaling(params: &MctsParams, node_stats: &ActionStats, node: &Node) -> f32 {
+        #[cfg(not(feature = "datagen"))]
+        {
+            let mut scale = Self::base_explore_scaling(params, node_stats);
+
+            let gini = node.gini_impurity();
+            scale *= (0.679 - 1.634 * (gini + 0.001).ln()).min(2.1);
+            scale
+        }
+
+        #[cfg(feature = "datagen")]
+        Self::base_explore_scaling(params, node_stats)
+    }
+
     /// First Play Urgency
     ///
     /// #### Note

src/tree.rs

@@ -72,6 +72,7 @@ impl Tree {
         let t = &mut *self[to].actions_mut();
 
         self[to].set_state(self[from].state());
+        self[to].set_gini_impurity(self[from].gini_impurity());
 
         if f.is_empty() {
             return;

src/tree/node.rs

@@ -1,5 +1,5 @@
 use std::sync::{
-    atomic::{AtomicU16, Ordering},
+    atomic::{AtomicU16, AtomicU32, Ordering},
     RwLock, RwLockReadGuard, RwLockWriteGuard,
 };
 
@@ -14,6 +14,9 @@ pub struct Node {
     actions: RwLock<Vec<Edge>>,
     state: AtomicU16,
     threads: AtomicU16,
+
+    // heuristics used in search
+    gini_impurity: AtomicU32,
 }
 
 impl Node {
@@ -22,12 +25,14 @@ impl Node {
             actions: RwLock::new(Vec::new()),
             state: AtomicU16::new(u16::from(state)),
             threads: AtomicU16::new(0),
+            gini_impurity: AtomicU32::new(0),
         }
     }
 
     pub fn set_new(&self, state: GameState) {
         *self.actions_mut() = Vec::new();
         self.set_state(state);
+        self.set_gini_impurity(0.0);
     }
 
     pub fn is_terminal(&self) -> bool {
@@ -74,9 +79,19 @@ impl Node {
         self.state() == GameState::Ongoing && !self.has_children()
     }
 
+    pub fn gini_impurity(&self) -> f32 {
+        f32::from_bits(self.gini_impurity.load(Ordering::Relaxed))
+    }
+
+    pub fn set_gini_impurity(&self, gini_impurity: f32) {
+        self.gini_impurity
+            .store(f32::to_bits(gini_impurity), Ordering::Relaxed);
+    }
+
     pub fn clear(&self) {
         *self.actions.write().unwrap() = Vec::new();
         self.set_state(GameState::Ongoing);
+        self.set_gini_impurity(0.0);
     }
 
     pub fn expand<const ROOT: bool>(
@@ -122,11 +137,17 @@ impl Node {
             total += policy;
         }
 
+        let mut sum_of_squares = 0.0;
+
         for action in actions.iter_mut() {
             let policy = f32::from_bits(action.ptr().inner()) / total;
             action.set_ptr(NodePtr::NULL);
             action.set_policy(policy);
+            sum_of_squares += policy * policy;
         }
+
+        let gini_impurity = (1.0 - sum_of_squares).clamp(0.0, 1.0);
+        self.set_gini_impurity(gini_impurity);
     }
 
     pub fn relabel_policy(&self, pos: &ChessState, params: &MctsParams, policy: &PolicyNetwork) {