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

feat(magnitude): advantage standardization + 5× entropy boost — prevents Q-value and distribution collapse

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 18:37:55 +02:00
parent e083e84809
commit 8736afb16c
4 changed files with 60 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -41,10 +41,13 @@ extern "C" __global__ void c51_grad_kernel(
* Float exp() handles full range — no bf16 overflow */
float d_combined = isw * (expf(lp) - proj);
/* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) */
/* Entropy regularization — magnitude branch (d==1) gets 5× boost to prevent
* atom distribution collapse. Standardized advantages prevent Q-value scale
* collapse; this prevents atom sharpness collapse. Two-pronged defense. */
if (entropy_coeff > 0.0f) {
float ent_scale = (d == 1) ? 5.0f : 1.0f;
float lp_clamped = fmaxf(lp, -10.0f);
d_combined += entropy_coeff * (1.0f + lp_clamped);
d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
}
/* d_value_logits is f32 — native atomicAdd, no overflow risk. */

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

@@ -335,14 +335,26 @@ 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])
* All 4 branches now have small action counts — mean subtraction is safe. */
/* Dueling: Q[j] = V[j] + A'[a_d,j] where A' is centered (all branches)
* and standardized to unit variance (magnitude branch d==1 only).
* Standardization prevents winner-take-all Q-value collapse where Small
* dominates due to lower PnL variance → tighter C51 distributions. */
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;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean;
float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
if (d == 1) {
float sq_sum = 0.0f;
for (int a = 0; a < n_d; a++) {
float diff = shmem_adv[a * num_atoms + j] - a_mean;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
centered /= (a_std + 1e-6f);
}
shmem_lp[j] = shmem_val[j] + centered;
}
__syncthreads();
@@ -352,17 +364,21 @@ extern "C" __global__ void c51_loss_batched(
* non-Flat direction, weighted at 0.3×. This gives the magnitude branch
* gradient signal on ALL samples, not just the 42% with directional actions. */
if (d == 1 && branch_action[0] == 1) {
/* Magnitude advantage is already in shmem_adv. Compute counterfactual
* Q = V + 0.3 * (A_mag[a_d] - mean(A_mag)) instead of Q = V only.
* The 0.3× weight prevents counterfactual from dominating real gradient. */
float cf_weight = 0.3f;
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;
shmem_lp[j] = shmem_val[j]
+ cf_weight * (shmem_adv[a_d * num_atoms + j] - a_mean);
float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
float sq_sum = 0.0f;
for (int a = 0; a < n_d; a++) {
float diff = shmem_adv[a * num_atoms + j] - a_mean;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
centered /= (a_std + 1e-6f);
shmem_lp[j] = shmem_val[j] + cf_weight * centered;
}
}
__syncthreads();

7
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
View File

@@ -45,6 +45,13 @@ extern "C" __global__ void mse_grad_kernel(
/* Gradient of MSE loss through softmax expectation (float arithmetic) */
float d_combined = isw * td_error * p_j * (z_j - e_q);
/* Magnitude entropy boost: prevent atom distribution collapse for d==1.
* MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */
if (d == 1) {
float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
d_combined += 0.005f * (1.0f + lp_approx);
}
/* Route through dueling: d_value[b,j] += d_combined.
* d_value_logits is f32 — native atomicAdd, no overflow risk. */
atomicAdd(&d_value_logits[b * num_atoms + j], d_combined);

32
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -283,14 +283,26 @@ extern "C" __global__ void mse_loss_batched(
shmem_adv[i] = on_adv_row[d][i];
__syncthreads();
/* Dueling: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j])
* All 4 branches now have small action counts — mean subtraction is safe. */
/* Dueling: Q[j] = V[j] + A'[a_d,j] where A' is centered (all branches)
* and standardized to unit variance (magnitude branch d==1 only).
* Standardization prevents winner-take-all Q-value collapse where Small
* dominates due to lower PnL variance → tighter C51 distributions. */
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;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean;
float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
if (d == 1) {
float sq_sum = 0.0f;
for (int a = 0; a < n_d; a++) {
float diff = shmem_adv[a * num_atoms + j] - a_mean;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
centered /= (a_std + 1e-6f);
}
shmem_lp[j] = shmem_val[j] + centered;
}
__syncthreads();
@@ -300,17 +312,21 @@ extern "C" __global__ void mse_loss_batched(
* non-Flat direction, weighted at 0.3×. This gives the magnitude branch
* gradient signal on ALL samples, not just the 42% with directional actions. */
if (d == 1 && branch_action[0] == 1) {
/* Magnitude advantage is already in shmem_adv. Compute counterfactual
* Q = V + 0.3 * (A_mag[a_d] - mean(A_mag)) instead of Q = V only.
* The 0.3× weight prevents counterfactual from dominating real gradient. */
float cf_weight = 0.3f;
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;
shmem_lp[j] = shmem_val[j]
+ cf_weight * (shmem_adv[a_d * num_atoms + j] - a_mean);
float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
float sq_sum = 0.0f;
for (int a = 0; a < n_d; a++) {
float diff = shmem_adv[a * num_atoms + j] - a_mean;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
centered /= (a_std + 1e-6f);
shmem_lp[j] = shmem_val[j] + cf_weight * centered;
}
}
__syncthreads();