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

// bellman_target_projection.cu — categorical (C51) projection of the
// Bellman target Z(s_{t+1}, a*) onto the discrete support [V_MIN, V_MAX]
// with Q_N_ATOMS atoms and Δ_z = (V_MAX - V_MIN) / (Q_N_ATOMS - 1) = 0.1.
//
// Standard C51 target projection (Bellemare et al. 2017):
//   1. Compute target atom values: T_z[j] = r + γ × (1 - done) × atom_value[j]
//      (clamped to [V_MIN, V_MAX])
//   2. For each target atom T_z[j], distribute its probability mass
//      across the two NEAREST support atoms (linear interpolation):
//        l = floor((T_z[j] - V_MIN) / Δ_z)   (clamp to [0, Q_N_ATOMS-1])
//        u = ceil (...)
//        target_dist[l] += target_probs[j] × (u - b_frac)
//        target_dist[u] += target_probs[j] × (b_frac - l)
//      where b_frac = (T_z - V_MIN) / Δ_z.
//
// γ is read from isv[RL_GAMMA_INDEX=400] per pearl_controller_anchors_isv_driven.
// γ is clamped to [0, 1] defensively — first-observation bootstrap (γ = 0)
// is OK (one-step bandit), the controller fills in the real anchor at the
// next emit.
//
// Inputs:
//   target_logits  [B × Q_N_ATOMS] — target-net's atom logits at the
//                                    argmax-Q action of s_{t+1}
//   rewards        [B]              — per-transition reward r_t
//   dones          [B]              — 0/1 done flag (multiplies γ × 0 at
//                                    terminal so V(s') is zeroed)
//   isv            [≥ 401]          — read γ at index 400
//   B                               — batch size
//
// Outputs:
//   target_dist    [B × Q_N_ATOMS]  — projected target distribution
//                                    (sums to 1 per batch row by construction)
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (Q_N_ATOMS, 1, 1)  — one thread per source/target atom
//
// Per `feedback_no_atomicadd.md`: no atomicAdd. Cross-thread accumulation
// uses shared memory + per-source-atom serial writes bracketed by
// __syncthreads — Q_N_ATOMS = 21 inner-loop syncs is negligible (single
// block per batch, small block size).

#define Q_N_ATOMS         21
#define N_ACTIONS          11
// V_MIN/V_MAX/DELTA_Z are now ISV-driven per audit 2026-05-24 followup
// so adaptive reward clamps also lift Q's distributional support.
// Slots RL_C51_V_MIN_INDEX/RL_C51_V_MAX_INDEX are ratchet (monotone-
// grow), so the C51 atom mapping never coarsens — only expands as
// wider trades are observed.
#define RL_C51_V_MAX_INDEX  484
#define RL_C51_V_MIN_INDEX  485
#define RL_GAMMA_INDEX   400


// ─────────────────────────────────────────────────────────────────────
// dqn_select_action_atoms — extract per-batch atom row at the requested
// action index from the full target-net logits tensor.
//
//   in:  full_logits  [B × N_ACTIONS × Q_N_ATOMS]
//        actions      [B]          (0..N_ACTIONS)
//   out: action_logits[B × Q_N_ATOMS]
//
// One block per batch element, Q_N_ATOMS threads per block. Each thread
// copies one atom of the selected action's row.
//
// Out-of-range action indices are clamped to 0 — defensive only; the
// integrated trainer enforces bound-checking at the loader boundary.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void dqn_select_action_atoms(
    const float* __restrict__ full_logits,    // [B × N_ACTIONS × Q_N_ATOMS]
    const int*   __restrict__ actions,        // [B]
    int          B,
    float*       __restrict__ action_logits   // [B × Q_N_ATOMS]
) {
    const int batch = blockIdx.x;
    const int atom  = threadIdx.x;
    if (batch >= B || atom >= Q_N_ATOMS) return;

    int a = actions[batch];
    if (a < 0)            a = 0;
    if (a >= N_ACTIONS)   a = 0;

    const long long src_idx =
        (long long)batch * N_ACTIONS * Q_N_ATOMS
        + (long long)a   * Q_N_ATOMS
        + (long long)atom;
    const long long dst_idx = (long long)batch * Q_N_ATOMS + (long long)atom;
    action_logits[dst_idx] = full_logits[src_idx];
}


