fix(dqn): close checkpoint validation gaps, atomic NormStats, clippy cleanup

- DQNAgent::load_from_safetensors: validate architecture hash before
  loading weights (was bypassing validation entirely - P0 production gap)
- DQNEnsemble::save_to_directory: embed architecture metadata via
  safetensors::serialize_to_file instead of bare VarMap::save (save/load
  round-trip was guaranteed to fail)
- architecture_hash: hash hidden_dims.len() before values to prevent
  theoretical collision between different-length configs
- train_baseline_rl: NormStats write now atomic (write .tmp then rename)
  with cleanup guard on rename failure; serialization error is now loud
  (error! + return None) instead of silently discarded
- train_baseline_rl: checkpoint rename failure cleans up orphaned .tmp
- hyperopt_baseline_rl: fix clippy single_match_else (match on equality
  check -> if/else)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml/examples/hyperopt_baseline_rl.rs

@@ -167,29 +167,26 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device)
     training_metrics::set_hyperopt_trial("dqn", 0.0, args.trials as f64);
 
     let start = Instant::now();
-    let result = match args.optimizer.as_str() {
-        "tpe" => {
-            info!("Using TPE (Tree-Parzen Estimator) optimizer");
-            ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed))
-                .context("DQN TPE hyperopt optimization failed")?
-        }
-        _ => {
-            let optimizer = ArgminOptimizer::builder()
-                .max_trials(args.trials)
-                .n_initial(args.n_initial)
-                .seed(args.seed)
-                .build();
+    let result = if args.optimizer.as_str() == "tpe" {
+        info!("Using TPE (Tree-Parzen Estimator) optimizer");
+        ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed))
+            .context("DQN TPE hyperopt optimization failed")?
+    } else {
+        let optimizer = ArgminOptimizer::builder()
+            .max_trials(args.trials)
+            .n_initial(args.n_initial)
+            .seed(args.seed)
+            .build();
 
-            if parallel > 1 {
-                info!("Using parallel PSO optimization ({} threads)", parallel);
-                optimizer
-                    .optimize_parallel(trainer)
-                    .context("DQN parallel hyperopt optimization failed")?
-            } else {
-                optimizer
-                    .optimize(trainer)
-                    .context("DQN hyperopt optimization failed")?
-            }
+        if parallel > 1 {
+            info!("Using parallel PSO optimization ({} threads)", parallel);
+            optimizer
+                .optimize_parallel(trainer)
+                .context("DQN parallel hyperopt optimization failed")?
+        } else {
+            optimizer
+                .optimize(trainer)
+                .context("DQN hyperopt optimization failed")?
         }
     };
     let elapsed = start.elapsed().as_secs_f64();
@@ -246,29 +243,26 @@ fn run_ppo_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device)
     training_metrics::set_hyperopt_trial("ppo", 0.0, args.trials as f64);
 
     let start = Instant::now();
