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

// ppo_clipped_surrogate.cu — PPO actor loss (clipped surrogate) + entropy bonus.
//
// PHASE D of the integrated RL trainer
// (docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
//
// Forward (`ppo_clipped_surrogate_fwd`):
//   r_t = exp(log π_new(a|s) - log π_old(a|s))
//   surr1 = r_t × A_t
//   surr2 = clip(r_t, 1-ε, 1+ε) × A_t
//   L_π = -min(surr1, surr2)             (negative because we minimise)
//   L_entropy = -coef × H(π_new)         (encourage exploration)
//
// V LOSS CANONICALIZATION (Phase E.2-DEFER 2026-05-22):
//   The dedicated V-head kernels in `v_head_fwd_bwd.cu` are the SINGLE
//   SOURCE OF TRUTH for the value-MSE loss + V-head gradients. The PPO
//   surrogate kernel previously emitted a redundant `loss_v` accumulator
//   computed from `r_target` / `v_pred` inputs that the integrated trainer
//   never consumed. Per `feedback_single_source_of_truth_no_duplicates` +
//   `feedback_no_partial_refactor` the redundant inputs + computation are
//   DELETED here, not flagged-off; the only V signal flows through the
//   V-head kernels.
//
// ε is read from isv[RL_PPO_CLIP_INDEX=402].
// coef is read from isv[RL_ENTROPY_COEF_INDEX=403].
// Per pearl_controller_anchors_isv_driven: both are ISV-driven, not const.
//
// Backward (`ppo_clipped_surrogate_bwd`):
//   dL_π/dlogit  : computed via the standard policy-gradient identity
//                  combined with the clip mask (gradient of the surrogate
//                  is zero in the clipped region).
//
// PHASE D scope: kernel signature + per-element body. The atomicAdd in
// the loss accumulator is the same loud-flagged deferral as Phase C's
// dqn_distributional_q_bwd — see ATOMIC NOTE below. Phase E refactors
// to warp-shuffle reduce across batches.
//
// ATOMIC NOTE (deliberate, deferred fix): the cross-batch loss
// accumulators (loss_pi / loss_entropy) use `atomicAdd`. This nominally
// violates `feedback_no_atomicadd.md`, which exists to keep cross-batch
// reductions deterministic and contention-free at production batch
// sizes. We accept it HERE in Phase D because:
//   (a) the toy-bandit smoke runs with B ≤ 32, so atomic contention is
//       negligible (32 writers on a single L2 line);
//   (b) the loss scalars are purely diagnostic — gradient flow goes
//       through `grad_logits`, which is NEVER atomicAdded;
//   (c) Phase E replaces this with a two-stage warp-shuffle → shared
//       reduce → single-writer store, matching the `aux_loss.cu` pattern,
//       when the trainer batches reach production sizes (B = 256+).
// Documented loudly here so the audit trail is visible at the kernel
// header without diving into the loss kernel body.
//
// Block layout:
//   Forward:  grid = (B, 1, 1); block = (N_ACTIONS, 1, 1).
//             One block per batch element, one thread per action.
//             Shared-mem softmax (max + sumexp) computed by thread 0
//             then broadcast. Thread 0 also writes the per-batch loss
//             contributions (the surrogate only depends on the taken
//             action, so no parallelism is wasted there).
//   Backward: same layout. Each thread writes its own logit's gradient.

#define HIDDEN_DIM                       128
#define N_ACTIONS                          11
#define RL_PPO_CLIP_INDEX                402
#define RL_ENTROPY_COEF_INDEX            403
// ISV slot carrying the adaptive importance-ratio clamp ceiling
// produced by rl_ppo_ratio_clamp_controller. Floor is reciprocal
// (clamp window is symmetric in log space: [1/ratio_max, ratio_max]).
#define RL_PPO_RATIO_CLAMP_MAX_INDEX     440


