Files
foxhunt/docs/superpowers/plans/2026-03-30-dqn-gems-implementation.md
jgrusewski 4eaac54700 feat: wire OFI features through TOML profile + CUDA kernel guard tests
Fix mbp10_data_dir/trades_data_dir TOML→DQNHyperparameters wiring gap.
Fields were deserialized but silently dropped — now applied in
training_profile.rs apply_to(). Production TOML activates OFI (8 features
from MBP-10 order book), smoketest leaves it off for fast iteration.

3 new OFI integration tests: state vector positioning, graceful
degradation (zero-fill without MBP-10), feature name ordering.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-31 01:35:04 +02:00

45 KiB
Raw Blame History

DQN Gems & Gaps Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Remove remaining feature-flag dead code from the DQN agent config, add three proven generalization techniques (spectral decoupling, manifold mixup, regime-aware replay), and redesign hyperopt to cover the full generalization parameter space without combinatorial explosion.

Architecture: Five phases, each independently testable. Phase 1 is pure cleanup (same pattern as the enable_/use_ removal from DQNHyperparameters). Phases 2-3 add new loss terms and replay logic at the host/kernel level. Phase 4 restructures hyperopt search. Phase 5 (multi-horizon returns) is scoped but deferred — it requires kernel rewrites and is tracked separately.

Tech Stack: Rust (crates/ml-dqn, crates/ml), CUDA C (*.cu kernels), TOML config, PSO hyperopt


Phase 1: DQNConfig Inner Flag Removal (Gap 2)

Task 1: Remove enable_q_value_clipping — always clip

This is the simplest flag: always true in every config variant. Remove the field, remove the conditional, always clamp.

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:121 (field def)

  • Modify: crates/ml-dqn/src/dqn.rs:283,594,661,738 (Default/constructors)

  • Modify: crates/ml-dqn/src/dqn.rs:1517 (usage in forward pass)

  • Modify: crates/ml-dqn/tests/gpu_smoketest.rs:67 (test config)

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs:300 (builder)

  • Modify: crates/ml/src/benchmark/dqn_benchmark.rs:453 (benchmark config)

  • Modify: crates/ml/src/trainers/dqn/config.rs:1711 (profile)

  • Step 1: Remove field from DQNConfig struct

In crates/ml-dqn/src/dqn.rs, delete the field definition and its doc comments:

// DELETE these 3 lines (~119-121):
// / BUG #37 FIX: Q-value clipping configuration (prevents step-level explosions)
// /// Whether to enable Q-value clipping in forward pass
// / Default: true (prevents explosions without breaking gradients)
// pub enable_q_value_clipping: bool,
  • Step 2: Remove from all Default/constructor impls

In crates/ml-dqn/src/dqn.rs, delete enable_q_value_clipping: true, from every impl block:

  • Default (~line 283)

  • new_production() (~line 594)

  • new_conservative() (~line 661)

  • new_emergency() (~line 738)

  • Step 3: Remove conditional from forward pass — always clamp

In crates/ml-dqn/src/dqn.rs around line 1517, change:

// BEFORE:
let q_values = if self.config.enable_q_value_clipping {
    let clamped = q_values.clamp(
        self.config.q_value_clip_min,
        self.config.q_value_clip_max,
    );
    clamped
} else {
    q_values
};

// AFTER:
let q_values = q_values.clamp(
    self.config.q_value_clip_min,
    self.config.q_value_clip_max,
);
  • Step 4: Remove from all external config sites

Delete enable_q_value_clipping: true, from:

  • crates/ml/src/trainers/dqn/trainer/constructor.rs:300

  • crates/ml/src/benchmark/dqn_benchmark.rs:453

  • crates/ml/src/trainers/dqn/config.rs:1711

  • crates/ml-dqn/tests/gpu_smoketest.rs:67

  • Step 5: Compile check

SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50

Expected: clean compile.

  • Step 6: Commit
git add crates/ml-dqn/src/dqn.rs crates/ml/src/trainers/dqn/trainer/constructor.rs \
  crates/ml/src/benchmark/dqn_benchmark.rs crates/ml/src/trainers/dqn/config.rs \
  crates/ml-dqn/tests/gpu_smoketest.rs
git commit -m "cleanup: remove enable_q_value_clipping — always clamp Q-values"

Task 2: Remove use_soft_updates — always soft EMA

Hard updates caused 10K-16K gradient norms (Wave 16H). Soft EMA is the only production-safe path.

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:71 (field def)

  • Modify: crates/ml-dqn/src/dqn.rs:266,577,642,713 (constructors)

  • Modify: crates/ml-dqn/src/dqn.rs:2783 (update_target_networks)

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs:271 (builder)

  • Modify: crates/ml/src/trainers/dqn/config.rs:1697 (profile)

  • Modify: crates/ml-dqn/tests/gpu_smoketest.rs:52 (test config)

  • Modify: crates/ml/src/benchmark/dqn_benchmark.rs:424 (benchmark — currently false, must fix)

  • Step 1: Remove field from DQNConfig struct

In crates/ml-dqn/src/dqn.rs, delete:

// DELETE (~line 70-71):
// /// Use soft (EMA) or hard target updates
// pub use_soft_updates: bool,
  • Step 2: Remove from all constructors

Delete use_soft_updates: true, (or false) from Default, new_production, new_conservative, new_emergency.

  • Step 3: Remove conditional from update_target_networks()

In crates/ml-dqn/src/dqn.rs around line 2783. The if self.config.use_soft_updates { ... } else { ... } block should collapse to ONLY the soft update path. Delete the entire hard-update else branch.

