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

fix(direction+magnitude): complete C51 variance bias elimination

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 01:02:20 +02:00
parent 5ba9f376c4
commit 5d5a2c1f3d
7 changed files with 162 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -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++) {

6
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -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];

120
crates/ml/src/cuda_pipeline/expected_q_kernel.cu
View File

@@ -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);

77
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -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. */

10
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
View File

@@ -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++) {

36
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -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;

26
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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