feat(dqn): add IQN network fields to DQN struct with initialization

Add iqn_network and iqn_target_network as Option<QuantileNetwork> fields
to the DQN struct. When use_iqn is enabled, DQN::new() creates both networks
with QuantileConfig derived from DQNConfig (embed_dim from last hidden layer).
Each network gets its own VarMap for independent parameter tracking.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-20 15:17:29 +01:00
parent 02393a4cd1
commit 33a3ee99ea

View File

@@ -730,6 +730,11 @@ pub struct DQN {
/// N-step buffer for multi-step TD learning (Wave 11.5)
nstep_buffer: Option<Arc<Mutex<super::NStepBuffer>>>,
/// IQN (Implicit Quantile Network) for distributional RL (Dabney et al. 2018b)
iqn_network: Option<super::quantile_regression::QuantileNetwork>,
/// IQN target network (frozen copy for stable target computation)
iqn_target_network: Option<super::quantile_regression::QuantileNetwork>,
// WAVE 23 P0 Fix #2: Early stopping on gradient collapse
/// Consecutive epochs with gradient collapse (resets when gradients recover)
gradient_collapse_counter: usize,
@@ -868,6 +873,34 @@ impl DQN {
None
};
// IQN initialization (replaces broken C51 — no scatter_add needed)
let (iqn_network, iqn_target_network) = if config.use_iqn {
let iqn_config = super::quantile_regression::QuantileConfig {
num_quantiles: config.iqn_num_quantiles,
quantile_embedding_dim: config.iqn_embedding_dim,
kappa: config.iqn_kappa,
num_actions: config.num_actions,
};
let embed_dim = config.hidden_dims.last().copied().unwrap_or(config.state_dim);
let iqn_vars = VarMap::new();
let iqn_vb = VarBuilder::from_varmap(&iqn_vars, DType::F32, &device);
let iqn_net = super::quantile_regression::QuantileNetwork::new(
&iqn_config, embed_dim, iqn_vb
)?;
let iqn_target_vars = VarMap::new();
let iqn_target_vb = VarBuilder::from_varmap(&iqn_target_vars, DType::F32, &device);
let iqn_target = super::quantile_regression::QuantileNetwork::new(
&iqn_config, embed_dim, iqn_target_vb
)?;
(Some(iqn_net), Some(iqn_target))
} else {
(None, None)
};
Ok(Self {
epsilon: config.epsilon_start,
q_network,
@@ -886,7 +919,8 @@ impl DQN {
device,
recent_actions: VecDeque::with_capacity(100),
nstep_buffer,
// WAVE 23 P0 Fix #2/#4: Initialize early stopping counters
iqn_network,
iqn_target_network,
gradient_collapse_counter: 0,
q_value_divergence_counter: 0,
})
@@ -2522,4 +2556,22 @@ mod tests {
assert!(loss.is_finite(), "CQL loss should be finite, got {}", loss);
Ok(())
}
#[test]
fn test_dqn_with_iqn_creation() {
let mut config = DQNConfig::default();
config.state_dim = 8;
config.num_actions = 3;
config.hidden_dims = vec![16, 16];
config.use_iqn = true;
config.iqn_num_quantiles = 32;
config.use_distributional = false;
config.use_dueling = false;
let dqn = DQN::new(config);
assert!(
dqn.is_ok(),
"DQN with IQN should create successfully: {:?}",
dqn.err()
);
}
}