gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: configurable gradient budget (IQN 75%→40% default, C51 10%→45%)

Browse Source 
Hardcoded IQN_GRAD_BUDGET=0.75 starved C51 directional learning:
grad_norm frozen at 0.655 for 20+ epochs, WinRate stuck at 44-46%.
IQN dominated the gradient, C51's contribution after budget clipping
was effectively zero.

Now configurable via DQNHyperparameters:
- iqn_grad_budget: 0.40 (was 0.75 const)
- cql_grad_budget: 0.10 (was 0.10 const)
- ens_grad_budget: 0.05 (was 0.05 const)
- C51 gets remainder: 1.0 - 0.40 - 0.10 - 0.05 = 0.45 (was 0.10)

Wired through GpuDqnTrainConfig, TOML profiles, hyperopt search space.
Hyperopt can now tune the gradient budget split.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-12 11:58:48 +02:00
parent 91e05990a9
commit fbe185cb92
5 changed files with 49 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.claude/scheduled_tasks.lock Normal file
View File

@@ -0,0 +1 @@
{"sessionId":"4d4aa47f-4eb8-44d0-9d38-840da6e33fc0","pid":4104094,"acquiredAt":1775984198815}

13
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -223,6 +223,12 @@ pub struct GpuDqnTrainConfig {
/// market features from portfolio+MTF+OFI features in the state vector.
/// OFI features (8 dims, when enabled) bypass the bottleneck via portfolio_dim.
pub market_dim: usize,
/// Gradient budget fraction for IQN. Default 0.40.
pub iqn_grad_budget: f32,
/// Gradient budget fraction for CQL. Default 0.10.
pub cql_grad_budget: f32,
/// Gradient budget fraction for ensemble. Default 0.05.
pub ens_grad_budget: f32,
}
impl Default for GpuDqnTrainConfig {
@@ -272,6 +278,9 @@ impl Default for GpuDqnTrainConfig {
enable_gradient_vaccine: true, // always on
bottleneck_dim: 16,
market_dim: 42, // Default: 42 base features. Overridden to 50 when OFI (MBP-10) enabled.
iqn_grad_budget: 0.40,
cql_grad_budget: 0.10,
ens_grad_budget: 0.05,
}
}
}
@@ -1289,7 +1298,7 @@ impl GpuDqnTrainer {
// Clipped SAXPY: grad_buf += iqn_lambda * clip(scratch, iqn_budget)
let grad_ptr = self.ptrs.grad_buf;
let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::IQN_GRAD_BUDGET;
let max_component_norm = self.config.max_grad_norm * self.config.iqn_grad_budget;
let scale = self.config.iqn_lambda;
unsafe {
self.stream
@@ -1536,7 +1545,7 @@ impl GpuDqnTrainer {
// Clipped SAXPY
let grad_ptr = self.ptrs.grad_buf;
let max_component_norm = self.config.max_grad_norm * crate::trainers::dqn::fused_training::ENS_GRAD_BUDGET;
let max_component_norm = self.config.max_grad_norm * self.config.ens_grad_budget;
unsafe {
self.stream
.launch_builder(&self.clipped_saxpy_kernel)

14
crates/ml/src/trainers/dqn/config.rs
View File

@@ -900,6 +900,17 @@ pub struct DQNHyperparameters {
/// L_total = L_c51 + iqn_lambda * L_iqn
/// Range [0.0, 2.0]: 0.0 = C51 only, 0.5 = balanced, 1.0 = equal weight
pub iqn_lambda: f64,
/// Gradient budget fraction for IQN auxiliary objective.
/// Controls how much of the total gradient norm IQN consumes.
/// Range [0.1, 0.8]. Default 0.40. Higher = IQN dominates (good for magnitude sizing).
/// At 0.75, IQN starves C51 — directional learning flatlines.
pub iqn_grad_budget: f64,
/// Gradient budget fraction for CQL regularization.
/// Range [0.0, 0.3]. Default 0.10.
pub cql_grad_budget: f64,
/// Gradient budget fraction for ensemble diversity.
/// Range [0.0, 0.2]. Default 0.05. C51 gets the remainder.
pub ens_grad_budget: f64,
/// Spectral norm σ_max — constrains ||W||_σ ≤ σ_max.
/// Range [1.0, 10.0]. Default 3.0 (permits Xavier scaling, prevents Q-explosion).
pub spectral_norm_sigma_max: f64,
@@ -1501,6 +1512,9 @@ impl DQNHyperparameters {
num_quantiles: 32, // Default: 32 quantiles
qr_kappa: 1.0, // Default: 1.0 (standard quantile Huber loss)
iqn_lambda: 0.25, // Default: mild IQN regularization alongside C51
iqn_grad_budget: 0.40, // Default: balanced — was 0.75 (starved C51 directional learning)
cql_grad_budget: 0.10, // Default: mild CQL regularization
ens_grad_budget: 0.05, // Default: ensemble diversity. C51 gets remainder (0.45)
spectral_norm_sigma_max: 3.0, // Default: permits Xavier scaling [1.0, 10.0]
spectral_decoupling_lambda: 0.01, // Default: mild logit magnitude penalty (Pezeshki 2021)

29
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -91,13 +91,14 @@ unsafe impl Sync for RawCudaGraph {}
/// Per-component gradient norm budget fractions for auxiliary objectives.
/// C51 gets whatever remains: `1.0 - sum(active_auxiliary_budgets)`.
/// When all auxiliaries are active: IQN=60%, CQL=25%, C51=10%, Ens=5%.
/// IQN is PRIMARY distributional loss (Huber — variance-neutral, no magnitude collapse).
/// C51 demoted to 10% — its cross-entropy structurally prefers low-variance actions,
/// causing irrecoverable magnitude collapse once the target network locks in Small.
pub(crate) const CQL_GRAD_BUDGET: f32 = 0.10; // was 0.25 — suppressed rare magnitude actions
pub(crate) const IQN_GRAD_BUDGET: f32 = 0.75; // was 0.60 — IQN is variance-neutral, better for magnitude
pub(crate) const ENS_GRAD_BUDGET: f32 = 0.05;
/// Configurable via DQNHyperparameters (iqn_grad_budget, cql_grad_budget, ens_grad_budget).
/// Default: IQN=40%, CQL=10%, ENS=5%, C51=45%.
/// Old hardcoded: IQN=75% — starved C51 directional learning (grad_norm frozen at 0.655).
pub(crate) struct GradBudget {
pub cql: f32,
pub iqn: f32,
pub ens: f32,
}
/// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets.
///
@@ -375,6 +376,9 @@ impl FusedTrainingCtx {
enable_gradient_vaccine: true,
bottleneck_dim: hyperparams.bottleneck_dim,
market_dim: 42, // Always 42 base market features — OFI features bypass bottleneck via portfolio_dim
iqn_grad_budget: hyperparams.iqn_grad_budget as f32,
cql_grad_budget: hyperparams.cql_grad_budget as f32,
ens_grad_budget: hyperparams.ens_grad_budget as f32,
};
// Create weight set pointer views AFTER GpuDqnTrainer is constructed below.
@@ -1141,10 +1145,11 @@ impl FusedTrainingCtx {
}
let alpha = self.trainer.c51_alpha();
let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 };
let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 };
let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 };
let c51_frac = (1.0 - cql_frac - iqn_frac - ens_frac).max(0.30);
let cfg = self.trainer.config();
let cql_frac = if self.trainer.has_cql() { cfg.cql_grad_budget } else { 0.0 };
let iqn_frac = if self.gpu_iqn.is_some() { cfg.iqn_grad_budget } else { 0.0 };
let ens_frac = if !self.ensemble_extra_heads.is_empty() { cfg.ens_grad_budget } else { 0.0 };
let c51_frac = (1.0_f32 - cql_frac - iqn_frac - ens_frac).max(0.30);
let mgn = self.trainer.config().max_grad_norm;
let primary_budget = mgn * (1.0 - alpha) + mgn * c51_frac * alpha;
self.trainer.clip_grad_buf_inplace(primary_budget)
@@ -1268,7 +1273,7 @@ impl FusedTrainingCtx {
if self.trainer.has_cql() {
match self.trainer.apply_cql_gradient() {
Ok(true) => {
let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET;
let cql_budget = self.trainer.config().max_grad_norm * self.trainer.config().cql_grad_budget;
self.trainer.apply_cql_clipped_saxpy(cql_budget)
.map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?;
}

6
crates/ml/src/training_profile.rs
View File

@@ -380,6 +380,9 @@ pub struct SearchSpaceSection {
pub branch_hidden_dim: Option<[f64; 2]>,
pub gradient_accumulation_steps: Option<[f64; 2]>,
pub iqn_lambda: Option<[f64; 2]>,
pub iqn_grad_budget: Option<[f64; 2]>,
pub cql_grad_budget: Option<[f64; 2]>,
pub ens_grad_budget: Option<[f64; 2]>,
/// Spectral norm sigma max bounds.
pub spectral_norm_sigma_max: Option<[f64; 2]>,
/// C51 warmup epochs bounds.
@@ -549,6 +552,9 @@ impl HyperoptProfile {
"branch_hidden_dim" => ss.branch_hidden_dim,
"gradient_accumulation_steps" => ss.gradient_accumulation_steps,
"iqn_lambda" => ss.iqn_lambda,
"iqn_grad_budget" => ss.iqn_grad_budget,
"cql_grad_budget" => ss.cql_grad_budget,
"ens_grad_budget" => ss.ens_grad_budget,
"spectral_norm_sigma_max" => ss.spectral_norm_sigma_max,
"c51_warmup_epochs" => ss.c51_warmup_epochs,
"c51_alpha_max" => ss.c51_alpha_max,