eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs

@@ -3,10 +3,20 @@
 //! Two algorithms ship here:
 //!
 //! - **LinUCB** (Li et al. 2010) — contextual bandit with linear reward
-//!   models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`. Pick
-//!   the arm that maximises `θ_a · x + α · √(x^T A^{-1} x)`.
+//!   models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`, where
+//!   `r` is 1 when the stage resolved the query and 0 when it did not.
 //! - **Thompson Sampling** — Beta-Bernoulli on per-arm success rates. Cheap
 //!   and surprisingly competitive when contexts are noisy.
+//!
+//! Both learn each stage's chance of resolving a query, and the cost of each
+//! stage from the joules `observe` reports. Both pick the stage that is
+//! cheapest to answer from in expectation
+//! ([`cheapest_in_expectation`](crate::router::cheapest_in_expectation)),
+//! not the arm with the best success rate: the neural stage always has the
+//! best rate, and picking it skips every cheaper stage that would have
+//! answered. Exploration enters through the success estimate, optimistic for
+//! LinUCB (`θ_a · x + α · √(x^T A^{-1} x)`) and sampled for Thompson, so an
+//! untried cheap stage looks likely to resolve and gets tried.
 
 use std::collections::{BTreeMap, HashMap};
 
@@ -18,7 +28,10 @@ use rand_distr::{Beta, Distribution};
 use serde::{Deserialize, Serialize};
 
 use crate::error::{Error, Result};
-use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction};
+use crate::router::{
+    ALL_STAGES, LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation,
+    default_stage_joules, stage_index,
+};
 
 /// Per-arm state for [`LinUcbRouter`].
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -52,7 +65,8 @@ impl LinUcbArm {
         }
     }
 
