feat(sp20): allocate trainer aux_conf_at_state_buf + wire PER direct-gather

Phase 5 plumbing — consumer-side wire-up. Allocates `aux_conf_at_state_buf:
CudaSlice<f32>` ([batch_size]) on `GpuDqnTrainer`, exposes the raw_ptr via
`aux_conf_at_state_buf_ptr()` and the `FusedTrainerCtx` delegating accessor
`trainer_aux_conf_at_state_buf_ptr()`, and invokes
`GpuReplayBuffer::set_trainer_aux_conf_ptr` at both fused_ctx init sites in
training_loop.rs (init + re-init, atomic per `feedback_no_partial_refactor`).

The PER `gather_f32_scalar` now writes the SAMPLED bar's per-batch aux_conf
directly into the trainer's f32 buffer on every step — same direct-to-trainer
pattern as the SP13 B1.1b `aux_nb_label_buf` (i32) wire-up immediately above
the new call. The c51_loss_batched reward gate (lands in the next commit)
reads this buffer to compute `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD])
/ ISV[AUX_GATE_TEMP])` and applies `r_used = gate * reward` at the Bellman
projection.

`alloc_zeros` cold-start: 0.0 sentinel → at threshold ≈ 0.10 and temp ≈ 0.05
the gate is `sigmoid(-2) ≈ 0.12` → reward is mostly suppressed pre-population.
This is the "graceful degradation" semantic from the Phase 3 Task 3.4 audit
doc spec §4.4. Once PER's direct-gather populates from the producer ring on
the first sample step, the per-bar aux_conf values drive the gate as designed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -6160,6 +6160,27 @@ pub struct GpuDqnTrainer {
     /// `[max_aux_param_len]` reusable per-tensor reduce target for
     /// `aux_param_grad_reduce`. Sized to the largest of the 8 aux tensors.
     aux_param_grad_final_buf: CudaSlice<f32>,
+    /// SP20 Phase 5 (2026-05-10): per-sample aux confidence at the SAMPLED
+    /// state — `[batch_size]` f32, range `[0, 0.5]` (K=2 peak-softmax-above-
+    /// uniform; 0.0 = uniform / no information sentinel; 0.5 = peak softmax
+    /// at 1.0). Populated by direct-to-trainer PER gather from the replay
+    /// buffer's `aux_conf` ring (filled by `experience_env_step` per (i, t)
+    /// in Phase 3 Task 3.4). Wired in at trainer init via
+    /// `GpuReplayBuffer::set_trainer_aux_conf_ptr(self.aux_conf_at_state_buf
+    /// .raw_ptr())` so PER's `gather_f32_scalar` writes the SAMPLED bar's
+    /// per-batch aux_conf directly into here on every step (mirrors the
+    /// SP13 B1.1b `aux_nb_label_buf` direct-to-trainer pattern). Consumed
+    /// by `c51_loss_batched`'s reward gate at the Bellman projection:
+    /// `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`,
+    /// `r_used = gate × reward`. NULL-tolerant (gate=1.0 = no-op) at the
+    /// kernel level so test scaffolds without a wired aux head still work.
+    /// `alloc_zeros` cold-start: 0.0 sentinel → at threshold ≈ 0.10 and
+    /// temp ≈ 0.05 the gate is `sigmoid(-2) ≈ 0.12` → reward is mostly
+    /// suppressed pre-population — the "graceful degradation" semantic
+    /// from the Phase 3 Task 3.4 audit doc spec §4.4. Once the PER
+    /// gather populates it on the first sample step, the per-bar values
+    /// from the producer drive the gate as designed.
+    aux_conf_at_state_buf: CudaSlice<f32>,
     /// Per-step aux-loss weight, baked into the SAXPY alpha at graph-capture
     /// time. ISV-driven via `set_aux_weight`:
     ///   `aux_weight = clamp(0.1 × LEARNING_HEALTH × (1 - tanh(SHARPE_SCALE × sharpe_ema)), 0.05, 0.3)`.
@@ -22564,6 +22585,12 @@ impl GpuDqnTrainer {
             .max(aux_kr)
             .max(aux_knb);
         let aux_param_grad_final_buf = alloc_f32(&stream, max_aux_tensor_len, "aux_param_grad_final_buf")?;
+        // SP20 Phase 5: per-sample aux_conf at sampled state ([B] f32). PER's
+        // direct-to-trainer `gather_f32_scalar` writes here on every sample
+        // step once `set_trainer_aux_conf_ptr` is wired; until then, the
+        // alloc_zeros 0.0 sentinel drives the c51_loss reward-gate to ~0
+        // (graceful degradation per Phase 3 Task 3.4 audit doc spec §4.4).
+        let aux_conf_at_state_buf  = alloc_f32(&stream, aux_b, "aux_conf_at_state_buf")?;
         // Initial aux_weight at the lower numerical-stability clamp (0.05).
         // training_loop's per-step `set_aux_weight` recomputes from ISV every
         // step. The first graph capture bakes 0.05 — an honest cold-start
@@ -24975,6 +25002,7 @@ impl GpuDqnTrainer {
             aux_partial_rg_w2,
             aux_partial_rg_b2,
             aux_param_grad_final_buf,
+            aux_conf_at_state_buf,
             aux_weight,
             stochastic_depth_scale_buf,
             stochastic_depth_kernel,
@@ -33800,6 +33828,17 @@ impl GpuDqnTrainer {
     /// in `gpu_experience_collector.rs`) into here on every step.
     pub(crate) fn aux_nb_label_buf_ptr(&self) -> u64 { self.aux_nb_label_buf.raw_ptr() }
 
+    /// SP20 Phase 5 (2026-05-10): raw pointer to `aux_conf_at_state_buf`
+    /// for direct-to-trainer PER gather. Mirrors the SP13 B1.1b
+    /// `aux_nb_label_buf_ptr` direct-to-trainer pattern. The buffer is
+    /// `CudaSlice<f32>` `[batch_size]`; the c51 reward-gate consumer
+    /// (`c51_loss_batched`) reads it as f32 to compute
+    /// `gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`.
+    /// `GpuReplayBuffer::set_trainer_aux_conf_ptr` is invoked once at
+    /// trainer init in `training_loop.rs` so PER's `gather_f32_scalar`
+    /// writes the SAMPLED bar's aux_conf into here on every step.
+    pub(crate) fn aux_conf_at_state_buf_ptr(&self) -> u64 { self.aux_conf_at_state_buf.raw_ptr() }
+
     /// Padded state dimension (128-byte aligned) for direct-to-trainer gather.
     pub(crate) fn state_dim_padded(&self) -> usize { ml_core::state_layout::STATE_DIM_PADDED }
 

crates/ml/src/trainers/dqn/fused_training.rs

@@ -4160,6 +4160,19 @@ impl FusedTrainingCtx {
         self.trainer.aux_nb_label_buf_ptr()
     }
 
+    /// SP20 Phase 5 (2026-05-10): raw pointer to trainer's
+    /// `aux_conf_at_state_buf` for direct-to-trainer PER gather. Wraps
+    /// the inner `GpuDqnTrainer::aux_conf_at_state_buf_ptr()` accessor —
+    /// same delegation pattern as `trainer_aux_sign_labels_buf_ptr`.
+    /// Caller (`training_loop.rs`) pipes this into
+    /// `GpuReplayBuffer::set_trainer_aux_conf_ptr` so PER's
+    /// `gather_f32_scalar` writes the per-batch sampled aux_conf
+    /// straight into the trainer's f32 buffer the c51_loss reward-gate
+    /// consumer reads.
+    pub(crate) fn trainer_aux_conf_at_state_buf_ptr(&self) -> u64 {
+        self.trainer.aux_conf_at_state_buf_ptr()
+    }
+
     /// Padded state dimension for direct-to-trainer gather.
     pub(crate) fn trainer_state_dim_padded(&self) -> usize {
         self.trainer.state_dim_padded()

crates/ml/src/trainers/dqn/trainer/training_loop.rs

@@ -822,6 +822,19 @@ impl DQNTrainer {
                         fused.trainer_aux_sign_labels_buf_ptr(),
                         fused.trainer_state_dim_padded(),
                     );
+                    // SP20 Phase 5 (2026-05-10): wire trainer's
+                    // `aux_conf_at_state_buf` raw_ptr so PER's
+                    // `gather_f32_scalar` writes the SAMPLED bar's
+                    // per-batch aux_conf directly into the trainer's
+                    // f32 buffer the c51_loss reward-gate consumer
+                    // reads. Same direct-to-trainer pattern as the
+                    // SP13 B1.1b `aux_nb_label_buf` (i32) wire-up
+                    // immediately above; the only difference is the
+                    // element type (f32 vs i32) and the consumer
+                    // (c51_loss_batched gate vs aux_next_bar CE).
+                    agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr(
+                        fused.trainer_aux_conf_at_state_buf_ptr(),
+                    );
                     fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr());
                 }
             }
@@ -2629,6 +2642,14 @@ impl DQNTrainer {
                         fused.trainer_aux_sign_labels_buf_ptr(),
                         fused.trainer_state_dim_padded(),
                     );
+                    // SP20 Phase 5 (2026-05-10): re-wire trainer's
+                    // `aux_conf_at_state_buf` raw_ptr after fused_ctx
+                    // re-init. Mirrors the init site above; both call
+                    // sites must stay in sync per
+                    // `feedback_no_partial_refactor`.
+                    agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr(
+                        fused.trainer_aux_conf_at_state_buf_ptr(),
+                    );
                     fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr());
                 }
             }

