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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user