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

// argmax_expected_q.cu — GPU-resident argmax over expected Q-values
// per action (Phase R4 of the integrated RL trainer rebuild;
// see docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md).
//
// Companion to `rl_action_kernel.cu`. Per
// `pearl_thompson_for_distributional_action_selection`: Thompson is
// used for ROLLOUT action selection (rl_action_kernel) and argmax over
// EXPECTED Q is used for the Bellman-target argmax (this kernel). The
// distinction matters for the canonical Double-DQN-style backup:
//
//   target Q = Q_target(s_{t+1}, argmax_a E[Q_online(s_{t+1}, a)])
//
// Replaces the host-side `argmax_f32(&action_expected)` loop the
// flawed Phase F shipped in step_with_lobsim — which violated
// `feedback_cpu_is_read_only` by DtoH-copying q_logits and reducing
// on CPU.
//
// Per-batch: softmax over each action's atoms, accumulate
// `expected[a] = Σ_i p_i · support[i]`, then argmax over the 9
// per-action expected values. Deterministic (no PRNG).
//
// Per `feedback_no_atomicadd`: thread 0 writes to `next_actions[b]`
// after __syncthreads.

#define N_ACTIONS    11
#define Q_N_ATOMS   21

// One block per batch (grid_dim.x = b_size). N_ACTIONS threads per
// block; each thread computes its action's expected Q and writes to
// shared mem; thread 0 argmaxes + writes next_actions[b].
//
// Inputs:
//   q_logits      [b_size, N_ACTIONS, Q_N_ATOMS] row-major
//   atom_supports [Q_N_ATOMS]
// Outputs:
//   next_actions  [b_size] argmax over expected Q per batch
extern "C" __global__ void argmax_expected_q(
    const float* __restrict__ q_logits,
    const float* __restrict__ atom_supports,
    int*         __restrict__ next_actions,
    int                       b_size
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= b_size) return;
    if (a >= N_ACTIONS) return;

    __shared__ float expected_q[N_ACTIONS];

    // Softmax + expected value over this action's atoms.
    const int row_off = (b * N_ACTIONS + a) * Q_N_ATOMS;
    float max_l = -INFINITY;
    #pragma unroll
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        const float l = q_logits[row_off + i];
        if (l > max_l) max_l = l;
    }
    float sum_exp = 0.0f;
    #pragma unroll
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        sum_exp += expf(q_logits[row_off + i] - max_l);
    }
    if (!isfinite(sum_exp) || sum_exp <= 0.0f) sum_exp = 1.0f;

    float ev = 0.0f;
    #pragma unroll
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        const float p = expf(q_logits[row_off + i] - max_l) / sum_exp;
        ev += p * atom_supports[i];
    }
    expected_q[a] = ev;
    __syncthreads();

    if (a == 0) {
        int   best_a = 0;
        float best_v = expected_q[0];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) {
            if (expected_q[i] > best_v) {
                best_v = expected_q[i];
                best_a = i;
            }
        }
        next_actions[b] = best_a;
    }
}