feat(dqn): add IQN action selection with optional CVaR risk-aware mode

Wire IQN into select_action() greedy branch. When use_iqn is enabled,
computes quantile-based Q-values via uniform quantile sampling. CVaR mode
(use_cvar_action_selection) optimizes for worst-case outcomes instead of
mean, providing risk-averse action selection for conservative trading.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-20 15:30:15 +01:00
parent 19f0de509c
commit bafe784dbf

View File

@@ -1116,14 +1116,47 @@ impl DQN {
)
.map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?;
let q_values = self.forward(&state_tensor)?;
let best_action_idx = q_values
.argmax(1)?
.get(0)?
.to_scalar::<u32>()
.map_err(|e| MLError::ModelError(format!("Failed to get best action: {}", e)))?;
if self.config.use_iqn && self.iqn_network.is_some() {
// IQN: Compute quantile-based Q-values
let iqn_net = self.iqn_network.as_ref().unwrap();
let state_embed = self.get_state_embedding(&state_tensor)?;
FactoredAction::from_index(best_action_idx as usize)?
// Use fixed uniform quantiles for deterministic action selection
let taus = iqn_net.sample_uniform_quantiles(1, &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to sample taus: {}", e)))?;
// Forward: [1, num_actions, num_quantiles]
let all_quantiles = iqn_net.forward(&state_embed, &taus)?;
// Score each action
let action_scores = if self.config.use_cvar_action_selection {
// CVaR: Optimize for worst-case outcomes (risk-averse)
iqn_net.compute_cvar(&all_quantiles, self.config.cvar_alpha)
.map_err(|e| MLError::ModelError(format!("CVaR computation failed: {}", e)))?
} else {
// Standard: Mean over quantiles (risk-neutral)
iqn_net.to_expected_q(&all_quantiles)
.map_err(|e| MLError::ModelError(format!("Expected Q computation failed: {}", e)))?
};
// Argmax over actions
let best_action_idx = action_scores
.argmax(1)?
.get(0)?
.to_scalar::<u32>()?;
FactoredAction::from_index(best_action_idx as usize)?
} else {
// Standard Q-network action selection
let q_values = self.forward(&state_tensor)?;
let best_action_idx = q_values
.argmax(1)?
.get(0)?
.to_scalar::<u32>()
.map_err(|e| MLError::ModelError(format!("Failed to get best action: {}", e)))?;
FactoredAction::from_index(best_action_idx as usize)?
}
};
// Track action for entropy penalty calculation
@@ -2704,6 +2737,50 @@ mod tests {
Ok(())
}
#[test]
fn test_iqn_action_selection() {
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 = 8;
config.use_distributional = false;
config.use_dueling = false;
config.epsilon_start = 0.0; // Force greedy for testing
config.use_noisy_nets = false;
config.warmup_steps = 0;
let mut dqn = DQN::new(config).unwrap();
let state = vec![0.5f32; 8];
let action = dqn.select_action(&state);
assert!(action.is_ok(), "IQN action selection should succeed: {:?}", action.err());
}
#[test]
fn test_iqn_cvar_action_selection() {
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 = 8;
config.use_distributional = false;
config.use_dueling = false;
config.epsilon_start = 0.0;
config.use_noisy_nets = false;
config.warmup_steps = 0;
config.use_cvar_action_selection = true;
config.cvar_alpha = 0.05;
let mut dqn = DQN::new(config).unwrap();
let state = vec![0.5f32; 8];
let action = dqn.select_action(&state);
assert!(action.is_ok(), "IQN CVaR action selection should succeed: {:?}", action.err());
}
#[test]
fn test_dqn_with_iqn_creation() {
let mut config = DQNConfig::default();