gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(dqn): per-branch independent epsilon for factored action diversity

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 14:12:27 +01:00
parent cbf81d293d
commit 2de3def317
1 changed files with 107 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

159
crates/ml-dqn/src/dqn.rs
View File

@@ -1619,21 +1619,58 @@ impl DQN {
};
// Epsilon-greedy exploration (forced to random during warmup)
let action = if in_warmup || rng.gen::<f32>() < 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::<f32>() < effective_epsilon;
let ord_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
let urg_random = in_warmup || rng.gen::<f32>() < 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::<f32>() < 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::<f32>() < 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::<f32>() < effective_epsilon;
let ord_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
let urg_random = in_warmup || rng.gen::<f32>() < 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::<f32>() < 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())