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

// rl_pi_action_kernel.cu — GPU-resident multinomial sampler over the
// PPO policy head's softmax(pi_logits). Replaces rl_action_kernel
// (which sampled from Q-Thompson) as the canonical actor-side
// decision per `pearl_q_thompson_actor_makes_pi_dead_weight`.
//
// The prior architecture had Q being BOTH the actor (via Thompson)
// AND the critic (via Bellman target). π trained as a third head
// via PPO surrogate against Q's actions but never drove any
// decision — `q_pi_agree_ema` decayed to 0 across every smoke
// because π converged to Q's Thompson SAMPLING distribution mode,
// not Q's argmax mode. The actor-critic design was inverted.
//
// New architecture (this kernel): π drives action selection via
// multinomial sampling from softmax(pi_logits). Q remains the
// critic (Bellman target via `argmax_expected_q` on h_{t+1}). The
// PPO importance-ratio surrogate now has its canonical semantics —
// π_new / π_old at the action that π_old sampled.
//
// PRNG: same xorshift32-per-batch pattern as rl_action_kernel.
// Single thread per block (N_ACTIONS=9 is small enough that the
// per-thread parallelism of the Thompson kernel is unjustified;
// multinomial CDF walk is sequential anyway).
//
// Anti-collapse: each action probability is floored at P_MIN before
// the CDF walk, then renormalised. This is a STRUCTURAL guarantee
// per `pearl_blend_formulas_must_have_permanent_floor` — H(π_sample)
// is bounded below regardless of what the logits say. P_MIN=0.02 →
// max-collapse entropy ≈ 0.77 nats (vs uniform ln(9)=2.197), well
// above the δ-function attractor observed in alpha-rl-9k9x6 where
// the policy collapsed to action 1 by step 2500 with action_entropy
// EMA at 0.001 and never recovered (controller pegged at COEF_MAX
// without effect because PPO updates couldn't escape the fixed point).
//
// Discrepancy with PPO importance ratio: log_pi_at_action records
// log(softmax(logits)[a]), not log(p_floored[a]). At P_MIN=0.02 the
// divergence is negligible when π is healthy (max abs Δlog_prob is
// ~0.01 for π(a)=0.99), and during recovery it nudges PPO toward
// conservative updates on rare-action samples — desirable.
//
// P_MIN is exempt from `feedback_isv_for_adaptive_bounds` per the
// user-stated "floors and clamp bounds" exemption; lives here as a
// design constant alongside N_ACTIONS.
//
// Per `feedback_no_atomicadd`: single-thread kernel, no atomics.
// Per `feedback_cpu_is_read_only`: pure device-side sampling.

#include <stdint.h>

#define N_ACTIONS    11
#define P_MIN        0.02f

__device__ static uint32_t xorshift32(uint32_t* state) {
    uint32_t x = *state;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *state = x;
    return x;
}

// One block per batch (grid_dim.x = b_size). Single thread per
// block (tid==0). Computes numerically-stable softmax over the
// batch's N_ACTIONS pi_logits, then multinomial-samples via CDF
// walk using the per-batch xorshift32 PRNG state.
//
// Inputs:
//   pi_logits     [b_size × N_ACTIONS]  row-major raw logits
//   prng_state    [b_size]  per-batch xorshift32 state (mutated)
// Outputs:
//   actions       [b_size]  multinomial-sampled action index per batch
extern "C" __global__ void rl_pi_action_kernel(
    const float* __restrict__ pi_logits,
    uint32_t*    __restrict__ prng_state,
    int*         __restrict__ actions,
    int                       b_size
) {
    const int b = blockIdx.x;
    if (b >= b_size) return;
    if (threadIdx.x != 0) return;

    // ── Numerically-stable softmax over pi_logits[b, :]. ──────────
    const int base = b * N_ACTIONS;
    float logits[N_ACTIONS];
    float max_logit = pi_logits[base];
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        logits[a] = pi_logits[base + a];
        if (logits[a] > max_logit) max_logit = logits[a];
    }
    float sum_exp = 0.0f;
    float probs[N_ACTIONS];
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        probs[a] = expf(logits[a] - max_logit);
        sum_exp += probs[a];
    }
    // Guard against pathological zero-sum (shouldn't happen with
    // expf of bounded values, but defensive).
    const float inv_sum = (sum_exp > 1e-9f) ? (1.0f / sum_exp) : (1.0f / (float)N_ACTIONS);
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) probs[a] *= inv_sum;

    // ── Anti-collapse floor: max(P_MIN, p[a]) then renormalise. ───
    // Structural guarantee that no action probability falls below
    // P_MIN. See header comment for rationale + alpha-rl-9k9x6
    // failure mode.
    float renorm_sum = 0.0f;
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        probs[a] = fmaxf(P_MIN, probs[a]);
        renorm_sum += probs[a];
    }
    const float inv_renorm = 1.0f / renorm_sum;
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) probs[a] *= inv_renorm;

    // ── Draw uniform u ∈ [0, 1) from per-batch PRNG. ──────────────
    uint32_t state = prng_state[b];
    // A few warmup advances dispel low-quality seeding correlations.
    #pragma unroll
    for (int w = 0; w < 4; ++w) xorshift32(&state);
    const uint32_t r = xorshift32(&state);
    // Convert to [0, 1): divide by 2^32. Use 24-bit mantissa precision
    // (good enough for multinomial — Q_N_ATOMS=21 actions max).
    const float u = (float)(r >> 8) * (1.0f / 16777216.0f);
    prng_state[b] = state;

    // ── CDF walk: pick the action whose cumulative prob crosses u. ──
    int chosen = N_ACTIONS - 1;  // defensive default if numerical drift loses the last bit
    float cdf = 0.0f;
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        cdf += probs[a];
        if (u < cdf) { chosen = a; break; }
    }

    actions[b] = chosen;
}