From c10340708f9ee9261faa7d3fa987c15472ee3f55 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 5 Mar 2026 09:57:48 +0100 Subject: [PATCH] feat(ml): replace greedy argmax with Gumbel-max softmax for hyperopt eval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DQN hyperopt walk-forward eval collapsed to 1 trade when Q-values were nearly uniform (greedy argmax always picked the same action). Replace batch_greedy_actions with batch_softmax_actions using the Gumbel-max trick (argmax(Q/T + Gumbel(0,1))) — fully GPU-resident, no CPU softmax/sampling. Temperature 0.1: nearly greedy when Q-values are separated, diverse when uniform. Logs unique_actions/45 per backtest. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/dqn/dqn.rs | 45 +++++++++++++++++++++ crates/ml/src/dqn/regime_conditional.rs | 53 +++++++++++++++++++++++++ crates/ml/src/hyperopt/adapters/dqn.rs | 15 +++++-- crates/ml/src/trainers/dqn/config.rs | 12 ++++++ 4 files changed, 122 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index bd8f2a56a..31602ceac 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -1813,6 +1813,51 @@ impl DQN { Ok(action_indices.into_iter().map(|a| a as usize).collect()) } + /// Batch softmax (Boltzmann) action selection via Gumbel-max trick — fully GPU-resident. + /// + /// Equivalent to sampling from softmax(Q/temperature) per row, but avoids + /// CPU transfer of Q-values. Uses: argmax(Q/T - log(-log(U))) where U~Uniform(0,1). + /// At low temperature (~0.1), nearly greedy; provides diversity when Q-values are uniform. + /// + /// # Arguments + /// * `states` - Tensor of shape `[N, state_dim]` + /// * `temperature` - Boltzmann temperature (clamped to >= 1e-6) + pub fn batch_softmax_actions( + &self, + states: &Tensor, + temperature: f64, + ) -> Result, MLError> { + let q_values = self.forward(states)?; + let temp = temperature.max(1e-6); + let temp_tensor = + Tensor::new(&[temp as f32], &self.device) + .and_then(|t| t.broadcast_as(q_values.shape())) + .map_err(|e| MLError::ModelError(format!("Temperature broadcast failed: {}", e)))?; + let scaled = q_values + .broadcast_div(&temp_tensor) + .map_err(|e| MLError::ModelError(format!("Q/T division failed: {}", e)))?; + // Gumbel noise: -log(-log(U)), U ~ Uniform(eps, 1-eps) to avoid log(0) + let uniform = Tensor::rand(0.001_f32, 0.999_f32, q_values.shape(), &self.device) + .map_err(|e| MLError::ModelError(format!("Gumbel uniform generation failed: {}", e)))?; + let gumbel = uniform + .log() + .and_then(|t| t.neg()) + .and_then(|t| t.log()) + .and_then(|t| t.neg()) + .map_err(|e| MLError::ModelError(format!("Gumbel noise computation failed: {}", e)))?; + let perturbed = scaled + .add(&gumbel) + .map_err(|e| MLError::ModelError(format!("Gumbel perturbation failed: {}", e)))?; + let action_indices = perturbed + .argmax(1) + .map_err(|e| MLError::ModelError(format!("Gumbel-max argmax failed: {}", e)))? + .to_vec1::() + .map_err(|e| { + MLError::ModelError(format!("Failed to transfer action indices to CPU: {}", e)) + })?; + Ok(action_indices.into_iter().map(|a| a as usize).collect()) + } + /// Track action for entropy penalty calculation /// /// This should be called after action selection (both single and batch modes) diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml/src/dqn/regime_conditional.rs index 3ccb4b19b..714a29302 100644 --- a/crates/ml/src/dqn/regime_conditional.rs +++ b/crates/ml/src/dqn/regime_conditional.rs @@ -345,6 +345,59 @@ impl RegimeConditionalDQN { Ok(results) } + /// Batch softmax action selection across regime heads (Gumbel-max, GPU-resident). + /// + /// Same regime-dispatch logic as `batch_greedy_actions` but delegates to + /// per-head `batch_softmax_actions` for Boltzmann sampling. + pub fn batch_softmax_actions( + &self, + states: &Tensor, + temperature: f64, + ) -> Result, MLError> { + let state_vecs = states.to_vec2::().map_err(|e| { + MLError::ModelError(format!("Failed to extract states for regime: {}", e)) + })?; + let n = state_vecs.len(); + let mut results = vec![0_usize; n]; + + let mut trending_idx = Vec::new(); + let mut ranging_idx = Vec::new(); + let mut volatile_idx = Vec::new(); + for (i, sv) in state_vecs.iter().enumerate() { + match RegimeType::classify_from_features(sv) { + RegimeType::Trending => trending_idx.push(i), + RegimeType::Ranging => ranging_idx.push(i), + RegimeType::Volatile => volatile_idx.push(i), + } + } + + let device = &self.device; + for (indices, head) in [ + (&trending_idx, &self.trending_head), + (&ranging_idx, &self.ranging_head), + (&volatile_idx, &self.volatile_head), + ] { + if indices.is_empty() { + continue; + } + let dim = state_vecs[0].len(); + let mut flat = Vec::with_capacity(indices.len() * dim); + for &i in indices { + flat.extend_from_slice(&state_vecs[i]); + } + let sub_tensor = + Tensor::from_vec(flat, (indices.len(), dim), device).map_err(|e| { + MLError::ModelError(format!("Regime sub-batch failed: {}", e)) + })?; + let actions = head.batch_softmax_actions(&sub_tensor, temperature)?; + for (j, &idx) in indices.iter().enumerate() { + results[idx] = actions[j]; + } + } + + Ok(results) + } + /// Store experience in all head buffers /// /// Since we no longer have shared memory, each head maintains its own buffer. diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 07eb4dd3a..6cedb5ab3 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -2728,13 +2728,14 @@ impl HyperparameterOptimizable for DQNTrainer { let num_chunks = total_bars.div_ceil(EVAL_CHUNK_SIZE); let mut ohlcv_bars = Vec::with_capacity(total_bars); let mut conversion_failures: usize = 0; + let mut unique_actions = std::collections::HashSet::new(); tracing::info!( "Batched backtest: {} bars in {} chunks of {} (device: {:?})", total_bars, num_chunks, EVAL_CHUNK_SIZE, device, ); - // Acquire read lock once — batch_greedy_actions is &self (no mutation) + // Acquire read lock once — batch_softmax_actions is &self (no mutation) let runtime = tokio::runtime::Runtime::new().map_err(|e| { MLError::TrainingError(format!("Failed to create runtime for backtest: {}", e)) })?; @@ -2780,7 +2781,12 @@ impl HyperparameterOptimizable for DQNTrainer { )) })?; - let action_indices = agent_guard.batch_greedy_actions(&batch_tensor)?; + // Softmax temperature for hyperopt eval — nearly greedy when Q-values + // are well-separated, but prevents degenerate single-action collapse + // when Q-values are uniform (Gumbel-max, fully GPU-resident). + const EVAL_SOFTMAX_TEMP: f64 = 0.1; + let action_indices = + agent_guard.batch_softmax_actions(&batch_tensor, EVAL_SOFTMAX_TEMP)?; // 3. Sequential trade simulation on CPU for (i, &action_idx) in action_indices.iter().enumerate() { @@ -2790,6 +2796,7 @@ impl HyperparameterOptimizable for DQNTrainer { let bar_idx = chunk_start + i; let close = val_close_prices[bar_idx] as f32; + unique_actions.insert(action_idx); let factored = crate::dqn::FactoredAction::from_index(action_idx)?; let bar = OHLCVBarF32 { @@ -2837,7 +2844,7 @@ impl HyperparameterOptimizable for DQNTrainer { tracing::info!( "Batched backtest complete: {} bars, {} chunks, {} trades, \ Sharpe {:.4}, Win Rate {:.2}%, Max DD {:.2}%, Return {:.2}%, \ - conversion_failures={}", + unique_actions={}/{}, conversion_failures={}", ohlcv_bars.len(), num_chunks, metrics.total_trades, @@ -2845,6 +2852,8 @@ impl HyperparameterOptimizable for DQNTrainer { metrics.win_rate, metrics.max_drawdown_pct, metrics.total_return_pct, + unique_actions.len(), + 45, // 5 exposure × 3 order × 3 urgency conversion_failures, ); diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index f4751f596..afd4f0101 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -71,6 +71,18 @@ impl DQNAgentType { } } + /// Batch softmax (Boltzmann) action selection — dispatches to underlying DQN variant. + pub fn batch_softmax_actions( + &self, + states: &Tensor, + temperature: f64, + ) -> Result, MLError> { + match self { + Self::Standard(agent) => agent.batch_softmax_actions(states, temperature), + Self::RegimeConditional(agent) => agent.batch_softmax_actions(states, temperature), + } + } + /// Reset count-based exploration bonus (call at epoch boundary) pub fn reset_count_bonus(&mut self) { match self {