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