From 7eccfc53c9f5a1596ad8d8b654e8da8667a79e69 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 21:07:23 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20Q-gap=20floor=20gradient=20=E2=80=94=20?= =?UTF-8?q?perpetual=20action=20differentiation=20pressure?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a spread gradient to the C51 advantage logits that pushes the taken action's distribution toward higher atoms and non-taken toward lower. Scale = inv_batch * delta_z (adaptive to per-sample atom resolution, zero hardcoded constants). This gradient is ORTHOGONAL to the Bellman equation — it depends on atom position, not target match. Active on ALL samples, providing perpetual pressure to differentiate Q-values even when the C51 cross-entropy gradient vanishes at convergence. Prevents the Q-gap plateau where all actions have identical Q-values. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 24 ++++++++++++++++++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 1 + 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index 67f7eb4c6..e3d4dd7d1 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -23,7 +23,8 @@ extern "C" __global__ void c51_grad_kernel( int b0_size, int b1_size, int b2_size, int b3_size, int total_branch_atoms, float entropy_coeff, - const float* __restrict__ branch_scales) /* [B, 4] per-sample per-branch gradient scale */ + const float* __restrict__ branch_scales, /* [B, 4] per-sample per-branch gradient scale */ + const float* __restrict__ per_sample_support) /* [B, 3] per-sample [v_min, v_max, delta_z] */ { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_elems = batch_size * num_atoms; @@ -72,9 +73,30 @@ extern "C" __global__ void c51_grad_kernel( int a_d = branch_actions[d]; /* d_adv: each (b,d,a,j) slot is written by exactly ONE thread — plain write */ + + /* Q-gap floor gradient: pushes taken action's distribution toward + * higher atoms and non-taken toward lower. Creates Q-value spread + * even when the C51 Bellman has converged (zero CE gradient). + * Scale: proportional to atom position (normalized to [-1, +1]). + * Weight: inv_batch * 0.01 (small relative to main gradient). + * The gradient is orthogonal to Bellman — it depends on atom + * position, not on the target match. Active on ALL samples, + * providing perpetual differentiation pressure. */ + float z_norm = 2.0f * (float)j / fmaxf((float)(num_atoms - 1), 1.0f) - 1.0f; + /* Scale proportional to delta_z — tighter atoms = smaller spread needed. + * No hardcoded constants: spread = inv_batch * delta_z (same order as CE grad). */ + float delta_z = per_sample_support[b * 3 + 2]; + float spread_scale = inv_batch * delta_z; + for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); float grad_val = branch_scale * d_combined * dueling_grad; + + /* Q-gap floor: taken action → push toward high z (positive Q), + * non-taken → push toward low z (negative Q). Creates spread. */ + float spread_grad = (a == a_d) ? (spread_scale * z_norm) : (-spread_scale * z_norm * inv_A); + grad_val += spread_grad; + int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j; d_adv_logits[adv_idx] = grad_val; } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7df515b89..c1eb4e58c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -5436,6 +5436,7 @@ impl GpuDqnTrainer { .arg(&total_branch_atoms_i32) .arg(&entropy_coeff) .arg(&self.branch_scales_ptr) + .arg(&self.per_sample_support_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1),