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

// rl_kurtosis_streaming.cu — EMA-streaming kurtosis `M4 / M2²` that
// folds ACROSS STEPS instead of across batch.
//
// Replaces `rl_kurtosis_b` for the b_size=1 hot path of the integrated
// RL trainer. The per-batch reduction was mathematically undefined at
// b_size=1 (kurtosis of a single sample is 0/0), so the td_kurtosis
// EMA stayed identically zero forever and the rl_per_alpha controller
// never received signal. Canonical incident: alpha-rl-mjzfk fold0
// (commit 53aeef099) — 100% of 50k steps had td_kurtosis = 0.
//
// Design: maintain three streaming Welford-EMA moments in ISV slots —
// streaming mean (μ_t), M2 (E[(x-μ_t)²]), M4 (E[(x-μ_t)⁴]). Each step:
//
//   1. Compute the current batch's mean x_t.
//   2. Update streaming mean (sentinel-zero first-obs replace):
//        μ_t = (1-α) μ_{t-1} + α x_t
//   3. Re-deviate against the UPDATED mean (per Welford):
//        δ = x_t − μ_t
//   4. Update streaming M2 and M4:
//        M2_t = (1-α) M2_{t-1} + α δ²
//        M4_t = (1-α) M4_{t-1} + α δ⁴
//   5. Write kurtosis directly to the controller's input slot:
//        out_slot = M4_t / (M2_t² + ε)
//
// For a Gaussian distribution, kurtosis = 3.0 (Pearson). Heavy-tailed
// distributions (e.g. TD errors after Q has started learning) have
// kurtosis > 3 — which is what the rl_per_alpha controller cranks up
// PER's priority exponent α to track. By streaming across STEPS the
// kernel sees the actual per-step variation in td_per_sample's batch
// mean, which is a meaningful kurtosis signal even when each batch
// has only one entry.
//
// α = 0.05 — matches rl_var_over_abs_mean_streaming and
// LR_LOSS_EMA_ALPHA for consistency. The downstream
// `ema_update_per_step` call is REMOVED for this signal — the
// streaming kernel writes the smoothed estimate directly to the
// controller's input ISV slot.
//
// Per `pearl_first_observation_bootstrap` + `pearl_blend_formulas_must_have_permanent_floor`:
// sentinel-zero state triggers replace-direct first-observation
// bootstrap; the M2² denominator is permanent-floored at 1e-12 to
// prevent division blow-up when streaming variance is genuinely zero.
//
// Per `pearl_no_atomicadd`: single-thread kernel.

#define STREAM_ALPHA   0.05f
#define M2_SQ_FLOOR   1e-12f

extern "C" __global__ void rl_kurtosis_streaming(
    const float* __restrict__ x,
    float*       __restrict__ isv,
    int                       b_size,
    int                       out_slot,
    int                       mean_slot,
    int                       m2_slot,
    int                       m4_slot
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    // Current batch mean (trivial at b_size=1).
    float batch_mean = 0.0f;
    for (int b = 0; b < b_size; ++b) batch_mean += x[b];
    batch_mean /= (float)b_size;

    // Streaming mean with first-observation bootstrap.
    const float mean_prev = isv[mean_slot];
    float mean_new;
    if (mean_prev == 0.0f) {
        mean_new = batch_mean;
    } else {
        mean_new = (1.0f - STREAM_ALPHA) * mean_prev
                 + STREAM_ALPHA * batch_mean;
    }
    isv[mean_slot] = mean_new;

    // Welford-EMA on δ² and δ⁴ against the UPDATED mean.
    const float dev  = batch_mean - mean_new;
    const float dev2 = dev * dev;
    const float dev4 = dev2 * dev2;

    const float m2_prev = isv[m2_slot];
    float m2_new;
    if (m2_prev == 0.0f && mean_prev == 0.0f) {
        m2_new = dev2;  // first observation
    } else {
        m2_new = (1.0f - STREAM_ALPHA) * m2_prev
               + STREAM_ALPHA * dev2;
    }
    isv[m2_slot] = m2_new;

    const float m4_prev = isv[m4_slot];
    float m4_new;
    if (m4_prev == 0.0f && mean_prev == 0.0f) {
        m4_new = dev4;  // first observation
    } else {
        m4_new = (1.0f - STREAM_ALPHA) * m4_prev
               + STREAM_ALPHA * dev4;
    }
    isv[m4_slot] = m4_new;

    // Kurtosis = M4 / M2² (Pearson). Permanent-floor on M2² prevents
    // div-by-zero before the streaming variance has accumulated.
    const float m2_sq = m2_new * m2_new;
    isv[out_slot] = m4_new / (m2_sq + M2_SQ_FLOOR);
}