Read the full function first to understand both branches, then keep only the soft EMA path.

  • Step 4: Remove from external sites

  • constructor.rs:271 — delete use_soft_updates: matches!(hyperparams.target_update_mode, TargetUpdateMode::Soft),

  • config.rs:1697 — delete use_soft_updates: true,

  • dqn_benchmark.rs:424 — delete use_soft_updates: false,

  • gpu_smoketest.rs:52 — delete use_soft_updates: true,

Also check if TargetUpdateMode enum is now unused after this removal. If so, delete it too from crates/ml/src/trainers/mod.rs or wherever it's defined. If it's still referenced elsewhere (e.g., hyperopt adapter target_update_mode field), leave it for now and note it as tech debt.

  • Step 5: Compile check
SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50
  • Step 6: Commit
git add -u crates/ml-dqn/ crates/ml/
git commit -m "cleanup: remove use_soft_updates — always soft EMA target updates"

Task 3: Remove use_cvar_action_selection — always use CVaR

Already hardcoded to true in constructor.rs:315. CVaR provides risk-aware action selection via IQN quantiles — no reason to ever disable it when IQN is active.

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:225 (field def)

  • Modify: crates/ml-dqn/src/dqn.rs:322,608,675,767 (constructors)

  • Modify: crates/ml-dqn/src/dqn.rs:1671,1878,2066,2134 (action selection conditionals)

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs:315

  • Modify: crates/ml-dqn/tests/gpu_smoketest.rs:92

  • Modify: crates/ml/src/benchmark/dqn_benchmark.rs:466

  • Step 1: Remove field from struct

Delete pub use_cvar_action_selection: bool, and its comment from DQNConfig.

  • Step 2: Remove from constructors

Delete use_cvar_action_selection: true/false, from all config constructors.

  • Step 3: Collapse all 4 action-selection conditionals

There are 4 if self.config.use_cvar_action_selection { CVaR } else { mean } patterns in dqn.rs at lines ~1671, ~1878, ~2066, ~2134. For each one, keep ONLY the CVaR path and delete the else branch.

Example pattern (applies to all 4):

// BEFORE:
let action_scores = if self.config.use_cvar_action_selection {
    iqn_net.compute_cvar(&all_quantiles, self.config.cvar_alpha)
        .map_err(|e| MLError::ModelError(format!("CVaR computation failed: {}", e)))?
} else {
    iqn_net.expected_q(&all_quantiles)?
};

// AFTER:
let action_scores = iqn_net.compute_cvar(&all_quantiles, self.config.cvar_alpha)
    .map_err(|e| MLError::ModelError(format!("CVaR computation failed: {}", e)))?;
  • Step 4: Remove from external sites

Delete from constructor.rs:315, gpu_smoketest.rs:92, dqn_benchmark.rs:466.

  • Step 5: Compile check
SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50
  • Step 6: Commit
git add -u crates/ml-dqn/ crates/ml/
git commit -m "cleanup: remove use_cvar_action_selection — always CVaR risk-aware selection"

Task 4: Remove use_count_bonus — always compute, coefficient gates magnitude

Currently parameter-gated: use_count_bonus: count_bonus_coefficient > 0.0. The UCB computation is trivial (a few f32 ops per action selection) so always running it with coefficient = 0.0 is functionally identical to disabling it — the bonus terms are all zero.

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:159 (field def)

  • Modify: crates/ml-dqn/src/dqn.rs:292,748 (constructors)

  • Modify: crates/ml-dqn/src/dqn.rs:1684,1709,1732,1887,1913,1938 (6 usage sites)

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs:324

  • Modify: crates/ml/src/trainers/dqn/config.rs:198-202 (use_count_bonus() accessor)

  • Modify: crates/ml-dqn/tests/gpu_smoketest.rs:74

  • Modify: crates/ml/tests/dqn_action_collapse_fix_test.rs:305 (test sets use_count_bonus)

  • Step 1: Remove field and accessor

Delete pub use_count_bonus: bool, from DQNConfig. Delete the use_count_bonus() method from crates/ml/src/trainers/dqn/config.rs:198-202.

  • Step 2: Remove from constructors

Delete use_count_bonus: true/false, from all constructors. Delete use_count_bonus: hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, from constructor.rs:324.

  • Step 3: Remove all 6 conditionals — always apply bonus

For each if self.config.use_count_bonus { ... } at lines ~1684, ~1709, ~1732, ~1887, ~1913, ~1938:

For the 4 action-selection conditionals (lines ~1684, ~1709, ~1887, ~1913), remove the if and always add the bonus:

// BEFORE:
let action_scores = if self.config.use_count_bonus {
    let bonuses = self.count_bonus.bonuses();
    // ... add bonus tensor ...
} else {
    action_scores
};

// AFTER:
let action_scores = {
    let bonuses = self.count_bonus.bonuses();
    // ... add bonus tensor ...
};

For the 2 record_action sites (lines ~1732, ~1938), remove the if guard:

// BEFORE:
if self.config.use_count_bonus {
    self.count_bonus.record_action(action.exposure as usize);
}

// AFTER:
self.count_bonus.record_action(action.exposure as usize);
  • Step 4: Compile check
SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50
  • Step 5: Commit
git add -u crates/ml-dqn/ crates/ml/
git commit -m "cleanup: remove use_count_bonus — always apply UCB exploration bonus"

Task 5: Remove use_iqn — IQN always active

