# SP22 H6 — Aux→Policy State Bridge Implementation Runbook

**Author**: jgrusewski + assistant pair
**Date**: 2026-05-12
**Branch**: `sp20-aux-h-fixed`
**Predecessor**: SP22 H1 (falsified — aux learns at H=200 but doesn't propagate)
**Parent plan**: `docs/plans/2026-05-12-sp22-wr-plateau-investigation.md`
**Goal**: Wire aux head's directional probability into policy STATE as input feature, preserving the trunk separation invariant.

## Why this design

Validated by SP22 H1 + memory pearls (see audit `## 2026-05-12 — SP22 H6 design` entry):

1. **Aux PROVES directional signal extractable** — aux_dir_acc 78% at H=200 within-fold (H1 HD[2])
2. **Policy can't extract direction** — WR=50.1% across all v5-v11 conditions
3. **Trunk separation walls aux from policy** — `pearl_separate_aux_trunk_when_shared_starves`: aux on independent trunk with stop-grad
4. **Existing infrastructure** — `exp_aux_nb_softmax_buf: [alloc_episodes × 2]` already exists in `GpuExperienceCollector` (per-env softmax tile written by `aux_next_bar_forward` during rollout)
5. **State slot available** — `PADDING_START=121` reserved 7 slots in current STATE_DIM=128 layout

## Sequence (chicken-and-egg resolution)

At step t in rollout:
1. State[env, t] built using **previous step's** aux signal (lag of 1 bar — fine for H=200 slow-moving signal)
2. Forward: encoder(state) → policy Q-head → action; AND encoder(state) → aux trunk → aux softmax (parallel)
3. After forward: copy `aux_softmax[env, 1]` (p_up) → `prev_aux_dir_prob[env]` buffer
4. Next step (t+1): state assembly reads buffer

Cold-start at step 0: buffer = 0.5 (neutral, no signal). Same at FoldReset.

## Implementation Steps (atomic order)

### Step 1: State layout constants

**File**: `crates/ml-core/src/state_layout.rs`

Add after line 56 (PADDING_START):
```rust
// SP22 H6 (2026-05-12) — aux directional probability from previous step.
// Lives in the first padding slot. PADDING_DIM stays 7 (no STATE_DIM change);
// `assemble_state` writes aux_dir_prob to [AUX_DIR_PROB_INDEX] and zeros the
// remaining 6 padding slots [PADDING_START+1..STATE_DIM).
pub const AUX_DIR_PROB_INDEX: usize = PADDING_START;  // = 121
```

No need to change PADDING_DIM (still 7) — slot is just repurposed for non-zero content.

**Sentinel value**: 0.5 means "neutral aux signal" (50/50 up/down). Used at:
- Cold start (no aux output yet)
- FoldReset (fresh state, no signal)
- Eval-time when aux pipeline missing (A3 fallback)

### Step 2: Modify `assemble_state` in state_layout.cuh

**File**: `crates/ml/src/cuda_pipeline/state_layout.cuh` (~line 625)

Current signature:
```c
__device__ __forceinline__ void assemble_state(
    float* __restrict__ out,
    const float* __restrict__ market,
    const float* __restrict__ ofi,
    const float  mtf[SL_MTF_DIM],
    const float  portfolio[SL_PORTFOLIO_BASE_DIM],
    const float  plan_isv[SL_PORTFOLIO_PLAN_DIM]
)
```

Add 7th parameter `float aux_dir_prob`. Change the padding block:
```c
// Padding [121..128) — slot 0 = aux directional probability (SP22 H6);
// slots 1..7 zero-filled for 8-alignment.
out[SL_PADDING_START + 0] = aux_dir_prob;
for (int k = 1; k < SL_PADDING_DIM; k++)
    out[SL_PADDING_START + k] = 0.0f;
```

### Step 3: Update training state-gather kernel

**File**: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (search for `experience_state_gather`)

Add kernel arg `const float* __restrict__ aux_dir_prob_per_env` (right after `plan_isv`). In the kernel body, read `aux_dir_prob_per_env[env_idx]` (or 0.5 if NULL — defensive) and pass to `assemble_state`:

```c
float aux_dir_prob = (aux_dir_prob_per_env != NULL)
    ? aux_dir_prob_per_env[env_idx]
    : 0.5f;
assemble_state(out, market, ofi, mtf, pf, pisv, aux_dir_prob);
```

### Step 4: Update eval state-gather kernel

**File**: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (search for `backtest_state_gather`)

Same change as Step 3 — add `aux_dir_prob_per_env` arg, pass to `assemble_state`. For eval-time (Phase 1, A3 fallback): callers pass NULL → kernel uses 0.5.

### Step 5: New copy kernel

**File**: `crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu` (new file)

Mirror pattern from `aux_pred_to_isv_tanh_kernel.cu`. Single-block, BLOCK=256, no atomics:

```c
extern "C" __global__ void aux_softmax_to_per_env_kernel(
    const float* __restrict__ aux_softmax,   // [n_envs * K=2]
    float*       __restrict__ prev_aux_dir_prob,  // [n_envs] — written
    int n_envs,
    int K  // = 2 for current K=2 softmax head
) {
    int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= n_envs) return;
    // Write p_up (softmax[1]). Slot used by next step's state assembly.
    prev_aux_dir_prob[env] = aux_softmax[env * K + 1];
}
```

Launch: grid = `((n_envs + 255) / 256)`, block = 256.

### Step 6: Wire per-env buffer in GpuExperienceCollector

**File**: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

Search for `exp_aux_nb_softmax_buf` field declaration (~line 1191).

Add new field after it:
```rust
/// SP22 H6 (2026-05-12) — per-env aux directional probability cache.
/// Sized [alloc_episodes] f32 (one slot per env). Written end-of-step by
/// `aux_softmax_to_per_env_kernel`. Read start-of-next-step by
/// `experience_state_gather` for `state[AUX_DIR_PROB_INDEX]`. Cold-start
/// and FoldReset both initialize to 0.5 (neutral; p_up = 50%).
prev_aux_dir_prob: cudarc::driver::CudaSlice<f32>,
```

Allocate in `new()`:
```rust
let prev_aux_dir_prob = stream
    .alloc_zeros::<f32>(alloc_episodes)
    .map_err(|e| MLError::ModelError(format!("prev_aux_dir_prob alloc: {e}")))?;
// Initialize to 0.5 (neutral) — host write since rare init path.
let init_vec = vec![0.5_f32; alloc_episodes];
stream.memcpy_htod(&init_vec, &mut prev_aux_dir_prob.clone()).map_err(...)?;
```

(Use `memcpy_htod` from the stream's host write helpers; check cudarc API.)

Add to `FoldReset` cleanup: re-init to 0.5 each fold (matches state_reset_registry pattern).

### Step 7: Launch copy kernel after aux forward

**File**: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

Find where `aux_pred_to_isv_tanh_kernel` is launched (search `exp_sp13_aux_pred_to_isv_tanh_kernel`). The new copy kernel launches on the same stream, RIGHT AFTER the aux softmax tile is available.

```rust
// SP22 H6: copy per-env aux p_up into prev_aux_dir_prob for next step.
unsafe {
    self.stream
        .launch_builder(&self.exp_aux_softmax_to_per_env_kernel)
        .arg(&self.exp_aux_nb_softmax_buf.dev_ptr)
        .arg(&self.prev_aux_dir_prob.dev_ptr)
        .arg(&n_envs_i32)
        .arg(&aux_k_i32)
        .launch(LaunchConfig { ... })
        .map_err(...)?;
}
```

Cubin loading: add `aux_softmax_to_per_env_kernel.cu` to the build script (search for how `aux_pred_to_isv_tanh_kernel.cu` is registered in `build.rs`).

### Step 8: Pass buffer to state-gather kernel

**File**: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

Find `experience_state_gather` launch call. Add `prev_aux_dir_prob.dev_ptr` as the new arg (matching kernel signature change from Step 3).

### Step 9: Eval-side (A3 fallback)

**File**: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs`

Find `backtest_state_gather` launch call (~line 1517 area, in `gather_states`). Pass `&0u64` (NULL) for the new `aux_dir_prob_per_env` arg — kernel will use 0.5 fallback per Step 4.

**For A2 (deferred follow-up)**: add aux trunk forward to the closure-path eval. Requires:
- Load aux trunk weights from checkpoint (already in safetensors metadata?)
- Run encoder → aux trunk → aux softmax per step (parallel with policy Q forward)
- Wire aux softmax → per-env buffer (same copy kernel)
- Pass buffer ptr to state-gather

### Step 10: Tests + Smoke

```bash
SQLX_OFFLINE=true cargo check -p ml --features cuda                # compile
SQLX_OFFLINE=true cargo test -p ml --test gpu_backtest_validation \
  --features cuda --release                                        # local validation
/usr/local/cuda/bin/compute-sanitizer --tool=memcheck <test_bin>   # OOB check
```

Then dispatch smoke:
```bash
./scripts/argo-train.sh dqn --baseline --branch sp20-aux-h-fixed \
  --epochs 3 --gpu-pool ci-training-l40s
```

**Watch**: cycle 1+ — `aux_dir_acc` (should still be ~78% within-fold per H1 baseline), and **WR** (the headline metric).

**Decision**:
- WR > 50.5% → H6 confirmed. Justify A2 (eval aux integration) for production parity.
- WR pinned at 50.1-50.2% → H6 falsified. Pivot to H3 (reward density) or audit V/A unidentifiability fix.

## Risks + mitigations

| Risk | Mitigation |
|---|---|
| Cross-fold aux overfit (78%→49% from H1) | Initial test with overfit signal is fine — we're testing if policy CAN use any aux input. Generalization fix is separate concern. |
| Kernel-arg ABI mismatch in launchers | After kernel sig change, every launcher must be updated atomically. Grep `experience_state_gather` + `backtest_state_gather` call sites; compile catches missing arg. |
| Stop-grad violation | CudaSlice<f32> write/read crosses no autograd graph (we're in raw CUDA, not Candle). Stop-grad is implicit. |
| Eval-time NULL → segfault | Step 4 kernel defensive check `(aux_dir_prob_per_env != NULL) ? ... : 0.5f`. |
| Replay buffer state augmentation | Replay buffer stores states as `[N, STATE_DIM=128]`. Aux slot is already part of STATE_DIM. Buffer storage unchanged; only the meaning of `state[121]` changes from zero-padding to aux signal. |

## Estimated effort

| Phase | Effort | Risk |
|---|---|---|
| Steps 1-2 (constants + assemble_state) | 30 min | Low |
| Steps 3-4 (state-gather kernels) | 1 hr | Medium (kernel ABI) |
| Steps 5-6-7 (copy kernel + buffer + launch) | 1.5 hr | Medium (cubin build, kernel registration) |
| Step 8 (launcher arg wireup) | 30 min | Low |
| Step 9 (eval-side fallback) | 30 min | Low |
| Step 10 (test + smoke) | 1 hr + smoke wait | Low |
| **Total Phase 1 (A3 only)** | **~5 hr** | |
| Step 9 A2 (eval aux integration) | +1 day | Higher (eval has no aux infra yet) |

## Files to touch (Phase 1)

- `crates/ml-core/src/state_layout.rs` (constants)
- `crates/ml/src/cuda_pipeline/state_layout.cuh` (assemble_state)
- `crates/ml/src/cuda_pipeline/experience_kernels.cu` (2 kernels)
- `crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu` (new)
- `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (buffer + kernel launch)
- `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` (eval launcher arg)
- `crates/ml/build.rs` or similar (register new cubin)
- `crates/ml/src/trainers/dqn/state_reset_registry.rs` (FoldReset entry for prev_aux_dir_prob)

## Verification protocol

After implementation, BEFORE dispatching smoke:

1. `cargo check -p ml --features cuda` → 0 errors
2. `cargo test -p ml --test gpu_backtest_validation --features cuda --release`
   → all 4 passing tests still pass (the 2 PnL-assertion failures are pre-existing)
3. `compute-sanitizer --tool=memcheck <test_bin>` → 0 CUDA errors
4. Smoke dispatch on l40s, 3 epochs

Kill criteria:
- Compile errors that can't be resolved in 30 min → checkpoint and ask for help
- compute-sanitizer reports any non-zero CUDA errors → fix before dispatch
- Smoke cycle 1 WR is still 50.1-50.2% → H6 falsified, do not run further folds

## Audit pointer

Append findings to `docs/dqn-wire-up-audit.md` under
`## 2026-05-XX — SP22 H6 implementation: ...` header.
