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

fix(critical): expected_q kernel per-atom dueling mean — was global mean destroying Q-value differentiation

Browse Source 
The expected_q kernel computed mean_adv as a SINGLE scalar averaged
across ALL atoms AND actions, then subtracted from every logit.
This destroyed per-atom advantage structure, making all actions
produce nearly identical expected Q-values.

Fixed to compute per-atom mean: mean_a(A[a,j]) separately for each
atom j, matching the C51/MSE loss kernels' correct implementation.

This bug affected: backtest evaluation action selection, Q-gap
conviction filter (always 0 → floor 0.25), and Q-value monitoring.
Training loss kernels were NOT affected (they had correct per-atom mean).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 19:15:32 +02:00
parent 8d2817293b
commit 652d78f549
1 changed files with 28 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -40,26 +40,32 @@ 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;
// Compute mean advantage logit sum for this branch (for dueling centering)
float mean_adv_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;
for (int j = 0; j < num_atoms; j++) {
mean_adv_sum += adv_aa[j];
}
}
float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms);
// Numerically stable log_softmax over combined = val[j] + adv[j] - mean_adv_per_atom
// Numerically stable log_softmax with PER-ATOM dueling mean subtraction.
// Q[j] = V[j] + A[a,j] - mean_a(A[*,j]) (mean over actions, separately per atom)
// BUG FIX: was computing a single global mean across ALL atoms AND actions,
// which destroyed per-atom structure and made all actions produce identical Q-values.
float max_logit = -1e30f;
for (int j = 0; j < num_atoms; j++) {
float combined = val[j] + adv[j] - mean_adv_per_atom;
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 combined = val[j] + adv[j] - a_mean_j;
max_logit = fmaxf(max_logit, combined);
}
float sum_exp = 0.0f;
for (int j = 0; j < num_atoms; j++) {
float combined = val[j] + adv[j] - mean_adv_per_atom;
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 combined = val[j] + adv[j] - a_mean_j;
sum_exp += expf(combined - max_logit);
}
float log_sum = logf(sum_exp + 1e-8f) + max_logit;
@@ -67,7 +73,14 @@ extern "C" __global__ void compute_expected_q(
// Expected Q = sum_j( softmax_j * z_j )
float eq = 0.0f;
for (int j = 0; j < num_atoms; j++) {
float combined = val[j] + adv[j] - mean_adv_per_atom;
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 combined = val[j] + adv[j] - a_mean_j;
float p = expf(combined - log_sum);
float z = v_min + (float)j * dz;
eq += p * z;