feat(A2): LearningHealth module with 7-component composition + EMA

crates/ml/src/cuda_pipeline/learning_health.rs

@@ -0,0 +1,175 @@
+//! LearningHealth signal — unified training health metric.
+//!
+//! Composes 7 normalized components into a scalar [0, 1] that senses training
+//! collapse and drives adaptive hyperparameters.
+
+/// Smoothstep function: smooth transition from 0 at `edge0` to 1 at `edge1`.
+/// Standard GLSL smoothstep formula.
+#[inline]
+pub fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 {
+    if edge1 == edge0 { return if x >= edge0 { 1.0 } else { 0.0 }; }
+    let t = ((x - edge0) / (edge1 - edge0)).clamp(0.0, 1.0);
+    t * t * (3.0 - 2.0 * t)
+}
+
+/// Raw input signals for LearningHealth composition.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct HealthComponents {
+    pub q_gap: f32,                // max - 2nd max of Q values across actions
+    pub q_var: f32,                // variance of Q across actions
+    pub atom_util: f32,            // C51 atom utilization fraction [0, 1]
+    pub grad_norm: f32,            // gradient L2 norm
+    pub ens_disagreement: f32,     // mean pairwise diff across ensemble heads
+    pub grad_consistency: f32,     // cosine similarity of successive gradients
+    pub spectral_gap: f32,         // sigma_1 / sigma_2 from SVD of Q matrix
+}
+
+/// Computed normalized signals in [0, 1].
+#[derive(Debug, Clone, Copy, Default)]
+pub struct NormalizedComponents {
+    pub q_gap_norm: f32,
+    pub q_var_norm: f32,
+    pub atom_util_norm: f32,
+    pub grad_stable: f32,
+    pub ens_agree: f32,
+    pub grad_consistency_norm: f32,
+    pub spectral_gap_norm: f32,
+}
+
+impl NormalizedComponents {
+    pub fn from_raw(raw: &HealthComponents) -> Self {
+        Self {
+            q_gap_norm: smoothstep(0.01, 0.5, raw.q_gap),
+            q_var_norm: smoothstep(0.001, 0.1, raw.q_var),
+            atom_util_norm: smoothstep(0.2, 0.7, raw.atom_util),
+            grad_stable: 1.0 - smoothstep(10.0, 100.0, raw.grad_norm),
+            ens_agree: (1.0 - raw.ens_disagreement).clamp(0.0, 1.0),
+            grad_consistency_norm: smoothstep(-0.2, 0.5, raw.grad_consistency),
+            spectral_gap_norm: 1.0 - smoothstep(2.0, 10.0, raw.spectral_gap),
+        }
+    }
+
+    /// Weighted sum per spec Layer 1 composition.
+    pub fn compose(&self) -> f32 {
+        0.25 * self.q_gap_norm
+        + 0.15 * self.q_var_norm
+        + 0.15 * self.atom_util_norm
+        + 0.15 * self.grad_stable
+        + 0.10 * self.ens_agree
+        + 0.10 * self.grad_consistency_norm
+        + 0.10 * self.spectral_gap_norm
+    }
+}
+
+/// LearningHealth state: EMA-smoothed health value with bounds and warmup.
+#[derive(Debug, Clone)]
+pub struct LearningHealth {
+    /// Current EMA-smoothed health [0.2, 0.95].
+    pub value: f32,
+    /// Epoch counter — enables warmup mode for first N epochs.
+    pub epoch: u32,
+    /// Latest normalized components (for logging).
+    pub components: NormalizedComponents,
+    /// EMA smoothing factor (new sample weight).
+    ema_alpha: f32,
+    /// Warmup epochs — health clamped to 0.5 during warmup.
+    warmup_epochs: u32,
+}
+
+impl LearningHealth {
+    pub fn new() -> Self {
+        Self {
+            value: 0.5, // neutral start
+            epoch: 0,
+            components: NormalizedComponents::default(),
+            ema_alpha: 0.3,
+            warmup_epochs: 3,
+        }
+    }
+
+    /// Update health from new raw components. Returns updated health value.
+    pub fn update(&mut self, raw: &HealthComponents) -> f32 {
+        self.components = NormalizedComponents::from_raw(raw);
+        let new_sample = self.components.compose();
+
+        if self.epoch < self.warmup_epochs {
+            // Warmup: pin to 0.5, allow components to accumulate
+            self.value = 0.5;
+        } else {
+            // EMA update
+            let ema = (1.0 - self.ema_alpha) * self.value + self.ema_alpha * new_sample;
+            // Clamp to safe bounds
+            self.value = ema.clamp(0.2, 0.95);
+        }
+
+        self.epoch += 1;
+        self.value
+    }
+}
+
+impl Default for LearningHealth {
+    fn default() -> Self { Self::new() }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_smoothstep_boundaries() {
+        assert_eq!(smoothstep(0.0, 1.0, -0.5), 0.0);
+        assert_eq!(smoothstep(0.0, 1.0, 1.5), 1.0);
+        assert!((smoothstep(0.0, 1.0, 0.5) - 0.5).abs() < 0.001);
+    }
+
+    #[test]
+    fn test_collapsed_inputs_give_low_health() {
+        let raw = HealthComponents {
+            q_gap: 0.01,          // collapsed
+            q_var: 0.001,         // collapsed
+            atom_util: 0.1,       // under-utilized
+            grad_norm: 50_000.0,  // exploding
+            ens_disagreement: 0.0,// everyone agrees (on uniform)
+            grad_consistency: -0.3, // gradients contradict
+            spectral_gap: 50.0,   // rank 1
+        };
+        let mut health = LearningHealth::new();
+        // Skip warmup
+        for _ in 0..5 { health.update(&raw); }
+        assert!(health.value <= 0.3, "Collapsed state should give health <= 0.3, got {}", health.value);
+    }
+
+    #[test]
+    fn test_healthy_inputs_give_high_health() {
+        let raw = HealthComponents {
+            q_gap: 1.0,            // strong differentiation
+            q_var: 0.5,            // good variance
+            atom_util: 0.9,        // atoms fully utilized
+            grad_norm: 5.0,        // stable
+            ens_disagreement: 0.1, // moderate diversity
+            grad_consistency: 0.9, // consistent direction
+            spectral_gap: 1.5,     // balanced rank
+        };
+        let mut health = LearningHealth::new();
+        for _ in 0..5 { health.update(&raw); }
+        assert!(health.value >= 0.7, "Healthy state should give health >= 0.7, got {}", health.value);
+    }
+
+    #[test]
+    fn test_warmup_clamps_to_neutral() {
+        let raw = HealthComponents {
+            q_gap: 0.01, q_var: 0.001, atom_util: 0.1,
+            grad_norm: 50_000.0, ens_disagreement: 0.0,
+            grad_consistency: -0.5, spectral_gap: 100.0,
+        };
+        let mut health = LearningHealth::new();
+        // First 3 updates are warmup — should stay at 0.5
+        for _ in 0..3 {
+            let v = health.update(&raw);
+            assert_eq!(v, 0.5, "Warmup should pin health to 0.5");
+        }
+        // 4th update starts EMA
+        health.update(&raw);
+        assert!(health.value < 0.5, "Post-warmup should start dropping with collapsed inputs");
+    }
+}

crates/ml/src/cuda_pipeline/mod.rs

@@ -37,6 +37,7 @@ pub mod gpu_iql_trainer;
 pub mod gpu_iqn_head;
 pub mod gpu_attention;
 pub mod decision_transformer;
+pub mod learning_health;
 // gpu_replay_buffer moved to ml-dqn crate
 
 /// Maximum bytes allowed for a single GPU upload (2 GB safety limit).