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;