From 22475efeb04f151781eaf28702fbfe7dde4bf4b5 Mon Sep 17 00:00:00 2001 From: Sonic Shih Date: Wed, 22 Jul 2026 09:34:54 +0800 Subject: [PATCH] refactor(research): extract shared UCT kernel --- rust_hft/ARCHITECTURE.md | 8 +- rust_hft/Cargo.lock | 8 + rust_hft/Cargo.toml | 1 + rust_hft/alpha-harness/engine/Cargo.toml | 1 + .../alpha-harness/engine/src/engines/mcts.rs | 237 ++++++---- .../alpha-harness/engine/src/engines/mod.rs | 23 +- .../research-core/search-kernel/Cargo.toml | 8 + .../research-core/search-kernel/src/lib.rs | 421 ++++++++++++++++++ 8 files changed, 588 insertions(+), 119 deletions(-) create mode 100644 rust_hft/research-core/search-kernel/Cargo.toml create mode 100644 rust_hft/research-core/search-kernel/src/lib.rs diff --git a/rust_hft/ARCHITECTURE.md b/rust_hft/ARCHITECTURE.md index 1b267b378..deac80f8d 100644 --- a/rust_hft/ARCHITECTURE.md +++ b/rust_hft/ARCHITECTURE.md @@ -44,7 +44,13 @@ flowchart TB - `alpha-harness/domain`: mission, LoopRun, candidate, evaluation, learning, approval, bundle, and signed deployment contracts. - `alpha-harness/store`: DuckDB migrations and append-only control-plane repositories. -- `alpha-harness/engine`: resumable search kernel, GP/MCTS/Bayesian engines, causal evaluator, bounded LLM client, offline-RL lab engine, and learning coordinator. +- `research-core/search-kernel`: domain-neutral deterministic UCT selection, + topology and reward statistics. It owns no candidate grammar, evaluator, + dataset, or storage contract. +- `alpha-harness/engine`: research-domain search adapters, GP/MCTS/Bayesian + engines, causal evaluator, bounded LLM client, offline-RL lab engine, and + learning coordinator. Domain checkpoints bind candidates and evaluator state + around the shared UCT mechanics. - `alpha-harness/app`: structured CLI for data, loops, missions, evaluation, promotion, approvals, policies, feedback, and signing. - `tools/collector`: streaming connectors plus governed public Binance OHLCV acquisition. - `market-core` and `data-pipelines`: events, runtime construction, venue market data, and replay. diff --git a/rust_hft/Cargo.lock b/rust_hft/Cargo.lock index 6d95cb8f4..8f48d9c39 100644 --- a/rust_hft/Cargo.lock +++ b/rust_hft/Cargo.lock @@ -642,6 +642,7 @@ dependencies = [ "hft-core", "hft-factor-dsl", "hft-research-manifest", + "hft-search-kernel", "hft-strategy-formula", "reqwest 0.12.28", "rust_decimal", @@ -6027,6 +6028,13 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "hft-search-kernel" +version = "0.1.0" +dependencies = [ + "serde", +] + [[package]] name = "hft-shared-config" version = "0.1.0" diff --git a/rust_hft/Cargo.toml b/rust_hft/Cargo.toml index e422b7ba8..b79173728 100644 --- a/rust_hft/Cargo.toml +++ b/rust_hft/Cargo.toml @@ -82,6 +82,7 @@ members = [ "research-core/factor-dsl", "research-core/manifest", "research-core/ml", + "research-core/search-kernel", # Bounded Loop Engineer research plane "alpha-harness/domain", diff --git a/rust_hft/alpha-harness/engine/Cargo.toml b/rust_hft/alpha-harness/engine/Cargo.toml index 3500633c1..67caa87b8 100644 --- a/rust_hft/alpha-harness/engine/Cargo.toml +++ b/rust_hft/alpha-harness/engine/Cargo.toml @@ -14,6 +14,7 @@ alpha-store = { path = "../store", optional = true } chrono = { workspace = true, features = ["serde"] } hex = { workspace = true } hft-factor-dsl = { path = "../../research-core/factor-dsl" } +hft-search-kernel = { path = "../../research-core/search-kernel" } reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } diff --git a/rust_hft/alpha-harness/engine/src/engines/mcts.rs b/rust_hft/alpha-harness/engine/src/engines/mcts.rs index 73268e989..c0a6b96f1 100644 --- a/rust_hft/alpha-harness/engine/src/engines/mcts.rs +++ b/rust_hft/alpha-harness/engine/src/engines/mcts.rs @@ -5,6 +5,7 @@ use crate::{ }; use alpha_domain::{CandidateArtifact, EngineKind}; use hft_factor_dsl::{validate_live_formula, FactorAst, FactorOperator, FactorTerminal}; +use hft_search_kernel::{backpropagate, select_expandable, validate_tree, UctNode, UctStats}; use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, BTreeSet}; @@ -31,6 +32,34 @@ struct Node { best_reward: Option, } +impl UctNode for Node { + fn parent(&self) -> Option { + self.parent + } + + fn children(&self) -> &[usize] { + &self.children + } + + fn is_expandable(&self) -> bool { + !self.unexpanded_actions.is_empty() + } + + fn depth(&self) -> usize { + self.depth + } + + fn stats(&self) -> Result { + UctStats::from_parts(self.visits, self.total_reward, self.best_reward) + } + + fn replace_stats(&mut self, stats: UctStats) { + self.visits = stats.visits(); + self.total_reward = stats.total_reward(); + self.best_reward = stats.best_reward(); + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] struct MctsConfigV2 { @@ -191,39 +220,6 @@ impl MctsEngine { }) } - fn select_expandable(&self, node_id: usize) -> Option { - let node = &self.nodes[node_id]; - if !node.unexpanded_actions.is_empty() { - return Some(node_id); - } - node.children - .iter() - .copied() - .filter(|child| self.has_expandable(*child)) - .max_by(|left, right| { - self.uct(*left, node.visits) - .total_cmp(&self.uct(*right, node.visits)) - }) - .and_then(|child| self.select_expandable(child)) - } - - fn has_expandable(&self, node_id: usize) -> bool { - !self.nodes[node_id].unexpanded_actions.is_empty() - || self.nodes[node_id] - .children - .iter() - .any(|child| self.has_expandable(*child)) - } - - fn uct(&self, node_id: usize, parent_visits: u64) -> f64 { - let node = &self.nodes[node_id]; - if node.visits == 0 { - return f64::INFINITY; - } - node.total_reward / node.visits as f64 - + self.exploration * ((parent_visits.max(1) as f64).ln() / node.visits as f64).sqrt() - } - fn expand(&mut self, parent_id: usize) -> Result { let action_index = self .rng @@ -283,8 +279,8 @@ impl ProposalEngine for MctsEngine { return Err("MCTS expansion budget is exhausted".to_string()); } for attempt in 1..=remaining.expansions.min(256) { - let parent = self - .select_expandable(0) + let parent = select_expandable(&self.nodes, 0, self.exploration) + .map_err(|error| error.to_string())? .ok_or_else(|| "MCTS tree has no expandable node".to_string())?; let node_id = self.expand(parent)?; let formula = self.nodes[node_id].ast.to_string(); @@ -305,25 +301,13 @@ impl ProposalEngine for MctsEngine { } fn observe(&mut self, proposal: &EngineProposal, evaluation: &CandidateEvaluation) { - let Some(mut node_id) = self.candidates.remove(&proposal.candidate_id) else { + let Some(node_id) = self.candidates.remove(&proposal.candidate_id) else { return; }; if !evaluation.score.is_finite() { return; } - loop { - let node = &mut self.nodes[node_id]; - node.visits += 1; - node.total_reward += evaluation.score; - node.best_reward = Some( - node.best_reward - .map_or(evaluation.score, |best| best.max(evaluation.score)), - ); - match node.parent { - Some(parent) => node_id = parent, - None => break, - } - } + let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score); } fn abandon(&mut self, proposal: &EngineProposal) { @@ -331,6 +315,8 @@ impl ProposalEngine for MctsEngine { } fn restore(&mut self, observations: &[HistoricalObservation]) -> Result<(), String> { + let mut nodes = self.nodes.clone(); + let mut seen = self.seen.clone(); for observation in observations { let CandidateArtifact::Formula(ast) = &observation.proposal.artifact else { return Err("MCTS history contains a non-formula artifact".to_string()); @@ -338,17 +324,12 @@ impl ProposalEngine for MctsEngine { if !observation.evaluation.score.is_finite() { return Err("MCTS history contains a non-finite score".to_string()); } - self.seen.insert(ast.to_string()); - let root = &mut self.nodes[0]; - root.visits += 1; - root.total_reward += observation.evaluation.score; - root.best_reward = Some( - root.best_reward - .map_or(observation.evaluation.score, |best| { - best.max(observation.evaluation.score) - }), - ); + seen.insert(ast.to_string()); + backpropagate(&mut nodes, 0, 0, observation.evaluation.score) + .map_err(|error| format!("invalid MCTS history: {error}"))?; } + self.nodes = nodes; + self.seen = seen; Ok(()) } @@ -427,6 +408,8 @@ impl MctsCheckpointV2 { if root.ast != self.config.root_ast || root.parent.is_some() || root.depth != 0 { return Err("MCTS checkpoint root does not match its configuration".to_string()); } + validate_tree(&self.nodes, 0, self.config.max_depth) + .map_err(|error| format!("invalid MCTS checkpoint tree: {error}"))?; for (node_id, node) in self.nodes.iter().enumerate() { node.ast @@ -437,18 +420,9 @@ impl MctsCheckpointV2 { format!("invalid live MCTS checkpoint node {node_id}: {error}") })?; } - if !node.total_reward.is_finite() - || node.best_reward.is_some_and(|reward| !reward.is_finite()) - || node.depth > self.config.max_depth - || (node.depth == self.config.max_depth && !node.unexpanded_actions.is_empty()) - { + if node.depth == self.config.max_depth && !node.unexpanded_actions.is_empty() { return Err(format!("invalid MCTS checkpoint node {node_id}")); } - if (node.visits == 0) != (node.total_reward == 0.0 && node.best_reward.is_none()) { - return Err(format!( - "invalid MCTS checkpoint rewards for node {node_id}" - )); - } let actions = node .unexpanded_actions .iter() @@ -464,33 +438,6 @@ impl MctsCheckpointV2 { "invalid MCTS checkpoint actions for node {node_id}" )); } - if node.children.windows(2).any(|pair| pair[0] >= pair[1]) { - return Err(format!( - "invalid MCTS checkpoint children for node {node_id}" - )); - } - for child_id in &node.children { - let Some(child) = self.nodes.get(*child_id) else { - return Err(format!("invalid MCTS checkpoint child for node {node_id}")); - }; - if *child_id <= node_id || child.parent != Some(node_id) { - return Err(format!("invalid MCTS checkpoint child for node {node_id}")); - } - } - if node_id > 0 { - let Some(parent_id) = node.parent else { - return Err(format!("MCTS checkpoint node {node_id} has no parent")); - }; - let Some(parent) = self.nodes.get(parent_id) else { - return Err(format!("MCTS checkpoint node {node_id} has invalid parent")); - }; - if parent_id >= node_id - || node.depth != parent.depth + 1 - || !parent.children.contains(&node_id) - { - return Err(format!("MCTS checkpoint node {node_id} is inconsistent")); - } - } } let mut pending_nodes = BTreeSet::new(); @@ -632,6 +579,100 @@ mod tests { assert!(restored.candidates.contains_key(&pending.candidate_id)); } + #[test] + fn shared_kernel_preserves_seeded_expansion_and_checkpoint_wire_shape() { + let mut engine = MctsEngine::new(3, "oi", "imbalance", 1.4, 3).unwrap(); + let dataset = super::super::test_dataset(); + let first = engine + .propose("mission", 0, &dataset.proposal_context(), &budget()) + .unwrap(); + let CandidateArtifact::Formula(first_ast) = &first.artifact else { + panic!("MCTS must emit a formula candidate"); + }; + assert_eq!(first_ast.to_string(), "mean(oi, 20)"); + engine.observe(&first, &evaluation(0.5)); + + let checkpoint = engine.checkpoint().unwrap(); + assert_eq!( + checkpoint.state, + serde_json::json!({ + "config": { + "seed": 3, + "root_ast": {"Terminal": {"Field": "oi"}}, + "secondary_field": "imbalance", + "exploration": 1.4, + "max_depth": 3, + "live_only": false + }, + "rng": 2_088_359_638_719_790_806_u64, + "nodes": [ + { + "ast": {"Terminal": {"Field": "oi"}}, + "parent": null, + "children": [1], + "unexpanded_actions": [0, 1, 3], + "depth": 0, + "visits": 1, + "total_reward": 0.5, + "best_reward": 0.5 + }, + { + "ast": {"Call": { + "operator": "Mean", + "args": [ + {"Terminal": {"Field": "oi"}}, + {"Terminal": {"Constant": "20"}} + ] + }}, + "parent": 0, + "children": [], + "unexpanded_actions": [0, 1, 2, 3], + "depth": 1, + "visits": 1, + "total_reward": 0.5, + "best_reward": 0.5 + } + ], + "candidates": {}, + "seen": ["mean(oi, 20)"] + }) + ); + let expected_trace = vec![ + MctsNodeSnapshot { + node_id: 0, + parent_id: None, + visits: 1, + total_reward: 0.5, + best_reward: Some(0.5), + formula: "oi".to_string(), + }, + MctsNodeSnapshot { + node_id: 1, + parent_id: Some(0), + visits: 1, + total_reward: 0.5, + best_reward: Some(0.5), + formula: "mean(oi, 20)".to_string(), + }, + ]; + assert_eq!(engine.trace(), expected_trace); + + let mut restored = MctsEngine::new(3, "oi", "imbalance", 1.4, 3).unwrap(); + restored.restore_checkpoint(&checkpoint, &[]).unwrap(); + assert_eq!(restored.trace(), expected_trace); + let expected = engine + .propose("mission", 1, &dataset.proposal_context(), &budget()) + .unwrap(); + let actual = restored + .propose("mission", 1, &dataset.proposal_context(), &budget()) + .unwrap(); + assert_eq!(actual, expected); + let CandidateArtifact::Formula(actual_ast) = &actual.artifact else { + panic!("MCTS must emit a formula candidate"); + }; + assert_eq!(actual_ast.to_string(), "rank(oi)"); + } + #[test] fn restored_search_continues_like_uninterrupted_search() { let mut uninterrupted = MctsEngine::new(9, "oi", "imbalance", 1.4, 4).unwrap(); @@ -681,6 +722,10 @@ mod tests { let mut malformed = checkpoint; malformed.state = serde_json::json!({"unexpected": true}); assert!(restored.restore_checkpoint(&malformed, &[]).is_err()); + + let mut cyclic = engine.checkpoint().unwrap(); + cyclic.state["nodes"][0]["children"] = serde_json::json!([0]); + assert!(restored.restore_checkpoint(&cyclic, &[]).is_err()); } #[test] diff --git a/rust_hft/alpha-harness/engine/src/engines/mod.rs b/rust_hft/alpha-harness/engine/src/engines/mod.rs index 6d8510bff..4c46783aa 100644 --- a/rust_hft/alpha-harness/engine/src/engines/mod.rs +++ b/rust_hft/alpha-harness/engine/src/engines/mod.rs @@ -3,34 +3,13 @@ mod gp; mod mcts; mod offline_rl; -use serde::{Deserialize, Serialize}; +use hft_search_kernel::DeterministicRng; pub use bayesian::BayesianOptimizerEngine; pub use gp::GeneticProgrammingEngine; pub use mcts::{MctsEngine, MctsNodeSnapshot}; pub use offline_rl::{OfflineRlEngine, OfflineTrace}; -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -struct DeterministicRng(u64); - -impl DeterministicRng { - fn new(seed: u64) -> Self { - Self(seed.max(1)) - } - - fn next_u64(&mut self) -> u64 { - self.0 = self - .0 - .wrapping_mul(6_364_136_223_846_793_005) - .wrapping_add(1_442_695_040_888_963_407); - self.0 - } - - fn index(&mut self, len: usize) -> usize { - (self.next_u64() as usize) % len - } -} - #[cfg(test)] fn test_dataset() -> crate::evaluation::PreparedDataset { use crate::evaluation::{prepare_dataset, ResearchRow}; diff --git a/rust_hft/research-core/search-kernel/Cargo.toml b/rust_hft/research-core/search-kernel/Cargo.toml new file mode 100644 index 000000000..9c07cc142 --- /dev/null +++ b/rust_hft/research-core/search-kernel/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "hft-search-kernel" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +serde = { workspace = true, features = ["derive"] } diff --git a/rust_hft/research-core/search-kernel/src/lib.rs b/rust_hft/research-core/search-kernel/src/lib.rs new file mode 100644 index 000000000..e769dc085 --- /dev/null +++ b/rust_hft/research-core/search-kernel/src/lib.rs @@ -0,0 +1,421 @@ +//! Domain-neutral deterministic UCT tree mechanics. +//! +//! Candidate payloads, expansion grammars, evaluators, and persistence remain +//! owned by domain adapters. This crate owns deterministic choice, UCT +//! selection, tree invariants, and reward-statistic updates. + +use serde::{Deserialize, Serialize}; +use std::fmt; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeterministicRng(u64); + +impl DeterministicRng { + pub fn new(seed: u64) -> Self { + Self(seed.max(1)) + } + + pub fn index(&mut self, len: usize) -> usize { + assert!(len > 0, "deterministic RNG requires a non-empty range"); + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (self.0 as usize) % len + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +pub struct UctStats { + visits: u64, + total_reward: f64, + best_reward: Option, +} + +impl Default for UctStats { + fn default() -> Self { + Self { + visits: 0, + total_reward: 0.0, + best_reward: None, + } + } +} + +impl UctStats { + pub fn from_parts( + visits: u64, + total_reward: f64, + best_reward: Option, + ) -> Result { + let stats = Self { + visits, + total_reward, + best_reward, + }; + stats.validate()?; + Ok(stats) + } + + pub fn visits(self) -> u64 { + self.visits + } + + pub fn total_reward(self) -> f64 { + self.total_reward + } + + pub fn best_reward(self) -> Option { + self.best_reward + } + + fn validate(self) -> Result<(), UctError> { + if !self.total_reward.is_finite() + || self.best_reward.is_some_and(|reward| !reward.is_finite()) + || (self.visits == 0) != (self.total_reward == 0.0 && self.best_reward.is_none()) + { + return Err(UctError::InvalidStats); + } + Ok(()) + } + + fn record(self, reward: f64) -> Result { + self.validate()?; + if !reward.is_finite() { + return Err(UctError::InvalidReward); + } + let visits = self.visits.checked_add(1).ok_or(UctError::StatsOverflow)?; + let total_reward = self.total_reward + reward; + if !total_reward.is_finite() { + return Err(UctError::StatsOverflow); + } + Ok(Self { + visits, + total_reward, + best_reward: Some(self.best_reward.map_or(reward, |best| best.max(reward))), + }) + } +} + +/// The tree facts required by UCT. Domain adapters may keep the three statistic +/// fields flat in an existing wire format; their update semantics live here. +pub trait UctNode { + fn parent(&self) -> Option; + fn children(&self) -> &[usize]; + fn is_expandable(&self) -> bool; + fn depth(&self) -> usize; + fn stats(&self) -> Result; + fn replace_stats(&mut self, stats: UctStats); +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UctError { + InvalidExploration, + InvalidNode(usize), + InvalidReward, + InvalidStats, + StatsOverflow, + InvalidTopology(usize), + CyclicTree, +} + +impl fmt::Display for UctError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidExploration => { + formatter.write_str("UCT exploration must be finite and non-negative") + } + Self::InvalidNode(node_id) => { + write!(formatter, "UCT tree references missing node {node_id}") + } + Self::InvalidReward => formatter.write_str("UCT reward must be finite"), + Self::InvalidStats => formatter.write_str("UCT reward statistics are invalid"), + Self::StatsOverflow => formatter.write_str("UCT reward statistics overflowed"), + Self::InvalidTopology(node_id) => { + write!(formatter, "UCT tree topology is invalid at node {node_id}") + } + Self::CyclicTree => formatter.write_str("UCT tree contains a cycle"), + } + } +} + +impl std::error::Error for UctError {} + +pub fn validate_tree( + nodes: &[N], + root: usize, + max_depth: usize, +) -> Result<(), UctError> { + let root_node = nodes.get(root).ok_or(UctError::InvalidNode(root))?; + if root_node.parent().is_some() || root_node.depth() != 0 { + return Err(UctError::InvalidTopology(root)); + } + + for (node_id, node) in nodes.iter().enumerate() { + node.stats()?; + if node.depth() > max_depth || node.children().windows(2).any(|pair| pair[0] >= pair[1]) { + return Err(UctError::InvalidTopology(node_id)); + } + for &child_id in node.children() { + if child_id <= node_id { + return Err(UctError::CyclicTree); + } + let child = nodes.get(child_id).ok_or(UctError::InvalidNode(child_id))?; + if child.parent() != Some(node_id) || child.depth() != node.depth() + 1 { + return Err(UctError::InvalidTopology(node_id)); + } + } + + if node_id == root { + continue; + } + let parent_id = node.parent().ok_or(UctError::InvalidTopology(node_id))?; + if parent_id >= node_id { + return Err(UctError::CyclicTree); + } + let parent = nodes + .get(parent_id) + .ok_or(UctError::InvalidNode(parent_id))?; + if !parent.children().contains(&node_id) || node.depth() != parent.depth() + 1 { + return Err(UctError::InvalidTopology(node_id)); + } + } + Ok(()) +} + +pub fn select_expandable( + nodes: &[N], + root: usize, + exploration: f64, +) -> Result, UctError> { + if !exploration.is_finite() || exploration < 0.0 { + return Err(UctError::InvalidExploration); + } + validate_tree(nodes, root, usize::MAX)?; + select_from(nodes, root, exploration) +} + +fn select_from( + nodes: &[N], + node_id: usize, + exploration: f64, +) -> Result, UctError> { + let node = nodes.get(node_id).ok_or(UctError::InvalidNode(node_id))?; + if node.is_expandable() { + return Ok(Some(node_id)); + } + + let parent_visits = node.stats()?.visits(); + let mut best = None; + for &child_id in node.children() { + if let Some(expandable_id) = select_from(nodes, child_id, exploration)? { + let score = uct_score(nodes, child_id, parent_visits, exploration)?; + if best.is_none_or(|(_, best_score)| { + score.total_cmp(&best_score) != std::cmp::Ordering::Less + }) { + best = Some((expandable_id, score)); + } + } + } + Ok(best.map(|(expandable_id, _)| expandable_id)) +} + +fn uct_score( + nodes: &[N], + node_id: usize, + parent_visits: u64, + exploration: f64, +) -> Result { + let stats = nodes + .get(node_id) + .ok_or(UctError::InvalidNode(node_id))? + .stats()?; + if stats.visits() == 0 { + return Ok(f64::INFINITY); + } + Ok(stats.total_reward() / stats.visits() as f64 + + exploration * ((parent_visits.max(1) as f64).ln() / stats.visits() as f64).sqrt()) +} + +pub fn backpropagate( + nodes: &mut [N], + root: usize, + leaf: usize, + reward: f64, +) -> Result<(), UctError> { + if !reward.is_finite() { + return Err(UctError::InvalidReward); + } + validate_tree(nodes, root, usize::MAX)?; + + let mut lineage = Vec::new(); + let mut current = leaf; + loop { + let node = nodes.get(current).ok_or(UctError::InvalidNode(current))?; + lineage.push(current); + if current == root { + break; + } + current = node.parent().ok_or(UctError::InvalidTopology(current))?; + } + + let updated = lineage + .iter() + .map(|&node_id| nodes[node_id].stats()?.record(reward)) + .collect::, UctError>>()?; + for (node_id, stats) in lineage.into_iter().zip(updated) { + nodes[node_id].replace_stats(stats); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone)] + struct Node { + parent: Option, + children: Vec, + expandable: bool, + depth: usize, + stats: UctStats, + } + + impl UctNode for Node { + fn parent(&self) -> Option { + self.parent + } + fn children(&self) -> &[usize] { + &self.children + } + fn is_expandable(&self) -> bool { + self.expandable + } + fn depth(&self) -> usize { + self.depth + } + fn stats(&self) -> Result { + Ok(self.stats) + } + fn replace_stats(&mut self, stats: UctStats) { + self.stats = stats; + } + } + + #[test] + fn deterministic_rng_preserves_the_existing_sequence() { + let mut rng = DeterministicRng::new(3); + assert_eq!(rng.index(4), 2); + } + + #[test] + fn selects_and_backpropagates_without_domain_knowledge() { + let mut nodes = vec![ + Node { + parent: None, + children: vec![1], + expandable: false, + depth: 0, + stats: UctStats::from_parts(1, 0.2, Some(0.2)).unwrap(), + }, + Node { + parent: Some(0), + children: vec![], + expandable: true, + depth: 1, + stats: UctStats::default(), + }, + ]; + assert_eq!(select_expandable(&nodes, 0, 1.4).unwrap(), Some(1)); + backpropagate(&mut nodes, 0, 1, 0.5).unwrap(); + assert_eq!( + nodes[0].stats, + UctStats::from_parts(2, 0.7, Some(0.5)).unwrap() + ); + assert_eq!( + nodes[1].stats, + UctStats::from_parts(1, 0.5, Some(0.5)).unwrap() + ); + } + + #[test] + fn selection_preserves_total_order_for_signed_zero() { + let nodes = vec![ + Node { + parent: None, + children: vec![1, 2], + expandable: false, + depth: 0, + stats: UctStats::from_parts(2, 0.0, Some(0.0)).unwrap(), + }, + Node { + parent: Some(0), + children: vec![], + expandable: true, + depth: 1, + stats: UctStats::from_parts(1, 0.0, Some(0.0)).unwrap(), + }, + Node { + parent: Some(0), + children: vec![], + expandable: true, + depth: 1, + stats: UctStats::from_parts(1, -0.0, Some(-0.0)).unwrap(), + }, + ]; + + assert_eq!(select_expandable(&nodes, 0, -0.0).unwrap(), Some(1)); + } + + #[test] + fn rejects_an_expandable_cycle_before_mutating_rewards() { + let mut nodes = vec![ + Node { + parent: None, + children: vec![1], + expandable: false, + depth: 0, + stats: UctStats::default(), + }, + Node { + parent: Some(0), + children: vec![2], + expandable: true, + depth: 1, + stats: UctStats::default(), + }, + Node { + parent: Some(1), + children: vec![1], + expandable: false, + depth: 2, + stats: UctStats::default(), + }, + ]; + assert_eq!( + backpropagate(&mut nodes, 0, 2, 0.5), + Err(UctError::CyclicTree) + ); + assert!(nodes.iter().all(|node| node.stats == UctStats::default())); + assert_eq!(select_expandable(&nodes, 0, 1.4), Err(UctError::CyclicTree)); + } + + #[test] + fn rejects_reward_overflow_without_partial_mutation() { + let original = UctStats::from_parts(1, f64::MAX, Some(f64::MAX)).unwrap(); + let mut nodes = vec![Node { + parent: None, + children: vec![], + expandable: true, + depth: 0, + stats: original, + }]; + + assert_eq!( + backpropagate(&mut nodes, 0, 0, f64::MAX), + Err(UctError::StatsOverflow) + ); + assert_eq!(nodes[0].stats, original); + } +}