gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add c51_alpha_max to cap MSE→C51 blend and prevent gradient starvation

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 15:42:41 +02:00
parent 30f46c04c1
commit a0aab3baa5
8 changed files with 153 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
config/training/dqn-hyperopt.toml
View File

@@ -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]

1
config/training/dqn-production.toml
View File

@@ -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

1
config/training/dqn-smoketest.toml
View File

@@ -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

225
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -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<Self, MLError> {
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;

7
crates/ml/src/trainers/dqn/config.rs
View File

@@ -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

4
crates/ml/src/trainers/dqn/trainer/tests.rs
View File

@@ -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)");

7
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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);

14
crates/ml/src/training_profile.rs
View File

@@ -157,6 +157,8 @@ pub struct AdvancedSection {
pub tau: Option<f64>,
/// C51 distributional loss warmup epochs (MSE before cross-entropy).
pub c51_warmup_epochs: Option<usize>,
/// Maximum C51 alpha blend (caps warmup ramp so MSE always contributes).
pub c51_alpha_max: Option<f32>,
/// Hindsight Experience Replay ratio (0.0-1.0, 0.0 = disabled).
pub her_ratio: Option<f64>,
/// 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);