This is the most impactful flag removal. use_iqn controls whether IQN networks are allocated at all. In production and conservative configs it's true. In Default and emergency it's false. Making it always-on means:

  • IQN networks always allocated (~2x quantile params)
  • Architecture hash changes (line 360 always includes IQN)
  • Checkpoint metadata always writes dqn.use_iqn = true

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:170 (field def)

  • Modify: crates/ml-dqn/src/dqn.rs:299,603,670,752 (constructors)

  • Modify: crates/ml-dqn/src/dqn.rs:347,360 (architecture_hash)

  • Modify: crates/ml-dqn/src/dqn.rs:383 (checkpoint metadata)

  • Modify: crates/ml-dqn/src/dqn.rs:463,488 (checkpoint loading)

  • Modify: crates/ml-dqn/src/dqn.rs:1309 (IQN initialization conditional)

  • Modify: crates/ml-dqn/src/dqn.rs:1656,1868,2017,2050,2123 (5 inference conditionals)

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs (builder)

  • Modify: crates/ml/src/validation/adapters.rs:225

  • Modify: crates/ml-dqn/tests/gpu_smoketest.rs:77

  • Modify: crates/ml/src/benchmark/dqn_benchmark.rs:462

  • Modify: crates/ml/tests/dqn_iqn_integration_test.rs:91,175 (test sets use_iqn)

  • Modify: crates/ml/examples/evaluate_baseline.rs:922,1091 (example configs)

  • Step 1: Remove field from struct

Delete pub use_iqn: bool, from DQNConfig.

  • Step 2: Remove from constructors

Delete use_iqn: true/false, from all constructors.

  • Step 3: Make IQN initialization unconditional

At crates/ml-dqn/src/dqn.rs:1309, change:

// BEFORE:
let (iqn_network, iqn_target_network) = if config.use_iqn {
    // ... allocate IQN ...
    (Some(iqn_net), Some(iqn_target))
} else {
    (None, None)
};

// AFTER:
let (iqn_network, iqn_target_network) = {
    // ... allocate IQN ...
    (iqn_net, iqn_target)  // No Option wrapping
};

IMPORTANT: This also requires changing iqn_network and iqn_target_network fields in the DQN struct from Option<IqnNetwork> to IqnNetwork. Then remove all .as_ref().unwrap() / .is_some() checks at the 5 inference sites.

  • Step 4: Collapse inference conditionals

At each if self.config.use_iqn && self.iqn_network.is_some() site (lines ~1656, ~1868, ~2017, ~2050, ~2123), the IQN path is now always taken. Remove the condition and the non-IQN else branch. Replace self.iqn_network.as_ref().ok_or_else(...) with direct &self.iqn_network.

  • Step 5: Fix architecture_hash and metadata

  • Line 360: Change hasher.update([self.use_iqn as u8]) to hasher.update([true as u8]) (backward compat with existing checkpoints, same pattern as use_dueling on line 357).

  • Line 383: Change to meta.insert("dqn.use_iqn".to_owned(), "true".to_owned()).

  • Line 463: Change checkpoint loading to ignore the parsed use_iqn value (same as use_dueling on line 462).

  • Step 6: Fix external sites

  • validation/adapters.rs:225 — delete config.use_iqn = false;

  • gpu_smoketest.rs:77 — delete use_iqn: false,

  • dqn_benchmark.rs:462 — delete use_iqn: false,

  • Step 7: Compile check

SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50

This may surface additional Option unwrapping errors. Fix each by converting Option<IqnNetwork> accesses to direct accesses.

  • Step 8: Commit
git add -u crates/ml-dqn/ crates/ml/
git commit -m "cleanup: remove use_iqn — IQN quantile networks always active"

Phase 2: Spectral Decoupling & Manifold Mixup

Task 6: Add Spectral Decoupling loss term

What: L2 penalty on Q-value logit magnitudes (not weights). Prevents overconfident Q-values more directly than weight decay. Paper: Pezeshki et al. 2021 "Gradient Starvation."

Where it plugs in: The fused CUDA training pipeline in gpu_dqn_trainer.rs assembles loss from C51 + CQL + IQN + ensemble components. Spectral decoupling is an additional scalar added to the total loss before the backward pass.

The implementation is host-side only — no kernel changes. The Q-value logits are already computed during forward; we just add λ_sd × mean(q_logits²) to the loss.

Files:

  • Modify: crates/ml/src/trainers/dqn/config.rs — add spectral_decoupling_lambda: f64 to DQNHyperparameters

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — add loss term after forward pass

  • Modify: config/training/dqn-production.toml — add parameter

  • Modify: config/training/dqn-smoketest.toml — add parameter

  • Modify: crates/ml/src/trainers/dqn/smoke_tests/helpers.rs — add to test config

  • Step 1: Add hyperparameter

In crates/ml/src/trainers/dqn/config.rs, add to DQNHyperparameters:

/// Spectral decoupling coefficient: L2 penalty on Q-value logit magnitudes.
/// Prevents overconfident Q-values more directly than weight decay.
/// Pezeshki et al. 2021 "Gradient Starvation". Range: [0.0, 0.1]. Default: 0.01.
pub spectral_decoupling_lambda: f64,

Add to Default impl: spectral_decoupling_lambda: 0.01,

  • Step 2: Wire into GpuDqnTrainConfig

In crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs, add pub spectral_decoupling_lambda: f32 to GpuDqnTrainConfig struct. Set from hyperparams in fused_training.rs where GpuDqnTrainConfig is constructed.

  • Step 3: Add loss term in training step

In gpu_dqn_trainer.rs, after the C51 loss kernel computes per_sample_loss and before the backward pass, add spectral decoupling. The Q-value logits are already available from the forward pass in the q_out buffer.

Find the loss assembly site (search for where total_loss or the combined loss is computed). Add:

// Spectral decoupling: L2 on Q-value logit magnitudes
// q_out shape: [batch_size, num_actions] — raw logits from forward pass
let sd_lambda = self.config.spectral_decoupling_lambda;
if sd_lambda > 0.0 {
    // Compute mean(q²) over the batch via cuBLAS dot product: ||q||² / n
    // q_out is already on GPU in f32 from the forward pass
    let q_sq_sum = cublas_dot(q_out, q_out);  // sum of squares
    let sd_loss = sd_lambda * q_sq_sum / (batch_size * num_actions) as f32;
    // Add to per-sample loss accumulator (or total loss scalar)
}