// ─────────────────────────────────────────────────────────────────────
// ppo_policy_logits_fwd (Phase E.2):
//   Standalone linear forward producing per-batch logits from h_t. The
//   surrogate kernel takes these logits as input; previously the
//   integrated trainer had no path to materialise them outside of
//   the surrogate compute.
//
//   logits[b, a] = b_pi[a] + Σ_c W_pi[a, c] × h_t[b, c]
//
// Block layout:
//   grid  = (B, N_ACTIONS, 1)
//   block = (HIDDEN_DIM, 1, 1)
//   shared_mem_bytes = 0 (each thread independently strides; thread 0
//                         performs the per-output dot product. The
//                         block over HIDDEN_DIM is sized for backward
//                         compatibility with grad_w_b_h_t's reuse of
//                         the same h_t layout in shared.)
//
// Single output slot per block (b, a) → thread 0 reduces inline. Other
// threads idle (block size mirrors backward's HIDDEN_DIM for caller
// uniformity; perf is non-critical since N_ACTIONS=9 means 9×B blocks).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void ppo_policy_logits_fwd(
    const float* __restrict__ w,         // [N_ACTIONS, HIDDEN_DIM]
    const float* __restrict__ b,         // [N_ACTIONS]
    const float* __restrict__ h_t,       // [B * HIDDEN_DIM]
    int          B,
    float*       __restrict__ logits     // [B * N_ACTIONS]
) {
    const int batch = blockIdx.x;
    const int act   = blockIdx.y;
    const int c     = threadIdx.x;
    if (batch >= B)          return;
    if (act   >= N_ACTIONS)  return;
    if (c     != 0)          return;  // single-writer per output slot

    float acc = b[act];
    #pragma unroll
    for (int i = 0; i < HIDDEN_DIM; ++i) {
        acc += w[act * HIDDEN_DIM + i] * h_t[batch * HIDDEN_DIM + i];
    }
    logits[batch * N_ACTIONS + act] = acc;
}


