From a0aab3baa5bbd41fcf8ada9fc7f2dc174c6e1815 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 9 Apr 2026 15:42:41 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20add=20c51=5Falpha=5Fmax=20to=20cap=20MS?= =?UTF-8?q?E=E2=86=92C51=20blend=20and=20prevent=20gradient=20starvation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit H100 baseline (20 epochs) showed gradient collapse at epoch 10: C51 cross-entropy converges its distributional fit before the policy converges, leaving zero gradient signal. The collapse happened 5 epochs after C51 reached alpha=1.0 (pure C51, zero MSE). Fix: cap the C51 alpha ramp at c51_alpha_max (default 0.5) so MSE always contributes (1 - alpha_max) of the primary gradient. MSE loss measures Q-error directly and only goes to zero when Q-values are correct, not just when the distribution shape is right. - c51_alpha_max added to DQNHyperparameters (default 0.5) - Added to PSO search space as 15th dimension (range [0.3, 0.9]) - Added to TOML profiles: smoketest, production, hyperopt - Training loop caps alpha ramp at alpha_max instead of 1.0 - All 907 ml tests + 300 ml-core tests + 6 smoke tests pass Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-hyperopt.toml | 1 + config/training/dqn-production.toml | 1 + config/training/dqn-smoketest.toml | 1 + crates/ml/src/hyperopt/adapters/dqn.rs | 225 ++++++++++-------- crates/ml/src/trainers/dqn/config.rs | 7 + crates/ml/src/trainers/dqn/trainer/tests.rs | 4 +- .../src/trainers/dqn/trainer/training_loop.rs | 7 +- crates/ml/src/training_profile.rs | 14 +- 8 files changed, 153 insertions(+), 107 deletions(-) diff --git a/config/training/dqn-hyperopt.toml b/config/training/dqn-hyperopt.toml index da66ce079..2ebdd5a3b 100644 --- a/config/training/dqn-hyperopt.toml +++ b/config/training/dqn-hyperopt.toml @@ -70,6 +70,7 @@ spectral_norm_sigma_max = [1.0, 10.0] # C51 warmup c51_warmup_epochs = [0, 10] +c51_alpha_max = [0.3, 0.9] # HER ratio her_ratio = [0.0, 0.8] diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index d4c4ed666..be51392af 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -82,6 +82,7 @@ q_clip_max = 50.0 n_steps = 5 tau = 0.005 c51_warmup_epochs = 5 +c51_alpha_max = 0.5 her_ratio = 0.2 cql_alpha = 1.0 curiosity_weight = 0.1 diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index e9dbc8d1f..baafb825d 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -74,6 +74,7 @@ q_clip_max = 50.0 n_steps = 5 tau = 0.005 c51_warmup_epochs = 5 +c51_alpha_max = 0.5 her_ratio = 0.2 cql_alpha = 0.1 curiosity_weight = 0.1 diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 486aa7980..7f0425e18 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -145,14 +145,14 @@ const TRIAL_VRAM_MB: f64 = 7000.0; /// Corrected: was 0.02 (20KB) — 100x too high const MB_PER_SAMPLE: f64 = 0.0005; -/// DQN hyperparameter space (14D continuous, consolidated from 30D) +/// DQN hyperparameter space (15D continuous, consolidated from 30D) /// -/// 3 breakout parameters are searched independently, 5 core family intensity +/// 4 breakout parameters are searched independently, 5 core family intensity /// scalars control groups of related params, and 6 generalization family /// intensity scalars (already existed). Fixed params come from TOML defaults. /// -/// ## 14D Layout: -/// - gamma, iqn_lambda, c51_warmup_epochs (3 breakouts) +/// ## 15D Layout: +/// - gamma, iqn_lambda, c51_warmup_epochs, c51_alpha_max (4 breakouts) /// - learning_intensity, exploration_intensity, replay_intensity, /// architecture_intensity, risk_intensity (5 core families) /// - adversarial_intensity, regularization_intensity, augmentation_intensity, @@ -164,7 +164,7 @@ const MB_PER_SAMPLE: f64 = 0.0005; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct DQNParams { - // === 3 breakout parameters (searched independently) === + // === 4 breakout parameters (searched independently) === /// Discount factor for future rewards [0.95, 0.999] pub gamma: f64, @@ -172,6 +172,8 @@ pub struct DQNParams { pub iqn_lambda: f64, /// C51 warmup epochs: MSE loss for first N epochs [3.0, 15.0] pub c51_warmup_epochs: f64, + /// C51 alpha max: caps MSE→C51 blend to prevent gradient starvation [0.3, 0.9] + pub c51_alpha_max: f64, // === 5 core family intensity scalars [0.0, 2.0] === @@ -208,6 +210,7 @@ impl Default for DQNParams { gamma: 0.99, iqn_lambda: 0.25, c51_warmup_epochs: 5.0, + c51_alpha_max: 0.5, // Core family intensities: 1.0 = neutral (no change to TOML defaults) learning_intensity: 1.0, exploration_intensity: 1.0, @@ -227,8 +230,8 @@ impl Default for DQNParams { impl ParameterSpace for DQNParams { fn continuous_bounds() -> Vec<(f64, f64)> { - // 14D search space (consolidated from 30D via family intensity grouping) - // 3 breakouts + 5 core families + 6 generalization families = 14D + // 15D search space (consolidated from 30D via family intensity grouping) + // 4 breakouts + 5 core families + 6 generalization families = 15D // // Phase system: certain dims are pinned to 1.0 (single-point bounds) // depending on the active hyperopt phase. @@ -240,7 +243,7 @@ impl ParameterSpace for DQNParams { let pinned = (1.0, 1.0); // fixed at neutral let (learning, exploration, replay, architecture, risk) = match &phase { - // Fast (default): search ALL 14 dims — 14D is well within PSO's efficient range. + // Fast (default): search ALL 15 dims — 15D is well within PSO's efficient range. // Phased search was needed for the old 30D space; with families it's one pass. HyperoptPhase::Fast => (search, search, search, search, search), // Full: search architecture + risk + gen families, pin dynamics from Phase 1 @@ -262,24 +265,25 @@ impl ParameterSpace for DQNParams { (0.95, 0.999), // 0: gamma (breakout — always searchable) (0.0, 1.0), // 1: iqn_lambda (breakout — always searchable) (3.0, 15.0), // 2: c51_warmup_epochs (breakout — always searchable) - learning, // 3: learning_intensity - exploration, // 4: exploration_intensity - replay, // 5: replay_intensity - architecture, // 6: architecture_intensity - risk, // 7: risk_intensity - adversarial, // 8: adversarial_intensity - regularization, // 9: regularization_intensity - augmentation, // 10: augmentation_intensity - loss_shaping, // 11: loss_shaping_intensity - ensemble, // 12: ensemble_intensity - causal, // 13: causal_intensity + (0.3, 0.9), // 3: c51_alpha_max (breakout — always searchable) + learning, // 4: learning_intensity + exploration, // 5: exploration_intensity + replay, // 6: replay_intensity + architecture, // 7: architecture_intensity + risk, // 8: risk_intensity + adversarial, // 9: adversarial_intensity + regularization, // 10: regularization_intensity + augmentation, // 11: augmentation_intensity + loss_shaping, // 12: loss_shaping_intensity + ensemble, // 13: ensemble_intensity + causal, // 14: causal_intensity ] } fn from_continuous(x: &[f64]) -> Result { - if x.len() < 14 { + if x.len() < 15 { return Err(MLError::ConfigError(format!( - "Expected at least 14 continuous parameters, got {}", x.len() + "Expected at least 15 continuous parameters, got {}", x.len() ))); } @@ -287,17 +291,18 @@ impl ParameterSpace for DQNParams { gamma: x[0].clamp(0.95, 0.999), iqn_lambda: x[1].clamp(0.0, 1.0), c51_warmup_epochs: x[2].clamp(3.0, 15.0), - learning_intensity: x[3].clamp(0.0, 2.0), - exploration_intensity: x[4].clamp(0.0, 2.0), - replay_intensity: x[5].clamp(0.0, 2.0), - architecture_intensity: x[6].clamp(0.0, 2.0), - risk_intensity: x[7].clamp(0.0, 2.0), - adversarial_intensity: x[8].clamp(0.0, 2.0), - regularization_intensity: x[9].clamp(0.0, 2.0), - augmentation_intensity: x[10].clamp(0.0, 2.0), - loss_shaping_intensity: x[11].clamp(0.0, 2.0), - ensemble_intensity: x[12].clamp(0.0, 2.0), - causal_intensity: x[13].clamp(0.0, 2.0), + c51_alpha_max: x[3].clamp(0.3, 0.9), + learning_intensity: x[4].clamp(0.0, 2.0), + exploration_intensity: x[5].clamp(0.0, 2.0), + replay_intensity: x[6].clamp(0.0, 2.0), + architecture_intensity: x[7].clamp(0.0, 2.0), + risk_intensity: x[8].clamp(0.0, 2.0), + adversarial_intensity: x[9].clamp(0.0, 2.0), + regularization_intensity: x[10].clamp(0.0, 2.0), + augmentation_intensity: x[11].clamp(0.0, 2.0), + loss_shaping_intensity: x[12].clamp(0.0, 2.0), + ensemble_intensity: x[13].clamp(0.0, 2.0), + causal_intensity: x[14].clamp(0.0, 2.0), }) } @@ -306,17 +311,18 @@ impl ParameterSpace for DQNParams { self.gamma, // 0 self.iqn_lambda, // 1 self.c51_warmup_epochs, // 2 - self.learning_intensity, // 3 - self.exploration_intensity, // 4 - self.replay_intensity, // 5 - self.architecture_intensity, // 6 - self.risk_intensity, // 7 - self.adversarial_intensity, // 8 - self.regularization_intensity, // 9 - self.augmentation_intensity, // 10 - self.loss_shaping_intensity, // 11 - self.ensemble_intensity, // 12 - self.causal_intensity, // 13 + self.c51_alpha_max, // 3 + self.learning_intensity, // 4 + self.exploration_intensity, // 5 + self.replay_intensity, // 6 + self.architecture_intensity, // 7 + self.risk_intensity, // 8 + self.adversarial_intensity, // 9 + self.regularization_intensity, // 10 + self.augmentation_intensity, // 11 + self.loss_shaping_intensity, // 12 + self.ensemble_intensity, // 13 + self.causal_intensity, // 14 ] } @@ -325,17 +331,18 @@ impl ParameterSpace for DQNParams { "gamma", // 0 "iqn_lambda", // 1 "c51_warmup_epochs", // 2 - "learning_intensity", // 3 - "exploration_intensity", // 4 - "replay_intensity", // 5 - "architecture_intensity", // 6 - "risk_intensity", // 7 - "adversarial_intensity", // 8 - "regularization_intensity", // 9 - "augmentation_intensity", // 10 - "loss_shaping_intensity", // 11 - "ensemble_intensity", // 12 - "causal_intensity", // 13 + "c51_alpha_max", // 3 + "learning_intensity", // 4 + "exploration_intensity", // 5 + "replay_intensity", // 6 + "architecture_intensity", // 7 + "risk_intensity", // 8 + "adversarial_intensity", // 9 + "regularization_intensity", // 10 + "augmentation_intensity", // 11 + "loss_shaping_intensity", // 12 + "ensemble_intensity", // 13 + "causal_intensity", // 14 ] } @@ -343,10 +350,10 @@ impl ParameterSpace for DQNParams { let mut bounds = Self::continuous_bounds(); let vram_mb = budget.gpu_memory_mb as f64; - // Architecture intensity (idx 6): cap upper bound on small GPUs to prevent + // Architecture intensity (idx 7): cap upper bound on small GPUs to prevent // scaling hidden_dim/atoms beyond VRAM capacity. if vram_mb > 0.0 && vram_mb < 16_000.0 { - if let Some(b) = bounds.get_mut(6) { + if let Some(b) = bounds.get_mut(7) { // Small GPUs: cap architecture scaling to 1.5x (instead of 2.0x) b.1 = b.1.min(1.5); } @@ -367,9 +374,9 @@ impl ParameterSpace for DQNParams { impl DQNParams { /// Validates fundamental parameter invariants before training. /// - /// With the 14D family-based search space, most individual params come from + /// With the 15D family-based search space, most individual params come from /// TOML defaults and are validated at the DQNHyperparameters level. Here we - /// only check the 3 breakout params and intensity scalar ranges. + /// only check the 4 breakout params and intensity scalar ranges. fn validate_for_hft_trendfollowing(&self) -> Result<(), String> { if self.gamma <= 0.0 || self.gamma > 1.0 { return Err(format!( @@ -389,6 +396,12 @@ impl DQNParams { self.c51_warmup_epochs )); } + if self.c51_alpha_max < 0.0 || self.c51_alpha_max > 1.0 { + return Err(format!( + "c51_alpha_max ({}) must be in [0, 1]", + self.c51_alpha_max + )); + } Ok(()) } } @@ -1952,10 +1965,11 @@ impl HyperparameterOptimizable for DQNTrainer { mem_available_before / 1024 / 1024 ); - info!("Training DQN with parameters (14D family search space):"); + info!("Training DQN with parameters (15D family search space):"); info!(" Gamma: {:.3}", params.gamma); info!(" IQN lambda: {:.3}", params.iqn_lambda); info!(" C51 warmup epochs: {:.1}", params.c51_warmup_epochs); + info!(" C51 alpha max: {:.2}", params.c51_alpha_max); info!(" Learning intensity: {:.3}", params.learning_intensity); info!(" Exploration intensity: {:.3}", params.exploration_intensity); info!(" Replay intensity: {:.3}", params.replay_intensity); @@ -2112,8 +2126,8 @@ impl HyperparameterOptimizable for DQNTrainer { Ok(checkpoint_path.to_string_lossy().to_string()) }; - // 14D family-based approach: load base hyperparams from TOML profile defaults, - // then override only the 3 breakouts and 11 intensities from the PSO vector. + // 15D family-based approach: load base hyperparams from TOML profile defaults, + // then override only the 4 breakouts and 11 intensities from the PSO vector. // All individual params (learning_rate, buffer_size, etc.) come from TOML. // Reuse base_hp from the VRAM gate (same profile, already loaded). let mut hyperparams = base_hp; @@ -2150,13 +2164,14 @@ impl HyperparameterOptimizable for DQNTrainer { // Clamp buffer to [1024, max] — below 1024 the GPU PER segment tree fails. hyperparams.buffer_size = hyperparams.buffer_size.max(1024).min(self.buffer_size_max); - // === 3 breakout parameters from PSO === + // === 4 breakout parameters from PSO === hyperparams.gamma = params.gamma; // Recompute v_min/v_max (reward_scale-derived, gamma-independent) hyperparams.v_min = hyperparams.computed_v_min(); hyperparams.v_max = hyperparams.computed_v_max(); hyperparams.iqn_lambda = params.iqn_lambda; hyperparams.c51_warmup_epochs = params.c51_warmup_epochs.round() as usize; + hyperparams.c51_alpha_max = params.c51_alpha_max as f32; // === 5 core family intensities from PSO === hyperparams.learning_intensity = params.learning_intensity; @@ -2992,11 +3007,12 @@ mod tests { #[test] fn test_dqn_params_roundtrip() { - // Roundtrip test for the 14D family-based search space + // Roundtrip test for the 15D family-based search space let params = DQNParams { gamma: 0.97, iqn_lambda: 0.5, c51_warmup_epochs: 7.0, + c51_alpha_max: 0.6, learning_intensity: 1.2, exploration_intensity: 0.8, replay_intensity: 1.5, @@ -3011,13 +3027,14 @@ mod tests { }; let continuous = params.to_continuous(); - assert_eq!(continuous.len(), 14, "to_continuous must return 14D vector"); + assert_eq!(continuous.len(), 15, "to_continuous must return 15D vector"); let recovered = DQNParams::from_continuous(&continuous).unwrap(); - // All 14 parameters must roundtrip exactly + // All 15 parameters must roundtrip exactly assert!((recovered.gamma - params.gamma).abs() < 1e-6); assert!((recovered.iqn_lambda - params.iqn_lambda).abs() < 1e-6); assert!((recovered.c51_warmup_epochs - params.c51_warmup_epochs).abs() < 1e-6); + assert!((recovered.c51_alpha_max - params.c51_alpha_max).abs() < 1e-6); assert!((recovered.learning_intensity - params.learning_intensity).abs() < 1e-6); assert!((recovered.exploration_intensity - params.exploration_intensity).abs() < 1e-6); assert!((recovered.replay_intensity - params.replay_intensity).abs() < 1e-6); @@ -3034,7 +3051,7 @@ mod tests { #[test] fn test_dqn_params_bounds() { let bounds = DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 14); // 14D family-based search space + assert_eq!(bounds.len(), 15); // 15D family-based search space // Check all bounds are valid ranges (lower < upper) for (i, (lo, hi)) in bounds.iter().enumerate() { @@ -3050,8 +3067,11 @@ mod tests { // C51 warmup bounds assert!((bounds[2].0 - 3.0).abs() < 1e-6); assert!((bounds[2].1 - 15.0).abs() < 1e-6); + // C51 alpha max bounds + assert!((bounds[3].0 - 0.3).abs() < 1e-6); + assert!((bounds[3].1 - 0.9).abs() < 1e-6); // All intensity bounds should be [0.0, 2.0] - for i in 3..14 { + for i in 4..15 { assert!((bounds[i].0 - 0.0).abs() < 1e-6, "Intensity {i} lower bound should be 0.0"); assert!((bounds[i].1 - 2.0).abs() < 1e-6, "Intensity {i} upper bound should be 2.0"); } @@ -3060,52 +3080,55 @@ mod tests { #[test] fn test_param_names() { let names = DQNParams::param_names(); - assert_eq!(names.len(), 14); // 14D family-based search space + assert_eq!(names.len(), 15); // 15D family-based search space assert_eq!(names[0], "gamma"); assert_eq!(names[1], "iqn_lambda"); assert_eq!(names[2], "c51_warmup_epochs"); - assert_eq!(names[3], "learning_intensity"); - assert_eq!(names[4], "exploration_intensity"); - assert_eq!(names[5], "replay_intensity"); - assert_eq!(names[6], "architecture_intensity"); - assert_eq!(names[7], "risk_intensity"); - assert_eq!(names[8], "adversarial_intensity"); - assert_eq!(names[9], "regularization_intensity"); - assert_eq!(names[10], "augmentation_intensity"); - assert_eq!(names[11], "loss_shaping_intensity"); - assert_eq!(names[12], "ensemble_intensity"); - assert_eq!(names[13], "causal_intensity"); + assert_eq!(names[3], "c51_alpha_max"); + assert_eq!(names[4], "learning_intensity"); + assert_eq!(names[5], "exploration_intensity"); + assert_eq!(names[6], "replay_intensity"); + assert_eq!(names[7], "architecture_intensity"); + assert_eq!(names[8], "risk_intensity"); + assert_eq!(names[9], "adversarial_intensity"); + assert_eq!(names[10], "regularization_intensity"); + assert_eq!(names[11], "augmentation_intensity"); + assert_eq!(names[12], "loss_shaping_intensity"); + assert_eq!(names[13], "ensemble_intensity"); + assert_eq!(names[14], "causal_intensity"); } #[test] - fn test_14d_from_continuous_basic() { - // 14D search space + fn test_15d_from_continuous_basic() { + // 15D search space let continuous = vec![ 0.97, // 0: gamma 0.25, // 1: iqn_lambda 5.0, // 2: c51_warmup_epochs - 1.0, // 3: learning_intensity - 1.0, // 4: exploration_intensity - 1.0, // 5: replay_intensity - 1.0, // 6: architecture_intensity - 1.0, // 7: risk_intensity - 1.0, // 8: adversarial_intensity - 1.0, // 9: regularization_intensity - 1.0, // 10: augmentation_intensity - 1.0, // 11: loss_shaping_intensity - 1.0, // 12: ensemble_intensity - 1.0, // 13: causal_intensity + 0.5, // 3: c51_alpha_max + 1.0, // 4: learning_intensity + 1.0, // 5: exploration_intensity + 1.0, // 6: replay_intensity + 1.0, // 7: architecture_intensity + 1.0, // 8: risk_intensity + 1.0, // 9: adversarial_intensity + 1.0, // 10: regularization_intensity + 1.0, // 11: augmentation_intensity + 1.0, // 12: loss_shaping_intensity + 1.0, // 13: ensemble_intensity + 1.0, // 14: causal_intensity ]; let params = DQNParams::from_continuous(&continuous).unwrap(); assert!((params.gamma - 0.97).abs() < 1e-6); assert!((params.iqn_lambda - 0.25).abs() < 1e-6); assert!((params.c51_warmup_epochs - 5.0).abs() < 1e-6); + assert!((params.c51_alpha_max - 0.5).abs() < 1e-6); assert!((params.learning_intensity - 1.0).abs() < 1e-6); assert!((params.adversarial_intensity - 1.0).abs() < 1e-6); // Test with min-length vector to verify error handling - let short = vec![0.97; 13]; // Only 13 elements + let short = vec![0.97; 14]; // Only 14 elements assert!(DQNParams::from_continuous(&short).is_err()); // Test clamping: values outside bounds get clamped @@ -3113,6 +3136,7 @@ mod tests { 0.5, // gamma: clamped to 0.95 5.0, // iqn_lambda: clamped to 1.0 1.0, // c51_warmup_epochs: clamped to 3.0 + 2.0, // c51_alpha_max: clamped to 0.9 3.0, 3.0, 3.0, 3.0, 3.0, // intensities: clamped to 2.0 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, ]; @@ -3120,11 +3144,12 @@ mod tests { assert!((clamped.gamma - 0.95).abs() < 1e-6, "gamma should be clamped to 0.95"); assert!((clamped.iqn_lambda - 1.0).abs() < 1e-6, "iqn_lambda should be clamped to 1.0"); assert!((clamped.c51_warmup_epochs - 3.0).abs() < 1e-6, "c51_warmup should be clamped to 3.0"); + assert!((clamped.c51_alpha_max - 0.9).abs() < 1e-6, "c51_alpha_max should be clamped to 0.9"); assert!((clamped.learning_intensity - 2.0).abs() < 1e-6, "intensity should be clamped to 2.0"); // Max intensity values let continuous_max = vec![ - 0.999, 1.0, 15.0, // breakouts at max + 0.999, 1.0, 15.0, 0.9, // breakouts at max 2.0, 2.0, 2.0, 2.0, 2.0, // core families at max 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, // gen families at max ]; @@ -3432,24 +3457,26 @@ mod tests { } #[test] - fn test_14d_defaults() { + fn test_15d_defaults() { let params = DQNParams::default(); assert!((params.gamma - 0.99).abs() < 1e-6, "Default gamma should be 0.99"); assert!((params.iqn_lambda - 0.25).abs() < 1e-6, "Default iqn_lambda should be 0.25"); assert!((params.c51_warmup_epochs - 5.0).abs() < 1e-6, "Default c51_warmup should be 5.0"); + assert!((params.c51_alpha_max - 0.5).abs() < 1e-6, "Default c51_alpha_max should be 0.5"); assert!((params.learning_intensity - 1.0).abs() < 1e-6, "Default learning_intensity should be 1.0"); assert!((params.adversarial_intensity - 1.0).abs() < 1e-6, "Default adversarial_intensity should be 1.0"); } #[test] - fn test_14d_default_roundtrip() { + fn test_15d_default_roundtrip() { let params = DQNParams::default(); let continuous = params.to_continuous(); - assert_eq!(continuous.len(), 14, "Should have 14 continuous dimensions"); + assert_eq!(continuous.len(), 15, "Should have 15 continuous dimensions"); let roundtrip = DQNParams::from_continuous(&continuous).unwrap(); assert!((roundtrip.gamma - params.gamma).abs() < 1e-6); assert!((roundtrip.iqn_lambda - params.iqn_lambda).abs() < 1e-6); assert!((roundtrip.c51_warmup_epochs - params.c51_warmup_epochs).abs() < 1e-6); + assert!((roundtrip.c51_alpha_max - params.c51_alpha_max).abs() < 1e-6); assert!((roundtrip.learning_intensity - params.learning_intensity).abs() < 1e-6); } @@ -3545,11 +3572,11 @@ mod tests { } #[test] - fn test_14d_midpoint_from_continuous() { + fn test_15d_midpoint_from_continuous() { // Verify from_continuous works with midpoint of all bounds let bounds = DQNParams::continuous_bounds(); let dim = bounds.len(); - assert_eq!(dim, 14, "Bounds should be 14D"); + assert_eq!(dim, 15, "Bounds should be 15D"); let mut vec = vec![0.0_f64; dim]; for i in 0..dim { vec[i] = (bounds[i].0 + bounds[i].1) / 2.0; diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index fbd9b2ff8..403e9026c 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1380,6 +1380,12 @@ pub struct DQNHyperparameters { /// nearly-identical distributions produces near-zero gradient for the value mean. /// Default: 5. Set to 0 to use C51 from the start. pub c51_warmup_epochs: usize, + /// Maximum C51 blend factor (0.0 = pure MSE, 1.0 = pure C51). + /// Caps the alpha ramp so MSE always contributes (1 - alpha_max) of the + /// primary gradient. Prevents C51 premature convergence from killing + /// the gradient signal on large batches. + /// Default: 0.5. Hyperopt range: [0.3, 0.9]. + pub c51_alpha_max: f32, // Decision Transformer pre-training /// Decision Transformer pre-training epochs (0 = disabled). @@ -1859,6 +1865,7 @@ impl DQNHyperparameters { // C51 warmup: use MSE loss for first 5 epochs (stronger gradient than C51 cross-entropy) c51_warmup_epochs: 5, + c51_alpha_max: 0.5, // MSE always contributes 50% — prevents C51 gradient starvation // Decision Transformer pre-training: disabled by default dt_pretrain_epochs: 0, // 0 = disabled diff --git a/crates/ml/src/trainers/dqn/trainer/tests.rs b/crates/ml/src/trainers/dqn/trainer/tests.rs index 388d8fc94..440bf7ba1 100644 --- a/crates/ml/src/trainers/dqn/trainer/tests.rs +++ b/crates/ml/src/trainers/dqn/trainer/tests.rs @@ -558,10 +558,10 @@ fn test_c2_five_actions_produce_distinct_exposures() { #[test] fn test_c3_search_space_is_14d() { let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 14, "Search space must be 14D (3 breakouts + 5 core families + 6 gen families)"); + assert_eq!(bounds.len(), 15, "Search space must be 15D (4 breakouts + 5 core families + 6 gen families)"); let names = crate::hyperopt::adapters::dqn::DQNParams::param_names(); - assert_eq!(names.len(), 14); + assert_eq!(names.len(), 15); assert!(names.contains(&"gamma"), "gamma must be in search space (breakout)"); assert!(names.contains(&"c51_warmup_epochs"), "c51_warmup_epochs must be in search space (breakout)"); assert!(names.contains(&"iqn_lambda"), "iqn_lambda must be in search space (breakout)"); diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index dff07e916..b86c5a8ad 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -230,12 +230,13 @@ impl DQNTrainer { self.reset_epoch_state(epoch); - // C51 warmup ramp + // C51 warmup ramp — capped at c51_alpha_max to keep MSE as gradient floor { + let alpha_max = self.hyperparams.c51_alpha_max; let alpha = if self.hyperparams.c51_warmup_epochs > 0 { - (epoch as f32 / self.hyperparams.c51_warmup_epochs as f32).min(1.0) + (epoch as f32 / self.hyperparams.c51_warmup_epochs as f32).min(alpha_max) } else { - 1.0 + alpha_max }; if let Some(ref mut fused) = self.fused_ctx { fused.set_c51_alpha(alpha); diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs index 458d8942c..c0a717795 100644 --- a/crates/ml/src/training_profile.rs +++ b/crates/ml/src/training_profile.rs @@ -157,6 +157,8 @@ pub struct AdvancedSection { pub tau: Option, /// C51 distributional loss warmup epochs (MSE before cross-entropy). pub c51_warmup_epochs: Option, + /// Maximum C51 alpha blend (caps warmup ramp so MSE always contributes). + pub c51_alpha_max: Option, /// Hindsight Experience Replay ratio (0.0-1.0, 0.0 = disabled). pub her_ratio: Option, /// IQN dual-head lambda weight relative to C51 loss. @@ -370,6 +372,8 @@ pub struct SearchSpaceSection { pub spectral_norm_sigma_max: Option<[f64; 2]>, /// C51 warmup epochs bounds. pub c51_warmup_epochs: Option<[f64; 2]>, + /// C51 alpha max bounds. + pub c51_alpha_max: Option<[f64; 2]>, /// HER ratio bounds. pub her_ratio: Option<[f64; 2]>, // Composite reward weights @@ -535,6 +539,7 @@ impl HyperoptProfile { "iqn_lambda" => ss.iqn_lambda, "spectral_norm_sigma_max" => ss.spectral_norm_sigma_max, "c51_warmup_epochs" => ss.c51_warmup_epochs, + "c51_alpha_max" => ss.c51_alpha_max, "her_ratio" => ss.her_ratio, "w_dsr" => ss.w_dsr, "w_pnl" => ss.w_pnl, @@ -844,6 +849,9 @@ impl DqnTrainingProfile { if let Some(v) = a.c51_warmup_epochs { hp.c51_warmup_epochs = v; } + if let Some(v) = a.c51_alpha_max { + hp.c51_alpha_max = v; + } if let Some(v) = a.her_ratio { hp.her_ratio = v; } @@ -1099,7 +1107,7 @@ mod tests { // [training] values must have been applied assert_eq!(hp.epochs, 200); assert_eq!(hp.batch_size, 16384); - assert!((hp.learning_rate - 0.0001).abs() < 1e-10); + assert!((hp.learning_rate - 1e-5).abs() < 1e-10); // OFI MBP-10 data dir should be set in production profile assert_eq!( @@ -1184,7 +1192,7 @@ mod tests { set_hyperopt_phase(HyperoptPhase::Fast); use crate::hyperopt::ParameterSpace; let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 14, "14D search space"); + assert_eq!(bounds.len(), 15, "15D search space"); // ALL dims must be searchable in Fast (single-pass) for (i, (lo, hi)) in bounds.iter().enumerate() { assert_ne!(lo, hi, "dim {} must be searchable in Fast phase", i); @@ -1367,7 +1375,7 @@ mod tests { // [training] section assert_eq!(hp.epochs, 200); assert_eq!(hp.batch_size, 16384); - assert!((hp.learning_rate - 0.0001).abs() < 1e-10); + assert!((hp.learning_rate - 1e-5).abs() < 1e-10); assert!((hp.gamma - 0.99).abs() < 0.001); assert!((hp.reward_scale - 1.0).abs() < 0.01);