From 6c59e16cd65a8a957adf24d27b705afb8913bc2c Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Wed, 25 Feb 2026 22:19:22 +0100
Subject: [PATCH] fix(ml): fix train_mamba2_dbn and train_dqn_es_fut compile
 errors

- Add missing Mamba2Config fields (early_stopping_enabled, patience, min_delta, min_epochs)
- Replace unwrap_or on f64 (not Option) with direct field access
- Replace .unwrap() on path.to_str() with safe .ok_or_else()
- Fix DQN closure signature: 3 args (epoch, data, is_final), not 2

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/examples/train_dqn_es_fut.rs |  4 +--
 crates/ml/examples/train_mamba2_dbn.rs | 37 +++++++++++++++++---------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/crates/ml/examples/train_dqn_es_fut.rs b/crates/ml/examples/train_dqn_es_fut.rs
index ac4152af3..fc50b6aaa 100644
--- a/crates/ml/examples/train_dqn_es_fut.rs
+++ b/crates/ml/examples/train_dqn_es_fut.rs
@@ -202,10 +202,10 @@ async fn main() -> Result<()> {
     let mut checkpoint_count = 0;
 
     let metrics = trainer
-        .train(&args.data_dir, |epoch, checkpoint_data| {
+        .train(&args.data_dir, |epoch, checkpoint_data, is_final| {
             checkpoint_count += 1;
 
-            let checkpoint_path = if epoch == args.epochs {
+            let checkpoint_path = if is_final {
                 // Final checkpoint
                 output_path.clone()
             } else {
diff --git a/crates/ml/examples/train_mamba2_dbn.rs b/crates/ml/examples/train_mamba2_dbn.rs
index 865ad91e7..78115f612 100644
--- a/crates/ml/examples/train_mamba2_dbn.rs
+++ b/crates/ml/examples/train_mamba2_dbn.rs
@@ -506,7 +506,11 @@ async fn main() -> Result<()> {
         seq_len: config.seq_len,
         shuffle_batches: false,              // Reproducibility
         sequence_stride: config.seq_len / 2, // P2: Overlapping windows
-        norm_eps: 1e-5,                      // P2: Layer norm epsilon
+        norm_eps: 1e-5,                       // P2: Layer norm epsilon
+        early_stopping_enabled: true,
+        early_stopping_patience: config.early_stopping_patience,
+        early_stopping_min_delta: 1e-6,
+        early_stopping_min_epochs: 10,
     };
 
     let mut model =
@@ -561,7 +565,7 @@ async fn main() -> Result<()> {
 
     // Process training history with early stopping
     for (epoch_idx, epoch) in training_history.iter().enumerate() {
-        let train_loss = epoch.loss.unwrap_or(f64::NAN);
+        let train_loss = epoch.loss;
         let should_save = monitor.update(
             epoch_idx,
             train_loss,
@@ -576,18 +580,17 @@ async fn main() -> Result<()> {
                 .checkpoint_dir
                 .join(format!("best_model_epoch_{}.ckpt", epoch_idx));
 
+            let path_str = checkpoint_path
+                .to_str()
+                .ok_or_else(|| anyhow::anyhow!("Invalid checkpoint path"))?;
             model
-                .save_checkpoint(checkpoint_path.to_str().unwrap())
+                .save_checkpoint(path_str)
                 .await
                 .context("Failed to save checkpoint")?;
 
             info!(
                 "✓ Saved best model at epoch {} (loss: {:.6})",
-                epoch_idx,
-                epoch
-                    .loss
-                    .map(|l| format!("{:.6}", l))
-                    .unwrap_or_else(|| "N/A".to_string())
+                epoch_idx, epoch.loss
             );
         }
 
@@ -597,8 +600,11 @@ async fn main() -> Result<()> {
                 .checkpoint_dir
                 .join(format!("checkpoint_epoch_{}.ckpt", epoch_idx));
 
+            let path_str = checkpoint_path
+                .to_str()
+                .ok_or_else(|| anyhow::anyhow!("Invalid checkpoint path"))?;
             model
-                .save_checkpoint(checkpoint_path.to_str().unwrap())
+                .save_checkpoint(path_str)
                 .await
                 .context("Failed to save checkpoint")?;
 
@@ -607,7 +613,7 @@ async fn main() -> Result<()> {
 
         // Log progress every 5 epochs
         if epoch_idx % 5 == 0 {
-            let perplexity = epoch.loss.unwrap_or(f64::NAN).exp();
+            let perplexity = epoch.loss.exp();
             let elapsed = monitor.start_time.elapsed();
             let epochs_per_min = (epoch_idx + 1) as f64 / elapsed.as_secs_f64() * 60.0;
 
@@ -638,8 +644,11 @@ async fn main() -> Result<()> {
 
     // Save final model
     let final_model_path = config.checkpoint_dir.join("final_model.ckpt");
+    let final_path_str = final_model_path
+        .to_str()
+        .ok_or_else(|| anyhow::anyhow!("Invalid final model path"))?;
     model
-        .save_checkpoint(final_model_path.to_str().unwrap())
+        .save_checkpoint(final_path_str)
         .await
         .context("Failed to save final model")?;
     info!("✓ Final model saved: {:?}", final_model_path);
@@ -706,7 +715,11 @@ async fn main() -> Result<()> {
     // Loss reduction
     if !monitor.epoch_losses.is_empty() {
         let initial_loss = monitor.epoch_losses[0];
-        let final_loss = *monitor.epoch_losses.last().unwrap();
+        let final_loss = monitor
+            .epoch_losses
+            .last()
+            .copied()
+            .unwrap_or(0.0);
         let reduction = ((initial_loss - final_loss) / initial_loss) * 100.0;
 
         info!("Loss Reduction:");