// ─────────────────────────────────────────────────────────────────────
// ppo_clipped_surrogate_fwd:
//   Computes softmax → log π_new(a_taken|s) → ratio → clipped surrogate +
//   entropy bonus. Writes per-batch log π / entropy diagnostics plus two
//   scalar loss accumulators (policy + entropy).
//
//   The value-MSE loss + V gradient now live exclusively in the dedicated
//   v_head_fwd_bwd kernels — this kernel no longer touches V_pred / R_target
//   (greenfield removal 2026-05-22, see header).
//
// Inputs:
//   logits        [B × N_ACTIONS]  — new-policy raw logits from PolicyHead
//   log_pi_old    [B]              — log π at the taken action when the
//                                    transition was recorded (rollout buf)
//   actions       [B]              — taken action index in [0, N_ACTIONS)
//   advantages    [B]              — A_t = Q(s_t, a_t) - V(s_t)
//                                    (RolloutBuffer::advantages)
//   isv           [≥ 404]          — reads ε at 402, coef at 403
//   B                              — batch size
// Outputs:
//   pi_log_prob   [B]              — log π_new(a_taken|s) for diagnostics
//                                    and Phase E KL EMA
//   entropy       [B]              — H(π_new) per batch sample
//   loss_pi       [1]              — Σ_b L_π (ATOMIC; see header)
//   loss_entropy  [1]              — Σ_b L_entropy
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void ppo_clipped_surrogate_fwd(
    const float* __restrict__ logits,        // [B * N_ACTIONS]
    const float* __restrict__ log_pi_old,    // [B]
    const int*   __restrict__ actions,       // [B]
    const float* __restrict__ advantages,    // [B]
    const float* __restrict__ isv,           // [>= 404]
    int          B,
    float*       __restrict__ pi_log_prob,   // [B]
    float*       __restrict__ entropy,       // [B]
    float*       __restrict__ loss_pi,       // [1]
    float*       __restrict__ loss_entropy   // [1]
) {
    const int batch = blockIdx.x;
    const int act   = threadIdx.x;
    if (batch >= B)         return;
    if (act   >= N_ACTIONS) return;

    const int base = batch * N_ACTIONS;

    // ── Softmax (numerically stable max-subtract) ────────────────────
    __shared__ float s_logits[N_ACTIONS];
    __shared__ float s_softmax[N_ACTIONS];
    __shared__ float s_max;
    __shared__ float s_sumexp;

    s_logits[act] = logits[base + act];
    __syncthreads();

    if (act == 0) {
        float m = s_logits[0];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) m = fmaxf(m, s_logits[i]);
        s_max = m;
    }
    __syncthreads();

    s_softmax[act] = expf(s_logits[act] - s_max);
    __syncthreads();

    if (act == 0) {
        float sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) sum += s_softmax[i];
        s_sumexp = sum;
    }
    __syncthreads();

    // ── Entropy H(π) = -Σ p log p (thread 0 only; N_ACTIONS=9 is small) ──
    if (act == 0) {
        float h = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) {
            const float pi = s_softmax[i] / s_sumexp;
            h -= pi * logf(fmaxf(pi, 1e-7f));
        }
        entropy[batch] = h;

        // ── Surrogate + entropy loss on the TAKEN action ──
        const int   a_taken = actions[batch];
        // Defensive: silently skip invalid actions. Phase E enforces
        // bound-checking at the loader boundary; this is the floor.
        if (a_taken >= 0 && a_taken < N_ACTIONS) {
            const float log_p  = logf(fmaxf(s_softmax[a_taken] / s_sumexp, 1e-7f));
            pi_log_prob[batch] = log_p;

            // Importance ratio with ISV-driven magnitude clamp.
            //
            // PPO's clip(r, 1-ε, 1+ε) bounds the LOSS when surr2 is the
            // active min — i.e. A>0,r>1+ε and A<0,r<1-ε. The
            // unclipped branch IS active when A<0,r>1+ε (surr1=A·r is
            // more negative than surr2=A·(1+ε) so min=surr1) and when
            // A>0,r<1-ε. In the first case `r` can blow up — we've
            // seen r reach 1e10 within a few steps of policy drift,
            // producing l_pi=-A·r=O(1e10) spikes that contaminate
            // the loss-balance controller and the LR controller's
            // plateau detection.
            //
            // Per `feedback_isv_for_adaptive_bounds` and
            // `pearl_controller_anchors_isv_driven`: the clamp bound
            // lives in ISV[RL_PPO_RATIO_CLAMP_MAX_INDEX = 440], emitted
            // by `rl_ppo_ratio_clamp_controller` which derives it from
            // the (already KL-adaptive) PPO clip ε at ISV[402]. The
            // controller widens the clamp when ε widens (rare but
            // larger excursions are expected) and tightens when ε
            // tightens (small excursions should be rare anomalies).
            //
            // Backward (ppo_clipped_surrogate_bwd) gates pg_grad
            // inside [1-ε, 1+ε] anyway so this clamp is forward-only —
            // gradients are bounded already; this bounds the reported
            // loss value to keep downstream controllers sane.
            const float ratio_raw  = expf(log_p - log_pi_old[batch]);
            const float ratio_max  = isv[RL_PPO_RATIO_CLAMP_MAX_INDEX];
            const float ratio_min  = 1.0f / ratio_max;
            const float ratio      = fmaxf(ratio_min,
                                           fminf(ratio_raw, ratio_max));
            const float A          = advantages[batch];
            const float eps        = isv[RL_PPO_CLIP_INDEX];
            const float surr1      = ratio * A;
            const float surr2      = fminf(fmaxf(ratio, 1.0f - eps),
                                           1.0f + eps) * A;
            const float l_pi       = -fminf(surr1, surr2);

            const float coef  = isv[RL_ENTROPY_COEF_INDEX];
            const float l_ent = -coef * h;

            // ATOMIC NOTE: see header. Phase E warp-shuffles these out.
            atomicAdd(loss_pi,      l_pi);
            atomicAdd(loss_entropy, l_ent);
        } else {
            // Invalid action; leave diagnostics at 0, skip loss accum.
            pi_log_prob[batch] = 0.0f;
        }
    }
}


