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

feat(ml): add compute_gradients and apply_accumulated_gradients to DQN

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-20 22:16:09 +01:00
parent 4f9092c0ac
commit b4bee985b1
1 changed files with 140 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

161
ml/src/dqn/dqn.rs
View File

@@ -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<Vec<Experience>>,
) -> Result<GradientResult, MLError> {
// 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)