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

fix: skip dueling mean subtraction for exposure branch in C51 loss kernel (i%3 fix)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 10:53:18 +02:00
parent 5fbc6deb66
commit 6243abd971
1 changed files with 28 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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