foxhunt/docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md

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):

// 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:

__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:

// 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:

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:

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:

/// 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():

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.

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

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:

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