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