extern "C" __global__ void bellman_target_projection(
    const float* __restrict__ target_logits, // [B × Q_N_ATOMS]
    const float* __restrict__ rewards,       // [B]
    const float* __restrict__ dones,         // [B]
    const float* __restrict__ isv,           // ≥ 401
    int          B,
    float*       __restrict__ target_dist    // [B × Q_N_ATOMS]
) {
    const int batch = blockIdx.x;
    const int atom  = threadIdx.x;
    if (batch >= B || atom >= Q_N_ATOMS) return;

    __shared__ float s_softmax[Q_N_ATOMS];
    __shared__ float s_max;
    __shared__ float s_sumexp;
    __shared__ float s_proj[Q_N_ATOMS];

    const int base = batch * Q_N_ATOMS;

    // ── Softmax over target logits (numerically-stable max-subtract) ─
    if (atom == 0) {
        float m = target_logits[base];
        #pragma unroll
        for (int z = 1; z < Q_N_ATOMS; ++z)
            m = fmaxf(m, target_logits[base + z]);
        s_max = m;
    }
    __syncthreads();

    const float e = expf(target_logits[base + atom] - s_max);
    s_softmax[atom] = e;
    s_proj[atom] = 0.0f;  // zero the projection accumulator
    __syncthreads();

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

    const float p = s_softmax[atom] / s_sumexp;

    // ── γ from ISV + done masking ────────────────────────────────────
    const float gamma     = fmaxf(0.0f, fminf(1.0f, isv[RL_GAMMA_INDEX]));
    const float r         = rewards[batch];
    const float gamma_eff = gamma * (1.0f - dones[batch]);

    // ── Adaptive atom span from ISV (audit follow-up). DELTA_Z is
    // recomputed every kernel launch — ratchet semantics in the
    // controller guarantee V_MAX > V_MIN strictly (floored at
    // [-1, +1] baseline) so this is always well-defined.
    const float V_MIN_eff = isv[RL_C51_V_MIN_INDEX];
    const float V_MAX_eff = isv[RL_C51_V_MAX_INDEX];
    const float DELTA_Z   = (V_MAX_eff - V_MIN_eff) / (float)(Q_N_ATOMS - 1);

    // ── Per-source-atom target value + clamp + interpolation indices ─
    const float atom_value = V_MIN_eff + (float)atom * DELTA_Z;
    const float t_z        = r + gamma_eff * atom_value;
    const float t_z_clamp  = fmaxf(V_MIN_eff, fminf(V_MAX_eff, t_z));
    const float b_frac     = (t_z_clamp - V_MIN_eff) / DELTA_Z;
    const int   l          = max(0, min(Q_N_ATOMS - 1, (int)floorf(b_frac)));
    const int   u          = max(0, min(Q_N_ATOMS - 1, (int)ceilf(b_frac)));
    const float frac       = b_frac - (float)l;

    // ── Distribute mass into the two nearest support atoms ───────────
    //
    // Each thread holds (l, u, frac, p) for its OWN source atom. We
    // walk source-atom-index `src` from 0..Q_N_ATOMS and only the
    // matching thread writes into s_proj[l] / s_proj[u], bracketed by
    // __syncthreads so writes are serialised. Q_N_ATOMS = 21 inner-loop
    // syncs per block is negligible and trivially correct (no
    // cross-thread race, no atomicAdd).
    for (int src = 0; src < Q_N_ATOMS; ++src) {
        __syncthreads();
        if (atom == src) {
            if (l == u) {
                s_proj[l] += p;
            } else {
                s_proj[l] += p * (1.0f - frac);
                s_proj[u] += p * frac;
            }
        }
    }
    __syncthreads();
    target_dist[base + atom] = s_proj[atom];
}