Skip to content

refactor(research): extract shared UCT kernel #190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
proerror77 merged 1 commit into from
Jul 22, 2026
Merged

refactor(research): extract shared UCT kernel #190

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion rust_hft/ARCHITECTURE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 8 additions & 0 deletions rust_hft/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions rust_hft/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 1 addition & 0 deletions rust_hft/alpha-harness/engine/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
237 changes: 141 additions & 96 deletions rust_hft/alpha-harness/engine/src/engines/mcts.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand All @@ -31,6 +32,34 @@ struct Node {
best_reward: Option<f64>,
}

impl UctNode for Node {
fn parent(&self) -> Option<usize> {
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, hft_search_kernel::UctError> {
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 {
Expand Down Expand Up @@ -191,39 +220,6 @@ impl MctsEngine {
})
}

fn select_expandable(&self, node_id: usize) -> Option<usize> {
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<usize, String> {
let action_index = self
.rng
Expand Down Expand Up @@ -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();
Expand All @@ -305,50 +301,35 @@ 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);

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Propagate backpropagation failures

When a restored, otherwise accepted checkpoint has a node at u64::MAX visits or a finite reward total near overflow, backpropagate returns StatsOverflow without mutating the tree. Discarding that error after removing the pending candidate lets AutoResearchKernel checkpoint and persist the evaluation as successful even though its reward was permanently omitted; before this extraction, the resulting invalid statistics caused checkpointing to fail. Preserve the failure so the run remains fail-closed instead of silently corrupting search state.

AGENTS.md reference: AGENTS.md:L59-L61

Useful? React with 👍 / 👎.

}
Comment on lines 303 to 311

@coderabbitai coderabbitai Bot Jul 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

observe() silently drops backpropagate failures.

let _ = backpropagate(...) discards any UctError (e.g. StatsOverflow, InvalidTopology). This is fail-closed for data integrity (no partial mutation), but a genuine observation is lost with zero observability — nothing signals that a proposed candidate's evaluation never got recorded.

🔍 Suggested minimal observability fix 
-        let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score);
+        if let Err(error) = backpropagate(&mut self.nodes, 0, node_id, evaluation.score) {
+            tracing::warn!(%error, node_id, "MCTS observe: backpropagate rejected evaluation");
+        }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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 observe(&mut self, proposal: &EngineProposal, evaluation: &CandidateEvaluation) {
let Some(node_id) = self.candidates.remove(&proposal.candidate_id) else {
return;
};
if !evaluation.score.is_finite() {
return;
}
if let Err(error) = backpropagate(&mut self.nodes, 0, node_id, evaluation.score) {
tracing::warn!(%error, node_id, "MCTS observe: backpropagate rejected evaluation");
}
}
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rust_hft/alpha-harness/engine/src/engines/mcts.rs` around lines 303 - 311,
Update MCTS::observe so failures returned by backpropagate are not silently
discarded: capture the Result and emit an appropriate error or warning with the
candidate/node context when it fails, while preserving the existing successful
backpropagation path and non-finite-score handling.


fn abandon(&mut self, proposal: &EngineProposal) {
self.candidates.remove(&proposal.candidate_id);
}

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());
};
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(())
}

Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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();
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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]
Expand Down
Loading
Loading