diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index f7baf1581..68c9748a3 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -22,6 +22,12 @@ */ /* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */ +#ifndef MAX_PER_SAMPLE_CE +#define MAX_PER_SAMPLE_CE 50.0f /* Layer 1: 10x converged loss — breaks PER feedback loop */ +#endif +#ifndef LABEL_SMOOTHING_EPS +#define LABEL_SMOOTHING_EPS 0.01f /* Layer 2: 1% uniform mix — prevents target collapse */ +#endif #ifndef NUM_ATOMS #define NUM_ATOMS 51 #endif @@ -567,7 +573,19 @@ extern "C" __global__ void c51_loss_batched( block_bellman_project(shmem_proj, reward, done, gamma, shmem_lp, shmem_support, shmem_reduce, tid); - /* shmem_lp[0..NUM_ATOMS] = projected target distribution */ + __syncthreads(); + + /* Layer 2: Label smoothing — mix ε uniform into projected target. + * Prevents any atom from having zero probability → no log(0) in CE. */ + { + const float ls_uniform = LABEL_SMOOTHING_EPS / (float)NUM_ATOMS; + const float ls_scale = 1.0f - LABEL_SMOOTHING_EPS; + for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) + shmem_lp[j] = ls_scale * shmem_lp[j] + ls_uniform; + __syncthreads(); + } + + /* shmem_lp[0..NUM_ATOMS] = smoothed projected target distribution */ /* Save projected for backward kernel */ for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) @@ -594,9 +612,15 @@ extern "C" __global__ void c51_loss_batched( /* ── Write outputs (thread 0 only to avoid races) ────────────── */ if (tid == 0) { - float weighted_loss = avg_ce * is_weight; + /* Layer 1: clamp per-sample CE to prevent PER feedback loop. + * High loss → high PER priority → high IS weight → amplified loss → repeat. + * Clamping at 50.0 (10x converged loss of ~5) breaks the loop. + * The gradient itself is bounded (softmax - target ∈ [0,1]), + * so clamping the LOSS only affects PER priority, not gradient direction. */ + float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE); + float weighted_loss = clamped_ce * is_weight; per_sample_loss[sample_id] = weighted_loss; - td_errors[sample_id] = avg_ce; /* unweighted — used by PER for priorities */ + td_errors[sample_id] = clamped_ce; /* PER sees clamped too */ atomicAdd(total_loss, weighted_loss / (float)batch_size); } }