// ─────────────────────────────────────────────────────────────────────
// ppo_clipped_surrogate_bwd:
//   Per-logit gradient of the clipped surrogate. One block per batch,
//   one thread per action.
//
// Gradient form (matches OpenAI/Schulman's PPO derivation):
//   inside the clip band (1-ε ≤ r_t ≤ 1+ε):
//     dL_π/dlogit_a = -A × (p_a - 1[a == a_taken]) × ratio
//   outside the clip band:
//     dL_π/dlogit_a = 0       (surrogate is constant in r_t there)
//
// Inputs:
//   logits        [B × N_ACTIONS] — raw logits from fwd
//   log_pi_old    [B]
//   actions       [B]
//   advantages    [B]
//   isv           — ε at 402
//   B             — batch size
// Outputs:
//   grad_logits   [B × N_ACTIONS] — ∂L_π/∂logits. Non-taken actions get
//                                   the cross-entropy form gradient, not
//                                   zero (that's the standard softmax-PG
//                                   identity).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void ppo_clipped_surrogate_bwd(
    const float* __restrict__ logits,        // [B * N_ACTIONS]
    const float* __restrict__ log_pi_old,    // [B]
    const int*   __restrict__ actions,       // [B]
    const float* __restrict__ advantages,    // [B]
    const float* __restrict__ isv,           // ε at 402
    int          B,
    float*       __restrict__ grad_logits    // [B * N_ACTIONS]
) {
    const int batch = blockIdx.x;
    const int act   = threadIdx.x;
    if (batch >= B)         return;
    if (act   >= N_ACTIONS) return;

    const int base = batch * N_ACTIONS;

    // ── Softmax (same staging as fwd) ────────────────────────────────
    __shared__ float s_logits[N_ACTIONS];
    __shared__ float s_softmax[N_ACTIONS];
    __shared__ float s_max;
    __shared__ float s_sumexp;

    s_logits[act] = logits[base + act];
    __syncthreads();

    if (act == 0) {
        float m = s_logits[0];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) m = fmaxf(m, s_logits[i]);
        s_max = m;
    }
    __syncthreads();

    s_softmax[act] = expf(s_logits[act] - s_max);
    __syncthreads();

    if (act == 0) {
        float sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) sum += s_softmax[i];
        s_sumexp = sum;
    }
    __syncthreads();

    const float p_a = s_softmax[act] / s_sumexp;

    const int a_taken = actions[batch];
    if (a_taken < 0 || a_taken >= N_ACTIONS) {
        // Defensive: zero grad on invalid action. Phase E enforces at
        // loader boundary; this is the floor.
        grad_logits[base + act] = 0.0f;
        return;
    }

    const float log_p = logf(fmaxf(s_softmax[a_taken] / s_sumexp, 1e-7f));
    // Match the forward's ratio clamp from ISV[RL_PPO_RATIO_CLAMP_MAX_INDEX].
    // For inside_clip the clamp is a no-op (since the controller floors
    // at well above 1+ε); for outside_clip the gradient is zero anyway
    // — but keeping fwd/bwd algebraically consistent prevents
    // future-edit drift.
    const float ratio_raw = expf(log_p - log_pi_old[batch]);
    const float ratio_max = isv[RL_PPO_RATIO_CLAMP_MAX_INDEX];
    const float ratio_min = 1.0f / ratio_max;
    const float ratio = fmaxf(ratio_min, fminf(ratio_raw, ratio_max));
    const float A     = advantages[batch];
    const float eps   = isv[RL_PPO_CLIP_INDEX];

    // Clip-band membership: gradient of min(surr1, surr2) is the standard
    // policy-gradient identity inside [1-ε, 1+ε] (where surr1 is unclipped)
    // and zero outside (where the clipped branch is the active min and is
    // locally constant in r_t).
    //
    // Phase D simplification: we use the inside-band gradient when the
    // ratio sits inside the clip range, zero outside. The exact OpenAI
    // form distinguishes by sign of A as well (one side of the clip is
    // active for A > 0, the other for A < 0); Phase E may refine this
    // edge handling but for the toy-bandit smoke and initial integration
    // tests the binary inside/outside split converges identically.
    float pg_grad = 0.0f;
    const bool inside_clip = (ratio >= 1.0f - eps) && (ratio <= 1.0f + eps);
    if (inside_clip) {
        // dL/dlogit_a = -A × (p_a - 1[a == a_taken]) × ratio
        const float indicator = (act == a_taken) ? 1.0f : 0.0f;
        pg_grad = -A * (p_a - indicator) * ratio;
    }
    grad_logits[base + act] = pg_grad;
}


