From 6687499cf155e5be14f6260b2faee3e3c51ae456 Mon Sep 17 00:00:00 2001 From: Kelechi Ukah Date: Sun, 12 Jul 2026 20:36:18 -0700 Subject: [PATCH] fix(mccfr): normalize sampling distribution used for importance reach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sampling_distribution/sampling() produced unnormalized weights: smoothing is added per-edge in the numerator but once in the denominator, temperature scales only the numerator, and the curiosity floor adds mass — so the per-infoset sum Z is > 1. The external sampler draws via WeightedIndex, which normalizes internally, so the true draw probability is weight(a)/Z, but sampling_reach treated the raw weight(a) as the sampling probability. That mis-scales counterfactual values by a path-dependent product of Z, compounding worst in deep (postflop) trees. Normalize by Z at the single source (sampling_distribution) and route the per-edge sampling() through the same, so sampling_reach / ancestor_reach / recursed_value all use the true draw probability. The draw itself is unchanged (WeightedIndex renormalizes regardless), so this corrects only the importance term; Kuhn still converges to the analytical Nash equilibrium. Fixes #53. --- crates/mccfr/src/strategy/flow.rs | 67 ++++++++++++++++++------------- 1 file changed, 39 insertions(+), 28 deletions(-) diff --git a/crates/mccfr/src/strategy/flow.rs b/crates/mccfr/src/strategy/flow.rs index 03f65f6e..dad75616 100644 --- a/crates/mccfr/src/strategy/flow.rs +++ b/crates/mccfr/src/strategy/flow.rs @@ -24,28 +24,38 @@ pub trait CfrFlow: RefProf + CfrSampling { fn weight_denom(&self, info: &Self::I) -> Probability { info.choices().map(|ref e| self.weight(info, e)).sum::() + self.smoothing() } - /// Compute sampling distribution for all edges of an info (single pass). - /// Returns exploration-adjusted probabilities for MCCFR sampling. + /// Unnormalized exploration weight for one edge: `max(ε, (σ(a)/τ + β) / (Σσ + β))`. + /// `denom` is [`Self::weight_denom`], hoisted so full-infoset callers compute it once. + /// Sums to `Z ≥ 1` over an infoset; divide by that sum to get a probability. + fn sampling_weight(&self, info: &Self::I, edge: &Self::E, denom: Probability) -> Probability { + ((self.weight(info, edge) / self.temperature() + self.smoothing()) / denom).max(self.curiosity()) + } + /// Sampling distribution over an infoset's edges: the normalized probabilities + /// the external sampler draws from. fn sampling_distribution(&self, info: &Self::I) -> Policy { - let raw = info.choices().map(|e| (e, self.weight(info, &e))).collect::>(); - let denom = raw.iter().map(|(_, w)| *w).sum::() + self.smoothing(); - raw.into_iter() - .map(|(e, w)| (e, w / self.temperature())) - .map(|(e, w)| (e, w + self.smoothing())) - .map(|(e, w)| (e, w / denom)) - .map(|(e, w)| (e, w.max(self.curiosity()))) - .collect() + let denom = self.weight_denom(info); + let raw = info + .choices() + .map(|e| (e, self.sampling_weight(info, &e, denom))) + .collect::>(); + let z = raw.iter().map(|(_, w)| *w).sum::(); + raw.into_iter().map(|(e, w)| (e, w / z)).collect() } /// Calculate immediate policy via regret matching for a single edge. /// Prefer `iterated_distribution` when multiple edges needed. fn instant_policy(&self, info: &Self::I, edge: &Self::E) -> Probability { self.regret(info, edge) / self.regret_denom(info) } - /// Calculate sampling probability for a single edge. + /// Normalized probability the external sampler draws `edge` — exactly what + /// `WeightedIndex` samples from, so the correct importance term for reaches. /// Prefer `sampling_distribution` when multiple edges needed. fn sampling(&self, info: &Self::I, edge: &Self::E) -> Probability { - ((self.weight(info, edge) / self.temperature() + self.smoothing()) / self.weight_denom(info)) - .max(self.curiosity()) + let denom = self.weight_denom(info); + let z = info + .choices() + .map(|ref e| self.sampling_weight(info, e, denom)) + .sum::(); + self.sampling_weight(info, edge, denom) / z } /// Fused regret + expected value computation for an information set. @@ -154,16 +164,12 @@ pub trait CfrFlow: RefProf + CfrSampling { /// Both products share the same path (root->tree_root) and filters /// (non-chance, non-walker), so we compute them in a single pass. fn ancestor_reach(&self, root: &Node) -> Utility { - let (cfactual, sampling) = root.decisions().filter(|(t, _, _)| *t != self.walker()).fold( - (1.0, 1.0), - |(cf, sm), (_, ref info, ref edge)| { - ( - cf * (self.regret(info, edge) / self.regret_denom(info)), - sm * ((self.weight(info, edge) / self.temperature() + self.smoothing()) / self.weight_denom(info)) - .max(self.curiosity()), - ) - }, - ); + let (cfactual, sampling) = root + .decisions() + .filter(|(t, _, _)| *t != self.walker()) + .fold((1.0, 1.0), |(cf, sm), (_, ref info, ref edge)| { + (cf * self.instant_policy(info, edge), sm * self.sampling(info, edge)) + }); cfactual / sampling } @@ -186,7 +192,15 @@ pub trait CfrFlow: RefProf + CfrSampling { let chance = node.game().turn() == Self::T::chance(); let walker = node.game().turn() == self.walker(); let regret_denom = (!chance).then(|| self.regret_denom(node.info())); - let weight_denom = (!chance && !walker).then(|| self.weight_denom(node.info())); + let sampling = (!chance && !walker).then(|| { + let denom = self.weight_denom(node.info()); + let z = node + .info() + .choices() + .map(|ref e| self.sampling_weight(node.info(), e, denom)) + .sum::(); + (denom, z) + }); node.edges() .map(|(child, edge)| (node.at(child), edge)) .map(|(ref child, edge)| { @@ -195,10 +209,7 @@ pub trait CfrFlow: RefProf + CfrSampling { child, relative_reach * regret_denom.map_or(1.0, |d| self.regret(node.info(), edge) / d), sampling_reach - * weight_denom.map_or(1.0, |d| { - ((self.weight(node.info(), edge) / self.temperature() + self.smoothing()) / d) - .max(self.curiosity()) - }), + * sampling.map_or(1.0, |(denom, z)| self.sampling_weight(node.info(), edge, denom) / z), ) }) .sum()