feat(sp20): c51_loss_batched aux_conf_at_state reward gate

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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-10 14:46:00 +02:00
parent d3a057af4f
commit 96b76d9298
2 changed files with 112 additions and 1 deletions

View File

@@ -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. */

View File

@@ -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