feat(dqn): IQN + Branching coexistence — risk-adjusted confidence scoring

Two changes enabling IQN and Branching DQN to complement each other:

1. Optimizer fix: IQN vars excluded from optimizer when branching is the
   primary loss path. Previously, IQN weights would silently decay to
   zero via weight_decay with zero gradients.

2. Confidence coexistence in select_action_with_confidence():
   When both use_iqn and use_branching are enabled, branching handles
   action decomposition (per-branch greedy selection) while IQN provides
   distributional risk assessment (CVaR-based confidence scoring).
   This is the "complement" design: branching = what to do, IQN = how
   risky.

Tests: ml-dqn=416, 0 failures

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 15:55:45 +01:00
parent a24c3d2ee0
commit 2a71f6bf4c

View File

@@ -2042,18 +2042,44 @@ impl DQN {
})?;
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)?;
// Confidence: IQN CVaR when available (risk-sensitive), else branching softmax.
// IQN + Branching coexistence: branching handles action decomposition,
// IQN provides distributional risk assessment for confidence scoring.
let conf = if self.config.use_iqn {
if let Some(ref iqn_net) = self.iqn_network {
let state_embed = self.get_state_embedding(&state_tensor)?;
let taus = iqn_net
.sample_uniform_quantiles(1, &self.device)
.map_err(|e| MLError::ModelError(format!("IQN taus: {}", e)))?;
let all_quantiles = iqn_net.forward(&state_embed, &taus)?;
let cvar_scores = iqn_net
.compute_cvar(&all_quantiles, self.config.cvar_alpha)
.map_err(|e| MLError::ModelError(format!("IQN CVaR: {}", e)))?;
// CVaR-based confidence: softmax over CVaR scores of exposure actions
Self::softmax_confidence(&cvar_scores)?
} else {
// IQN enabled but network not initialized — fall through to branching
let mut total_conf = 0.0_f32;
for adv in &output.advantages {
total_conf += Self::softmax_confidence(adv)?;
}
total_conf / output.advantages.len() as f32
}
} else {
// Per-branch confidence: average of per-branch softmax max probs
let mut total_conf = 0.0_f32;
for adv in &output.advantages {
total_conf += Self::softmax_confidence(adv)?;
}
total_conf / output.advantages.len() as f32
};
let conf = conf.clamp(0.5, 0.95);
return Ok((FactoredAction { exposure, order, urgency }, conf));
}
@@ -2422,9 +2448,13 @@ impl DQN {
self.q_network.vars().all_vars()
};
// IQN: Add IQN network vars (loss flows through both IQN and base q_network)
if let Some(ref iqn_net) = self.iqn_network {
vars.extend(iqn_net.vars().all_vars());
// IQN: Add IQN network vars only when IQN is the primary loss path.
// When branching is active, IQN gets zero gradients but weight_decay
// would silently decay its weights to zero — skip to prevent that.
if !self.config.use_branching {
if let Some(ref iqn_net) = self.iqn_network {
vars.extend(iqn_net.vars().all_vars());
}
}
self.optimizer = Some(