diff --git a/crates/ml/src/cuda_pipeline/learning_health.rs b/crates/ml/src/cuda_pipeline/learning_health.rs new file mode 100644 index 000000000..2dd02195c --- /dev/null +++ b/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"); + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index c714a72f1..a8df655e9 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/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).