From 652d78f549a293684d7dba753aef5df2dba08781 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 8 Apr 2026 19:15:32 +0200 Subject: [PATCH] =?UTF-8?q?fix(critical):=20expected=5Fq=20kernel=20per-at?= =?UTF-8?q?om=20dueling=20mean=20=E2=80=94=20was=20global=20mean=20destroy?= =?UTF-8?q?ing=20Q-value=20differentiation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../ml/src/cuda_pipeline/expected_q_kernel.cu | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/expected_q_kernel.cu b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu index 88efb286c..4be24193c 100644 --- a/crates/ml/src/cuda_pipeline/expected_q_kernel.cu +++ b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu @@ -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;