69d8038a80
Root cause: confidence gate decouples policy from actions. Policy softmax stays high-entropy (~2.4) while the gate forces Hold, producing action entropy ~0.7. SAC read policy entropy EMA (slot 420), saw "above target", and kept LOWERING τ — the exact opposite of what was needed. New kernel `action_entropy_per_step` computes H(action_histogram) from the POST-gate actions buffer and writes to ISV slot 583 (RL_ACTION_ENTROPY_EMA_INDEX). SAC co-tuning in rl_q_pi_distill_grad now reads this slot. Launched OUTSIDE CUDA Graph capture (after confidence gate + FRD gate) in both training and prefill paths. Kernel design: 11 threads (N_ACTIONS), each thread counts its action across all b_size elements. Thread 0 computes entropy from the histogram and updates the EMA. No atomics. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
67 lines
2.3 KiB
Plaintext
67 lines
2.3 KiB
Plaintext
// action_entropy_per_step.cu — compute batch-level action entropy from
|
||
// the actions buffer and update an ISV-resident EMA.
|
||
//
|
||
// The confidence gate decouples policy entropy from action entropy:
|
||
// policy softmax can be high-entropy while the gate forces most actions
|
||
// to Hold. SAC α/τ co-tuning must respond to ACTION entropy (what the
|
||
// agent actually does) not policy entropy (what the model predicts).
|
||
//
|
||
// Single block, N_ACTIONS threads. Each thread counts its action in
|
||
// shared memory via tree reduction over the batch, then thread 0
|
||
// computes H(histogram) and updates the EMA.
|
||
//
|
||
// Per feedback_no_atomicadd: uses tree reduction, no atomics.
|
||
|
||
#include <stdint.h>
|
||
|
||
#define N_ACTIONS 11
|
||
|
||
extern "C" __global__ void action_entropy_per_step(
|
||
float* __restrict__ isv, // ISV bus (RW)
|
||
int slot_index,// ISV slot to write
|
||
float alpha, // EMA blend α (≥ 0.4)
|
||
const int* __restrict__ actions, // [b_size] per-batch action indices
|
||
int b_size
|
||
) {
|
||
// Phase 1: each of N_ACTIONS threads counts how many batch elements
|
||
// chose that action. Single-pass scan over the actions buffer.
|
||
__shared__ int counts[N_ACTIONS];
|
||
const int a = threadIdx.x;
|
||
if (a >= N_ACTIONS) return;
|
||
|
||
counts[a] = 0;
|
||
__syncthreads();
|
||
|
||
// Sequential scan — N_ACTIONS threads each check all b_size elements.
|
||
// At b=1024, N_ACTIONS=11: 1024/11 ≈ 93 iterations per thread.
|
||
// No atomics needed since each thread owns its histogram bin.
|
||
for (int b = 0; b < b_size; ++b) {
|
||
if (actions[b] == a) {
|
||
counts[a] += 1;
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
// Phase 2: thread 0 computes entropy from the histogram.
|
||
if (a == 0) {
|
||
float h = 0.0f;
|
||
const float inv_b = 1.0f / (float)b_size;
|
||
for (int i = 0; i < N_ACTIONS; ++i) {
|
||
if (counts[i] > 0) {
|
||
float p = (float)counts[i] * inv_b;
|
||
h -= p * logf(p);
|
||
}
|
||
}
|
||
|
||
// EMA update with first-observation bootstrap.
|
||
const float prev = isv[slot_index];
|
||
if (prev == 0.0f) {
|
||
if (h != 0.0f) {
|
||
isv[slot_index] = h;
|
||
}
|
||
} else {
|
||
isv[slot_index] = (1.0f - alpha) * prev + alpha * h;
|
||
}
|
||
}
|
||
}
|