foxhunt/crates/ml-alpha/cuda/rl_q_pi_distill_grad.cu

// rl_q_pi_distill_grad.cu — Q→π KL distillation gradient kernel.
//
// Audit 2026-05-24 (vj5f6 follow-up): the C51 V_MAX lift made Q
// distributional learning calibrated to actual reward magnitudes
// (l_q dropped 100×), but reward economics didn't change because
// per Option B (`pearl_q_thompson_actor_makes_pi_dead_weight`), π
// drives action selection via multinomial sample of softmax(pi_logits)
// and is trained by PPO surrogate using advantage = returns - V.
// V regression doesn't benefit from C51 atom-span calibration, so
// Q's improved knowledge stays trapped in the critic.
//
// Fix: add a KL distillation term to the policy loss that pulls π
// toward a Boltzmann distribution over Q's expected action values:
//
//   E_Q[s,a]   = sum_z(softmax(q_logits[s,a,*])[z] × atom_supports[z])
//   π_target   = softmax(E_Q[s,*] / τ)              over actions
//   L_distill  = KL(π_target || π_new)             (forward KL)
//   ∂L/∂logits = λ × (π_new(a) - π_target(a))      (xent-like)
//
// The kernel ADDS this gradient to pi_grad_logits (additive — does
// not overwrite the existing PPO surrogate gradient). PPO retains
// the dominant learning signal; distill is a soft bias toward Q's
// argmax. Forward KL chosen over reverse so π_target's high-prob
// actions (Q-preferred) dominate the gradient; reverse KL would
// just keep π broad.
//
// Block layout: one block per batch, N_ACTIONS threads. Each thread
// handles one action's E_Q + softmax row + final gradient write.
// Shared mem holds the E_Q vector + π_target + π_new for the
// per-block softmax reductions.
//
// Per `feedback_no_atomicadd`: shared-mem reductions only, no atomics.
// Per `feedback_cpu_is_read_only`: all state in ISV / device buffers.

#include <stdint.h>

#define N_ACTIONS                            11
#define Q_N_ATOMS                            21
#define RL_Q_DISTILL_LAMBDA_INDEX            486
#define RL_Q_DISTILL_TEMPERATURE_INDEX       487
#define RL_Q_DISTILL_KL_EMA_INDEX            488
// audit-isv 2026-05-24: KL_EMA_ALPHA was hardcoded 0.05f — caught
// by scripts/audit-isv.sh on the first manifest-driven dogfood
// run. ISV-resident now per SP20 §0.1.
#define RL_Q_DISTILL_KL_EMA_ALPHA_INDEX      493