-    fn ucb(&self, x: &DVector<f32>, alpha: f32) -> f32 {
+    /// The mean reward estimate `θ · x` and its upper confidence bound.
+    fn estimate(&self, x: &DVector<f32>, alpha: f32) -> (f32, f32) {
         let theta = self.theta();
         let mean = theta.dot(x);
         let a_inv = self
@@ -61,7 +75,7 @@ impl LinUcbArm {
             .try_inverse()
             .unwrap_or_else(|| DMatrix::<f32>::identity(self.dim, self.dim));
         let bonus_sq = (x.transpose() * &a_inv * x)[(0, 0)].max(0.0);
-        mean + alpha * bonus_sq.sqrt()
+        (mean, mean + alpha * bonus_sq.sqrt())
     }
 
     fn update(&mut self, x: &DVector<f32>, reward: f32) {
@@ -71,6 +85,10 @@ impl LinUcbArm {
 }
 
 /// LinUCB contextual bandit over the four cascade stages.
+///
+/// The reward model is linear in the context with no intercept, as in Li et
+/// al.: put a constant feature in the context for a stage's base rate to be
+/// learnable.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct LinUcbRouter {
     /// Context (embedding) dimension.
@@ -81,6 +99,9 @@ pub struct LinUcbRouter {
     pub ridge: f32,
     /// Per-stage arm state.
     pub arms: HashMap<Stage, LinUcbArm>,
+    /// Per-stage cost estimate in microjoules, learned from `observe`.
+    #[serde(default = "default_stage_joules")]
+    pub joules: HashMap<Stage, u64>,
 }
 
 impl LinUcbRouter {
@@ -95,11 +116,14 @@ impl LinUcbRouter {
             alpha,
             ridge,
             arms,
+            joules: default_stage_joules(),
         }
     }
 
-    /// Argmax UCB across stages.
-    pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> {
+    /// Each stage's mean estimate and upper confidence bound that it resolves
+    /// the query, in [`ALL_STAGES`] order. A missing arm counts as a certain
+    /// miss.
+    fn estimates(&self, embedding: &[f32]) -> Result<([f32; 4], [f32; 4])> {
         if embedding.len() != self.dim {
             return Err(Error::DimensionMismatch {
                 expected: self.dim,
@@ -107,20 +131,27 @@ impl LinUcbRouter {
             });
         }
         let x = DVector::<f32>::from_row_slice(embedding);
-        let mut best = Stage::Cache;
-        let mut best_score = f32::NEG_INFINITY;
-        for s in ALL_STAGES {
-            let arm = match self.arms.get(&s) {
-                Some(a) => a,
-                None => continue,
-            };
-            let u = arm.ucb(&x, self.alpha);
-            if u > best_score {
-                best_score = u;
-                best = s;
+        let mut means = [0.0_f32; 4];
+        let mut bounds = [0.0_f32; 4];
+        for (i, s) in ALL_STAGES.iter().enumerate() {
+            if let Some(arm) = self.arms.get(s) {
+                (means[i], bounds[i]) = arm.estimate(&x, self.alpha);
             }
         }
-        Ok((best, best_score))
+        Ok((means, bounds))
+    }
+
+    /// The stage cheapest to answer from in expectation, with every stage's
+    /// chance taken at its upper confidence bound, and that stage's bound.
+    pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> {
+        let (_, bounds) = self.estimates(embedding)?;
+        let stage = cheapest_in_expectation(&bounds, &self.joules);
+        Ok((stage, bounds[stage_index(stage)]))
+    }
+
+    /// Record an observed stage cost — EMA with coefficient 0.1.
+    pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) {
+        blend_joules(&mut self.joules, stage, microjoules);
     }
 
     /// Update the chosen arm with the observed reward.
@@ -145,13 +176,14 @@ impl LearnedRouter for LinUcbRouter {
         let Some(emb) = query.embedding.as_ref() else {
             return StagePrediction::new(Stage::Cache, 0.0);
         };
-        match self.pick(emb) {
-            Ok((stage, score)) => {
-                // Squash UCB score into [0,1] via logistic for "confidence".
-                let conf = 1.0 / (1.0 + (-score).exp());
-                StagePrediction::new(stage, conf)
+        match (self.pick(emb), self.estimates(emb)) {
+            // The confidence is the mean estimate, not the optimistic bound
+            // the pick used: the bound is how much the router is willing to
+            // explore, not how likely the stage is to resolve.
+            (Ok((stage, _)), Ok((means, _))) => {
+                StagePrediction::new(stage, means[stage_index(stage)])
             }
-            Err(_) => StagePrediction::new(Stage::Cache, 0.0),
+            _ => StagePrediction::new(Stage::Cache, 0.0),
         }
     }
 
@@ -160,8 +192,9 @@ impl LearnedRouter for LinUcbRouter {
         query: &Query,
         chosen_stage: Stage,
         was_correct: bool,
-        _joule_cost: u64,
+        joule_cost: u64,
     ) {
+        self.observe_cost(chosen_stage, joule_cost);
         if let Some(emb) = query.embedding.as_ref() {
             let reward = if was_correct { 1.0 } else { 0.0 };
             let _ = self.update(emb, chosen_stage, reward);
@@ -224,6 +257,9 @@ pub struct ThompsonSamplingRouter {
     pub arms: HashMap<Stage, BetaArm>,
     /// RNG seed for reproducibility — incremented every sample.
     pub seed: u64,
+    /// Per-stage cost estimate in microjoules, learned from `observe`.
+    #[serde(default = "default_stage_joules")]
+    pub joules: HashMap<Stage, u64>,
 }
 
 impl ThompsonSamplingRouter {
@@ -233,30 +269,43 @@ impl ThompsonSamplingRouter {
         for s in ALL_STAGES {
             arms.insert(s, BetaArm::new());
         }
-        Self { arms, seed }
+        Self {
+            arms,
+            seed,
+            joules: default_stage_joules(),
+        }
     }
 
-    /// Sample once from each arm; return the argmax. Iteration order is the
-    /// fixed cheapest-first stage order, so the result is deterministic for
-    /// a given `seed`.
+    /// Sample each stage's chance of resolving once, and return the stage
+    /// cheapest to answer from in expectation under those samples, with its
+    /// sample. Sampling runs in the fixed cheapest-first stage order, so the
+    /// result is deterministic for a given `seed`. A missing arm counts as a
+    /// certain miss.
     pub fn pick(&self) -> (Stage, f32) {
         let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed);
-        let mut best = Stage::Cache;
-        let mut best_v = f32::NEG_INFINITY;
-        for s in ALL_STAGES {
-            let arm = match self.arms.get(&s) {
-                Some(a) => a,
-                None => continue,
+        let mut samples = [0.0_f32; 4];
+        for (i, s) in ALL_STAGES.iter().enumerate() {
+            let Some(arm) = self.arms.get(s) else {
+                continue;
             };
             let dist = Beta::new(arm.alpha as f64, arm.beta as f64)
                 .unwrap_or_else(|_| Beta::new(1.0, 1.0).expect("beta(1,1) ok"));
-            let v = dist.sample(&mut rng) as f32;
-            if v > best_v {
-                best_v = v;
-                best = s;
-            }
+            samples[i] = dist.sample(&mut rng) as f32;
         }
-        (best, best_v)
+        let stage = cheapest_in_expectation(&samples, &self.joules);
+        (stage, samples[stage_index(stage)])
+    }
+
+    /// The posterior mean chance that `stage` resolves a query.
+    pub fn mean(&self, stage: Stage) -> f32 {
+        self.arms
+            .get(&stage)
+            .map_or(0.0, |a| a.alpha / (a.alpha + a.beta))
+    }
+
+    /// Record an observed stage cost — EMA with coefficient 0.1.
+    pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) {
+        blend_joules(&mut self.joules, stage, microjoules);
     }
 
     /// Update Beta posterior for `stage`.
@@ -275,8 +324,11 @@ impl ThompsonSamplingRouter {
 #[async_trait]
 impl LearnedRouter for ThompsonSamplingRouter {
     async fn route(&self, _query: &Query) -> StagePrediction {
-        let (stage, v) = self.pick();
-        StagePrediction::new(stage, v.clamp(0.0, 1.0))
+        // The confidence is the posterior mean, not the sample the pick used:
+        // a sample is a draw for exploration, and would make the threshold
+        // policy's skip decision a coin toss.
+        let (stage, _) = self.pick();
+        StagePrediction::new(stage, self.mean(stage))
     }
 
     async fn observe(
@@ -284,8 +336,9 @@ impl LearnedRouter for ThompsonSamplingRouter {
         _query: &Query,
         chosen_stage: Stage,
         was_correct: bool,
-        _joule_cost: u64,
+        joule_cost: u64,
     ) {
+        self.observe_cost(chosen_stage, joule_cost);
         self.update(chosen_stage, was_correct);
     }
 
@@ -332,4 +385,39 @@ mod tests {
         assert_eq!(arm.alpha, 2.0);
         assert_eq!(arm.beta, 1.0);
     }
+
+    /// `observe` is how a stage's cost reaches the router; it was ignored.
+    #[tokio::test]
+    async fn observing_a_stage_learns_its_cost() {
+        let q = Query::new("q").with_embedding(vec![1.0, 0.0]);
+        let mut lin = LinUcbRouter::new(2, 0.1, 1.0);
+        let mut ts = ThompsonSamplingRouter::new(0);
+        lin.observe(&q, Stage::Kv, false, 1_050_000).await;
+        ts.observe(&q, Stage::Kv, false, 1_050_000).await;
+        // A tenth of the way from the default 50 000.
+        assert_eq!(lin.joules[&Stage::Kv], 150_000);
+        assert_eq!(ts.joules[&Stage::Kv], 150_000);
+        assert_eq!(ts.arms[&Stage::Kv].beta, 2.0);
+    }
+
+    /// The confidence handed to the threshold policy is the router's belief
+    /// that the stage resolves, not the bound or sample it explored with: an
+    /// untried router believes nothing yet.
+    #[tokio::test]
+    async fn confidence_is_the_mean_belief_not_the_exploration_bound() {
+        let q = Query::new("q").with_embedding(vec![1.0, 0.0]);
+        let lin = LinUcbRouter::new(2, 1.0, 1.0);
+        let (stage, bound) = lin.pick(&[1.0, 0.0]).unwrap();
+        assert!(bound > 0.5, "an untried arm explores optimistically");
+        let pred = lin.route(&q).await;
+        assert_eq!((pred.recommended, pred.confidence), (stage, 0.0));
+
+        let mut ts = ThompsonSamplingRouter::new(3);
+        for _ in 0..3 {
+            ts.update(Stage::Cache, true);
+        }
+        let pred = ts.route(&q).await;
+        assert_eq!(pred.confidence, ts.mean(pred.recommended));
+        assert_eq!(ts.mean(Stage::Cache), 0.8);
+    }
 }

eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs

@@ -21,19 +21,10 @@ use rand_distr::{Distribution, Normal};
 use serde::{Deserialize, Serialize};
 
 use crate::error::{Error, Result};
-use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction};
+use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction, stage_index};
 
 const NUM_STAGES: usize = 4;
 
-fn stage_index(s: Stage) -> usize {
-    match s {
-        Stage::Cache => 0,
-        Stage::Kv => 1,
-        Stage::Graph => 2,
-        Stage::Neural => 3,
-    }
-}
-
 fn stage_from_index(i: usize) -> Stage {
     ALL_STAGES[i.min(3)]
 }
@@ -179,8 +170,16 @@ impl LearnedRouter for LogRegRouter {
         was_correct: bool,
         _joule_cost: u64,
     ) {
-        // Only train on successful resolutions — we want P(stage = correct | x).
-        if was_correct && let Some(emb) = query.embedding.as_ref() {
+        // The classifier predicts the cheapest stage that resolves a query.
+        // A stage that resolved it shows that stage is enough, not that it is
+        // the cheapest that is, so it is evidence only when it is no costlier
+        // than the stage already predicted. Otherwise a correct Neural answer,
+        // which every query gets, would teach it to predict Neural for all.
+        if was_correct
+            && let Some(emb) = query.embedding.as_ref()
+            && let Ok((predicted, _, _)) = self.predict(emb)
+            && stage_index(chosen_stage) <= stage_index(predicted)
+        {
             let _ = self.sgd_step(emb, chosen_stage, 0.01);
         }
     }
@@ -228,4 +227,57 @@ mod tests {
         let r = LogRegRouter::new(8, 1e-3, 1);
         assert!(r.predict_probs(&[0.0; 7]).is_err());
     }
+
+    /// Neural answers every query correctly. Reinforcing it whenever it did
+    /// drifted the classifier to predict Neural for everything; only a
+    /// correct stage no costlier than the prediction may move it.
+    #[tokio::test]
+    async fn a_correct_costlier_stage_does_not_move_the_prediction() {
+        let emb = vec![1.0, -0.5, 0.25];
+        let q = Query::new("q").with_embedding(emb.clone());
+        let mut r = LogRegRouter::new(3, 0.0, 9);
+        for _ in 0..200 {
+            r.sgd_step(&emb, Stage::Kv, 0.1).unwrap();
+        }
+        assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv);
+
+        for _ in 0..500 {
+            r.observe(&q, Stage::Neural, true, 50_000_000).await;
+        }
+        assert_eq!(
+            r.predict(&emb).unwrap().0,
+            Stage::Kv,
+            "Neural being enough is no evidence Kv is not"
+        );
+
+        for _ in 0..500 {
+            r.observe(&q, Stage::Cache, true, 1_000).await;
+        }
+        assert_eq!(
+            r.predict(&emb).unwrap().0,
+            Stage::Cache,
+            "a cheaper stage that resolved it is"
+        );
+    }
+
+    /// The predicted stage resolving the query is evidence for it.
+    #[tokio::test]
+    async fn the_predicted_stage_resolving_reinforces_it() {
+        let emb = vec![1.0, -0.5, 0.25];
+        let q = Query::new("q").with_embedding(emb.clone());
+        let mut r = LogRegRouter::new(3, 0.0, 9);
+        for _ in 0..20 {
+            r.sgd_step(&emb, Stage::Kv, 0.1).unwrap();
+        }
+        let (stage, before, _) = r.predict(&emb).unwrap();
+        assert_eq!(stage, Stage::Kv);
+        for _ in 0..50 {
+            r.observe(&q, Stage::Kv, true, 50_000).await;
+        }
+        let (_, after, _) = r.predict(&emb).unwrap();
+        assert!(
+            after > before,
+            "confidence in Kv went from {before} to {after}"
+        );
+    }
 }

eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs

@@ -1,13 +1,25 @@
 //! Routing-quality metrics.
 //!
-//! `evaluate` consumes a `LearnedRouter` and a held-out test set, returning
-//! `RouterMetrics`:
+//! `evaluate` replays a router's recommendations over a held-out set of
+//! per-stage outcomes, returning `RouterMetrics`:
 //!
-//! - **accuracy** — fraction of examples where the predicted stage matched
-//!   the recorded successful stage.
-//! - **joules_per_correct** — total observed joules / correct predictions.
-//! - **lift** — joules saved vs. always running the full cascade (sum of all
-//!   stage costs).
+//! - **accuracy** — fraction of scored queries the dispatch answered
+//!   correctly.
+//! - **total_microjoules** — what the dispatches the router chose spend.
+//! - **full_cascade_microjoules** — what the plain cascade walk spends on the
+//!   same queries.
+//! - **lift** — the walk's joules minus the router's: what routing saved, or
+//!   cost when negative.
+//! - **joules_per_correct** — total joules / correct answers, absent when
+//!   none was correct.
+//!
+//! Every joule figure is priced from the recorded per-stage costs, run the way
+//! [`LearnedCascade`](crate::inference::LearnedCascade) dispatches: Kv or
+//! Graph, falling through to Neural on a miss; Neural alone; or, for Cache,
+//! the full walk. A row that did not succeed is a miss. A query whose
+//! dispatch or walk would run a stage the set has no row for is not scored:
+//! it is counted in `unscored` and enters no figure, rather than being priced
+//! at zero.
 //!
 //! Latency metrics are deferred to the caller (we don't run the cascade
 //! here, just predictions).
@@ -17,76 +29,238 @@ use std::collections::HashMap;
 use eoc_core::{Query, Stage};
 use serde::{Deserialize, Serialize};
 
-use crate::router::LearnedRouter;
-use crate::training::Example;
+use crate::router::{LearnedRouter, stage_index};
+use crate::training::{Example, query_key};
 
 /// Summary statistics produced by [`evaluate`].
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct RouterMetrics {
-    /// Number of test examples scored.
+    /// Number of queries scored.
     pub n: usize,
-    /// Number of correct (predicted == actual successful) predictions.
+    /// Queries left out because their dispatch or the walk would run a stage
+    /// the test set has no row for.
+    pub unscored: usize,
+    /// Scored queries whose dispatch ended on a correct answer.
     pub correct: usize,
     /// `correct / n`.
     pub accuracy: f32,
-    /// Sum of observed joule cost.
+    /// Joules the router's dispatches spend on the scored queries.
     pub total_microjoules: u64,
-    /// `total_microjoules / max(correct, 1)`.
-    pub joules_per_correct: u64,
-    /// Per-stage prediction count.
+    /// `total_microjoules / correct`; `None` when no answer was correct.
+    pub joules_per_correct: Option<u64>,
+    /// Per-stage recommendation count, over every query routed.
     pub stage_counts: HashMap<Stage, usize>,
-    /// Joules saved vs. running the full cascade for every example
-    /// (sum of all stage costs in the dataset).
+    /// Joules the plain cascade walk spends on the scored queries.
+    pub full_cascade_microjoules: u64,
+    /// `full_cascade_microjoules − total_microjoules`: what routing saved
+    /// over the walk, negative when it cost more.
     pub lift_microjoules: i64,
 }
 
-/// Evaluate a router on a labelled test set.
+/// One stage's recorded outcome on a query: whether it resolved it, and what
+/// it spent.
+type Outcome = Option<(bool, u64)>;
+
+/// What skipping to `stage` spends on a query and whether it ends correct,
+/// or `None` when it would run a stage with no row.
+fn dispatch(stage: Stage, rows: &[Outcome; 4]) -> Option<(bool, u64)> {
+    let neural = rows[stage_index(Stage::Neural)];
+    match stage {
+        Stage::Cache => walk(rows),
+        Stage::Neural => neural,
+        Stage::Kv | Stage::Graph => match rows[stage_index(stage)]? {
+            (true, uj) => Some((true, uj)),
+            (false, uj) => neural.map(|(ok, n)| (ok, uj.saturating_add(n))),
+        },
+    }
+}
+
+/// The plain cascade: every stage in order until one resolves the query,
+/// Neural last.
+fn walk(rows: &[Outcome; 4]) -> Option<(bool, u64)> {
+    let mut spent = 0u64;
+    for (i, row) in rows.iter().enumerate() {
+        let (ok, uj) = (*row)?;
+        spent = spent.saturating_add(uj);
+        if ok || i == rows.len() - 1 {
+            return Some((ok, spent));
+        }
+    }
+    unreachable!("the walk ends at Neural")
+}
+
+/// Evaluate a router on a labelled test set of per-stage outcomes.
 ///
-/// Each example is treated as a (embedding → routed_stage) prediction. If the
-/// example's `success` is true and the router picks the same stage, we count
-/// it correct. If the example was a failure, we ignore correctness but still
-/// account for joules.
+/// Rows with the same embedding are one query, and each row is one stage's
+/// outcome on it; when a stage has more than one row for a query, the first
+/// counts. See the module docs for how each recommendation is priced.
 pub async fn evaluate<R: LearnedRouter>(router: &R, test_set: &[Example]) -> RouterMetrics {
+    let mut index: HashMap<Vec<u32>, usize> = HashMap::new();
+    let mut queries: Vec<(&[f32], [Outcome; 4])> = Vec::new();
+    for ex in test_set {
+        let i = *index.entry(query_key(&ex.embedding)).or_insert_with(|| {
+            queries.push((&ex.embedding, [None; 4]));
+            queries.len() - 1
+        });
+        let slot = &mut queries[i].1[stage_index(ex.stage)];
+        if slot.is_none() {
+            *slot = Some((ex.success, ex.cost_microjoules));
+        }
+    }
+
+    let mut n = 0usize;
+    let mut unscored = 0usize;
     let mut correct = 0usize;
     let mut total_microjoules: u64 = 0;
     let mut full_cascade_microjoules: u64 = 0;
     let mut stage_counts: HashMap<Stage, usize> = HashMap::new();
 
-    for ex in test_set {
-        let q = Query::new("").with_embedding(ex.embedding.clone());
-        let prediction = router.route(&q).await;
-        *stage_counts.entry(prediction.recommended).or_insert(0) += 1;
-        if ex.success && prediction.recommended == ex.stage {
-            correct += 1;
-        }
-        // Saved joules: only pay for the recommended stage, not all four.
-        total_microjoules = total_microjoules.saturating_add(ex.cost_microjoules);
-        // "Full-cascade" baseline pays every stage's average cost — we use
-        // 4x the recorded cost as a conservative proxy when per-stage costs
-        // aren't broken out in the test set.
-        full_cascade_microjoules = full_cascade_microjoules.saturating_add(ex.cost_microjoules * 4);
+    for (embedding, rows) in &queries {
+        let q = Query::new("").with_embedding(embedding.to_vec());
+        let recommended = router.route(&q).await.recommended;
+        *stage_counts.entry(recommended).or_insert(0) += 1;
+        let (Some((ok, uj)), Some((_, walk_uj))) = (dispatch(recommended, rows), walk(rows)) else {
+            unscored += 1;
+            continue;
+        };
+        n += 1;
+        correct += usize::from(ok);
+        total_microjoules = total_microjoules.saturating_add(uj);
+        full_cascade_microjoules = full_cascade_microjoules.saturating_add(walk_uj);
     }
 
-    let n = test_set.len();
     let accuracy = if n == 0 {
         0.0
     } else {
         correct as f32 / n as f32
     };
-    let joules_per_correct = if correct == 0 {
-        0
-    } else {
-        total_microjoules / correct as u64
-    };
+    let joules_per_correct = (correct > 0).then(|| total_microjoules / correct as u64);
     let lift_microjoules = full_cascade_microjoules as i64 - total_microjoules as i64;
 
     RouterMetrics {
         n,
+        unscored,
         correct,
         accuracy,
         total_microjoules,
         joules_per_correct,
         stage_counts,
+        full_cascade_microjoules,
         lift_microjoules,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::router::{RouterState, StagePrediction};
+    use async_trait::async_trait;
+
+    /// Recommends one stage for every query.
+    struct Always(Stage);
+
+    #[async_trait]
+    impl LearnedRouter for Always {
+        async fn route(&self, _q: &Query) -> StagePrediction {
+            StagePrediction::new(self.0, 1.0)
+        }
+        async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {}
+        fn export_state(&self) -> RouterState {
+            unimplemented!()
+        }
+        fn import_state(_s: RouterState) -> crate::Result<Self> {
+            unimplemented!()
+        }
+    }
+
+    fn row(q: f32, stage: Stage, success: bool, uj: u64) -> Example {
+        Example::new(vec![q], stage, success, uj)
+    }
+
+    /// Query 1 is resolved first by Graph. Query 2 has rows only for Cache
+    /// (a hit) and Neural.
+    fn set() -> Vec<Example> {
+        vec![
+            row(1.0, Stage::Cache, false, 1),
+            row(1.0, Stage::Kv, false, 10),
+            row(1.0, Stage::Graph, true, 100),
+            row(1.0, Stage::Neural, true, 1_000),
+            row(2.0, Stage::Cache, true, 1),
+            row(2.0, Stage::Neural, true, 1_000),
+        ]
+    }
+
+    #[tokio::test]
+    async fn the_routers_choice_decides_the_joules() {
+        let m = evaluate(&Always(Stage::Graph), &set()).await;
+        // Query 1: Graph resolves it for 100; the walk spends 1 + 10 + 100.
+        // Query 2: no Graph row, so it is not scored.
+        assert_eq!((m.n, m.unscored, m.correct), (1, 1, 1));
+        assert_eq!(
+            (m.total_microjoules, m.full_cascade_microjoules),
+            (100, 111)
+        );
+        assert_eq!(m.lift_microjoules, 11);
+        assert_eq!(m.joules_per_correct, Some(100));
+        assert_eq!(m.stage_counts[&Stage::Graph], 2);
+
+        let m = evaluate(&Always(Stage::Kv), &set()).await;
+        // Query 1: Kv misses for 10 and Neural answers for 1 000.
+        assert_eq!((m.n, m.total_microjoules), (1, 1_010));
+        assert_eq!(m.lift_microjoules, 111 - 1_010);
+
+        let m = evaluate(&Always(Stage::Cache), &set()).await;
+        // Both walks: 111, and 1 for query 2's cache hit.
+        assert_eq!((m.n, m.total_microjoules, m.lift_microjoules), (2, 112, 0));
+
+        let m = evaluate(&Always(Stage::Neural), &set()).await;
+        assert_eq!((m.n, m.total_microjoules), (2, 2_000));
+    }
+
+    /// The old metric charged each row's own cost whatever the router chose,
+    /// and priced the walk at four times it, so every router showed a lift of
+    /// three times its spend. One row per query prices no walk, and no
+    /// dispatch except to the stage the row is for.
+    #[tokio::test]
+    async fn one_row_per_query_scores_only_what_it_can_price() {
+        let rows = vec![
+            row(1.0, Stage::Kv, true, 50_000),
+            row(2.0, Stage::Neural, true, 50_000_000),
+        ];
+        let m = evaluate(&Always(Stage::Kv), &rows).await;
+        assert_eq!((m.n, m.unscored), (0, 2));
+        assert_eq!(
+            (
+                m.total_microjoules,
+                m.full_cascade_microjoules,
+                m.lift_microjoules
+            ),
+            (0, 0, 0)
+        );
+        assert_eq!(m.joules_per_correct, None);
+    }
+
+    #[tokio::test]
+    async fn no_correct_answer_has_no_joules_per_correct() {
+        let rows = vec![
+            row(1.0, Stage::Cache, false, 1),
+            row(1.0, Stage::Kv, false, 10),
+            row(1.0, Stage::Graph, false, 100),
+            row(1.0, Stage::Neural, false, 1_000),
+        ];
+        let m = evaluate(&Always(Stage::Neural), &rows).await;
+        assert_eq!((m.n, m.correct, m.total_microjoules), (1, 0, 1_000));
+        assert_eq!(
+            m.joules_per_correct, None,
+            "zero would read as the best possible figure"
+        );
+    }
+
+    #[tokio::test]
+    async fn the_first_row_for_a_stage_counts() {
+        let mut rows = set();
+        rows.push(row(1.0, Stage::Graph, false, 7));
+        let m = evaluate(&Always(Stage::Graph), &rows).await;
+        assert_eq!((m.correct, m.total_microjoules), (1, 100));
+    }
+}

eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs

@@ -4,8 +4,11 @@
 //! strict order. That's optimal when most queries land in the cheap stages.
 //! For workloads where many queries unavoidably need the neural stage, walking
 //! the cheaper stages first is wasted joules. A *learned router* predicts the
-//! cheapest stage that is likely to satisfy the query and lets us **skip
-//! ahead** when the prediction is confident.
+//! stage the cascade answers from most cheaply in expectation
+//! ([`router::expected_joules`]) and lets us **skip ahead** when the
+//! prediction is confident. Cheapest in expectation is not likeliest to
+//! succeed: Neural succeeds on everything and costs the most, and a router
+//! that ranked stages by success alone learned to send every query there.
 //!
 //! This crate ships three orthogonal router families plus a threshold policy
 //! and an end-to-end cascade extension:

eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs

@@ -30,24 +30,14 @@ use rand_distr::{Distribution, Normal};
 use serde::{Deserialize, Serialize};
 
 use crate::error::{Error, Result};
-use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction};
+use crate::router::{
+    LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation,
+    default_stage_joules, stage_index,
+};
 
 /// Default latent dimension.
 pub const DEFAULT_LATENT_DIM: usize = 32;
 
-fn stage_index(s: Stage) -> usize {
-    match s {
-        Stage::Cache => 0,
-        Stage::Kv => 1,
-        Stage::Graph => 2,
-        Stage::Neural => 3,
-    }
-}
-
-fn stage_from_index(i: usize) -> Stage {
-    ALL_STAGES[i.min(3)]
-}
-
 /// Matrix-factorisation router (embedding-conditioned).
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MfRouter {
@@ -91,7 +81,7 @@ impl MfRouter {
             stage_factors,
             stage_bias,
             global_bias: 0.0,
-            joule_ema: default_joule_ema(),
+            joule_ema: default_stage_joules(),
         }
     }
 
@@ -105,11 +95,11 @@ impl MfRouter {
         self.embedding_dim
     }
 
-    /// Forward pass: predict success probability per stage.
+    /// Forward pass: the probability that each stage resolves the query.
     ///
-    /// Returns logits *and* softmax-normalised probabilities. The softmax is
-    /// taken over per-stage sigmoid outputs so the recommended-stage
-    /// probability stays in [0, 1] when the router is uncertain.
+    /// One sigmoid per stage, not a softmax: the stages are not alternatives
+    /// to one another (a query Kv resolves, Neural resolves too), so the
+    /// probabilities need not sum to one.
     pub fn predict_probs(&self, embedding: &[f32]) -> Result<[f32; 4]> {
         if embedding.len() != self.embedding_dim {
             return Err(Error::DimensionMismatch {
@@ -129,11 +119,13 @@ impl MfRouter {
         Ok(probs)
     }
 
-    /// Pick the recommended stage and its confidence.
+    /// Pick the stage cheapest to answer from in expectation
+    /// ([`cheapest_in_expectation`], priced at [`Self::joule_estimates`]).
+    /// The confidence is that stage's probability of resolving the query.
     pub fn predict(&self, embedding: &[f32]) -> Result<(Stage, f32, [f32; 4])> {
         let probs = self.predict_probs(embedding)?;
-        let (idx, conf) = argmax(&probs);
-        Ok((stage_from_index(idx), conf, probs))
+        let stage = cheapest_in_expectation(&probs, &self.joule_ema);
+        Ok((stage, probs[stage_index(stage)], probs))
     }
 
     /// One SGD step on a single example.
@@ -196,9 +188,7 @@ impl MfRouter {
 
     /// Record an observed joule cost — EMA with coefficient 0.1.
     pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) {
-        let prev = self.joule_ema.get(&stage).copied().unwrap_or(microjoules) as f64;
-        let blended = 0.9 * prev + 0.1 * microjoules as f64;
-        self.joule_ema.insert(stage, blended as u64);
+        blend_joules(&mut self.joule_ema, stage, microjoules);
     }
 
     /// Current per-stage joule estimate (EMA).
@@ -207,16 +197,6 @@ impl MfRouter {
     }
 }
 
-fn default_joule_ema() -> HashMap<Stage, u64> {
-    let mut m = HashMap::new();
-    // Conservative defaults, ordered cheap → expensive.
-    m.insert(Stage::Cache, 1_000);
-    m.insert(Stage::Kv, 50_000);
-    m.insert(Stage::Graph, 500_000);
-    m.insert(Stage::Neural, 50_000_000);
-    m
-}
-
 fn sigmoid(x: f32) -> f32 {
     if x >= 0.0 {
         let z = (-x).exp();
@@ -227,18 +207,6 @@ fn sigmoid(x: f32) -> f32 {
     }
 }
 
-fn argmax(xs: &[f32; 4]) -> (usize, f32) {
-    let mut best = 0;
-    let mut bv = xs[0];
-    for (i, v) in xs.iter().enumerate().skip(1) {
-        if *v > bv {
-            best = i;
-            bv = *v;
-        }
-    }
-    (best, bv)
-}
-
 #[async_trait]
 impl LearnedRouter for MfRouter {
     async fn route(&self, query: &Query) -> StagePrediction {
@@ -323,6 +291,15 @@ mod tests {
         }
     }
 
+    #[test]
+    fn the_confidence_is_the_chosen_stages_chance() {
+        let r = MfRouter::new(4, 4, 7);
+        let emb = [0.5, -0.5, 0.25, -0.25];
+        let (stage, conf, probs) = r.predict(&emb).unwrap();
+        assert_eq!(conf, probs[stage_index(stage)]);
+        assert!(conf < 1.0);
+    }
+
     #[test]
     fn sgd_reduces_loss() {
         let mut r = MfRouter::new(4, 4, 7);

eoc/eoc-rs/crates/eoc-route-learned/src/router.rs

@@ -63,10 +63,20 @@ pub struct RouterState {
 /// A learned router predicts the best cascade stage from a query.
 #[async_trait]
 pub trait LearnedRouter: Send + Sync {
-    /// Predict the best stage for `query`.
+    /// Predict the best stage for `query`: the one [`LearnedCascade`] answers
+    /// from most cheaply in expectation (see [`expected_joules`]), not the one
+    /// most likely to succeed. The neural stage is the most likely to succeed
+    /// and the most expensive, so a router ranking by success alone learns to
+    /// skip every cheap stage.
+    ///
+    /// [`LearnedCascade`]: crate::inference::LearnedCascade
     async fn route(&self, query: &Query) -> StagePrediction;
 
-    /// Online update — observe what actually happened.
+    /// Online update with the outcome of one stage that ran: whether
+    /// `chosen_stage` resolved `query` correctly, and the microjoules that
+    /// stage itself spent. When a dispatch ran several stages (a full walk,
+    /// or a miss that fell through to Neural), report each of them: a stage
+    /// that missed is evidence too.
     async fn observe(
         &mut self,
         query: &Query,
@@ -86,3 +96,186 @@ pub trait LearnedRouter: Send + Sync {
 
 /// The four cascade stages, in cheapest-first order, useful for iteration.
 pub const ALL_STAGES: [Stage; 4] = [Stage::Cache, Stage::Kv, Stage::Graph, Stage::Neural];
+
+/// A stage's position in [`ALL_STAGES`].
+pub(crate) fn stage_index(s: Stage) -> usize {
+    match s {
+        Stage::Cache => 0,
+        Stage::Kv => 1,
+        Stage::Graph => 2,
+        Stage::Neural => 3,
+    }
+}
+
+/// Per-stage cost assumed before any is observed, in microjoules. Conservative
+/// defaults, ordered cheap → expensive.
+pub fn default_stage_joules() -> HashMap<Stage, u64> {
+    ALL_STAGES.iter().map(|s| (*s, default_cost(*s))).collect()
+}
+
+fn default_cost(stage: Stage) -> u64 {
+    match stage {
+        Stage::Cache => 1_000,
+        Stage::Kv => 50_000,
+        Stage::Graph => 500_000,
+        Stage::Neural => 50_000_000,
+    }
+}
+
+/// Fold one observed stage cost into a per-stage estimate — EMA with
+/// coefficient 0.1.
+pub fn blend_joules(joules: &mut HashMap<Stage, u64>, stage: Stage, microjoules: u64) {
+    let prev = joules.get(&stage).copied().unwrap_or(microjoules) as f64;
+    let blended = 0.9 * prev + 0.1 * microjoules as f64;
+    joules.insert(stage, blended as u64);
+}
+
+/// Expected microjoules for [`LearnedCascade`] to answer a query it was told
+/// to skip to `stage`, given each stage's probability of resolving the query
+/// (`resolves`, in [`ALL_STAGES`] order) and each stage's cost (`joules`).
+///
+/// It follows the dispatch exactly:
+///
+/// ```text
+///     Neural      c_N
+///     Kv, Graph   c_s + (1 − p_s) · c_N         a miss falls through to Neural
+///     Cache       c_C + (1 − p_C)(c_K + (1 − p_K)(c_G + (1 − p_G) · c_N))
+///                                               the full walk
+/// ```
+///
+/// Neural always answers, so its own probability does not enter. The walk
+/// treats the stages' misses as independent, because the routers estimate
+/// each stage's chance on its own rather than given that a cheaper one
+/// missed. A probability outside [0, 1] is clamped and NaN counts as a miss,
+/// so a broken estimate prices a stage dearer, never cheaper. A stage with no
+/// cost estimate is priced at [`default_stage_joules`].
+///
+/// [`LearnedCascade`]: crate::inference::LearnedCascade
+pub fn expected_joules(stage: Stage, resolves: &[f32; 4], joules: &HashMap<Stage, u64>) -> f64 {
+    let cost = |s: Stage| joules.get(&s).copied().unwrap_or_else(|| default_cost(s)) as f64;
+    let miss = |i: usize| {
+        let p = resolves[i];
+        if p.is_nan() {
+            1.0
+        } else {
+            1.0 - f64::from(p.clamp(0.0, 1.0))
+        }
+    };
+    let neural = cost(Stage::Neural);
+    match stage {
+        Stage::Neural => neural,
+        Stage::Kv => cost(Stage::Kv) + miss(1) * neural,
+        Stage::Graph => cost(Stage::Graph) + miss(2) * neural,
+        Stage::Cache => {
+            let graph_on = cost(Stage::Graph) + miss(2) * neural;
+            let kv_on = cost(Stage::Kv) + miss(1) * graph_on;
+            cost(Stage::Cache) + miss(0) * kv_on
+        }
+    }
+}
+
+/// The stage [`LearnedCascade`] answers from most cheaply in expectation
+/// ([`expected_joules`]). A tie goes to the cheaper stage.
+///
+/// [`LearnedCascade`]: crate::inference::LearnedCascade
+pub fn cheapest_in_expectation(resolves: &[f32; 4], joules: &HashMap<Stage, u64>) -> Stage {
+    let mut best = ALL_STAGES[0];
+    let mut best_uj = expected_joules(best, resolves, joules);
+    for s in &ALL_STAGES[1..] {
+        let uj = expected_joules(*s, resolves, joules);
+        if uj < best_uj {
+            best = *s;
+            best_uj = uj;
+        }
+    }
+    best
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn costs() -> HashMap<Stage, u64> {
+        let mut m = HashMap::new();
+        m.insert(Stage::Cache, 1);
+        m.insert(Stage::Kv, 10);
+        m.insert(Stage::Graph, 100);
+        m.insert(Stage::Neural, 1_000);
+        m
+    }
+
+    #[test]
+    fn each_dispatch_is_priced_as_the_cascade_runs_it() {
+        let p = [0.5, 0.25, 0.75, 0.0];
+        let c = costs();
+        assert_eq!(expected_joules(Stage::Neural, &p, &c), 1_000.0);
+        assert_eq!(expected_joules(Stage::Kv, &p, &c), 10.0 + 0.75 * 1_000.0);
+        assert_eq!(
+            expected_joules(Stage::Graph, &p, &c),
+            100.0 + 0.25 * 1_000.0
+        );
+        let walk = 1.0 + 0.5 * (10.0 + 0.75 * (100.0 + 0.25 * 1_000.0));
+        assert_eq!(expected_joules(Stage::Cache, &p, &c), walk);
+    }
+
+    /// The defect this rule replaces: Neural is the likeliest to succeed, so
+    /// an argmax over success picked it, at a thousand times Graph's cost.
+    #[test]
+    fn the_likeliest_stage_is_not_the_cheapest_to_answer_from() {
+        let c = costs();
+        let p = [0.0, 0.0, 0.95, 1.0];
+        assert_eq!(cheapest_in_expectation(&p, &c), Stage::Graph);
+        // Hopeless cheap stages: skipping them all is cheapest.
+        assert_eq!(
+            cheapest_in_expectation(&[0.0, 0.0, 0.0, 1.0], &c),
+            Stage::Neural
+        );
+        // A likely cache hit: walk from the top.
+        assert_eq!(
+            cheapest_in_expectation(&[0.9, 0.0, 0.0, 1.0], &c),
+            Stage::Cache
+        );
+    }
+
+    #[test]
+    fn a_broken_estimate_never_prices_a_stage_cheaper() {
+        let c = costs();
+        let nan = [f32::NAN; 4];
+        assert_eq!(expected_joules(Stage::Kv, &nan, &c), 1_010.0);
+        assert_eq!(expected_joules(Stage::Kv, &[0.0, 7.0, 0.0, 0.0], &c), 10.0);
+        assert_eq!(
+            expected_joules(Stage::Kv, &[0.0, -3.0, 0.0, 0.0], &c),
+            1_010.0
+        );
+    }
+
+    #[test]
+    fn a_tie_goes_to_the_cheaper_stage() {
+        let mut c = costs();
+        c.insert(Stage::Graph, 10);
+        // Kv and Graph both cost 10 and both certainly resolve.
+        assert_eq!(
+            cheapest_in_expectation(&[0.0, 1.0, 1.0, 1.0], &c),
+            Stage::Kv
+        );
+    }
+
+    #[test]
+    fn a_stage_without_an_estimate_is_priced_at_the_default() {
+        let p = [0.0; 4];
+        assert_eq!(
+            expected_joules(Stage::Neural, &p, &HashMap::new()),
+            50_000_000.0
+        );
+    }
+
+    #[test]
+    fn blending_moves_a_tenth_of_the_way() {
+        let mut j = costs();
+        blend_joules(&mut j, Stage::Neural, 2_000);
+        assert_eq!(j[&Stage::Neural], 1_100);
+        let mut empty = HashMap::new();
+        blend_joules(&mut empty, Stage::Kv, 40);
+        assert_eq!(empty[&Stage::Kv], 40);
+    }
+}

eoc/eoc-rs/crates/eoc-route-learned/src/training.rs

@@ -5,6 +5,8 @@
 //! `train_logreg` (logistic regression). Both consume `&[Example]` and emit a
 //! ready-to-serve router.
 
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::path::Path;
 
 use eoc_core::Stage;
@@ -15,6 +17,7 @@ use serde::{Deserialize, Serialize};
 use crate::classifier::LogRegRouter;
 use crate::error::{Error, Result};
 use crate::matrix_factorization::{DEFAULT_LATENT_DIM, MfRouter};
+use crate::router::stage_index;
 
 /// One labeled training example.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -141,16 +144,46 @@ pub fn train_mf_with_dim(
     router
 }
 
-/// Train a logistic-regression router. Only successful examples shape the
-/// classifier — failed-stage examples are skipped because we want
-/// `P(stage = correct | x)`.
+/// The key that makes rows one query: their embedding, bit for bit.
+pub(crate) fn query_key(embedding: &[f32]) -> Vec<u32> {
+    embedding.iter().map(|v| v.to_bits()).collect()
+}
+
+/// Each query's cheapest successful row, in the order queries first appear.
+/// Rows with the same embedding are one query; a query no stage resolved has
+/// no row here.
+pub(crate) fn cheapest_successes(examples: &[Example]) -> Vec<&Example> {
+    let mut index: HashMap<Vec<u32>, usize> = HashMap::new();
+    let mut out: Vec<&Example> = Vec::new();
+    for e in examples.iter().filter(|e| e.success) {
+        match index.entry(query_key(&e.embedding)) {
+            Entry::Occupied(slot) => {
+                let kept = &mut out[*slot.get()];
+                if stage_index(e.stage) < stage_index(kept.stage) {
+                    *kept = e;
+                }
+            }
+            Entry::Vacant(slot) => {
+                slot.insert(out.len());
+                out.push(e);
+            }
+        }
+    }
+    out
+}
+
+/// Train a logistic-regression router to predict the cheapest stage that
+/// resolves a query. Each query is labelled with its cheapest successful row
+/// ([`Example`]s with the same embedding are one query); failed rows label
+/// nothing. Labelling every successful row instead would teach it the stage
+/// that most often succeeds, and Neural succeeds on everything.
 pub fn train_logreg(examples: &[Example], config: TrainingConfig) -> LogRegRouter {
     let dim = first_embedding_dim(examples).unwrap_or(1);
     let mut router = LogRegRouter::new(dim, config.l2, config.seed);
     if examples.is_empty() {
         return router;
     }
-    let positive: Vec<&Example> = examples.iter().filter(|e| e.success).collect();
+    let positive = cheapest_successes(examples);
     if positive.is_empty() {
         return router;
     }
@@ -182,4 +215,50 @@ mod tests {
         let r = train_logreg(&[], TrainingConfig::default());
         assert_eq!(r.embedding_dim(), 1);
     }
+
+    #[test]
+    fn a_query_is_labelled_with_its_cheapest_success() {
+        let a = vec![1.0, 0.0];
+        let b = vec![0.0, 1.0];
+        let rows = vec![
+            Example::new(a.clone(), Stage::Neural, true, 50_000_000),
+            Example::new(a.clone(), Stage::Cache, false, 1_000),
+            Example::new(b.clone(), Stage::Graph, false, 500_000),
+            Example::new(a.clone(), Stage::Graph, true, 500_000),
+            Example::new(a.clone(), Stage::Kv, true, 50_000),
+            Example::new(b.clone(), Stage::Neural, true, 50_000_000),
+            Example::new(vec![9.0, 9.0], Stage::Kv, false, 50_000),
+        ];
+        let kept: Vec<(Vec<f32>, Stage)> = cheapest_successes(&rows)
+            .into_iter()
+            .map(|e| (e.embedding.clone(), e.stage))
+            .collect();
+        assert_eq!(kept, vec![(a, Stage::Kv), (b, Stage::Neural)]);
+    }
+
+    /// Every query here is resolved by Kv, its cheapest stage, and by Graph
+    /// and Neural; Neural ran twice. Trained on every successful row, the
+    /// classifier learned Neural, the stage that most often succeeds.
+    #[test]
+    fn logreg_learns_the_cheapest_stage_not_the_most_frequent_success() {
+        let mut rows = Vec::new();
+        for i in 0..40 {
+            let emb = vec![1.0, i as f32 / 40.0];
+            rows.push(Example::new(emb.clone(), Stage::Kv, true, 50_000));
+            rows.push(Example::new(emb.clone(), Stage::Graph, true, 500_000));
+            rows.push(Example::new(emb.clone(), Stage::Neural, true, 50_000_000));
+            rows.push(Example::new(emb, Stage::Neural, true, 50_000_000));
+        }
+        let cfg = TrainingConfig {
+            epochs: 30,
+            learning_rate: 0.1,
+            l2: 0.0,
+            ..TrainingConfig::default()
+        };
+        let r = train_logreg(&rows, cfg);
+        for i in 0..40 {
+            let emb = [1.0, i as f32 / 40.0];
+            assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv);
+        }
+    }
 }

eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs

@@ -1,68 +1,193 @@
-//! 100-round LinUCB synthetic bandit test.
+//! Synthetic bandit tests: LinUCB and Thompson sampling.
 //!
-//! Four arms with known mean rewards. Context is a fixed-length embedding
-//! that *informs* but doesn't determine reward (Gaussian noise). After ~50
-//! rounds the router should pull the best arm more than any other.
+//! Each stage resolves a query with a known probability. Each round the
+//! router picks a stage, the pick is run the way `LearnedCascade` dispatches
+//! it (Kv or Graph falling through to Neural on a miss, Neural alone, or for
+//! Cache the full walk), and every stage that ran reports its outcome and
+//! cost. The router must learn the dispatch that spends least, not the stage
+//! likeliest to succeed: before, both bandits converged on Neural.
 
 use eoc_core::Stage;
-use eoc_route_learned::bandit::LinUcbRouter;
-use rand::SeedableRng;
+use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter};
+use eoc_route_learned::router::ALL_STAGES;
+use rand::{RngExt, SeedableRng};
 use rand_distr::{Distribution, Normal};
 use std::collections::HashMap;
 
 const DIM: usize = 4;
 
-fn arm_mean(stage: Stage) -> f32 {
+fn cost(stage: Stage) -> u64 {
     match stage {
-        Stage::Cache => 0.10,
-        Stage::Kv => 0.20,
-        Stage::Graph => 0.30,
-        // Best arm: Neural.
-        Stage::Neural => 0.80,
+        Stage::Cache => 1_000,
+        Stage::Kv => 50_000,
+        Stage::Graph => 500_000,
+        Stage::Neural => 50_000_000,
     }
 }
 
-#[test]
-fn linucb_converges_to_best_arm() {
-    // Lower alpha → less exploration, more exploitation. Ridge=0.1 so the
-    // estimated theta moves quickly.
-    let mut router = LinUcbRouter::new(DIM, 0.1, 0.1);
-    let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF);
-    let ctx_dist = Normal::new(0.0_f32, 1.0).expect("valid normal");
-    let reward_noise = Normal::new(0.0_f32, 0.02).expect("valid normal");
-
-    let mut pulls: HashMap<Stage, usize> = HashMap::new();
-    let mut tail_pulls: HashMap<Stage, usize> = HashMap::new();
-
-    for round in 0..100 {
-        let ctx: Vec<f32> = (0..DIM).map(|_| ctx_dist.sample(&mut rng)).collect();
-        let (pick, _score) = router.pick(&ctx).expect("pick");
-        *pulls.entry(pick).or_insert(0) += 1;
-        if round >= 50 {
-            *tail_pulls.entry(pick).or_insert(0) += 1;
+/// The stages a dispatch to `stage` runs, each with whether it resolved.
+fn dispatch(stage: Stage, resolves: &HashMap<Stage, bool>) -> Vec<(Stage, bool)> {
+    let order: &[Stage] = match stage {
+        Stage::Cache => &ALL_STAGES,
+        Stage::Kv => &[Stage::Kv, Stage::Neural],
+        Stage::Graph => &[Stage::Graph, Stage::Neural],
+        Stage::Neural => &[Stage::Neural],
+    };
+    let mut ran = Vec::new();
+    for s in order {
+        ran.push((*s, resolves[s]));
+        if resolves[s] {
+            break;
         }
-        let reward = (arm_mean(pick) + reward_noise.sample(&mut rng)).clamp(0.0, 1.0);
-        router.update(&ctx, pick, reward).expect("update");
     }
+    ran
+}
 
-    println!("LinUCB pulls (all 100 rounds): {pulls:?}");
-    println!("LinUCB pulls (rounds 50-99): {tail_pulls:?}");
-
-    let neural_tail = *tail_pulls.get(&Stage::Neural).unwrap_or(&0);
-    let best_other = [Stage::Cache, Stage::Kv, Stage::Graph]
-        .iter()
-        .map(|s| *tail_pulls.get(s).unwrap_or(&0))
-        .max()
-        .unwrap_or(0);
-    assert!(
-        neural_tail >= best_other,
-        "neural ({neural_tail}) should be pulled at least as often as any other ({best_other}) after round 50"
-    );
-    // And neural should dominate over the full 100 rounds, since its mean
-    // reward is 0.8 vs ≤0.3 for every other arm.
-    let neural_total = *pulls.get(&Stage::Neural).unwrap_or(&0);
-    assert!(
-        neural_total >= 40,
-        "neural should be pulled ≥40 times over 100 rounds, got {neural_total}"
-    );
+/// Pick counts over the last `tail` of `rounds`, and the joules spent there.
+struct Run {
+    tail_picks: HashMap<Stage, usize>,
+    tail_microjoules: u64,
+}
+
+fn run(
+    rounds: usize,
+    tail: usize,
+    chance: [f64; 4],
+    mut pick: impl FnMut(&[f32]) -> Stage,
+    mut learn: impl FnMut(&[f32], Stage, bool, u64),
+) -> Run {
+    let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF);
+    let noise = Normal::new(0.0_f32, 1.0).expect("valid normal");
+    let mut out = Run {
+        tail_picks: HashMap::new(),
+        tail_microjoules: 0,
+    };
+    for round in 0..rounds {
+        // A constant feature first, so LinUCB can learn each stage's base rate.
+        let mut ctx = vec![1.0_f32];
+        ctx.extend((1..DIM).map(|_| 0.1 * noise.sample(&mut rng)));
+        let resolves: HashMap<Stage, bool> = ALL_STAGES
+            .iter()
+            .zip(chance)
+            .map(|(s, p)| (*s, rng.random_bool(p)))
+            .collect();
+        let stage = pick(&ctx);
+        for (s, ok) in dispatch(stage, &resolves) {
+            learn(&ctx, s, ok, cost(s));
+            if round >= rounds - tail {
+                out.tail_microjoules += cost(s);
+            }
+        }
+        if round >= rounds - tail {
+            *out.tail_picks.entry(stage).or_insert(0) += 1;
+        }
+    }
+    out
+}
+
+fn linucb(rounds: usize, tail: usize, chance: [f64; 4]) -> Run {
+    let router = std::cell::RefCell::new(LinUcbRouter::new(DIM, 0.1, 1.0));
+    run(
+        rounds,
+        tail,
+        chance,
+        |ctx| router.borrow().pick(ctx).expect("pick").0,
+        |ctx, s, ok, uj| {
+            let mut r = router.borrow_mut();
+            r.update(ctx, s, if ok { 1.0 } else { 0.0 })
+                .expect("update");
+            r.observe_cost(s, uj);
+        },
+    )
+}
+
+fn thompson(rounds: usize, tail: usize, chance: [f64; 4]) -> Run {
+    let router = std::cell::RefCell::new(ThompsonSamplingRouter::new(7));
+    run(
+        rounds,
+        tail,
+        chance,
+        |_| router.borrow().pick().0,
+        |_, s, ok, uj| {
+            let mut r = router.borrow_mut();
+            r.update(s, ok);
+            r.observe_cost(s, uj);
+        },
+    )
+}
+
+fn share(run: &Run, stage: Stage, tail: usize) -> f32 {
+    *run.tail_picks.get(&stage).unwrap_or(&0) as f32 / tail as f32
+}
+
+/// Graph resolves 95% of queries and nothing cheaper resolves any: skip to
+/// Graph, at about 3 million µJ a query, not Neural at 50 million.
+#[test]
+fn bandits_skip_to_the_stage_that_resolves() {
+    let chance = [0.0, 0.0, 0.95, 1.0];
+    for (name, r) in [
+        ("linucb", linucb(3000, 500, chance)),
+        ("thompson", thompson(3000, 500, chance)),
+    ] {
+        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
+        assert!(
+            share(&r, Stage::Graph, 500) > 0.8,
+            "{name} picks Graph only {}",
+            share(&r, Stage::Graph, 500)
+        );
+        assert!(
+            share(&r, Stage::Neural, 500) < 0.05,
+            "{name} still skips to Neural"
+        );
+        assert!(
+            r.tail_microjoules / 500 < 5_000_000,
+            "{name} spends {} µJ a query",
+            r.tail_microjoules / 500
+        );
+    }
+}
+
+/// Nothing short of Neural resolves anything, so going straight to Neural,
+/// at 50 million µJ, is cheapest; the full walk costs 50.551 million. The
+/// bandits come within 1% of Neural. They still probe Kv now and then: a
+/// probe costs a thousandth of a Neural call and would save the whole call
+/// if it hit, so it is worth trying until the bound on Kv's chance falls
+/// below 0.1%.
+#[test]
+fn bandits_come_within_a_percent_of_neural_when_nothing_cheaper_resolves() {
+    let chance = [0.0, 0.0, 0.0, 1.0];
+    for (name, r) in [
+        ("linucb", linucb(4000, 500, chance)),
+        ("thompson", thompson(4000, 500, chance)),
+    ] {
+        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
+        let per_query = r.tail_microjoules / 500;
+        assert!(
+            per_query < 50_500_000,
+            "{name} spends {per_query} µJ a query"
+        );
+    }
+}
+
+/// The old test's arms. Neural is likeliest to succeed (0.8), and the old
+/// bandits learned to pick it every time. But the cheap stages resolve often
+/// enough that walking them first costs least, so the walk is the pick.
+#[test]
+fn bandits_walk_when_cheap_stages_often_resolve() {
+    let chance = [0.1, 0.2, 0.3, 0.8];
+    for (name, r) in [
+        ("linucb", linucb(1000, 500, chance)),
+        ("thompson", thompson(1000, 500, chance)),
+    ] {
+        println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules);
+        assert!(
+            share(&r, Stage::Cache, 500) > 0.8,
+            "{name} walks only {}",
+            share(&r, Stage::Cache, 500)
+        );
+        assert!(
+            share(&r, Stage::Neural, 500) < 0.05,
+            "{name} skips to Neural"
+        );
+    }
 }

eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs

@@ -1,73 +1,204 @@
-//! Synthetic-data accuracy test for the MF router.
+//! Synthetic test for the MF router.
 //!
-//! Generates 1000 examples where the "best stage" is a deterministic
-//! function of the embedding (first two dims define quadrants → stage). The
-//! MF router should learn this assignment with >80% accuracy.
+//! Each query has one row per stage. Each cheap stage has its own coverage:
+//! the chance it resolves a query is a logistic function of one feature of
+//! the embedding (Cache's centred at 0.5, Kv's at 0, Graph's at −0.5), and
+//! Neural resolves everything. The router must learn to skip to the stage
+//! that is cheapest to answer from in expectation, not to the one likeliest
+//! to succeed, which is Neural.
+//!
+//! The coverage is soft, as real coverage is, and of the family the router
+//! models, so its learned chances can be right. They have to be: with the
+//! default costs a wrong skip pays a Neural call, a thousand times Graph,
+//! while a right one saves at most a hundredth of one. A router sure of
+//! itself on noise-free boundaries skips wrongly at their edges and spends
+//! more than the plain walk.
+//!
+//! With the default costs the plain walk is itself within a fraction of a
+//! percent of the best choices, so there is little for routing to win, and
+//! the router lands within 1% of the best. Where Kv and Graph cost a real
+//! share of a Neural call, the walk is 4–5% above the best, but at this
+//! training budget MF's chances are off by 5–9% a stage, and the router lands
+//! at about the walk's cost rather than capturing that. Training longer made
+//! it beat the walk by 2% while its chances got worse, so that is a property
+//! of the estimator, not of the rule; the test holds it to not losing.
 
-use eoc_core::Stage;
+use std::collections::HashMap;
+
+use async_trait::async_trait;
+use eoc_core::{Query, Stage};
+use eoc_route_learned::evaluation::evaluate;
 use eoc_route_learned::matrix_factorization::MfRouter;
+use eoc_route_learned::router::{
+    ALL_STAGES, LearnedRouter, RouterState, StagePrediction, cheapest_in_expectation,
+    default_stage_joules, expected_joules,
+};
 use eoc_route_learned::training::{Example, TrainingConfig, train_mf};
-use rand::SeedableRng;
+use rand::{RngExt, SeedableRng};
 use rand_distr::{Distribution, Normal};
 
 const DIM: usize = 8;
 
-fn quadrant_stage(emb: &[f32]) -> Stage {
-    let (a, b) = (emb[0], emb[1]);
-    match (a >= 0.0, b >= 0.0) {
-        (true, true) => Stage::Cache,
-        (true, false) => Stage::Kv,
-        (false, true) => Stage::Graph,
-        (false, false) => Stage::Neural,
-    }
+/// Kv and Graph at a real share of a Neural call.
+fn dear_cheap_stages() -> HashMap<Stage, u64> {
+    let mut c = default_stage_joules();
+    c.insert(Stage::Kv, 2_000_000);
+    c.insert(Stage::Graph, 10_000_000);
+    c
 }
 
-fn synthesise(n: usize, seed: u64) -> Vec<Example> {
+/// The true chance that each stage resolves the query with this embedding,
+/// in `ALL_STAGES` order.
+fn chances(emb: &[f32]) -> [f32; 4] {
+    let s = |x: f32| 1.0 / (1.0 + (-4.0 * x).exp());
+    [s(emb[0] - 0.5), s(emb[1]), s(emb[2] + 0.5), 1.0]
+}
+
+fn synthesise(n: usize, seed: u64, costs: &HashMap<Stage, u64>) -> Vec<Example> {
     let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
     let dist = Normal::new(0.0_f32, 1.0).expect("valid normal");
-    let mut out = Vec::with_capacity(n);
+    let mut out = Vec::with_capacity(n * 4);
     for _ in 0..n {
-        let mut emb = vec![0.0_f32; DIM];
-        for v in emb.iter_mut() {
-            *v = dist.sample(&mut rng);
+        let emb: Vec<f32> = (0..DIM).map(|_| dist.sample(&mut rng)).collect();
+        let p = chances(&emb);
+        for (i, stage) in ALL_STAGES.into_iter().enumerate() {
+            let success = rng.random_bool(f64::from(p[i]));
+            out.push(Example::new(emb.clone(), stage, success, costs[&stage]));
         }
-        let stage = quadrant_stage(&emb);
-        // Pure labels — the synthetic test focuses on classification accuracy,
-        // not noise-handling.
-        let success = true;
-        // cost varies per stage so cost estimates are meaningful.
-        let cost = match stage {
-            Stage::Cache => 1_000,
-            Stage::Kv => 50_000,
-            Stage::Graph => 500_000,
-            Stage::Neural => 50_000_000,
-        };
-        out.push(Example::new(emb, stage, success, cost));
     }
     out
 }
 
-#[test]
-fn mf_router_accuracy_above_80() {
-    let train = synthesise(1000, 17);
-    let test = synthesise(200, 42);
-
+fn router(costs: &HashMap<Stage, u64>) -> MfRouter {
     let cfg = TrainingConfig {
         epochs: 30,
         learning_rate: 0.1,
         l2: 1e-5,
         ..TrainingConfig::default()
     };
-    let router: MfRouter = train_mf(&train, cfg);
+    train_mf(&synthesise(1000, 17, costs), cfg)
+}
 
-    let mut correct = 0usize;
-    for ex in &test {
-        let (stage, _conf, _probs) = router.predict(&ex.embedding).expect("predict");
-        if stage == ex.stage {
-            correct += 1;
+/// The rule the routers used before: the stage likeliest to succeed.
+struct LikeliestToSucceed<'a>(&'a MfRouter);
+
+#[async_trait]
+impl LearnedRouter for LikeliestToSucceed<'_> {
+    async fn route(&self, q: &Query) -> StagePrediction {
+        let probs = self.0.predict_probs(q.embedding.as_ref().unwrap()).unwrap();
+        let best = (0..4).fold(0, |b, i| if probs[i] > probs[b] { i } else { b });
+        StagePrediction::new(ALL_STAGES[best], probs[best])
+    }
+    async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {}
+    fn export_state(&self) -> RouterState {
+        unimplemented!()
+    }
+    fn import_state(_s: RouterState) -> eoc_route_learned::Result<Self> {
+        unimplemented!()
+    }
+}
+
+/// Expected µJ over held-out queries, priced at their true chances: the best
+/// choices, the router's, the plain walk's and the likeliest-stage rule's;
+/// and the share of queries where the router's choice is the best one.
+struct Priced {
+    best: f64,
+    routed: f64,
+    walk: f64,
+    likeliest: f64,
+    agree: f64,
+}
+
+fn priced(costs: &HashMap<Stage, u64>) -> Priced {
+    let router = router(costs);
+    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+    let dist = Normal::new(0.0_f32, 1.0).expect("valid normal");
+    let mut out = Priced {
+        best: 0.0,
+        routed: 0.0,
+        walk: 0.0,
+        likeliest: 0.0,
+        agree: 0.0,
+    };
+    let mut fit = [0.0_f64; 4];
+    let n = 2000;
+    for _ in 0..n {
+        let emb: Vec<f32> = (0..DIM).map(|_| dist.sample(&mut rng)).collect();
+        let p = chances(&emb);
+        let (choice, _, learned) = router.predict(&emb).expect("predict");
+        let ideal = cheapest_in_expectation(&p, costs);
+        let top = (0..4).fold(0, |b, i| if learned[i] > learned[b] { i } else { b });
+        for i in 0..4 {
+            fit[i] += f64::from((learned[i] - p[i]).abs()) / f64::from(n);
         }
+        out.agree += f64::from(u8::from(choice == ideal)) / f64::from(n);
+        out.routed += expected_joules(choice, &p, costs);
+        out.walk += expected_joules(Stage::Cache, &p, costs);
+        out.best += expected_joules(ideal, &p, costs);
+        out.likeliest += expected_joules(ALL_STAGES[top], &p, costs);
     }
-    let acc = correct as f32 / test.len() as f32;
-    println!("MF synthetic accuracy: {:.3}", acc);
-    assert!(acc > 0.8, "MF accuracy {acc} below 0.8 threshold");
+    println!(
+        "expected µJ: best {:.0}, routed {:.0}, walk {:.0}, likeliest {:.0}; agrees with best {:.3}",
+        out.best, out.routed, out.walk, out.likeliest, out.agree
+    );
+    println!("mean |learned − true| chance per stage: {fit:.3?}");
+    out
+}
+
+#[test]
+fn with_the_default_costs_the_router_lands_within_a_percent_of_the_best() {
+    let r = priced(&default_stage_joules());
+    assert!(
+        r.agree > 0.9,
+        "the router's choice is the best for only {} of queries",
+        r.agree
+    );
+    assert!(
+        r.routed <= r.best * 1.01,
+        "routed {:.0} µJ is over 1% above the best {:.0}",
+        r.routed,
+        r.best
+    );
+    assert!(
+        r.routed * 2.0 < r.likeliest,
+        "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}",
+        r.routed,
+        r.likeliest
+    );
+}
+
+#[test]
+fn where_cheap_stages_are_dear_the_router_does_not_lose_to_the_walk() {
+    let r = priced(&dear_cheap_stages());
+    assert!(
+        r.routed <= r.walk * 1.01,
+        "routed {:.0} µJ is over 1% above the walk's {:.0}",
+        r.routed,
+        r.walk
+    );
+    assert!(
+        r.routed * 2.0 < r.likeliest,
+        "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}",
+        r.routed,
+        r.likeliest
+    );
+}
+
+/// The same, replayed over sampled outcomes by `evaluate`.
+#[tokio::test]
+async fn mf_router_spends_far_less_than_the_likeliest_stage_rule() {
+    let costs = default_stage_joules();
+    let router = router(&costs);
+    let test = synthesise(500, 42, &costs);
+    let routed = evaluate(&router, &test).await;
+    let likeliest = evaluate(&LikeliestToSucceed(&router), &test).await;
+    println!("routed:     {routed:?}");
+    println!("likeliest:  {likeliest:?}");
+    assert_eq!((routed.n, routed.unscored), (500, 0));
+    assert!(
+        routed.total_microjoules * 2 < likeliest.total_microjoules,
+        "routed {} µJ is not far below the likeliest-stage rule's {}",
+        routed.total_microjoules,
+        likeliest.total_microjoules
+    );
 }

eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs

@@ -46,3 +46,45 @@ fn rejects_wrong_algorithm_tag() {
     s.algorithm = "wrong".into();
     assert!(MfRouter::import_state(s).is_err());
 }
+
+/// The bandits learn each stage's cost. A learned cost survives export, and
+/// a state saved before they learned costs imports with the defaults rather
+/// than failing.
+#[test]
+fn bandit_stage_costs_roundtrip_and_older_states_import() {
+    use eoc_core::Stage;
+    use eoc_route_learned::router::default_stage_joules;
+
+    fn without_joules(state: &mut eoc_route_learned::router::RouterState) {
+        let mut v: serde_json::Value = serde_json::from_slice(&state.weights_blob).unwrap();
+        v.as_object_mut()
+            .unwrap()
+            .remove("joules")
+            .expect("the field is exported");
+        state.weights_blob = serde_json::to_vec(&v).unwrap();
+    }
+
+    let mut lin = LinUcbRouter::new(4, 1.0, 1.0);
+    lin.observe_cost(Stage::Neural, 10_000);
+    let restored = LinUcbRouter::import_state(lin.export_state()).expect("import");
+    assert_eq!(restored.joules, lin.joules);
+    let mut old = lin.export_state();
+    without_joules(&mut old);
+    assert_eq!(
+        LinUcbRouter::import_state(old).expect("import").joules,
+        default_stage_joules()
+    );
+
+    let mut ts = ThompsonSamplingRouter::new(3);
+    ts.observe_cost(Stage::Kv, 7);
+    let restored = ThompsonSamplingRouter::import_state(ts.export_state()).expect("import");
+    assert_eq!(restored.joules, ts.joules);
+    let mut old = ts.export_state();
+    without_joules(&mut old);
+    assert_eq!(
+        ThompsonSamplingRouter::import_state(old)
+            .expect("import")
+            .joules,
+        default_stage_joules()
+    );
+}