diff --git a/crates/ml-dqn/src/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs index 4b459b80e..db5e3077f 100644 --- a/crates/ml-dqn/src/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -71,6 +71,12 @@ pub struct NoisyLinear { in_features: usize, out_features: usize, + // Task 2.5 Bug #4: effective sigma scale applied by the scheduler via + // `reset_noise_with_sigma`. Retained so `sigma_mean()` can report the + // *effective* σ (scaled), not the raw σ tensor constant. Default 1.0 + // means "no scheduler attenuation", matching reset_noise()'s behaviour. + current_sigma_scale: f32, + // CUDA stream for GPU operations stream: Arc, } @@ -145,6 +151,7 @@ impl NoisyLinear { bias_epsilon, in_features, out_features, + current_sigma_scale: 1.0, stream, }) } @@ -163,6 +170,11 @@ impl NoisyLinear { // where f(x) = sign(x) * sqrt(|x|) let scale = sigma_scale as f32; + // Task 2.5 Bug #4: record scheduler attenuation so HEALTH_DIAG + // `sigma_mean` reports effective σ (= raw σ × current scale), + // not the constant raw weight_sigma tensor mean. + self.current_sigma_scale = scale; + // Generate raw Gaussian noise on host, apply f(x) = sign(x) * sqrt(|x|) let apply_f = |x: f32| -> f32 { let sign = if x >= 0.0 { 1.0_f32 } else { -1.0_f32 }; @@ -339,9 +351,16 @@ impl NoisyLinear { [&self.weight_sigma, &self.bias_sigma] } - /// Task 0.6 — mean of |sigma| across weight + bias σ tensors. Used by - /// HEALTH_DIAG to track NoisyNets exploration pressure per branch - /// (H7 detection signal). Sync stream-readback; epoch-boundary only. + /// Task 0.6 — mean of |sigma| across weight + bias σ tensors, scaled by + /// the scheduler's current sigma_scale. Used by HEALTH_DIAG to track + /// NoisyNets exploration pressure per branch (H7 detection signal). + /// Sync stream-readback; epoch-boundary only. + /// + /// Task 2.5 Bug #4: result is multiplied by `current_sigma_scale` so it + /// reports effective σ (raw σ × scheduler scale), not the constant raw + /// weight_sigma tensor mean. `reset_noise_with_sigma` scales the noise + /// epsilon samples; the σ tensor itself stays fixed, so without this + /// multiplier the diag reads constant 0.0320 regardless of schedule. pub fn sigma_mean(&self) -> Result { let n_w = self.weight_sigma.len(); let n_b = self.bias_sigma.len(); @@ -358,7 +377,8 @@ impl NoisyLinear { .map_err(|e| MLError::ModelError(format!("sigma_mean bias dtoh: {e}")))?; let sum_w: f64 = h_w.iter().map(|x| x.abs() as f64).sum(); let sum_b: f64 = h_b.iter().map(|x| x.abs() as f64).sum(); - Ok(((sum_w + sum_b) / ((n_w + n_b) as f64)) as f32) + let raw_mean = ((sum_w + sum_b) / ((n_w + n_b) as f64)) as f32; + Ok(raw_mean * self.current_sigma_scale) } /// Disable noise for evaluation (use mean parameters only) @@ -387,6 +407,9 @@ impl NoisyLinear { Self::dtod_copy_async(&src.bias_sigma, &mut self.bias_sigma, "bias_sigma", &self.stream)?; Self::dtod_copy_async(&src.weight_epsilon, &mut self.weight_epsilon, "weight_epsilon", &self.stream)?; Self::dtod_copy_async(&src.bias_epsilon, &mut self.bias_epsilon, "bias_epsilon", &self.stream)?; + // Task 2.5 Bug #4: propagate scheduler sigma scale so a target-net + // sync doesn't reset the effective σ reporting to 1.0. + self.current_sigma_scale = src.current_sigma_scale; Ok(()) }