NOTE: The exact integration point depends on how the loss pipeline works in the fused training. Read gpu_dqn_trainer.rs around the loss computation to find where auxiliary losses (CQL, ensemble) are added. Spectral decoupling goes in the same place.

  • Step 4: Add to TOML configs

In config/training/dqn-production.toml under [advanced]:

spectral_decoupling_lambda = 0.01

In config/training/dqn-smoketest.toml under [advanced]:

spectral_decoupling_lambda = 0.01
  • Step 5: Compile and smoke test
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -50
FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -30
  • Step 6: Commit
git add -u crates/ml/ config/training/
git commit -m "feat: add spectral decoupling — L2 penalty on Q-value logit magnitudes"

Task 7: Add Manifold Mixup at bottleneck

What: During training, interpolate between pairs of bottleneck embeddings (the 2D output of the Temporal Causal Bottleneck). This smooths the learned manifold and dramatically improves generalization in RL. Paper: Verma et al. 2019 "Manifold Mixup."

Where it plugs in: The bottleneck output is 2D (after [42 → 2] linear + tanh). Before it's concatenated with portfolio features and fed into the trunk, we mix pairs: h_mix = λ·h_i + (1-λ)·h_j where λ ~ Beta(α, α) and j is a random permutation of the batch.

This is a training-only modification. At inference, no mixup.

Files:

  • Modify: crates/ml/src/trainers/dqn/config.rs — add mixup_alpha: f64

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — add mixup in forward pass (training only)

  • Modify: config/training/dqn-production.toml — add parameter

  • Modify: config/training/dqn-smoketest.toml — add parameter

  • Step 1: Add hyperparameter

In crates/ml/src/trainers/dqn/config.rs, add to DQNHyperparameters:

/// Manifold Mixup alpha for Beta distribution at bottleneck layer.
/// Higher = more aggressive interpolation. 0.0 = disabled.
/// Applied training-only: mixes pairs of bottleneck embeddings.
/// Verma et al. 2019. Range: [0.0, 2.0]. Default: 0.2.
pub mixup_alpha: f64,

Add to Default: mixup_alpha: 0.2,

  • Step 2: Implement mixup in fused forward pass

In gpu_dqn_trainer.rs, find where the bottleneck output (h_bottleneck, the 2D tensor) is computed. The bottleneck is the first layer: [batch, 42] × W_bottleneck → [batch, 2] → tanh.

After the bottleneck tanh and before concatenation with portfolio features, add:

// Manifold Mixup (training only)
let mixup_alpha = self.config.mixup_alpha;
if mixup_alpha > 0.0 && is_training {
    // Sample lambda from Beta(alpha, alpha) — use rand crate
    let lambda: f32 = Beta::new(mixup_alpha, mixup_alpha).sample(&mut rng);
    // Shuffle batch indices for pairing
    let perm = random_permutation(batch_size);
    // h_mix[i] = lambda * h_bottleneck[i] + (1-lambda) * h_bottleneck[perm[i]]
    // This is a simple CUDA kernel: element-wise linear combination
    mixup_kernel(h_bottleneck, h_bottleneck_permuted, lambda, batch_size, bottleneck_dim);
    // Targets must also be mixed: target_mix[i] = lambda * target[i] + (1-lambda) * target[perm[i]]
    mixup_kernel(targets, targets_permuted, lambda, batch_size, target_dim);
}

CRITICAL: Mixup requires mixing BOTH the inputs AND the targets/labels. For C51 distributional loss, the target distributions must also be interpolated. This means:

  1. Mix bottleneck embeddings
  2. Mix the reward + done + next_state used to compute Bellman targets
  3. OR simpler: mix the projected target distributions after Bellman projection

The simpler approach is to do mixup AFTER the full forward pass but BEFORE loss computation:

  • Compute C51 projected distributions for all samples
  • Mix: projected_mix[i] = λ·projected[i] + (1-λ)·projected[perm[i]]
  • Compute cross-entropy loss against mixed targets

This avoids mixing at the bottleneck level (which would require re-running the forward) and instead mixes at the distribution level, which is mathematically cleaner for distributional RL.

  • Step 3: Write mixup CUDA kernel

Create crates/ml/src/cuda_pipeline/mixup_kernel.cu:

extern "C" __global__ void mixup_distributions(
    const float* __restrict__ dist_a,     // [batch, num_atoms]
    const float* __restrict__ dist_b,     // [batch, num_atoms] (permuted)
    float* __restrict__ dist_out,          // [batch, num_atoms] (mixed)
    const int* __restrict__ perm,          // [batch] permutation indices
    float lambda,
    int batch_size,
    int num_atoms
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size * num_atoms) return;
    int sample = idx / num_atoms;
    int atom = idx % num_atoms;
    int partner = perm[sample];
    dist_out[idx] = lambda * dist_a[idx] + (1.0f - lambda) * dist_b[partner * num_atoms + atom];
}
  • Step 4: Add to TOML configs

In dqn-production.toml under [generalization]:

mixup_alpha = 0.2

In dqn-smoketest.toml under [generalization]:

mixup_alpha = 0.2
  • Step 5: Compile and smoke test
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -50
FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -30
  • Step 6: Commit
git add crates/ml/src/cuda_pipeline/mixup_kernel.cu
git add -u crates/ml/ config/training/
git commit -m "feat: add manifold mixup on C51 target distributions"

Phase 3: Regime-Aware Replay Sampling

Task 8: Tag transitions with regime ID

What: Add a regime: u8 field to the Experience struct so transitions know which market regime they came from. Currently regime is recomputed on-the-fly from features[40:41] every time experiences are sampled — wasteful and prevents regime-biased sampling.