-    let result = match args.optimizer.as_str() {
-        "tpe" => {
-            info!("Using TPE (Tree-Parzen Estimator) optimizer");
-            ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed))
-                .context("PPO TPE hyperopt optimization failed")?
-        }
-        _ => {
-            let optimizer = ArgminOptimizer::builder()
-                .max_trials(args.trials)
-                .n_initial(args.n_initial)
-                .seed(args.seed)
-                .build();
+    let result = if args.optimizer.as_str() == "tpe" {
+        info!("Using TPE (Tree-Parzen Estimator) optimizer");
+        ml::hyperopt::optimize_with_tpe(trainer, args.trials, args.n_initial, Some(args.seed))
+            .context("PPO TPE hyperopt optimization failed")?
+    } else {
+        let optimizer = ArgminOptimizer::builder()
+            .max_trials(args.trials)
+            .n_initial(args.n_initial)
+            .seed(args.seed)
+            .build();
 
-            if parallel > 1 {
-                info!("Using parallel PSO optimization ({} threads)", parallel);
-                optimizer
-                    .optimize_parallel(trainer)
-                    .context("PPO parallel hyperopt optimization failed")?
-            } else {
-                optimizer
-                    .optimize(trainer)
-                    .context("PPO hyperopt optimization failed")?
-            }
+        if parallel > 1 {
+            info!("Using parallel PSO optimization ({} threads)", parallel);
+            optimizer
+                .optimize_parallel(trainer)
+                .context("PPO parallel hyperopt optimization failed")?
+        } else {
+            optimizer
+                .optimize(trainer)
+                .context("PPO hyperopt optimization failed")?
         }
     };
     let elapsed = start.elapsed().as_secs_f64();

crates/ml/examples/train_baseline_rl.rs

@@ -224,12 +224,24 @@ fn prepare_fold_data(
 
     // Save NormStats for this fold
     let norm_path = output_dir.join(format!("norm_stats_fold{}.json", window.fold));
-    if let Ok(json) = serde_json::to_string_pretty(&norm_stats) {
-        if let Err(e) = std::fs::write(&norm_path, json) {
-            warn!("  Failed to write NormStats: {}", e);
-        } else {
+    match serde_json::to_string_pretty(&norm_stats) {
+        Ok(json) => {
+            let norm_tmp = norm_path.with_extension("json.tmp");
+            if let Err(e) = std::fs::write(&norm_tmp, &json) {
+                error!("  Failed to write NormStats tmp {}: {}", norm_tmp.display(), e);
+                return None;
+            }
+            if let Err(e) = std::fs::rename(&norm_tmp, &norm_path) {
+                error!("  Failed to rename NormStats {} -> {}: {}", norm_tmp.display(), norm_path.display(), e);
+                drop(std::fs::remove_file(&norm_tmp));
+                return None;
+            }
             info!("  Saved NormStats to {}", norm_path.display());
         }
+        Err(e) => {
+            error!("  Failed to serialize NormStats: {}", e);
+            return None;
+        }
     }
 
     let train_warmup = window.train.len().saturating_sub(train_norm.len());
@@ -369,8 +381,10 @@ fn train_dqn_fold(
         let tmp_path = ckpt_path.with_extension("safetensors.tmp");
         std::fs::write(&tmp_path, &data)
             .with_context(|| format!("Failed to write checkpoint tmp: {}", tmp_path.display()))?;
-        std::fs::rename(&tmp_path, &ckpt_path)
-            .with_context(|| format!("Failed to rename checkpoint: {} -> {}", tmp_path.display(), ckpt_path.display()))?;
+        if let Err(e) = std::fs::rename(&tmp_path, &ckpt_path) {
+            drop(std::fs::remove_file(&tmp_path));
+            return Err(e).with_context(|| format!("Failed to rename checkpoint: {} -> {}", tmp_path.display(), ckpt_path.display()));
+        }
         info!("  [DQN] Fold {} saved checkpoint: {}", fold, ckpt_path.display());
         Ok(ckpt_path.to_string_lossy().into_owned())
     };

crates/ml/src/dqn/agent.rs

@@ -729,6 +729,18 @@ impl DQNAgent {
             )));
         }
 
+        // Validate architecture hash before loading weights
+        let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| {
+            MLError::CheckpointError(format!("Failed to read safetensors file: {}", e))
+        })?;
+        let (_header_size, st_metadata) =
+            safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
+                MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
+            })?;
+        self.config
+            .validate_checkpoint_metadata(st_metadata.metadata())?;
+        drop(raw_bytes);
+
         let mut vars_clone = self.q_network.vars().clone();
         vars_clone.load(&safetensors_path).map_err(|e| {
             MLError::CheckpointError(format!("Failed to load safetensors via VarMap: {}", e))

crates/ml/src/dqn/dqn.rs

@@ -283,14 +283,15 @@ impl DQNConfig {
     /// or panic on shape mismatch. This hash is embedded in safetensors metadata
     /// on save and validated on load.
     ///
-    /// Hashed fields: state_dim, num_actions, hidden_dims, use_dueling,
-    /// dueling_hidden_dim, use_distributional, num_atoms, use_noisy_nets,
-    /// use_iqn, iqn_embedding_dim, iqn_num_quantiles.
+    /// Hashed fields: state_dim, num_actions, hidden_dims (length + values),
+    /// use_dueling, dueling_hidden_dim, use_distributional, num_atoms,
+    /// use_noisy_nets, use_iqn, iqn_embedding_dim, iqn_num_quantiles.
     pub fn architecture_hash(&self) -> String {
         use sha2::{Digest, Sha256};
         let mut hasher = Sha256::new();
         hasher.update(self.state_dim.to_le_bytes());
         hasher.update(self.num_actions.to_le_bytes());
+        hasher.update(self.hidden_dims.len().to_le_bytes());
         for &dim in &self.hidden_dims {
             hasher.update(dim.to_le_bytes());
         }

crates/ml/src/dqn/ensemble.rs

@@ -750,7 +750,7 @@ impl DQNEnsemble {
         self.config.voting_strategy
     }
 
-    /// Save ensemble to directory (saves all agent weights)
+    /// Save ensemble to directory (saves all agent weights with architecture metadata)
     pub fn save_to_directory(&self, dir: &str) -> Result<(), MLError> {
         std::fs::create_dir_all(dir).map_err(|e| {
             MLError::CheckpointError(format!("Failed to create ensemble directory: {}", e))
@@ -759,12 +759,28 @@ impl DQNEnsemble {
         for (idx, agent) in self.agents.iter().enumerate() {
             let path = format!("{}/agent_{}.safetensors", dir, idx);
             let vars = agent.get_q_network_vars();
-            vars.save(&path).map_err(|e| {
-                MLError::CheckpointError(format!("Failed to save agent {}: {}", idx, e))
+            let vars_data = vars.data().lock().map_err(|e| {
+                MLError::LockError(format!("Failed to lock vars for ensemble agent {}: {}", idx, e))
+            })?;
+
+            let mut tensors: std::collections::HashMap<String, candle_core::Tensor> =
+                std::collections::HashMap::new();
+            for (name, var) in vars_data.iter() {
+                tensors.insert(name.clone(), var.as_tensor().clone());
+            }
+
+            let arch_metadata = Some(agent.config.checkpoint_metadata());
+            safetensors::serialize_to_file(
+                &tensors,
+                &arch_metadata,
+                std::path::Path::new(&path),
+            )
+            .map_err(|e| {
+                MLError::CheckpointError(format!("Failed to save ensemble agent {}: {}", idx, e))
             })?;
         }
 
-        info!("✓ Ensemble saved to {} ({} agents)", dir, self.config.num_agents);
+        info!("Ensemble saved to {} ({} agents)", dir, self.config.num_agents);
         Ok(())
     }