diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 4841ae4ea..476587a8a 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -22,6 +22,7 @@ use candle_nn::ops::leaky_relu; use candle_nn::Module; use candle_nn::{Linear, VarBuilder, VarMap}; use candle_optimisers::adam::ParamsAdam; +use candle_optimisers::Decay; use rand::{thread_rng, Rng}; use serde::{Deserialize, Serialize}; use tracing::debug; @@ -1497,7 +1498,8 @@ impl DQN { beta_1: 0.9, beta_2: 0.999, eps: 1.5e-4, // Rainbow DQN standard (was 1e-8) - weight_decay: None, + weight_decay: (self.config.weight_decay > 0.0) + .then_some(Decay::DecoupledWeightDecay(self.config.weight_decay)), amsgrad: false, }; @@ -3249,4 +3251,73 @@ mod tests { dqn.err() ); } + + /// Verify that DQNConfig.weight_decay is wired through to the Adam optimizer. + /// With weight_decay > 0, training a few steps should produce different final + /// weights than weight_decay == 0, proving L2 regularization is active. + #[test] + fn test_weight_decay_wired_to_optimizer() -> anyhow::Result<()> { + // Helper: run a few training steps and return the first Q-network parameter tensor + fn train_and_get_params(weight_decay: f64) -> anyhow::Result> { + let mut config = DQNConfig::emergency_safe_defaults(); + config.state_dim = 8; + config.num_actions = 3; + config.hidden_dims = vec![16, 16]; + config.use_iqn = false; + config.use_distributional = false; + config.use_dueling = false; + config.use_cql = false; + config.batch_size = 4; + config.min_replay_size = 4; + config.learning_rate = 1e-3; // Higher LR to make weight_decay effect visible + config.weight_decay = weight_decay; + + let mut dqn = DQN::new(config)?; + + // Fill replay buffer + for i in 0..10 { + let exp = Experience::new( + vec![0.1 * i as f32; 8], + (i % 3) as u8, + 0.5, + vec![0.2 * i as f32; 8], + false, + ); + dqn.store_experience(exp)?; + } + + // Run several training steps so weight_decay has a measurable effect + for _ in 0..20 { + let _ = dqn.train_step(None)?; + } + + // Extract first parameter tensor values + let vars = dqn.q_network.vars().all_vars(); + let first_var = vars.first().ok_or_else(|| { + anyhow::anyhow!("Q-network has no parameters") + })?; + // Cast to F32 in case mixed-precision uses BF16 weights on GPU + let values = first_var + .flatten_all()? + .to_dtype(candle_core::DType::F32)? + .to_vec1::()?; + Ok(values) + } + + let params_no_decay = train_and_get_params(0.0)?; + let params_with_decay = train_and_get_params(1e-2)?; // Strong decay to make difference clear + + // The parameters MUST differ — proving weight_decay is actually applied + let diff: f32 = params_no_decay + .iter() + .zip(params_with_decay.iter()) + .map(|(a, b)| (a - b).abs()) + .sum(); + assert!( + diff > 1e-6, + "Weight decay should produce different parameters (total abs diff = {diff})" + ); + + Ok(()) + } }