The regime tag is computed once during experience collection using the existing ADX/CUSUM classification: ADX > 0.25 → Trending (0), |CUSUM| > 0.7 → Volatile (2), else Ranging (1).

Files:

  • Modify: crates/ml-dqn/src/experience.rs:10-52 — add regime: u8 field

  • Modify: crates/ml/src/cuda_pipeline/gpu_experience_collector.rs — compute regime during collection

  • Modify: crates/ml-dqn/src/prioritized_replay.rs — store regime with transition

  • Step 1: Add regime field to Experience struct

In crates/ml-dqn/src/experience.rs:

pub struct Experience {
    pub state: Vec<f32>,
    pub action: u8,
    pub reward: i32,
    pub next_state: Vec<f32>,
    pub done: bool,
    pub timestamp: u64,
    /// Market regime at time of transition: 0=Trending, 1=Ranging, 2=Volatile
    pub regime: u8,
}
  • Step 2: Fix all Experience construction sites

Grep for Experience { and Experience{ across the workspace. Add regime: 1 (Ranging = safe default) to all construction sites that don't have access to the feature vector. For construction sites that DO have the state vector, compute regime:

fn classify_regime(state: &[f32]) -> u8 {
    let adx = state.get(40).copied().unwrap_or(0.0);
    let cusum = state.get(41).copied().unwrap_or(0.0);
    if adx > 0.25 { 0 }           // Trending
    else if cusum.abs() > 0.7 { 2 } // Volatile
    else { 1 }                      // Ranging
}
  • Step 3: Tag during GPU experience collection

In gpu_experience_collector.rs, after the state features are gathered into states_out, classify regime for each timestep. This can be done as a simple CUDA kernel or host-side post-processing (the state is already read back for PER priority computation anyway).

Add regime output buffer: regimes_out: CudaSlice<u8> with shape [alloc_episodes × alloc_timesteps].

  • Step 4: Compile check
SQLX_OFFLINE=true cargo check -p ml-dqn -p ml 2>&1 | head -50

This WILL produce many errors from Experience construction sites missing regime. Fix them all.

  • Step 5: Commit
git add -u crates/ml-dqn/ crates/ml/
git commit -m "feat: tag Experience transitions with market regime ID"

Task 9: Regime-biased PER sampling

What: Modify PER sampling to bias toward transitions from the current detected regime. This prevents the trending head from training on ranging data and vice versa. Implementation: multiply each transition's priority by a regime affinity factor before sampling.

Regime affinity model:

  • If current regime matches transition regime: factor = 1.0 (full weight)
  • If one step away (Trending ↔ Ranging, Ranging ↔ Volatile): factor = regime_decay (default 0.3)
  • If two steps away (Trending ↔ Volatile): factor = regime_decay² (default 0.09)

This is applied as a multiplicative modifier to the PER priority, not a hard filter — all transitions remain reachable to prevent catastrophic forgetting.

Files:

  • Modify: crates/ml-dqn/src/prioritized_replay.rs — add regime-aware sampling mode

  • Modify: crates/ml/src/trainers/dqn/config.rs — add regime_replay_decay: f64

  • Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs — pass current regime to sampler

  • Modify: config/training/dqn-production.toml — add parameter

  • Step 1: Add hyperparameter

In crates/ml/src/trainers/dqn/config.rs, add to DQNHyperparameters:

/// Regime-aware replay decay factor. Multiplied into PER priority for
/// transitions from non-matching regimes. 1.0 = no regime bias (uniform).
/// 0.0 = only sample from current regime. Default: 0.3.
pub regime_replay_decay: f64,

Default: regime_replay_decay: 0.3,

  • Step 2: Add regime-biased sampling to PER

In crates/ml-dqn/src/prioritized_replay.rs, add a method:

/// Sample proportionally with regime affinity weighting.
/// current_regime: 0=Trending, 1=Ranging, 2=Volatile
/// decay: multiplicative factor per regime distance (0.0-1.0)
pub fn sample_regime_biased(
    &self,
    batch_size: usize,
    beta: f32,
    current_regime: u8,
    decay: f32,
) -> (Vec<Experience>, Vec<f32>, Vec<usize>) {
    // For each stored experience, compute effective priority:
    // eff_priority = base_priority * decay^|regime_distance|
    // Then sample proportionally from effective priorities.
    // IS weights must account for the modified distribution.
}

Implementation strategy: Rather than modifying the segment tree (which would require O(n) rebuild every time regime changes), sample from the existing tree and apply rejection sampling or importance-weighted correction:

  1. Sample K candidates from standard PER (K = 2× batch_size for safety margin)
  2. For each candidate, compute regime factor: decay^|candidate.regime - current_regime|
  3. Accept with probability proportional to regime factor
  4. Apply corrected IS weight: w_corrected = w_per / regime_factor

This is simple, correct, and doesn't touch the segment tree structure.

  • Step 3: Wire into training loop

In training_loop.rs, detect current regime from the most recent feature vector (features[40:41]) and pass to the PER sampler. Replace the standard sample_proportional() call with sample_regime_biased().

  • Step 4: Add to TOML config

In dqn-production.toml under [replay_buffer]:

regime_replay_decay = 0.3
  • Step 5: Compile and smoke test
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -50
FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -30
  • Step 6: Commit
git add -u crates/ml-dqn/ crates/ml/ config/training/
git commit -m "feat: regime-biased PER sampling — decay cross-regime transitions"

Phase 4: Hyperopt Redesign (Gap 3)

Task 10: Implement family-based hyperopt scaling

What: Instead of adding all 34 generalization params individually to PSO (creating an intractable 50+ dimensional space), group related techniques into families and add one intensity scalar per family to the search space. Each family's individual params scale proportionally from their defaults.

Families (6 new PSO dimensions):

Family Scalar Params Scaled
Adversarial adversarial_intensity [0, 2] saboteur perturbation, warmup, interval, dd_threshold
Regularization regularization_intensity [0, 2] dropout, spectral norm, stochastic depth, spectral decoupling
Data Augmentation augmentation_intensity [0, 2] feature masking fraction, feature noise, domain rand, mirror, time reversal
Loss Shaping loss_shaping_intensity [0, 2] asymmetric DD weight, regret blend, trade clustering penalty, position entropy
Ensemble ensemble_intensity [0, 2] ensemble disagreement penalty, bottleneck dim scaling
Causal causal_intensity [0, 2] causal weight, intervention interval

An intensity of 1.0 = current defaults. 0.5 = halved from defaults. 2.0 = doubled.

This takes the search from 22D → 28D (manageable) while covering all 34 techniques.

Files:

  • Modify: crates/ml/src/hyperopt/adapters/dqn.rs — add 6 family scalars to search space

  • Modify: crates/ml/src/trainers/dqn/config.rs — add family fields + scaling logic

  • Step 1: Add 6 family intensity fields to DQNHyperparameters

In crates/ml/src/trainers/dqn/config.rs:

/// Family intensity scalars for hyperopt (1.0 = defaults, 0-2 range).
/// Each scales a group of related generalization params proportionally.
pub adversarial_intensity: f64,
pub regularization_intensity: f64,
pub augmentation_intensity: f64,
pub loss_shaping_intensity: f64,
pub ensemble_intensity: f64,
pub causal_intensity: f64,

All default to 1.0.

  • Step 2: Add scaling method

Add a method apply_family_scaling(&mut self) to DQNHyperparameters that multiplies each technique's param by its family intensity:

pub fn apply_family_scaling(&mut self) {
    // Adversarial family
    let ai = self.adversarial_intensity;
    self.adversarial_dd_threshold *= ai;
    self.adversarial_warmup_epochs = (self.adversarial_warmup_epochs as f64 / ai.max(0.01)) as usize;
    self.adversarial_checkpoint_interval = (self.adversarial_checkpoint_interval as f64 / ai.max(0.01)) as usize;

    // Regularization family
    let ri = self.regularization_intensity;
    self.dropout_initial = (self.dropout_initial as f64 * ri).clamp(0.0, 0.5) as f32;
    self.spectral_norm_sigma_max *= ri;
    self.stochastic_depth_prob = (self.stochastic_depth_prob * ri).clamp(0.0, 0.5);
    self.spectral_decoupling_lambda *= ri;

    // Data augmentation family
    let di = self.augmentation_intensity;
    self.feature_mask_fraction = (self.feature_mask_fraction * di).clamp(0.0, 0.8);
    self.feature_noise_scale *= di;

    // Loss shaping family
    let li = self.loss_shaping_intensity;
    self.asymmetric_dd_weight *= li;
    self.regret_blend = (self.regret_blend * li).clamp(0.0, 1.0);
    self.trade_clustering_penalty *= li;
    self.position_entropy_weight *= li;

    // Ensemble family
    let ei = self.ensemble_intensity;
    self.ensemble_disagreement_penalty *= ei;

    // Causal family
    let ci = self.causal_intensity;
    self.causal_weight *= ci;
    self.causal_intervention_interval = (self.causal_intervention_interval as f64 / ci.max(0.01)) as usize;
}
  • Step 3: Add to PSO search space

In crates/ml/src/hyperopt/adapters/dqn.rs, add 6 new dimensions to the PSO parameter vector:

// In build_search_space() or equivalent:
params.push(SearchParam::new("adversarial_intensity", 0.0, 2.0, 1.0));
params.push(SearchParam::new("regularization_intensity", 0.0, 2.0, 1.0));
params.push(SearchParam::new("augmentation_intensity", 0.0, 2.0, 1.0));
params.push(SearchParam::new("loss_shaping_intensity", 0.0, 2.0, 1.0));
params.push(SearchParam::new("ensemble_intensity", 0.0, 2.0, 1.0));
params.push(SearchParam::new("causal_intensity", 0.0, 2.0, 1.0));
  • Step 4: Call apply_family_scaling() in hyperopt trial setup

In the hyperopt adapter's build_hyperparams() method, call apply_family_scaling() after constructing the base hyperparameters from the PSO vector.

  • Step 5: Compile check
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -50
  • Step 6: Commit
git add -u crates/ml/
git commit -m "feat: family-based hyperopt — 6 intensity scalars for 34 generalization techniques"

Phase 5: Multi-Horizon Returns (Scoped, Deferred)

Not implemented in this plan. This phase requires kernel rewrites to c51_loss_kernel.cu, nstep_kernel.cu, and iqn_dual_head_kernel.cu. The current n-step mechanism (γ^n Bellman discount) provides partial multi-horizon coverage.

When to attempt: After Phases 1-4 are validated on the H100 with walk-forward results showing the generalization improvements from spectral decoupling, mixup, and regime-aware replay.

Key design decisions for future plan:

  1. Kernel approach: Triple the C51 projection to compute 3 projected distributions (γ=0.9, 0.95, 0.99), then average the cross-entropy losses. This requires 3× the save_projected buffer (3 × batch × 3 × 51 atoms = ~46KB per batch at batch=1024 — trivial on H100).
  2. IQN: Separate Bellman targets per gamma. The IQN quantile loss becomes: L_iqn = Σ_m w_m × L_huber(τ, Q_pred - (r + γ_m × Q_target)).
  3. N-step: Each gamma needs its own n-step accumulation. The nstep_kernel.cu already takes gamma as a parameter — launch it 3× with different gammas.
  4. Gradient budget: Spectral norm per-component budgets (C51=70%, etc.) stay the same, but applied to the averaged multi-horizon loss.

Phase 6: Core Hyperopt Families + Breakout Dimensions

Motivation

Phase 4 added 6 family intensity scalars for the 34 generalization techniques. But the original 24 individual PSO dimensions (learning_rate, gamma, batch_size, etc.) are still independently tuned. Many of these correlate — high lr with low gradient clip makes no physical sense. Grouping them into families reduces the effective search space and ensures every PSO particle represents a "sensible" configuration.

Design: Family + Breakout

5 core families (replace 21 of the 24 original dimensions):

Family Scalar Range Params Scaled
Learning learning_intensity [0, 2] learning_rate, tau, weight_decay
Exploration exploration_intensity [0, 2] entropy_coefficient, noisy_sigma_init
Replay replay_intensity [0, 2] buffer_size, per_alpha, per_beta_start, n_steps
Architecture architecture_intensity [0, 2] hidden_dim_base, dueling_hidden_dim, num_atoms
Risk risk_intensity [0, 2] max_position_absolute, huber_delta, minimum_profit_factor, w_dd, dd_threshold

3 breakout dimensions (stay independent — physically meaningful, exotic values matter):

Param Why independent
gamma Time horizon — doesn't correlate with other params, wrong value ruins everything
iqn_lambda Controls C51/IQN blend — nonlinear interaction with distributional loss
c51_warmup_epochs Transition timing — wrong value causes early training collapse

Fixed params (removed from search, always use defaults):

Param Default Why fixed
batch_size From TOML config Hardware-dependent, not a taste parameter
transaction_cost_multiplier 1.0 Physical constant (IBKR fees), not tunable
v_max Computed from gamma + reward_scale Derived quantity, not independent
min_hold_bars 5 Domain constraint, not a hyperparameter

Result: 30D → 14D (5 core families + 6 generalization families + 3 breakouts)

Task 11: Implement core hyperopt families

Files:

  • Modify: crates/ml/src/trainers/dqn/config.rs — add 5 core family fields + extend apply_family_scaling()

  • Modify: crates/ml/src/hyperopt/adapters/dqn.rs — restructure PSO vector from 30D to 14D

  • Modify: config/training/dqn-production.toml — add 5 core family scalars

  • Modify: config/training/dqn-smoketest.toml — add 5 core family scalars

  • Modify: config/training/dqn-localdev.toml — add 5 core family scalars

  • Step 1: Add 5 core family fields to DQNHyperparameters

In crates/ml/src/trainers/dqn/config.rs, add:

// Core hyperopt families (1.0 = defaults, range [0.0, 2.0])
/// Learning family: scales learning_rate, tau, weight_decay together.
pub learning_intensity: f64,
/// Exploration family: scales entropy_coefficient, noisy_sigma_init.
pub exploration_intensity: f64,
/// Replay family: scales buffer_size, per_alpha, per_beta_start, n_steps.
pub replay_intensity: f64,
/// Architecture family: scales hidden_dim_base, dueling_hidden_dim, num_atoms.
pub architecture_intensity: f64,
/// Risk family: scales max_position, huber_delta, min_profit_factor, w_dd, dd_threshold.
pub risk_intensity: f64,

All default to 1.0.

  • Step 2: Extend apply_family_scaling() with core families

Add to the existing method:

// ── Core families ──

// Learning family
let li = self.learning_intensity;
self.learning_rate *= li;
self.tau *= li;
self.weight_decay *= li;

// Exploration family
let ei = self.exploration_intensity;
self.entropy_coefficient = (self.entropy_coefficient * ei).clamp(0.001, 1.0);
self.noisy_sigma_init = (self.noisy_sigma_init as f64 * ei).clamp(0.01, 2.0) as f32;

// Replay family
let rpi = self.replay_intensity;
self.buffer_size = (self.buffer_size as f64 * rpi).max(1000.0) as usize;
self.per_alpha = (self.per_alpha * rpi).clamp(0.0, 1.0);
self.per_beta_start = (self.per_beta_start * rpi).clamp(0.0, 1.0);
self.n_steps = (self.n_steps as f64 * rpi).clamp(1.0, 20.0) as usize;

// Architecture family
let ari = self.architecture_intensity;
self.hidden_dim_base = ((self.hidden_dim_base as f64 * ari) as usize).max(64);
self.dueling_hidden_dim = ((self.dueling_hidden_dim as f64 * ari) as usize).max(32);
// num_atoms must stay odd for C51 symmetry
let scaled_atoms = (self.num_atoms as f64 * ari) as usize;
self.num_atoms = (scaled_atoms | 1).max(11); // ensure odd, minimum 11

// Risk family
let rki = self.risk_intensity;
self.max_position_absolute *= rki;
self.huber_delta *= rki;
self.minimum_profit_factor = 1.0 + (self.minimum_profit_factor - 1.0) * rki; // scale above 1.0
self.w_dd *= rki;
self.dd_threshold *= rki;
  • Step 3: Restructure PSO vector from 30D to 14D

In crates/ml/src/hyperopt/adapters/dqn.rs:

Replace the current 30D continuous_bounds() with 14D:

0: gamma                       [0.95, 0.999]      — breakout
1: iqn_lambda                  [0.0, 1.0]         — breakout
2: c51_warmup_epochs           [3.0, 15.0]        — breakout
3: learning_intensity          [0.0, 2.0]         — core family
4: exploration_intensity       [0.0, 2.0]         — core family
5: replay_intensity            [0.0, 2.0]         — core family
6: architecture_intensity      [0.0, 2.0]         — core family
7: risk_intensity              [0.0, 2.0]         — core family
8: adversarial_intensity       [0.0, 2.0]         — generalization family
9: regularization_intensity    [0.0, 2.0]         — generalization family
10: augmentation_intensity     [0.0, 2.0]         — generalization family
11: loss_shaping_intensity     [0.0, 2.0]         — generalization family
12: ensemble_intensity         [0.0, 2.0]         — generalization family
13: causal_intensity           [0.0, 2.0]         — generalization family

Update from_continuous() to construct DQNParams from 14D vector: set the 3 breakout params directly, set all 11 family intensities, and leave all other params at their TOML defaults (loaded from the training profile).

Update to_continuous() and param_names() to match the 14D layout.

BACKWARD COMPATIBILITY: The current hyperopt results JSON files store 30D vectors. Either:

  • Add a version field to detect old vs new format

  • Or read old results as 30D with the extra dims ignored

  • Step 4: Update train_with_params() to use family scaling

The call to hyperparams.apply_family_scaling() is already in place from Task 10. Verify it still runs after the restructured PSO vector is mapped to DQNHyperparameters.

  • Step 5: Add to TOML configs

Add under [advanced] or a new [hyperopt_families] section:

learning_intensity = 1.0
exploration_intensity = 1.0
replay_intensity = 1.0
architecture_intensity = 1.0
risk_intensity = 1.0
  • Step 6: Compile check
cargo check --workspace 2>&1 | head -80
  • Step 7: Run existing hyperopt tests
SQLX_OFFLINE=true cargo test -p ml --lib -- hyperopt --nocapture 2>&1 | tail -30

Verify the search space test (test_hyperopt_search_space_22d or similar) passes with the new 14D layout. Update the test if it hardcodes the old dimension count.

  • Step 8: Commit
git add -u crates/ml/ config/training/
git commit -m "feat: core hyperopt families — 30D → 14D PSO with breakout dimensions"

Phase 7: OFI + Trade Features — Wire Existing Infrastructure into DQN

Motivation

The crates/ml-features/ crate already has a complete OFI calculator (879 lines, 8 features from Cont/Xu/Easley/Kyle research) and MBP-10 loader (554 lines). Trade data loader (159 lines) also exists. But these features aren't wired into the DQN training pipeline — the 42-dim feature vector doesn't include them.

MBP-10 data exists locally at test_data/mbp10/ (2.6G, 10 days) and test_data/futures-baseline-mbp10/ (14G, ES Q2 2024).

Design

Expand the DQN feature vector from 42 → 50 dimensions:

  • Dims 0-39: existing OHLCV-derived technical features
  • Dim 40: ADX (regime indicator)
  • Dim 41: CUSUM (regime indicator)
  • Dims 42-49: 8 OFI features (OFI L1, OFI L5, depth imbalance, VPIN, Kyle's lambda, bid slope, ask slope, trade imbalance)

The bottleneck [42→2] becomes [50→2] — just a wider input, no architecture change. state_dim is already configurable via DQNConfig.state_dim.

Task 13: Wire OFI features into DQN data loading

Files:

  • Modify: crates/ml/src/trainers/dqn/data_loading.rs — add OFI extraction step after OHLCV features
  • Modify: crates/ml/src/features/extraction.rs — extend feature extraction to include OFI
  • Modify: crates/ml/src/trainers/dqn/config.rs — update state_dim default from 42 to 50
  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — update market_dim from 42 to 50

Steps:

  • Read data_loading.rs to understand how the 42-dim features are extracted from OHLCV bars
  • Find where OFICalculator is already referenced in the codebase (it has 5 files that import it)
  • In the data loading pipeline, after OHLCV feature extraction, check if MBP-10 data is available for the same time range
  • If available: compute 8 OFI features per bar using OFICalculator, append to feature vector
  • If not available: append 8 zeros (graceful degradation — model learns to ignore zero features)
  • Update state_dim and market_dim constants/defaults
  • Update the regime classification indices (ADX stays at 40, CUSUM at 41, OFI starts at 42)
  • Compile check and verify smoke tests still pass with OHLCV-only data (zeros for OFI)

Task 14: Regime-normalized Sharpe + gap detection

Steps:

  • After regime distribution is logged per fold (Task 12), compute regime difficulty weight: difficulty = 1.0 + 0.5 * volatile_pct/100 - 0.3 * trending_pct/100
  • Normalized Sharpe = raw Sharpe * difficulty (volatile folds score higher, trending folds lower)
  • Compute R² between regime distribution and OOS Sharpe across folds
  • Log verdict: "IS→OOS gap is regime-driven (R²={:.2})" or "IS→OOS gap is model-driven (R²={:.2})"

Task 15: Per-fold regime_replay_decay

Steps:

  • Before each fold's training, compute that fold's validation regime distribution
  • Set regime_replay_decay proportional to regime concentration: decay = base_decay * (1.0 - max_regime_pct/100)
  • If fold is 80% trending: decay = 0.3 * 0.2 = 0.06 (heavily bias toward trending)
  • If fold is balanced (33/33/33): decay = 0.3 * 0.67 = 0.2 (moderate bias)
  • Pass per-fold decay to the PER sampler

Verification Checklist

After all phases:

  • SQLX_OFFLINE=true cargo check --workspace — full workspace compiles
  • SQLX_OFFLINE=true cargo test -p ml --lib — unit tests pass
  • SQLX_OFFLINE=true cargo test -p ml-dqn --lib — agent tests pass
  • FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture — smoke tests pass
  • grep -r 'enable_\|use_' config/training/ — no feature flags in TOML configs
  • grep -rn 'use_soft_updates\|use_iqn\|enable_q_value_clipping\|use_cvar_action_selection\|use_count_bonus' crates/ml-dqn/src/dqn.rs — zero matches
  • Production TOML has spectral_decoupling_lambda, mixup_alpha, regime_replay_decay, and all 11 family intensity scalars
  • PSO search space is 14D (3 breakouts + 11 families) — verify with cargo test -p ml --lib -- hyperopt
  • apply_family_scaling() covers all 11 families (6 generalization + 5 core)