extern "C" __global__ void rl_q_pi_distill_grad(
    const float* __restrict__ q_logits,        // [B × N_ACTIONS × Q_N_ATOMS]
    const float* __restrict__ pi_logits,       // [B × N_ACTIONS]
    const float* __restrict__ atom_supports,   // [Q_N_ATOMS]
    float*       __restrict__ isv,             // ISV bus (RW — KL_ema)
    float*       __restrict__ pi_grad_logits,  // [B × N_ACTIONS] — ADD to
    int                       B
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= B || a >= N_ACTIONS) return;

    const float lambda = isv[RL_Q_DISTILL_LAMBDA_INDEX];
    const float tau    = isv[RL_Q_DISTILL_TEMPERATURE_INDEX];

    // Skip work if distillation is disabled (λ=0). Still need to
    // syncthreads to keep the block lockstep — but a single read +
    // early return on all threads is safe since they all read the
    // same λ. (Actually, returning here means subsequent
    // __syncthreads in the kernel hang; safer to keep going with
    // grad=0 contribution.)
    const bool active = (lambda > 0.0f);

    // ── Step 1: each thread computes E_Q for its OWN action ─────
    //
    //   E_Q[a] = sum_z(softmax(q_logits[b,a,*])[z] × atom_supports[z])
    //
    // Numerically-stable softmax: subtract max before exp.
    __shared__ float s_eq[N_ACTIONS];

    const int q_base = (b * N_ACTIONS + a) * Q_N_ATOMS;
    float max_q = q_logits[q_base];
    #pragma unroll
    for (int z = 1; z < Q_N_ATOMS; ++z) {
        max_q = fmaxf(max_q, q_logits[q_base + z]);
    }
    float sum_exp = 0.0f;
    float probs_local[Q_N_ATOMS];
    #pragma unroll
    for (int z = 0; z < Q_N_ATOMS; ++z) {
        probs_local[z] = expf(q_logits[q_base + z] - max_q);
        sum_exp += probs_local[z];
    }
    const float inv_sum = (sum_exp > 1e-9f) ? (1.0f / sum_exp)
                                            : (1.0f / (float)Q_N_ATOMS);
    float e_q = 0.0f;
    #pragma unroll
    for (int z = 0; z < Q_N_ATOMS; ++z) {
        e_q += probs_local[z] * inv_sum * atom_supports[z];
    }
    s_eq[a] = e_q;
    __syncthreads();

    // ── Step 2: thread 0 builds π_target = softmax(E_Q / τ) ─────
    __shared__ float s_pi_target[N_ACTIONS];
    if (a == 0) {
        float max_eq = s_eq[0];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) max_eq = fmaxf(max_eq, s_eq[i]);
        const float tau_safe = fmaxf(tau, 1e-3f);  // guard against τ=0
        float total = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) {
            s_pi_target[i] = expf((s_eq[i] - max_eq) / tau_safe);
            total += s_pi_target[i];
        }
        const float inv = (total > 1e-9f) ? (1.0f / total)
                                          : (1.0f / (float)N_ACTIONS);
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) s_pi_target[i] *= inv;
    }
    __syncthreads();

    // ── Step 3: thread 0 builds π_new = softmax(pi_logits) ──────
    __shared__ float s_pi_new[N_ACTIONS];
    if (a == 0) {
        const int pi_base = b * N_ACTIONS;
        float max_pi = pi_logits[pi_base];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) {
            max_pi = fmaxf(max_pi, pi_logits[pi_base + i]);
        }
        float total = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) {
            s_pi_new[i] = expf(pi_logits[pi_base + i] - max_pi);
            total += s_pi_new[i];
        }
        const float inv = (total > 1e-9f) ? (1.0f / total)
                                          : (1.0f / (float)N_ACTIONS);
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) s_pi_new[i] *= inv;
    }
    __syncthreads();

    // ── Step 4: ADD distill gradient. ───────────────────────────
    //
    //   ∂KL/∂logits[a] = λ × (π_new(a) - π_target(a))
    //
    // Additive so PPO's surrogate gradient (already written) keeps
    // its contribution. Each thread writes one action's grad slot.
    if (active) {
        const int pi_idx = b * N_ACTIONS + a;
        pi_grad_logits[pi_idx] += lambda * (s_pi_new[a] - s_pi_target[a]);
    }

    // ── Step 5 (diag): thread 0 of block 0 EMAs the per-batch
    // mean KL(π_target || π_new) into ISV. Done by block 0 only so
    // we don't race across the grid; we reduce within this block
    // first then commit.
    __shared__ float s_kl;
    if (a == 0) {
        float kl = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) {
            const float pt = s_pi_target[i];
            const float pn = fmaxf(s_pi_new[i], 1e-12f);
            if (pt > 1e-12f) kl += pt * logf(pt / pn);
        }
        s_kl = kl;
    }
    __syncthreads();

    if (b == 0 && a == 0) {
        const float kl_prev = isv[RL_Q_DISTILL_KL_EMA_INDEX];
        const float alpha   = isv[RL_Q_DISTILL_KL_EMA_ALPHA_INDEX];
        const float kl_new  = (kl_prev == 0.0f) ? s_kl
            : (1.0f - alpha) * kl_prev + alpha * s_kl;
        isv[RL_Q_DISTILL_KL_EMA_INDEX] = kl_new;
    }
}