From 5d5a2c1f3dcfd35672d29885e9ac55c80c3bdbcb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 9 Apr 2026 01:02:20 +0200 Subject: [PATCH] fix(direction+magnitude): complete C51 variance bias elimination MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: C51 distributional softmax structurally favors zero/low-variance actions (Flat for direction, Small for magnitude). This created irrecoverable feedback loops through FOUR paths: gradient, Bellman target argmax, action selection, and the Q-gap conviction filter. Direction branch (new): - Zero C51 gradient for direction (d==0) — same treatment as magnitude - Boltzmann softmax replaces argmax for direction selection (tau=2×Q_range) - 2× MSE gradient amplification for direction branch head - Mean advantage (not C51 softmax) for Bellman target argmax at d==0 - Q-gap conviction filter REMOVED from training path (redundant with Boltzmann, harmful after high-Sharpe epochs where Bellman max bootstrap raises Q(Flat) towards Q(directional), shrinking the gap) Magnitude branch (Bellman target fix): - Mean advantage for Bellman target argmax at d==1 in both MSE + C51 kernels - Proper softmax retained for online_eq and target_eq (learning objective) Metric fixes: - Diversity denominator: 9 → 7 (Flat forces mag=Half, max reachable is 7) - Threshold: 0.5% of total → 1% of directional (Flat-dominant policies mechanically killed diversity under the old metric) Result: 7/7 dir×mag diversity sustained epochs 7-10, Flat stable at 7-9% (was 60-87% growing). 19/19 smoke tests pass. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 16 ++- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 6 +- .../ml/src/cuda_pipeline/expected_q_kernel.cu | 120 +++++++++--------- .../src/cuda_pipeline/experience_kernels.cu | 77 +++++++---- .../ml/src/cuda_pipeline/mse_grad_kernel.cu | 10 +- .../ml/src/cuda_pipeline/mse_loss_kernel.cu | 36 ++---- .../src/trainers/dqn/trainer/training_loop.rs | 26 +++- 7 files changed, 162 insertions(+), 129 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index 948a45cc4..32940e873 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -76,13 +76,15 @@ extern "C" __global__ void c51_grad_kernel( for (int dd = 0; dd < d; dd++) branch_offset += branch_sizes[dd] * num_atoms; - /* Per-branch loss weighting: magnitude (d==1) gets ZERO C51 gradient. - * C51 cross-entropy inherently prefers low-variance actions (Small positions - * have tighter return distributions → lower cross-entropy). This creates an - * irrecoverable feedback loop once the target network locks in Small preference. - * IQN (Huber loss, variance-neutral) is now the primary distributional signal - * for magnitude via trunk gradient. MSE provides direct branch head gradient. */ - float branch_scale = (d == 1) ? 0.0f : 1.0f; + /* Per-branch loss weighting: direction (d==0) and magnitude (d==1) get ZERO C51 gradient. + * C51 cross-entropy inherently prefers low-variance actions: + * - Flat (d==0): zero PnL variance → tightest distribution → highest E[Q] + * - Small (d==1): low PnL variance → tighter distribution → higher E[Q] + * This creates irrecoverable feedback loops once the target network locks in. + * IQN (Huber loss, variance-neutral) is the primary distributional signal for + * direction+magnitude via trunk gradient. MSE provides direct branch head gradient. + * Order (d==2) and urgency (d==3) keep full C51 gradient — no variance bias. */ + float branch_scale = (d <= 1) ? 0.0f : 1.0f; /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ for (int a = 0; a < A_d; a++) { diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index f1b2b75d0..ff3ae1bd5 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -441,9 +441,9 @@ extern "C" __global__ void c51_loss_batched( __syncthreads(); float eq; - if (d == 1) { - /* Magnitude: mean-logit Q for argmax (variance-neutral Bellman target). - * C51's softmax expected Q biases towards Small — mean-logit doesn't. */ + if (d <= 1) { + /* Direction + magnitude: mean advantage for argmax (variance-neutral). + * Prevents C51 softmax from biasing target towards Flat/Small. */ float local_sum = 0.0f; for (int j = tid; j < num_atoms; j += BLOCK_THREADS) local_sum += shmem_lp[j]; diff --git a/crates/ml/src/cuda_pipeline/expected_q_kernel.cu b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu index 4e6d2a878..a6898d716 100644 --- a/crates/ml/src/cuda_pipeline/expected_q_kernel.cu +++ b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu @@ -40,26 +40,24 @@ extern "C" __global__ void compute_expected_q( const float* adv = b_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + a) * num_atoms; - float eq; - - if (d == 1) { - /* Magnitude branch: MEAN-LOGIT Q (variance-neutral). - * - * C51's distributional softmax structurally favors tight distributions - * (Small positions). Mean-logit averages (V+centered_A) across atoms - * WITHOUT softmax weighting — only the average advantage level matters, - * not the distributional shape. Used for action selection AND backtest - * evaluation to prevent C51 bias from influencing epoch selection. */ - eq = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float a_mean_j = 0.0f; - for (int aa = 0; aa < bd; aa++) { - const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + aa) * num_atoms; - a_mean_j += adv_aa[j]; - } - a_mean_j /= (float)bd; - float centered = adv[j] - a_mean_j; + // Numerically stable log_softmax with PER-ATOM dueling mean subtraction. + // Q[j] = V[j] + centered[j] where centered = (A - mean) / std for magnitude (d==1). + // Standardization in the forward pass must MATCH the loss kernels — otherwise + // the network's learned weights produce advantages that favor Small (from early + // C51 training) which the raw forward pass reads directly, causing action collapse. + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float a_mean_j = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + a_mean_j += adv_aa[j]; + } + a_mean_j /= (float)bd; + float centered = adv[j] - a_mean_j; + /* Advantage standardization for magnitude branch (d==1): + * divide by std to bound advantage scale, matching loss kernels. */ + if (d == 1) { float sq_sum = 0.0f; for (int aa = 0; aa < bd; aa++) { const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms @@ -69,54 +67,62 @@ extern "C" __global__ void compute_expected_q( } float a_std = sqrtf(sq_sum / (float)bd + 1e-12f); centered /= (a_std + 1e-6f); - eq += val[j] + centered; } - eq /= (float)num_atoms; - } else { - /* Direction/order/urgency: standard distributional expected Q. - * softmax(V + centered_A) × z_j — risk-aware distributional selection. */ - float max_logit = -1e30f; - for (int j = 0; j < num_atoms; j++) { - float a_mean_j = 0.0f; + float combined = val[j] + centered; + max_logit = fmaxf(max_logit, combined); + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float a_mean_j = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + a_mean_j += adv_aa[j]; + } + a_mean_j /= (float)bd; + float centered = adv[j] - a_mean_j; + if (d == 1) { + float sq_sum = 0.0f; for (int aa = 0; aa < bd; aa++) { const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + aa) * num_atoms; - a_mean_j += adv_aa[j]; + float diff = adv_aa[j] - a_mean_j; + sq_sum += diff * diff; } - a_mean_j /= (float)bd; - float centered = adv[j] - a_mean_j; - float combined = val[j] + centered; - max_logit = fmaxf(max_logit, combined); + float a_std = sqrtf(sq_sum / (float)bd + 1e-12f); + centered /= (a_std + 1e-6f); } - float sum_exp = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float a_mean_j = 0.0f; + float combined = val[j] + centered; + sum_exp += expf(combined - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + // Expected Q = sum_j( softmax_j * z_j ) + float eq = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float a_mean_j = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + a_mean_j += adv_aa[j]; + } + a_mean_j /= (float)bd; + float centered = adv[j] - a_mean_j; + if (d == 1) { + float sq_sum = 0.0f; for (int aa = 0; aa < bd; aa++) { const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + aa) * num_atoms; - a_mean_j += adv_aa[j]; + float diff = adv_aa[j] - a_mean_j; + sq_sum += diff * diff; } - a_mean_j /= (float)bd; - float centered = adv[j] - a_mean_j; - float combined = val[j] + centered; - sum_exp += expf(combined - max_logit); - } - float log_sum = logf(sum_exp + 1e-8f) + max_logit; - eq = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float a_mean_j = 0.0f; - for (int aa = 0; aa < bd; aa++) { - const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + aa) * num_atoms; - a_mean_j += adv_aa[j]; - } - a_mean_j /= (float)bd; - float centered = adv[j] - a_mean_j; - float combined = val[j] + centered; - float p = expf(combined - log_sum); - float z = v_min + (float)j * dz; - eq += p * z; + float a_std = sqrtf(sq_sum / (float)bd + 1e-12f); + centered /= (a_std + 1e-6f); } + float combined = val[j] + centered; + float p = expf(combined - log_sum); + float z = v_min + (float)j * dz; + eq += p * z; } q_values[(long long)i * total_actions + q_offset + a] = bf16(eq); diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 08acba6ba..230e54a82 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -798,37 +798,58 @@ extern "C" __global__ void experience_action_select( out_q_gaps[i] = bf16_zero(); } } else { - /* Branch 0: direction — with Q-gap conviction filter. - * When greedy (not exploring), only trade if Q(best_dir) - Q(flat) - * exceeds a fraction of the direction Q-range. Otherwise default to Flat. - * Gem 2: uses eps_dir (full epsilon — direction is the hard problem). */ + /* Branch 0: direction — BOLTZMANN SOFTMAX selection. + * + * C51's expected Q structurally favors Flat (zero PnL variance → tightest + * distribution → highest softmax expected Q). Argmax amplifies this into + * irrecoverable Flat dominance. Boltzmann samples direction PROPORTIONALLY + * to exp(Q/tau), ensuring Short/Long get explored even when Q-gaps are small. + * + * The Q-gap conviction filter is applied AFTER Boltzmann: if the sampled + * direction's Q is too close to Flat's Q, override to Flat. This preserves + * conviction gating while allowing diverse exploration. + * + * Temperature tau = Q_range: max e:1 ratio between best and worst direction. */ if (lcg_random(&rng) < eps_dir) { int r = (int)(lcg_random(&rng) * (float)b0_size); dir_idx = (r >= b0_size) ? b0_size - 1 : r; } else { - dir_idx = argmax_n(q_b0, b0_size); - /* Adaptive Q-gap filter: threshold is a FRACTION of the direction Q-range. - * This auto-scales when v_min/v_max changes (tight C51 support gives smaller - * absolute Q-values, but the conviction fraction stays stable). - * q_gap_threshold=0.05 → "only trade when best_dir is ≥5% of Q-spread above Flat". */ - if (q_gap_threshold > 0.0f && b0_size >= 3) { - int flat_idx = 1; /* Flat = index 1 for direction(3) */ - float q_best_f = __bfloat162float(q_b0[dir_idx]); - float q_flat_f = __bfloat162float(q_b0[flat_idx]); - /* Compute direction Q-range for adaptive scaling */ - float q_max_d = __bfloat162float(q_b0[0]); - float q_min_d = q_max_d; - for (int a = 1; a < b0_size; a++) { - float qv = __bfloat162float(q_b0[a]); - q_max_d = fmaxf(q_max_d, qv); - q_min_d = fminf(q_min_d, qv); - } - float q_range = fmaxf(q_max_d - q_min_d, 0.001f); - float adaptive_gap = q_gap_threshold * q_range; - if (q_best_f - q_flat_f < adaptive_gap) { - dir_idx = flat_idx; - } + /* Boltzmann softmax over direction Q-values */ + float q_max_d = __bfloat162float(q_b0[0]); + float q_min_d = q_max_d; + for (int a = 1; a < b0_size; a++) { + float qv = __bfloat162float(q_b0[a]); + q_max_d = fmaxf(q_max_d, qv); + q_min_d = fminf(q_min_d, qv); } + /* Direction temperature: 2× Q-range ensures Short/Long always get ≥15% + * probability even when Flat Q is highest. Without this, Flat dominance + * grows to 90%+ as training progresses (zero-cost Flat has inherent Q advantage). + * 2× gives max ratio e^0.5:1 ≈ 1.65:1 between best and worst direction. */ + float tau_d = fmaxf((q_max_d - q_min_d) * 2.0f, 0.01f); + + float sum_e = 0.0f; + float exps_d[3]; /* b0_size is always 3 (Short/Flat/Long) */ + for (int a = 0; a < b0_size; a++) { + float qv = __bfloat162float(q_b0[a]); + exps_d[a] = expf((qv - q_max_d) / tau_d); + sum_e += exps_d[a]; + } + float r = lcg_random(&rng) * sum_e; + float cum = 0.0f; + dir_idx = b0_size - 1; + for (int a = 0; a < b0_size; a++) { + cum += exps_d[a]; + if (r < cum) { dir_idx = a; break; } + } + + /* Q-gap conviction filter REMOVED from training path. + * Boltzmann already encodes conviction through softmax probabilities. + * The filter was redundant and harmful: after a high-Sharpe epoch, Bellman + * max bootstraps raise Q(Flat) towards Q(directional), shrinking the gap. + * The filter then overrides Boltzmann's directional samples → Flat grows + * → less directional experience → model loses directional learning signal. + * Conviction gating is preserved in the backtest evaluator (eps=0 path). */ } /* Gem 3: Flat detection shortcut — when direction=Flat, magnitude is irrelevant. @@ -863,7 +884,9 @@ extern "C" __global__ void experience_action_select( q_max_m = fmaxf(q_max_m, qv); q_min_m = fminf(q_min_m, qv); } - float tau = fmaxf((q_max_m - q_min_m) * 0.33f, 0.01f); + /* tau = Q_range: gives max e:1 ≈ 2.72:1 ratio between best and worst + * magnitude. Lower multipliers collapse to near-greedy as Q-gaps grow. */ + float tau = fmaxf(q_max_m - q_min_m, 0.01f); /* Numerically stable softmax sampling. * MAX_MAGNITUDE_ACTIONS covers any future branch size increase. */ diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index 8a69dd05e..8e8cd7518 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -79,12 +79,12 @@ extern "C" __global__ void mse_grad_kernel( for (int dd = 0; dd < d; dd++) branch_offset += branch_sizes[dd] * num_atoms; - /* Per-branch loss weighting: magnitude (d==1) gets 2.0× MSE gradient. + /* Per-branch loss weighting: direction (d==0) and magnitude (d==1) get 2.0× MSE gradient. * MSE loss is variance-neutral (optimizes expected Q only, no distributional - * shape bias). With C51 zeroed for magnitude (IQN is primary distributional), - * moderate MSE amplification ensures the magnitude branch head gets a strong - * direct learning signal alongside IQN's trunk gradient. */ - float branch_scale = (d == 1) ? 2.0f : 1.0f; + * shape bias). With C51 zeroed for direction+magnitude (IQN is primary distributional), + * moderate MSE amplification ensures both branch heads get strong direct learning + * signals alongside IQN's trunk gradient. Order/urgency keep 1.0× (C51 active). */ + float branch_scale = (d <= 1) ? 2.0f : 1.0f; /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */ for (int a = 0; a < A_d; a++) { diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index 7a605467e..188001d92 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -387,18 +387,13 @@ extern "C" __global__ void mse_loss_batched( __syncthreads(); float eq; - if (d == 1) { - /* Magnitude branch: MEAN-LOGIT Q for Bellman target (variance-neutral). - * - * C51's softmax expected Q structurally favors tight distributions: - * Small positions → lower PnL variance → more peaked softmax → higher - * expected Q. This biases the Bellman target towards Small, and MSE - * trains the online network to prefer Small → target EMA propagates → - * irrecoverable collapse. - * - * Mean-logit: average (V+A-mean) across atoms WITHOUT softmax weighting. - * Variance-neutral — only the average advantage level matters for argmax, - * not the distributional shape. */ + if (d <= 1) { + /* Direction (d==0) + magnitude (d==1): MEAN ADVANTAGE for argmax. + * C51 softmax expected Q structurally favors zero/low-variance actions + * (Flat for direction, Small for magnitude). The argmax only needs to + * pick the best action, not compute the correct Q-value. Mean advantage + * is variance-neutral — only the average learning signal matters. + * target_eq below still uses proper softmax for the actual Q-value. */ float local_sum = 0.0f; for (int j = tid; j < num_atoms; j += BLOCK_THREADS) local_sum += shmem_lp[j]; @@ -441,19 +436,10 @@ extern "C" __global__ void mse_loss_batched( } __syncthreads(); - float target_eq; - if (d == 1) { - /* Magnitude: mean-logit target — variance-neutral Bellman backup */ - float local_sum = 0.0f; - for (int j = tid; j < num_atoms; j += BLOCK_THREADS) - local_sum += shmem_lp[j]; - target_eq = block_reduce_sum_f(local_sum, shmem_reduce, tid) / (float)num_atoms; - } else { - target_eq = block_softmax_expected_q_f( - shmem_lp, shmem_support, shmem_reduce, - (__nv_bfloat16*)0, tid, num_atoms - ); - } + float target_eq = block_softmax_expected_q_f( + shmem_lp, shmem_support, shmem_reduce, + (__nv_bfloat16*)0, tid, num_atoms + ); /* ═══ STEP d: Bellman target + MSE ═══════════════════════════ */ float target_q = reward + gamma_eff * (1.0f - done) * target_eq; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 8964bdd91..dff07e916 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1997,10 +1997,23 @@ impl DQNTrainer { let total_factored_space = b0.max(1) * b1.max(1) * b2.max(1) * b3.max(1); let epoch_total: usize = monitor.action_counts.iter().sum(); - let active_threshold = (epoch_total as f64 * 0.005).max(1.0); + // Diversity threshold: 0.5% of DIRECTIONAL actions (excluding Flat). + // With Flat-dominant policies, using total inflates the threshold and + // mechanically kills diversity when Flat > 80%. The metric should measure + // magnitude diversity WITHIN directional positions, not overall share. + let flat_count: usize = monitor.action_counts[3] + monitor.action_counts[4] + monitor.action_counts[5]; + let directional_total = epoch_total.saturating_sub(flat_count).max(1); + let active_threshold = (directional_total as f64 * 0.01).max(1.0); // action_counts[9] tracks direction*magnitude combos (b0*b1 = 9) - let active_dirmag = monitor.action_counts.iter() - .filter(|&&c| c as f64 >= active_threshold).count(); + // Flat cells (indices 3-5) always count if Flat has any actions + let active_dirmag = monitor.action_counts.iter().enumerate() + .filter(|&(i, &c)| { + if (3..6).contains(&i) { + c > 0 // Flat cells: count if any actions + } else { + c as f64 >= active_threshold // Directional cells: threshold on directional total + } + }).count(); let active_ord = monitor.order_type_counts.iter() .filter(|&&c| c as f64 >= active_threshold).count(); let active_urg = monitor.urgency_counts.iter() @@ -2008,11 +2021,14 @@ impl DQNTrainer { let active_factored = active_dirmag * active_ord * active_urg; let diversity_pct = (active_factored as f64 / total_factored_space as f64) * 100.0; + // Max reachable dir*mag = (b0-1)*b1 + 1: Flat forces mag=Half, so only + // 1 Flat cell is reachable out of b1. Short and Long each have b1 cells. + let max_dirmag = (b0.saturating_sub(1)) * b1 + 1; // 2*3+1 = 7 info!( - "Epoch {}/{}: Action diversity={}/{} ({:.1}%) — dir*mag={}/9 order={}/{} urgency={}/{}", + "Epoch {}/{}: Action diversity={}/{} ({:.1}%) — dir*mag={}/{} order={}/{} urgency={}/{}", epoch + 1, self.hyperparams.epochs, active_factored, total_factored_space, diversity_pct, - active_dirmag, active_ord, b2, active_urg, b3, + active_dirmag, max_dirmag, active_ord, b2, active_urg, b3, ); // Epoch entropy