diff --git a/crates/ml/src/trainers/dqn/early_stopping.rs b/crates/ml/src/trainers/dqn/early_stopping.rs index a0e7bc6fd..52dd94ef6 100644 --- a/crates/ml/src/trainers/dqn/early_stopping.rs +++ b/crates/ml/src/trainers/dqn/early_stopping.rs @@ -21,7 +21,7 @@ use tracing::{debug, info}; /// for epoch in 0..100 { /// let val_loss = train_epoch(); /// -/// if early_stop.should_stop(val_loss) { +/// if early_stop.should_stop(val_loss, epoch) { /// println!("Early stopping at epoch {}", epoch); /// break; /// } @@ -37,10 +37,12 @@ pub struct EarlyStopping { best_val_loss: f64, /// Counter for epochs without improvement counter: usize, - /// Epoch when best validation loss was achieved + /// Training epoch (0-indexed) when best validation loss was achieved. + /// This is the epoch number passed by the caller, NOT an internal call + /// counter — outer callers may gate `should_stop` invocations (e.g. + /// `min_epochs_before_stopping`), so an internal counter would drift + /// from real training epochs and produce misleading log messages. best_epoch: usize, - /// Current epoch number - current_epoch: usize, } impl EarlyStopping { @@ -61,35 +63,37 @@ impl EarlyStopping { best_val_loss: f64::INFINITY, counter: 0, best_epoch: 0, - current_epoch: 0, } } /// Check if training should stop /// - /// Returns `true` if no improvement for `patience` consecutive epochs. - /// Automatically tracks epochs and best loss internally. + /// Returns `true` if no improvement for `patience` consecutive calls. /// /// # Arguments /// * `val_loss` - Current epoch validation loss + /// * `epoch` - Current training epoch (0-indexed). Used for accurate + /// log messages and for tracking which epoch achieved the best loss. + /// Pass the epoch number from the outer training loop — this struct + /// deliberately does NOT track its own epoch counter because callers + /// may gate invocations (e.g. skip until `min_epochs_before_stopping`), + /// and an internal counter would drift from real epochs. /// /// # Returns /// `true` if early stopping triggered, `false` to continue training - pub fn should_stop(&mut self, val_loss: f64) -> bool { - self.current_epoch += 1; - + pub fn should_stop(&mut self, val_loss: f64, epoch: usize) -> bool { // Check if validation loss improved by more than min_delta let improvement = self.best_val_loss - val_loss; if improvement > self.min_delta { // Significant improvement detected debug!( - "Early stopping: Improvement detected: {:.6} (val_loss: {:.6} -> {:.6})", - improvement, self.best_val_loss, val_loss + "Early stopping: Improvement detected: {:.6} (val_loss: {:.6} -> {:.6}) at epoch {}", + improvement, self.best_val_loss, val_loss, epoch ); self.best_val_loss = val_loss; - self.best_epoch = self.current_epoch; + self.best_epoch = epoch; self.counter = 0; false } else { @@ -97,14 +101,14 @@ impl EarlyStopping { self.counter += 1; debug!( - "Early stopping: No improvement for {} epoch(s) (improvement: {:.6} < min_delta: {:.6})", - self.counter, improvement, self.min_delta + "Early stopping: No improvement for {} call(s) (improvement: {:.6} < min_delta: {:.6}) at epoch {}", + self.counter, improvement, self.min_delta, epoch ); if self.counter >= self.patience { info!( - "🛑 Early stopping triggered at epoch {}! No improvement for {} epochs (best: {:.6} at epoch {})", - self.current_epoch, self.patience, self.best_val_loss, self.best_epoch + "🛑 Early stopping triggered at epoch {}! No improvement for {} consecutive evaluations (best val_loss: {:.6} at epoch {})", + epoch, self.patience, self.best_val_loss, self.best_epoch ); true } else { @@ -133,16 +137,14 @@ impl EarlyStopping { self.best_val_loss = f64::INFINITY; self.counter = 0; self.best_epoch = 0; - self.current_epoch = 0; } /// Restore early stopping state from checkpoint /// /// Useful when resuming training from a saved checkpoint. - pub fn restore(&mut self, best_val_loss: f64, best_epoch: usize, current_epoch: usize) { + pub fn restore(&mut self, best_val_loss: f64, best_epoch: usize) { self.best_val_loss = best_val_loss; self.best_epoch = best_epoch; - self.current_epoch = current_epoch; self.counter = 0; // Reset counter on restore } } @@ -155,24 +157,24 @@ mod tests { fn test_early_stopping_triggers_after_patience() { let mut early_stop = EarlyStopping::new(3, 0.001); - // Epoch 1: Initial loss (improves from infinity) - assert!(!early_stop.should_stop(1.0)); + // Epoch 0: Initial loss (improves from infinity) + assert!(!early_stop.should_stop(1.0, 0)); assert_eq!(early_stop.counter, 0); - // Epoch 2: Improves significantly - assert!(!early_stop.should_stop(0.5)); + // Epoch 1: Improves significantly + assert!(!early_stop.should_stop(0.5, 1)); assert_eq!(early_stop.counter, 0); - // Epoch 3: No improvement - assert!(!early_stop.should_stop(0.51)); + // Epoch 2: No improvement + assert!(!early_stop.should_stop(0.51, 2)); assert_eq!(early_stop.counter, 1); - // Epoch 4: No improvement - assert!(!early_stop.should_stop(0.52)); + // Epoch 3: No improvement + assert!(!early_stop.should_stop(0.52, 3)); assert_eq!(early_stop.counter, 2); - // Epoch 5: No improvement - triggers early stop - assert!(early_stop.should_stop(0.53)); + // Epoch 4: No improvement - triggers early stop + assert!(early_stop.should_stop(0.53, 4)); assert_eq!(early_stop.counter, 3); } @@ -181,15 +183,15 @@ mod tests { let mut early_stop = EarlyStopping::new(3, 0.001); // Initial loss - assert!(!early_stop.should_stop(1.0)); + assert!(!early_stop.should_stop(1.0, 0)); // No improvement for 2 epochs - assert!(!early_stop.should_stop(1.01)); - assert!(!early_stop.should_stop(1.02)); + assert!(!early_stop.should_stop(1.01, 1)); + assert!(!early_stop.should_stop(1.02, 2)); assert_eq!(early_stop.counter, 2); // Significant improvement - resets counter - assert!(!early_stop.should_stop(0.5)); + assert!(!early_stop.should_stop(0.5, 3)); assert_eq!(early_stop.counter, 0); } @@ -198,14 +200,14 @@ mod tests { let mut early_stop = EarlyStopping::new(3, 0.01); // 1% min improvement // Initial loss - assert!(!early_stop.should_stop(1.0)); + assert!(!early_stop.should_stop(1.0, 0)); // Small improvement (0.005 < 0.01) - doesn't reset counter - assert!(!early_stop.should_stop(0.995)); + assert!(!early_stop.should_stop(0.995, 1)); assert_eq!(early_stop.counter, 1); // Large improvement (0.02 > 0.01) - resets counter - assert!(!early_stop.should_stop(0.975)); + assert!(!early_stop.should_stop(0.975, 2)); assert_eq!(early_stop.counter, 0); } @@ -213,44 +215,45 @@ mod tests { fn test_best_loss_tracking() { let mut early_stop = EarlyStopping::new(5, 0.001); - early_stop.should_stop(1.0); + // Caller passes real training epoch — best_epoch tracks those, + // NOT an internal call counter that would drift if the outer + // loop gates calls (e.g. min_epochs_before_stopping). + early_stop.should_stop(1.0, 10); assert_eq!(early_stop.get_best_val_loss(), 1.0); - assert_eq!(early_stop.get_best_epoch(), 1); + assert_eq!(early_stop.get_best_epoch(), 10); - early_stop.should_stop(0.5); + early_stop.should_stop(0.5, 11); assert_eq!(early_stop.get_best_val_loss(), 0.5); - assert_eq!(early_stop.get_best_epoch(), 2); + assert_eq!(early_stop.get_best_epoch(), 11); - early_stop.should_stop(0.6); // Worse - doesn't update + early_stop.should_stop(0.6, 12); // Worse - doesn't update assert_eq!(early_stop.get_best_val_loss(), 0.5); - assert_eq!(early_stop.get_best_epoch(), 2); + assert_eq!(early_stop.get_best_epoch(), 11); } #[test] fn test_reset() { let mut early_stop = EarlyStopping::new(3, 0.001); - early_stop.should_stop(1.0); - early_stop.should_stop(1.01); + early_stop.should_stop(1.0, 0); + early_stop.should_stop(1.01, 1); early_stop.reset(); assert_eq!(early_stop.best_val_loss, f64::INFINITY); assert_eq!(early_stop.counter, 0); assert_eq!(early_stop.best_epoch, 0); - assert_eq!(early_stop.current_epoch, 0); } #[test] fn test_restore() { let mut early_stop = EarlyStopping::new(3, 0.001); - // Simulate resuming from checkpoint at epoch 50 with best loss 0.5 at epoch 45 - early_stop.restore(0.5, 45, 50); + // Simulate resuming from checkpoint — best loss 0.5 at epoch 45 + early_stop.restore(0.5, 45); assert_eq!(early_stop.best_val_loss, 0.5); assert_eq!(early_stop.best_epoch, 45); - assert_eq!(early_stop.current_epoch, 50); assert_eq!(early_stop.counter, 0); // Reset on restore } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 6ee9cfbc5..ec6c76944 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2955,7 +2955,7 @@ impl DQNTrainer { let old_should_stop = self.check_early_stopping(log_output.avg_q_value, epoch); let patience_should_stop = if self.hyperparams.early_stopping_enabled && epoch + 1 >= self.hyperparams.min_epochs_before_stopping { - self.early_stopping.should_stop(-log_output.epoch_sharpe) + self.early_stopping.should_stop(-log_output.epoch_sharpe, epoch) } else { false };