From 2de3def3172d5a6993b9c2e0717442a07efe4f08 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 12 Mar 2026 14:12:27 +0100 Subject: [PATCH] feat(dqn): per-branch independent epsilon for factored action diversity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The single global epsilon coin flip caused 90% of actions to use greedy argmax across ALL 3 branches simultaneously, collapsing to 1 composite action. Only 6/45 factored actions were used (13.3% diversity). Fix: Each branch (exposure, order, urgency) now flips its own epsilon coin independently: - P(all greedy) = (1-ε)³ ≈ 72.9% at ε=0.10 - P(at least one random) ≈ 27.1% vs previous 10% - Expected unique actions per epoch: significantly higher Applied to both select_action() and select_action_with_confidence(). The forward pass is skipped when all 3 branches happen to be random (0.1% chance), preserving the optimization. 1331 tests pass (416 ml-dqn + 915 ml), 0 failures. Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/dqn.rs | 159 ++++++++++++++++++++++++++------------- 1 file changed, 107 insertions(+), 52 deletions(-) diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 752d7a90d..9eca91797 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -1619,21 +1619,58 @@ impl DQN { }; // Epsilon-greedy exploration (forced to random during warmup) - let action = if in_warmup || rng.gen::() < effective_epsilon { - if self.config.use_branching { - // Branching: sample each branch independently for uniform factored coverage + let action = if self.config.use_branching { + // Per-branch independent epsilon: each branch flips its own coin. + // P(all greedy) = (1-ε)^3 ≈ 72.9% at ε=0.10 + // P(at least one random) ≈ 27.1% — much better factored coverage + let exp_random = in_warmup || rng.gen::() < effective_epsilon; + let ord_random = in_warmup || rng.gen::() < effective_epsilon; + let urg_random = in_warmup || rng.gen::() < effective_epsilon; + + if exp_random && ord_random && urg_random { + // All branches random — skip forward pass let exposure = ExposureLevel::from_index(rng.gen_range(0..5_usize))?; let order = OrderType::from_index(rng.gen_range(0..3_usize))?; let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?; FactoredAction { exposure, order, urgency } } else { - // Standard: random exposure index, deterministic order routing - let action_idx = rng.gen_range(0..self.config.num_actions); - let exposure = ExposureLevel::from_index(action_idx)?; - OrderRouter::route_default(exposure) + // At least one branch is greedy — need forward pass + let state_tensor = Tensor::from_vec( // gpu-ok: state upload + state.to_vec(), + (1, self.config.state_dim), + self.q_network.device(), + ) + .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; + let branching_net = self.branching_q_network.as_ref().ok_or_else(|| { + MLError::ModelError("Branching enabled but network not initialized".into()) + })?; + let output = branching_net.forward_branches_eval(&state_tensor)?; + let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output)?; + + let exposure = if exp_random { + ExposureLevel::from_index(rng.gen_range(0..5_usize))? + } else { + ExposureLevel::from_index(greedy.first().copied().unwrap_or(2) as usize)? + }; + let order = if ord_random { + OrderType::from_index(rng.gen_range(0..3_usize))? + } else { + OrderType::from_index(greedy.get(1).copied().unwrap_or(0) as usize)? + }; + let urgency = if urg_random { + Urgency::from_index(rng.gen_range(0..3_usize))? + } else { + Urgency::from_index(greedy.get(2).copied().unwrap_or(1) as usize)? + }; + FactoredAction { exposure, order, urgency } } + } else if in_warmup || rng.gen::() < effective_epsilon { + // Standard (non-branching): random exposure index, deterministic order routing + let action_idx = rng.gen_range(0..self.config.num_actions); + let exposure = ExposureLevel::from_index(action_idx)?; + OrderRouter::route_default(exposure) } else { - // Greedy action selection + // Greedy action selection (non-branching) let state_tensor = Tensor::from_vec( // gpu-ok: state upload state.to_vec(), (1, self.config.state_dim), @@ -1641,20 +1678,7 @@ impl DQN { ) .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; - if self.config.use_branching { - // Branching DQN: greedy argmax per branch independently - let branching_net = self.branching_q_network.as_ref().ok_or_else(|| { - MLError::ModelError("Branching enabled but network not initialized".into()) - })?; - let output = branching_net.forward_branches_eval(&state_tensor)?; - let branch_actions = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output)?; - - // branch_actions: [exposure_idx, order_idx, urgency_idx] - let exposure = ExposureLevel::from_index(branch_actions.first().copied().unwrap_or(2) as usize)?; - let order = OrderType::from_index(branch_actions.get(1).copied().unwrap_or(0) as usize)?; - let urgency = Urgency::from_index(branch_actions.get(2).copied().unwrap_or(1) as usize)?; - FactoredAction { exposure, order, urgency } - } else if self.config.use_iqn && self.iqn_network.is_some() { + if self.config.use_iqn && self.iqn_network.is_some() { // IQN: Compute quantile-based Q-values let iqn_net = self.iqn_network.as_ref().ok_or_else(|| { MLError::ModelError("IQN network not initialized despite use_iqn=true".into()) @@ -1795,23 +1819,74 @@ impl DQN { }; // Epsilon-greedy exploration (forced to random during warmup) - let (action, confidence) = if in_warmup || rng.gen::() < effective_epsilon { - if self.config.use_branching { - // Branching: sample each branch independently for uniform factored coverage + let (action, confidence) = if self.config.use_branching { + // Per-branch independent epsilon: each branch flips its own coin. + let exp_random = in_warmup || rng.gen::() < effective_epsilon; + let ord_random = in_warmup || rng.gen::() < effective_epsilon; + let urg_random = in_warmup || rng.gen::() < effective_epsilon; + + if exp_random && ord_random && urg_random { + // All branches random — skip forward pass let exposure = ExposureLevel::from_index(rng.gen_range(0..5_usize))?; let order = OrderType::from_index(rng.gen_range(0..3_usize))?; let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?; let uniform_conf = (1.0_f32 / 45.0).clamp(0.5, 0.95); (FactoredAction { exposure, order, urgency }, uniform_conf) } else { - // Standard: random exposure index, deterministic order routing - let action_idx = rng.gen_range(0..self.config.num_actions); - let uniform_conf = (1.0_f32 / self.config.num_actions as f32).clamp(0.5, 0.95); - let exposure = ExposureLevel::from_index(action_idx)?; - (OrderRouter::route_default(exposure), uniform_conf) + // At least one branch is greedy — need forward pass + let state_tensor = Tensor::from_vec( // gpu-ok: state upload + state.to_vec(), + (1, self.config.state_dim), + self.q_network.device(), + ) + .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; + let branching_net = self.branching_q_network.as_ref().ok_or_else(|| { + MLError::ModelError("Branching enabled but network not initialized".into()) + })?; + let output = branching_net.forward_branches_eval(&state_tensor)?; + let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output)?; + + // Per-branch confidence: average of per-branch softmax max probs (greedy branches only) + let mut total_conf = 0.0_f32; + let mut greedy_count = 0_u32; + for (i, adv) in output.advantages.iter().enumerate() { + let is_random = match i { 0 => exp_random, 1 => ord_random, _ => urg_random }; + if !is_random { + total_conf += Self::softmax_confidence(adv)?; + greedy_count += 1; + } + } + let conf = if greedy_count > 0 { + (total_conf / greedy_count as f32).clamp(0.5, 0.95) + } else { + 0.5 + }; + + let exposure = if exp_random { + ExposureLevel::from_index(rng.gen_range(0..5_usize))? + } else { + ExposureLevel::from_index(greedy.first().copied().unwrap_or(2) as usize)? + }; + let order = if ord_random { + OrderType::from_index(rng.gen_range(0..3_usize))? + } else { + OrderType::from_index(greedy.get(1).copied().unwrap_or(0) as usize)? + }; + let urgency = if urg_random { + Urgency::from_index(rng.gen_range(0..3_usize))? + } else { + Urgency::from_index(greedy.get(2).copied().unwrap_or(1) as usize)? + }; + (FactoredAction { exposure, order, urgency }, conf) } + } else if in_warmup || rng.gen::() < effective_epsilon { + // Standard (non-branching): random exposure index, deterministic order routing + let action_idx = rng.gen_range(0..self.config.num_actions); + let uniform_conf = (1.0_f32 / self.config.num_actions as f32).clamp(0.5, 0.95); + let exposure = ExposureLevel::from_index(action_idx)?; + (OrderRouter::route_default(exposure), uniform_conf) } else { - // Greedy action selection with confidence from Q-values + // Greedy action selection (non-branching) with confidence from Q-values let state_tensor = Tensor::from_vec( // gpu-ok: state upload state.to_vec(), (1, self.config.state_dim), @@ -1819,27 +1894,7 @@ impl DQN { ) .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?; - if self.config.use_branching { - // Branching DQN: greedy argmax per branch with per-branch confidence - let branching_net = self.branching_q_network.as_ref().ok_or_else(|| { - MLError::ModelError("Branching enabled but network not initialized".into()) - })?; - let output = branching_net.forward_branches_eval(&state_tensor)?; - - // Per-branch confidence: average of per-branch softmax max probs - let mut total_conf = 0.0_f32; - for adv in &output.advantages { - let branch_conf = Self::softmax_confidence(adv)?; - total_conf += branch_conf; - } - let conf = (total_conf / output.advantages.len() as f32).clamp(0.5, 0.95); - - let branch_actions = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output)?; - let exposure = ExposureLevel::from_index(branch_actions.first().copied().unwrap_or(2) as usize)?; - let order = OrderType::from_index(branch_actions.get(1).copied().unwrap_or(0) as usize)?; - let urgency = Urgency::from_index(branch_actions.get(2).copied().unwrap_or(1) as usize)?; - (FactoredAction { exposure, order, urgency }, conf) - } else if self.config.use_iqn && self.iqn_network.is_some() { + if self.config.use_iqn && self.iqn_network.is_some() { // IQN path let iqn_net = self.iqn_network.as_ref().ok_or_else(|| { MLError::ModelError("IQN network not initialized despite use_iqn=true".into())