From 96b76d92984a84475063f2be19a2f10fa2adf3f7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 10 May 2026 14:46:00 +0200 Subject: [PATCH] feat(sp20): c51_loss_batched aux_conf_at_state reward gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the Phase 5 consumer kernel-side gate. New kernel arg `const float* __restrict__ aux_conf_at_state` appended to `c51_loss_batched`'s signature. Gate computation runs once per sample at the kernel-entry reward-setup site (after the #27 ensemble- disagreement adjustment), then the gated `reward` propagates through every branch's `block_bellman_project_f` call without per-branch changes. Formula: gate = sigmoid((aux_conf - threshold) / temp) reward = gate * reward where: threshold = ISV[AUX_CONF_THRESHOLD_INDEX=518] temp = max(ISV[AUX_GATE_TEMP_INDEX=519], 1e-3) Mathematical interpretation: at low aux confidence (gate→0), `r_used → 0`, so the Bellman target becomes `gamma * Q(s', a')`. The Q value at the current state collapses toward `gamma * Q(s', a')` — model gets no reward feedback on uncertain transitions. Effectively "don't update Q on uncertain transitions" — the "uncertain-state neutralizer" semantic from the Phase 3 Task 3.4 audit doc spec §4.4. NULL-tolerant: `aux_conf_at_state == NULL` OR `isv_signals == NULL` ⇒ gate skipped (identity, no-op = pre-Phase-5 behaviour). Test scaffolds without a wired aux head still work. Out of scope: `iqn_dual_head_kernel.cu` — IQN is the auxiliary loss, C51 is production. Gating IQN is more complexity for marginal gain. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 59 ++++++++++++++++++- docs/dqn-wire-up-audit.md | 54 +++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 3047b5168..f5bbec251 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -584,7 +584,21 @@ extern "C" __global__ void c51_loss_batched( const float* __restrict__ cvar_alpha_buf, /* [B] per-sample CVaR alpha from risk branch. NULL = use iqn_readiness. */ /* ── B3/G4: ISV signals for health-scaled Expected SARSA temperature ── */ - const float* __restrict__ isv_signals /* [ISV_DIM=23] pinned device-mapped. isv_signals[12] = learning_health. NULL = 0.5 fallback. */ + const float* __restrict__ isv_signals, /* [ISV_DIM=23] pinned device-mapped. isv_signals[12] = learning_health. NULL = 0.5 fallback. */ + + /* ── SP20 Phase 5: per-sample aux confidence at sampled state ── + * `[B]` f32 written by PER's direct-to-trainer `gather_f32_scalar` + * from the replay buffer's `aux_conf` ring (filled by the + * SP20 Phase 3 Task 3.4 `experience_env_step` producer per (i, t)). + * Range `[0, 0.5]` (K=2 peak-softmax-above-uniform); 0.0 = uniform / + * no info sentinel. The reward gate at the Bellman target reads + * `aux_conf[sample_id]` and computes + * `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD_INDEX=518]) + * / max(ISV[AUX_GATE_TEMP_INDEX=519], 1e-3))`, + * then `r_used = gate * reward` before the Bellman projection. + * NULL-tolerant: when NULL (test scaffolds without a wired aux head), + * gate = 1.0 → behaviour is identical to pre-Phase-5 (no-op). */ + const float* __restrict__ aux_conf_at_state /* [B] f32 per-sample aux confidence at sampled state. NULL = no-op (gate = 1.0). */ ) { extern __shared__ float shmem_f[]; @@ -665,6 +679,49 @@ extern "C" __global__ void c51_loss_batched( reward -= ensemble_disagreement_weight * ensemble_std[sample_id]; } + /* ── SP20 Phase 5: aux→Q reward gate ──────────────────────────────── + * Gate the reward (NOT the Bellman target) by aux head confidence at + * the sampled state. At low aux confidence (gate→0), `r_used → 0`, + * so the Bellman target collapses toward `gamma * Q(s', a')` and the + * model gets no reward feedback on uncertain transitions — + * effectively "don't update Q on uncertain transitions" (the + * "uncertain-state neutralizer" semantic per Phase 3 Task 3.4 audit + * doc spec §4.4). + * + * Formula: + * gate = sigmoid((aux_conf - threshold) / temp) + * r_gated = gate * reward + * + * where threshold = ISV[AUX_CONF_THRESHOLD_INDEX=518] + * = clamp(aux_dir_acc_ema - 0.50, 0.01, 0.20) + * temp = max(ISV[AUX_GATE_TEMP_INDEX=519], 1e-3) + * = max(aux_conf_std_ema, 0.01) + * + * (both produced by `sp20_controllers_compute_kernel` per Phase 1.3). + * + * NULL-tolerance: + * - aux_conf_at_state == NULL ⇒ gate = 1.0 (identity, no-op). + * Test scaffolds without a wired aux head get pre-Phase-5 + * behaviour. + * - isv_signals == NULL ⇒ gate = 1.0 (cannot read threshold/temp, + * so degrade to identity rather than producing garbage). + * + * Applied AFTER the #27 ensemble-disagreement adjustment so the + * cascade stays well-defined: ensemble subtracts a magnitude penalty + * from raw reward (semantic: "reduce effective reward magnitude + * when uncertain"); then this gate scales the resulting reward + * toward zero when the aux head can't predict the next bar + * (semantic: "don't trust the reward signal at all when aux is + * lost"). The two mechanisms compose: ensemble damps magnitude, + * aux damps the entire signal. */ + if (aux_conf_at_state != NULL && isv_signals != NULL) { + float aux_conf = aux_conf_at_state[sample_id]; + float threshold = isv_signals[518]; /* AUX_CONF_THRESHOLD_INDEX */ + float temp = fmaxf(isv_signals[519], 1e-3f); /* AUX_GATE_TEMP_INDEX */ + float gate = 1.0f / (1.0f + __expf(-(aux_conf - threshold) / temp)); + reward = gate * reward; + } + /* ── Curiosity Q-penalty base (per-sample; factored per-branch below) ── * Shrinks gamma for novel/uncertain states where the curiosity forward model * has high MSE. Applied as a multiplicative factor per branch below. */ diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index e7e1c3fae..66ebb4f95 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,60 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-10 — SP20 Phase 5 (commit 2/4): c51_loss_kernel reward gate + +Adds the consumer kernel-side gate. New kernel arg `const float* +__restrict__ aux_conf_at_state` appended to `c51_loss_batched`'s signature +(after the existing `isv_signals` arg). Gate computation runs once per +sample at the kernel-entry reward-setup site (immediately after the +existing `#27 ensemble_disagreement` adjustment, so the gate composes +cleanly with the existing reward-shaping cascade): + +```cuda +if (aux_conf_at_state != NULL && isv_signals != NULL) { + float aux_conf = aux_conf_at_state[sample_id]; + float threshold = isv_signals[518]; /* AUX_CONF_THRESHOLD_INDEX */ + float temp = fmaxf(isv_signals[519], 1e-3f); /* AUX_GATE_TEMP_INDEX */ + float gate = 1.0f / (1.0f + __expf(-(aux_conf - threshold) / temp)); + reward = gate * reward; +} +``` + +The gated reward propagates through every branch's `block_bellman_project_f` +call without any per-branch changes — the gate is semantically a per-sample +modifier on the reward signal, which is per-sample by construction. + +### NULL-tolerance contract + +Two NULL paths, both → `gate = 1.0` (identity, no-op = pre-Phase-5 +behaviour): + + - `aux_conf_at_state == NULL` ⇒ test scaffolds without a wired aux head + or pre-Phase-5 launchers still work. + - `isv_signals == NULL` ⇒ cannot read threshold/temp, degrade to + identity rather than producing garbage. + +### IQN: explicitly out of scope + +`iqn_dual_head_kernel.cu` is the auxiliary distributional loss in this +codebase; C51 is the production loss. The Phase 5 gate lives in C51 only +per the user-supplied design rationale — adding it to IQN is more +complexity for marginal gain. + +### Files modified + +| File | Status | Purpose | +|------|--------|---------| +| `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` | +1 kernel arg, +gate computation | Reward-gate consumer | + +### Verification + +``` +SQLX_OFFLINE=true cargo check -p ml --lib # green +``` + +GPU verification deferred to Commit 3/4 (launcher + tests). + ## 2026-05-10 — SP20 Phase 5 (commit 1/4): trainer aux_conf_at_state_buf + PER direct-gather wire-up Phase 5 consumer-side plumbing — atomic across all struct boundaries per