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

feat(magnitude): advantage standardization in expected_q forward pass — matches loss kernels

Browse Source 
The forward pass (action selection) must apply the same advantage
standardization as the loss kernels. Otherwise the network's raw
advantages can encode stale C51 preferences that don't match the
standardized training signal.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 22:32:59 +02:00
parent 49bedfcd69
commit 10e445f2c7
1 changed files with 45 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -41,9 +41,10 @@ extern "C" __global__ void compute_expected_q(
+ (long long)(adv_offset + a) * num_atoms;
// 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.
// 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;
@@ -53,7 +54,21 @@ extern "C" __global__ void compute_expected_q(
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float combined = val[j] + adv[j] - a_mean_j;
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
+ (long long)(adv_offset + aa) * num_atoms;
float diff = adv_aa[j] - a_mean_j;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)bd + 1e-12f);
centered /= (a_std + 1e-6f);
}
float combined = val[j] + centered;
max_logit = fmaxf(max_logit, combined);
}
float sum_exp = 0.0f;
@@ -65,7 +80,19 @@ extern "C" __global__ void compute_expected_q(
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float combined = val[j] + adv[j] - a_mean_j;
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;
float diff = adv_aa[j] - a_mean_j;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)bd + 1e-12f);
centered /= (a_std + 1e-6f);
}
float combined = val[j] + centered;
sum_exp += expf(combined - max_logit);
}
float log_sum = logf(sum_exp + 1e-8f) + max_logit;
@@ -80,7 +107,19 @@ extern "C" __global__ void compute_expected_q(
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float combined = val[j] + adv[j] - a_mean_j;
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;
float diff = adv_aa[j] - a_mean_j;
sq_sum += diff * diff;
}
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;