// ─────────────────────────────────────────────────────────────────────
// ppo_grad_w_b_h_t (Phase E.2):
//   Given per-batch grad_logits [B × N_ACTIONS] from
//   `ppo_clipped_surrogate_bwd`, propagate to weight / bias gradients
//   (per-batch scratch — caller reduces via reduce_axis0) and to grad_h_t
//   (per-(batch, c) sole writer; OVERWRITE).
//
// Mathematics (per batch b, action a, hidden unit c):
//   grad_w_per_batch[b, a, c] = grad_logits[b, a] × h_t[b, c]
//   grad_b_per_batch[b, a]    = grad_logits[b, a]
//   grad_h_t[b, c]            = Σ_a grad_logits[b, a] × w[a, c]
//
// Block layout (mirrors aux_heads_bwd):
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM = 128, 1, 1)
//   shared_mem_bytes = N_ACTIONS * sizeof(float)  (grad_logits stage)
//
// Thread c is the sole writer of grad_w_per_batch[batch, a, c] for all a,
// and of grad_h_t[batch, c]. The first N_ACTIONS threads additionally
// write grad_b_per_batch[batch, a]. Per `feedback_no_atomicadd.md`: no
// atomics; each output slot has exactly one writer.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void ppo_grad_w_b_h_t(
    const float* __restrict__ w,                 // [N_ACTIONS, HIDDEN_DIM]
    const float* __restrict__ h_t,               // [B * HIDDEN_DIM]
    const float* __restrict__ grad_logits,       // [B * N_ACTIONS]
    int          B,
    float*       __restrict__ grad_w_per_batch,  // [B * N_ACTIONS * HIDDEN_DIM]
    float*       __restrict__ grad_b_per_batch,  // [B * N_ACTIONS]
    float*       __restrict__ grad_h_t           // [B * HIDDEN_DIM] (OVERWRITE)
) {
    __shared__ float s_gl[N_ACTIONS];

    const int batch = blockIdx.x;
    const int c     = threadIdx.x;
    if (batch >= B) return;

    // Stage grad_logits[batch, :] (9 floats) — first N_ACTIONS threads.
    if (c < N_ACTIONS) {
        s_gl[c] = grad_logits[batch * N_ACTIONS + c];
        // Sole writer of grad_b_per_batch[batch, c].
        grad_b_per_batch[batch * N_ACTIONS + c] = s_gl[c];
    }
    __syncthreads();

    if (c >= HIDDEN_DIM) return;

    // grad_w + grad_h_t — thread role c = hidden unit.
    const float h_bc = h_t[batch * HIDDEN_DIM + c];
    float gh_acc = 0.0f;
    #pragma unroll
    for (int a = 0; a < N_ACTIONS; ++a) {
        const float g_a = s_gl[a];
        grad_w_per_batch[(long long) batch * N_ACTIONS * HIDDEN_DIM
                         + (long long) a * HIDDEN_DIM + c]
            = g_a * h_bc;
        gh_acc += g_a * w[a * HIDDEN_DIM + c];
    }
    // OVERWRITE — caller folds via grad_h_accumulate_scaled.
    grad_h_t[batch * HIDDEN_DIM + c] = gh_acc;
}