From b4bee985b17d73b1bceaa6ab27cacdb919249ad5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Feb 2026 22:16:09 +0100 Subject: [PATCH] feat(ml): add compute_gradients and apply_accumulated_gradients to DQN Extract target network update logic into private update_target_networks() helper to avoid duplication between train_step() and the new gradient accumulation methods. Add three new public methods: - compute_gradients(): forward+backward without optimizer step - apply_accumulated_gradients(): apply pre-computed grads + target update - optimizer_vars(): expose optimizer variables for accumulation utilities Co-Authored-By: Claude Opus 4.6 --- ml/src/dqn/dqn.rs | 161 ++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 140 insertions(+), 21 deletions(-) diff --git a/ml/src/dqn/dqn.rs b/ml/src/dqn/dqn.rs index a1540f082..b8c7e1ca2 100644 --- a/ml/src/dqn/dqn.rs +++ b/ml/src/dqn/dqn.rs @@ -1956,28 +1956,139 @@ impl DQN { } // Update target network with cosine-annealed EMA or hard updates + self.update_target_networks()?; + + Ok((result.loss_value, grad_norm)) + } + + /// Compute clipped gradients for a batch WITHOUT applying an optimizer step. + /// + /// This is used in gradient accumulation workflows where multiple micro-batches + /// are processed before a single optimizer update via [`apply_accumulated_gradients`]. + /// + /// # Errors + /// + /// Returns an error if called during warmup, if the optimizer is not initialised, + /// or if the forward/backward pass fails. + pub fn compute_gradients( + &mut self, + batch: Option>, + ) -> Result { + // Skip during warmup + if self.total_steps < self.config.warmup_steps as u64 { + return Err(MLError::TrainingError( + "Cannot compute gradients during warmup period".to_string(), + )); + } + + let result = self.compute_loss_internal(batch)?; + + // Backward + clip WITHOUT optimizer step + let (grads, grad_norm) = if let Some(ref optimizer) = self.optimizer { + optimizer.backward_and_clip(&result.loss_tensor, self.gradient_clip_norm)? + } else { + return Err(MLError::TrainingError( + "Optimizer not initialized".to_string(), + )); + }; + + Ok(GradientResult { + loss: result.loss_value, + grad_norm: grad_norm as f32, + grads, + td_errors: result.td_errors, + indices: result.indices, + }) + } + + /// Apply pre-computed (accumulated) gradients and perform a single optimizer step. + /// + /// After calling this, `training_steps` is incremented and target networks are + /// updated according to the configured schedule (soft EMA or hard copy). + /// + /// # Errors + /// + /// Returns an error if the optimizer is not initialised or if the target network + /// update fails. + pub fn apply_accumulated_gradients( + &mut self, + grads: &GradStore, + ) -> Result<(), MLError> { + // Apply accumulated gradients via a single optimizer step + if let Some(ref mut optimizer) = self.optimizer { + optimizer.apply_grads(grads)?; + } else { + return Err(MLError::TrainingError( + "Optimizer not initialized".to_string(), + )); + } + + // Increment training steps (once per effective batch, not per mini-batch) + self.training_steps += 1; + + // Target network update (same schedule as train_step) + self.update_target_networks()?; + + Ok(()) + } + + /// Get a reference to the optimizer's tracked variables. + /// + /// This is needed by the gradient accumulation utilities to iterate over + /// variable tensors when accumulating or scaling gradient stores. + /// + /// # Errors + /// + /// Returns an error if the optimizer has not been initialised yet. + pub fn optimizer_vars(&self) -> Result<&[Var], MLError> { + if let Some(ref optimizer) = self.optimizer { + Ok(optimizer.vars()) + } else { + Err(MLError::TrainingError( + "Optimizer not initialized — call train_step or compute_gradients first" + .to_string(), + )) + } + } + + /// Update target networks using cosine-annealed EMA (soft) or hard copy. + /// + /// This is the shared implementation used by both [`train_step`] and + /// [`apply_accumulated_gradients`] to keep target network update logic + /// in a single place. + fn update_target_networks(&mut self) -> Result<(), MLError> { if self.config.use_soft_updates { // Cosine-annealed EMA (BYOL/MoCo v3 style) - // τ(t) = τ_final - (τ_final - τ_base) * (cos(π·t/T) + 1) / 2 - // Early: τ ≈ τ_base (fast adaptation) → Late: τ ≈ τ_final (stability) + // tau(t) = tau_final - (tau_final - tau_base) * (cos(pi*t/T) + 1) / 2 + // Early: tau ~ tau_base (fast adaptation) -> Late: tau ~ tau_final (stability) let current_tau = if self.config.tau_anneal_steps > 0 { - let progress = (self.training_steps as f64 / self.config.tau_anneal_steps as f64).min(1.0); + let progress = + (self.training_steps as f64 / self.config.tau_anneal_steps as f64).min(1.0); let cosine_factor = (std::f64::consts::PI * progress).cos(); - self.config.tau_final - (self.config.tau_final - self.config.tau) * (cosine_factor + 1.0) / 2.0 + self.config.tau_final + - (self.config.tau_final - self.config.tau) * (cosine_factor + 1.0) / 2.0 } else { - self.config.tau // Fixed τ when annealing disabled + self.config.tau // Fixed tau when annealing disabled }; if let (Some(ref dist_dueling_net), Some(ref mut dist_dueling_target)) = (&self.dist_dueling_q_network, &mut self.dist_dueling_target_network) { - polyak_update(dist_dueling_net.vars(), dist_dueling_target.vars(), current_tau) - .map_err(|e| MLError::TrainingError(format!("Hybrid EMA update failed: {}", e)))?; + polyak_update( + dist_dueling_net.vars(), + dist_dueling_target.vars(), + current_tau, + ) + .map_err(|e| { + MLError::TrainingError(format!("Hybrid EMA update failed: {}", e)) + })?; } else if let (Some(ref dueling_net), Some(ref mut dueling_target)) = (&self.dueling_q_network, &mut self.dueling_target_network) { polyak_update(dueling_net.vars(), dueling_target.vars(), current_tau) - .map_err(|e| MLError::TrainingError(format!("Dueling EMA update failed: {}", e)))?; + .map_err(|e| { + MLError::TrainingError(format!("Dueling EMA update failed: {}", e)) + })?; } else { polyak_update( self.q_network.vars(), @@ -1992,7 +2103,9 @@ impl DQN { (&self.iqn_network, &self.iqn_target_network) { polyak_update(iqn_net.vars(), iqn_target.vars(), current_tau) - .map_err(|e| MLError::TrainingError(format!("IQN EMA update failed: {}", e)))?; + .map_err(|e| { + MLError::TrainingError(format!("IQN EMA update failed: {}", e)) + })?; } // Log EMA update every 1000 steps @@ -2009,27 +2122,33 @@ impl DQN { if let (Some(ref dist_dueling_net), Some(ref mut dist_dueling_target)) = (&self.dist_dueling_q_network, &mut self.dist_dueling_target_network) { - // Update hybrid distributional dueling target network - hard_update(dist_dueling_net.vars(), dist_dueling_target.vars()) - .map_err(|e| MLError::TrainingError(format!("Hybrid hard update failed: {}", e)))?; + hard_update(dist_dueling_net.vars(), dist_dueling_target.vars()).map_err( + |e| { + MLError::TrainingError(format!( + "Hybrid hard update failed: {}", + e + )) + }, + )?; } else if let (Some(ref dueling_net), Some(ref mut dueling_target)) = (&self.dueling_q_network, &mut self.dueling_target_network) { - // Update dueling target network - hard_update(dueling_net.vars(), dueling_target.vars()) - .map_err(|e| MLError::TrainingError(format!("Dueling hard update failed: {}", e)))?; + hard_update(dueling_net.vars(), dueling_target.vars()).map_err(|e| { + MLError::TrainingError(format!("Dueling hard update failed: {}", e)) + })?; } else { - // Update standard target network - hard_update(self.q_network.vars(), self.target_network.vars()) - .map_err(|e| MLError::TrainingError(format!("Hard update failed: {}", e)))?; + hard_update(self.q_network.vars(), self.target_network.vars()).map_err( + |e| MLError::TrainingError(format!("Hard update failed: {}", e)), + )?; } // Also hard update IQN target network if present if let (Some(ref iqn_net), Some(ref mut iqn_target)) = (&self.iqn_network, &mut self.iqn_target_network) { - iqn_target.copy_weights_from(iqn_net) - .map_err(|e| MLError::TrainingError(format!("IQN hard update failed: {}", e)))?; + iqn_target.copy_weights_from(iqn_net).map_err(|e| { + MLError::TrainingError(format!("IQN hard update failed: {}", e)) + })?; } debug!( @@ -2039,7 +2158,7 @@ impl DQN { } } - Ok((result.loss_value, grad_norm)) + Ok(()) } /// Log Q-values for the first state in batch (Wave 10-A4 diagnostic monitoring)