From 502fc469fe80494d4ef4e04626e14613d1f5ce8e Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sun, 22 Mar 2026 14:48:40 +0100
Subject: [PATCH] fix: TPE evaluate_point catches training errors as penalty,
 not abort
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

NaN/divergence during training is a valid PSO/TPE outcome — it means
the sampled hyperparameters are unstable. Score failed trials with 1e6
penalty so the optimizer learns to avoid that region, instead of
aborting all remaining trials.

The PSO/argmin paths already handled this (lines 1118, 1261). Only
the shared evaluate_point (TPE path) propagated errors as fatal.

Tested: 6 trials × 20 epochs on RTX 3050 — 2 NaN trials scored as
penalty, optimizer completed all 6 trials in 597s.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml-hyperopt/src/optimizer.rs | 31 +++++++++++++++++------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/crates/ml-hyperopt/src/optimizer.rs b/crates/ml-hyperopt/src/optimizer.rs
index 840dd7ca5..72e0dd1fc 100644
--- a/crates/ml-hyperopt/src/optimizer.rs
+++ b/crates/ml-hyperopt/src/optimizer.rs
@@ -480,20 +480,25 @@ impl ArgminOptimizer {
         // Log CONVERTED parameters (shows actual values: learning_rate ~1e-4, not -11)
         info!("  Parameters (converted): {:?}", params);
 
-        // Train model with parameters
-        let metrics = model
-            .train_with_params(params.clone())
-            .context(format!("Training failed for trial {}", trial_num))?;
-
-        // Extract objective (replace NaN/Inf with large penalty)
-        let raw_objective = M::extract_objective(&metrics);
-        let objective = if raw_objective.is_finite() {
-            raw_objective
-        } else {
-            warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
-            1e6
+        // Train model with parameters — NaN/divergence is a valid PSO outcome,
+        // not a fatal error. Score failed trials with maximum penalty so PSO
+        // learns to avoid that region of the search space.
+        let (objective, trial_metrics) = match model.train_with_params(params.clone()) {
+            Ok(metrics) => {
+                let raw_objective = M::extract_objective(&metrics);
+                let obj = if raw_objective.is_finite() {
+                    raw_objective
+                } else {
+                    warn!("Trial {} produced non-finite objective ({:?}), using penalty 1e6", trial_num, raw_objective);
+                    1e6
+                };
+                (obj, M::extract_metrics(&metrics))
+            }
+            Err(e) => {
+                warn!("Trial {} training failed (scored as penalty): {:#}", trial_num, e);
+                (1e6, None)
+            }
         };
-        let trial_metrics = M::extract_metrics(&metrics);
         let duration_secs = start_time.elapsed().as_secs_f64();
 
         info!("✓ Trial {} completed in {:.1}s", trial_num, duration_secs);