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

cleanup(fflag): remove use_soft_updates — always EMA — [FFLAG-006]

Browse Source 
DQNConfig.use_soft_updates was true at all 4 construction sites; the
else branch (legacy hard update) was unreachable. Deleted field + 4
setters + collapsed update_target_networks() to the unconditional
cosine-annealed EMA path. Hard-copy mode removed per
feedback_no_feature_flags.md.
This commit is contained in:
jgrusewski
2026-04-20 22:23:22 +02:00
parent 9c9c7c6d8c
commit a05cae1e40
1 changed files with 34 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
crates/ml-dqn/src/dqn.rs
View File

@@ -65,8 +65,6 @@ pub struct DQNConfig {
pub tau_final: f64,
/// Total training steps for cosine schedule (T). 0 = use fixed tau (no annealing).
pub tau_anneal_steps: u64,
/// Use soft (EMA) or hard target updates
pub use_soft_updates: bool,
// Rainbow DQN warmup period
/// Number of steps to collect experiences with random exploration before training begins
@@ -253,7 +251,6 @@ impl Default for DQNConfig {
tau: 0.005,
tau_final: 0.0005, // 10x slower at convergence for offline RL stability
tau_anneal_steps: 100_000, // cosine anneal over 100K gradient steps
use_soft_updates: true,
warmup_steps: 1000,
n_steps: 1,
initial_capital: 100_000.0,
@@ -566,7 +563,6 @@ impl DQNConfig {
tau: 0.005,
tau_final: 0.0005,
tau_anneal_steps: 100_000,
use_soft_updates: true,
warmup_steps: 1000,
n_steps: 3,
initial_capital: 100000.0,
@@ -630,7 +626,6 @@ impl DQNConfig {
tau: 0.005,
tau_final: 0.0005,
tau_anneal_steps: 100_000,
use_soft_updates: true,
warmup_steps: 0,
n_steps: 1, // Default to single-step TD (most stable)
initial_capital: 100000.0,
@@ -700,7 +695,6 @@ impl DQNConfig {
tau: 0.005,
tau_final: 0.001, // Very conservative annealing for emergency mode
tau_anneal_steps: 50_000,
use_soft_updates: true,
// Rainbow DQN warmup period
warmup_steps: 0, // No warmup for emergency mode (safety first)
@@ -2055,63 +2049,45 @@ impl DQN {
Ok(())
}
/// Update target networks using cosine-annealed EMA (soft) or hard copy.
/// Update target networks using cosine-annealed EMA (soft updates).
///
/// This is the shared implementation used by both [`train_step`] and
/// Shared implementation used by both [`train_step`] and
/// [`apply_accumulated_gradients`] to keep target network update logic
/// in a single place.
/// in a single place. Hard-copy legacy mode was removed — all DQN
/// training uses cosine-annealed Polyak averaging.
fn update_target_networks(&mut self) -> Result<(), MLError> {
if self.config.use_soft_updates {
// Cosine-annealed EMA (BYOL/MoCo v3 style)
// 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 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
} else {
self.config.tau // Fixed tau when annealing disabled
};
// Cold-path fallback: copy_weights_from as hard-copy approximation.
// The real training path uses the fused CUDA EMA kernel on flat buffers.
// DistributionalDueling and IQN have no per-layer copy — fused trainer only.
if let (Some(ref branching_net), Some(ref mut branching_target)) =
(&self.branching_q_network, &mut self.branching_target_network)
{
branching_target.copy_weights_from(branching_net)
.map_err(|e| {
MLError::TrainingError(format!("Branching EMA update failed: {}", e))
})?;
}
// Log EMA update every 1000 steps
if self.training_steps % 1000 == 0 {
let half_life = convergence_half_life(current_tau);
debug!(
"EMA target update at step {} (τ={:.6}, half-life={:.0} steps)",
self.training_steps, current_tau, half_life
);
}
// Cosine-annealed EMA (BYOL/MoCo v3 style)
// 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 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
} else {
// Hard update: Full copy every N steps (legacy mode)
// DistributionalDueling and IQN have no per-layer copy — fused trainer only.
if self.training_steps % self.config.target_update_freq as u64 == 0 {
if let (Some(ref branching_net), Some(ref mut branching_target)) =
(&self.branching_q_network, &mut self.branching_target_network)
{
branching_target.copy_weights_from(branching_net).map_err(|e| {
MLError::TrainingError(format!("Branching hard update failed: {}", e))
})?;
}
self.config.tau // Fixed tau when annealing disabled
};
debug!(
"Hard target update at step {} (every {} steps)",
self.training_steps, self.config.target_update_freq
);
}
// Cold-path fallback: copy_weights_from as hard-copy approximation.
// The real training path uses the fused CUDA EMA kernel on flat buffers.
// DistributionalDueling and IQN have no per-layer copy — fused trainer only.
if let (Some(ref branching_net), Some(ref mut branching_target)) =
(&self.branching_q_network, &mut self.branching_target_network)
{
branching_target.copy_weights_from(branching_net)
.map_err(|e| {
MLError::TrainingError(format!("Branching EMA update failed: {}", e))
})?;
}
// Log EMA update every 1000 steps
if self.training_steps % 1000 == 0 {
let half_life = convergence_half_life(current_tau);
debug!(
"EMA target update at step {} (τ={:.6}, half-life={:.0} steps)",
self.training_steps, current_tau, half_life
);
}
Ok(())