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

feat: Q-gap floor gradient — perpetual action differentiation pressure

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 21:07:23 +02:00
parent 6271322b06
commit 7eccfc53c9
2 changed files with 24 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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;
}

1
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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),