diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 22ad1a259..d88adcae9 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1477,6 +1477,14 @@ impl FusedTrainingCtx { /// The training guard reads this pointer for NaN detection, so we must return /// `mse_loss_buf` during warmup to avoid false NaN halts. /// Once C51 is active (alpha > 0), returns `total_loss_buf` (C51 accumulator). + pub(crate) fn online_dueling_ref(&self) -> &DuelingWeightSet { + &self.online_dueling + } + + pub(crate) fn online_branching_ref(&self) -> &BranchingWeightSet { + &self.online_branching + } + pub(crate) fn loss_gpu_buf(&self) -> &cudarc::driver::CudaSlice { if self.trainer.c51_alpha() < 1e-6 { &self.trainer.mse_loss_buf diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs index a12091d6c..9225fe323 100644 --- a/crates/ml/src/trainers/dqn/trainer/metrics.rs +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -472,36 +472,25 @@ impl DQNTrainer { self.gpu_evaluator = Some(evaluator); } - // ── Extract current weights from agent ── + // ── Extract current weights from fused CUDA trainer (GPU-resident) ── + // The fused_ctx owns the up-to-date weights (updated by Adam each step). + // Do NOT use agent.get_q_network_vars() — that reads the Candle VarMap + // which is NOT synced during fused CUDA training (stale epoch-0 weights). + let fused = self.fused_ctx.as_ref() + .ok_or_else(|| anyhow::anyhow!("fused_ctx required for GPU validation"))?; + let online_weights = fused.online_dueling_ref(); + let branching_weights = fused.online_branching_ref(); + let agent = self.agent.read().await; - let is_branching = agent.is_using_branching(); let network_dims = agent.network_dims(); - let vars = agent.get_q_network_vars(); + let (b0, b1, b2) = agent.branch_sizes(); + let is_branching = agent.is_using_branching(); + drop(agent); let evaluator = self.gpu_evaluator.as_mut() .ok_or_else(|| anyhow::anyhow!("gpu_evaluator must be initialized"))?; - let eval_stream = evaluator.stream(); - - let online_weights = if is_branching { - crate::cuda_pipeline::gpu_weights::extract_dueling_weights_branching( - vars, eval_stream, - ).map_err(|e| anyhow::anyhow!("extract dueling weights (branching): {e}"))? - } else { - crate::cuda_pipeline::gpu_weights::extract_dueling_weights( - vars, eval_stream, - ).map_err(|e| anyhow::anyhow!("extract dueling weights: {e}"))? - }; - - let branching_weights = if is_branching { - Some(crate::cuda_pipeline::gpu_weights::extract_branching_weights( - vars, eval_stream, - ).map_err(|e| anyhow::anyhow!("extract branching weights: {e}"))?) - } else { - None - }; let hp = &self.hyperparams; - let (b0, b1, b2) = agent.branch_sizes(); #[allow(clippy::cast_possible_truncation)] let dqn_cfg = DqnBacktestConfig { shared_h1: network_dims.0, @@ -516,16 +505,13 @@ impl DQNTrainer { v_max: hp.v_max as f32, }; - // Drop agent read guard before mutable evaluator call - drop(agent); - // Weights changed since last epoch — invalidate the cached CUDA graph // so evaluate_dqn_graphed re-captures with current weight pointers. evaluator.invalidate_dqn_graph(); let metrics = evaluator.evaluate_dqn_graphed( - &online_weights, - branching_weights.as_ref(), + online_weights, + Some(branching_weights), &dqn_cfg, ).map_err(|e| anyhow::anyhow!("GPU backtest evaluate_dqn_graphed: {e}"))?;