docs/dqn-wire-up-audit.md

@@ -2,6 +2,89 @@
 
 **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 1/4): trainer aux_conf_at_state_buf + PER direct-gather wire-up
+
+Phase 5 consumer-side plumbing — atomic across all struct boundaries per
+`feedback_no_partial_refactor`. Lands the trainer-side buffer the c51_loss
+reward-gate kernel will consume in commit 2/4.
+
+### Trainer-side allocation
+
+`GpuDqnTrainer.aux_conf_at_state_buf: CudaSlice<f32>` — `[batch_size]` f32.
+Allocated via `alloc_f32(&stream, aux_b, "aux_conf_at_state_buf")` in the
+constructor, alongside the existing aux-heads buffers. Sentinel: `0.0`
+(alloc_zeros) — matches the K=2 uniform-prior default the Phase 3 Task 3.4
+producer writes when `aux_logits_per_env == NULL`.
+
+### Trainer-side accessor
+
+`pub(crate) fn aux_conf_at_state_buf_ptr(&self) -> u64` — mirrors the
+`aux_nb_label_buf_ptr` accessor used by SP13 B1.1b's i32 direct-gather. The
+wrapping `FusedTrainerCtx::trainer_aux_conf_at_state_buf_ptr` delegates to
+the inner trainer accessor (same delegation pattern as the 6 sibling getters).
+
+### Replay-buffer wire-up (already plumbed in Phase 3 Task 3.4)
+
+`GpuReplayBuffer::set_trainer_aux_conf_ptr(u64)` was added in Phase 3 Task 3.4
+(comment at `gpu_replay_buffer.rs:1163-1165` explicitly says: "Phase 5 will
+call this once after the trainer's `aux_conf_at_state_buf` is allocated").
+This commit calls it.
+
+### Training loop wire-up (atomic — both fused_ctx init sites)
+
+`crates/ml/src/trainers/dqn/trainer/training_loop.rs` — two call sites mirror
+the existing 6-arg `set_trainer_buffers` block:
+
+  - Line ~826 (init site): `agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr(
+    fused.trainer_aux_conf_at_state_buf_ptr())` immediately after
+    `set_trainer_buffers(...)`.
+  - Line ~2645 (re-init site after fused_ctx tear-down): same call. Both
+    sites stay in sync per `feedback_no_partial_refactor`.
+
+### Default-state semantics ("graceful degradation")
+
+`alloc_zeros` cold-start: 0.0 sentinel. The Phase 5 gate kernel (commit 2/4)
+will compute:
+
+```
+gate = sigmoid((aux_conf - threshold) / temp)
+     = sigmoid((0 - 0.10) / 0.05)         // typical ISV values per Phase 1.3
+     = sigmoid(-2)
+     ≈ 0.12
+```
+
+so `r_used = 0.12 × reward` ≈ "reward mostly suppressed". This is the
+"uncertain-state neutralizer" semantic from Phase 3 Task 3.4 audit doc
+spec §4.4 — pre-population (no aux signal yet), the model gets minimal
+reward feedback at the Bellman target → Q collapses to `γ × Q(s', a')`,
+i.e., "don't update Q on uncertain transitions." Once PER's
+`gather_f32_scalar` populates from the producer ring on the first sample
+step, the per-bar aux_conf values drive the gate as designed.
+
+### Files modified
+
+| File | Status | Purpose |
+|------|--------|---------|
+| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | +field, +alloc, +accessor | Trainer-side buffer ownership |
+| `crates/ml/src/trainers/dqn/fused_training.rs` | +delegating accessor | FusedTrainerCtx wrapper |
+| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | +2 call sites | PER direct-gather wire-up |
+
+### Verification
+
+```
+SQLX_OFFLINE=true cargo check --workspace --examples              # green
+```
+
+Forward-references (commits 2-4 of this Phase 5 lands in this branch):
+
+  - Commit 2/4: `c51_loss_kernel.cu` adds `aux_conf_at_state` arg + gate
+    computation at the Bellman projection. NULL-tolerant (gate=1.0
+    no-op).
+  - Commit 3/4: launcher in `gpu_dqn_trainer.rs::launch_c51_loss` threads
+    `self.aux_conf_at_state_buf.raw_ptr()` as the new last arg; unit
+    tests for the gate formula.
+  - Commit 4/4: audit-doc consolidation + close-out.
+
 ## 2026-05-10 — MAX_UPLOAD_BYTES 2GB → 8GB
 
 `crates/ml/src/cuda_pipeline/mod.rs:136`. L40S (48GB) and H100 (80GB) have plenty of headroom; the 2GB cap was a conservative leftover that tripped on workflow `zgjgc` with 17.8M imbalance bars (3.2GB needed). 8GB accommodates richer bar densities while leaving ~28GB free on L40S after model + activations + workspace. Updates both `DqnGpuData::upload_slices` (training data) and `PpoGpuData::upload` (market data) — same constant, same error message format.