From cf6e36c8a1730a2d7f914f1e038d3124c8f05946 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 12 Mar 2026 13:21:28 +0100
Subject: [PATCH] perf(dqn): restore best-epoch checkpoint before walk-forward
 evaluation

Walk-forward evaluation was using the final epoch model, which may have
overfit. Now loads the best per-epoch Sharpe checkpoint (saved during
training) back into the agent before running the walk-forward backtest.

Trial 0 showed Sharpe +3.30 at epoch 2 vs +1.84 at epoch 8 (final).
The walk-forward should evaluate the peak model, not the final one.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/src/hyperopt/adapters/dqn.rs | 28 ++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs
index 0266df1b8..9f204f354 100644
--- a/crates/ml/src/hyperopt/adapters/dqn.rs
+++ b/crates/ml/src/hyperopt/adapters/dqn.rs
@@ -3005,6 +3005,34 @@ impl HyperparameterOptimizable for DQNTrainer {
         tracing::debug!("Reached backtest decision point");
         tracing::debug!("enable_backtest = {}", self.enable_backtest);
 
+        // Restore best checkpoint before walk-forward evaluation.
+        // Training may overfit in later epochs; the best per-epoch Sharpe checkpoint
+        // (saved during training) gives a more accurate walk-forward evaluation.
+        let best_ckpt_path = self.training_paths.checkpoints_dir()
+            .join(format!("trial_{}_best.safetensors", current_trial));
+        if best_ckpt_path.exists() {
+            let best_epoch = internal_trainer.get_best_epoch();
+            let agent_arc = internal_trainer.get_agent().clone();
+            let bt_handle = self.runtime_handle.as_ref().ok_or_else(|| {
+                MLError::ConfigError("runtime_handle is None".to_owned())
+            })?;
+            let mut agent_guard = bt_handle.block_on(agent_arc.write());
+            let ckpt_str = best_ckpt_path.to_string_lossy().to_string();
+            match &mut *agent_guard {
+                crate::trainers::dqn::DQNAgentType::Standard(ref mut dqn) => {
+                    if let Err(e) = dqn.load_from_safetensors(&ckpt_str) {
+                        tracing::warn!("Failed to restore best checkpoint for walk-forward: {}", e);
+                    } else {
+                        info!("Restored best checkpoint (epoch {}) for walk-forward evaluation", best_epoch);
+                    }
+                }
+                crate::trainers::dqn::DQNAgentType::RegimeConditional(_) => {
+                    tracing::warn!("Best-checkpoint restore not yet implemented for RegimeConditional");
+                }
+            }
+            drop(agent_guard);
+        }
+
         // Run backtest if enabled
         let backtest_metrics = if self.enable_backtest {
             tracing::debug!("Backtest is enabled, starting backtest...");