From 3ac6daf3defd54ad1dc438109bdabff4e926bda2 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Tue, 7 Apr 2026 22:38:18 +0200
Subject: [PATCH] feat(reward-v7.1): add exposure_aux_grad_kernel for
 per-branch gradient signal

---
 .../src/cuda_pipeline/dqn_utility_kernels.cu  | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
index 84144dba4..c13dd3ff5 100644
--- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
+++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
@@ -1375,3 +1375,59 @@ extern "C" __global__ void regime_prefix_sum_kernel(
         prefix[r * stride + (i + 1)] = count;
     }
 }
+
+/* ══════════════════════════════════════════════════════════════════════
+ * EXPOSURE AUXILIARY GRADIENT KERNEL (v7.1)
+ *
+ * Breaks the i%3 Q-value degeneracy in branching DQN by giving each
+ * exposure output a UNIQUE gradient via cross-entropy against the
+ * counterfactual best exposure from the experience kernel's CEA loop.
+ *
+ * Follows IQN/CQL pattern: accumulates auxiliary gradients into grad_buf
+ * between C51 backward and Adam. The C51 gradient only differentiates
+ * selected vs non-selected (-1/9 for all 8 non-selected). This kernel
+ * gives each exposure output a target-specific gradient.
+ *
+ * Launch: grid=(ceil(batch_size/256)), block=(256).
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void exposure_aux_grad_kernel(
+    const float* __restrict__ adv_logits,  /* [batch_size, b0_size * num_atoms] exposure branch logits (f32) */
+    const int*   __restrict__ targets,     /* [batch_size] best exposure idx (-1 = skip) */
+    float*       __restrict__ grad_buf,    /* [total_params] gradient accumulator */
+    int          adv_offset,               /* offset into grad_buf for exposure adv weights */
+    int          b0_size,                  /* 9 */
+    int          num_atoms,                /* 51 or 101 */
+    float        aux_weight,               /* scaled weight (decays over epochs) */
+    int          batch_size
+) {
+    int sample = blockIdx.x * blockDim.x + threadIdx.x;
+    if (sample >= batch_size) return;
+
+    int target = targets[sample];
+    if (target < 0 || target >= b0_size) return;  /* skip non-exit steps */
+
+    const float* logits = adv_logits + (long long)sample * b0_size * num_atoms;
+    float inv_batch = aux_weight / (float)batch_size;
+
+    for (int z = 0; z < num_atoms; z++) {
+        /* Stable softmax over b0_size for atom z */
+        float max_l = -1e30f;
+        for (int a = 0; a < b0_size; a++) {
+            float l = logits[a * num_atoms + z];
+            if (l > max_l) max_l = l;
+        }
+        float sum_exp = 0.0f;
+        for (int a = 0; a < b0_size; a++) {
+            sum_exp += expf(logits[a * num_atoms + z] - max_l);
+        }
+        float inv_sum = 1.0f / fmaxf(sum_exp, 1e-10f);
+
+        /* Cross-entropy gradient: softmax(logit) - one_hot(target) */
+        for (int a = 0; a < b0_size; a++) {
+            float prob = expf(logits[a * num_atoms + z] - max_l) * inv_sum;
+            float grad = (prob - ((a == target) ? 1.0f : 0.0f)) * inv_batch;
+            int param_idx = adv_offset + a * num_atoms + z;
+            atomicAdd(&grad_buf[param_idx], grad);
+        }
+    }
+}