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

// rl_q_pi_agree_b.cu — per-batch fraction of (argmax_a Q(s, a) ==
// argmax_a π(a|s)) folded into a streaming EMA at
// ISV[RL_Q_ARG_VS_PI_AGREE_INDEX = 407].
//
// Wires the previously-dead diagnostic slot 407. The slot's docstring
// promised "alarm fires if value < 0.5" but no producer ever wrote to
// it across all 16+ smokes through commit fe58b9f60 — slot stayed at
// sentinel 0.0 forever.
//
// What it measures: how often the DQN-derived greedy action agrees
// with the PPO-derived mode action. High agreement (≈1.0) = Q and π
// have learned the same policy ranking → they're consistent. Low
// agreement (<0.5) = Q and π disagree → either Q is wrong (PPO is
// outracing it) or π is wrong (DQN is more grounded). Either way an
// alarm-worthy diagnostic.
//
// Design: per-batch argmax over the q_logits / pi_logits tensors,
// then count matches, divide by b_size, EMA into slot 407 with
// LR_LOSS_EMA_ALPHA = 0.05 (slow smoothing, same as the LR
// controller's loss EMA — agreement signal is slow-moving).
//
// At b_size=1 this is trivial: one comparison per step, EMA'd over
// time. At b_size>1 it gives the fraction across the batch.
//
// Per `pearl_first_observation_bootstrap`: sentinel-zero state means
// "first observation"; replace directly with the current step's
// agreement (don't blend against zero).
//
// Per `pearl_no_atomicadd`: single-block kernel; tree reduce over
// per-thread argmax-match indicators.

#define RL_Q_ARG_VS_PI_AGREE_INDEX 407
#define Q_N_ATOMS                  21
#define N_ACTIONS                   11
#define AGREE_EMA_ALPHA            0.05f

extern "C" __global__ void rl_q_pi_agree_b(
    const float* __restrict__ q_logits,   // [b_size × N_ACTIONS × Q_N_ATOMS]
    const float* __restrict__ pi_logits,  // [b_size × N_ACTIONS]
    float*       __restrict__ isv,        // [≥ 408]
    int                       b_size
) {
    extern __shared__ float s_match[];
    const int tid = threadIdx.x;

    // Per-thread: compare argmax(Q[tid]) to argmax(π[tid]) if tid < b_size,
    // else contribute 0.0 (mean is taken over b_size real entries below).
    float match = 0.0f;
    if (tid < b_size) {
        // Q argmax = argmax_a Σ_z atom[z] · softmax(Q_logits[a, z])
        // Simpler approximation: argmax_a max_z Q_logits[a, z] (atom-mode).
        // For the agreement diagnostic this is sufficient — we only need
        // ranking consistency, not calibrated expectations.
        int   q_arg = 0;
        float q_max = -INFINITY;
        for (int a = 0; a < N_ACTIONS; ++a) {
            float ma = -INFINITY;
            for (int z = 0; z < Q_N_ATOMS; ++z) {
                const float v = q_logits[tid * N_ACTIONS * Q_N_ATOMS
                                          + a * Q_N_ATOMS + z];
                if (v > ma) ma = v;
            }
            if (ma > q_max) { q_max = ma; q_arg = a; }
        }

        int   pi_arg = 0;
        float pi_max = pi_logits[tid * N_ACTIONS + 0];
        for (int a = 1; a < N_ACTIONS; ++a) {
            const float v = pi_logits[tid * N_ACTIONS + a];
            if (v > pi_max) { pi_max = v; pi_arg = a; }
        }

        match = (q_arg == pi_arg) ? 1.0f : 0.0f;
    }
    s_match[tid] = match;
    __syncthreads();

    // Tree reduce sum (caller enforces block_dim is power of 2 ≥ b_size).
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_match[tid] += s_match[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float agree_frac = s_match[0] / (float)b_size;
        const float prev       = isv[RL_Q_ARG_VS_PI_AGREE_INDEX];
        float out;
        if (prev == 0.0f) {
            // First-observation replace per pearl_first_observation_bootstrap.
            out = agree_frac;
        } else {
            out = (1.0f - AGREE_EMA_ALPHA) * prev
                + AGREE_EMA_ALPHA * agree_frac;
        }
        isv[RL_Q_ARG_VS_PI_AGREE_INDEX] = out;
    }
}