From 6243abd9717988763aa430db98e31fa8f464cb9f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 8 Apr 2026 10:53:18 +0200 Subject: [PATCH] fix: skip dueling mean subtraction for exposure branch in C51 loss kernel (i%3 fix) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dueling A'[i] = A[i] - mean(A) gives identical -1/9 gradient to all 8 non-selected exposure outputs, causing i%3 Q-value degeneracy. For d==0 (exposure, 9 outputs): skip mean subtraction entirely. Non-selected outputs get ZERO C51 gradient — only the separate aux optimizer updates them with unique per-output directional signal. Branches d=1,2 (order=3, urgency=3) keep mean subtraction — the -1/2 gradient (33%) is strong enough to differentiate 3 outputs. 4 locations patched in the C51 loss kernel. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 41 +++++++++++++------ 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 7fb8d4027..3f2b5b4d9 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -326,12 +326,21 @@ extern "C" __global__ void c51_loss_batched( shmem_adv[i] = on_adv_row[d][i]; /* already f32 */ __syncthreads(); - /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) */ + /* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) + * v8: Skip mean subtraction for branch 0 (exposure, 9 outputs). + * Mean subtraction gives identical -1/9 gradient to all 8 non-selected + * outputs, causing i%3 Q-value degeneracy. Without it, non-selected + * outputs get ZERO C51 gradient — only the separate aux optimizer + * updates them, providing unique per-output directional signal. + * Branches 1+2 (order=3, urgency=3) keep mean subtraction — the -1/2 + * gradient is strong enough (33%) to differentiate 3 outputs. */ for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; - for (int a = 0; a < n_d; a++) - a_mean += shmem_adv[a * num_atoms + j]; - a_mean /= (float)n_d; + if (d > 0) { /* skip mean for exposure branch (d==0) */ + for (int a = 0; a < n_d; a++) + a_mean += shmem_adv[a * num_atoms + j]; + a_mean /= (float)n_d; + } shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean; } __syncthreads(); @@ -345,9 +354,11 @@ extern "C" __global__ void c51_loss_batched( for (int a = 0; a < n_d; a++) { for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; - for (int aa = 0; aa < n_d; aa++) - a_mean += shmem_adv[aa * num_atoms + j]; - a_mean /= (float)n_d; + if (d > 0) { + for (int aa = 0; aa < n_d; aa++) + a_mean += shmem_adv[aa * num_atoms + j]; + a_mean /= (float)n_d; + } float logit = shmem_val[j] + shmem_adv[a * num_atoms + j] - a_mean; local_sq += logit * logit; } @@ -378,9 +389,11 @@ extern "C" __global__ void c51_loss_batched( for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; - for (int a = 0; a < n_d; a++) - a_mean += shmem_adv[a * num_atoms + j]; - a_mean /= (float)n_d; + if (d > 0) { + for (int a = 0; a < n_d; a++) + a_mean += shmem_adv[a * num_atoms + j]; + a_mean /= (float)n_d; + } shmem_proj[j] = a_mean; } __syncthreads(); @@ -420,9 +433,11 @@ extern "C" __global__ void c51_loss_batched( for (int j = tid; j < num_atoms; j += BLOCK_THREADS) { float a_mean = 0.0f; - for (int a = 0; a < n_d; a++) - a_mean += shmem_adv[a * num_atoms + j]; - a_mean /= (float)n_d; + if (d > 0) { + for (int a = 0; a < n_d; a++) + a_mean += shmem_adv[a * num_atoms + j]; + a_mean /= (float)n_d; + } shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * num_atoms + j] - a_mean; } __syncthreads();