diff --git a/docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md b/docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md new file mode 100644 index 000000000..b6db0c878 --- /dev/null +++ b/docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md @@ -0,0 +1,261 @@ +# 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, +``` + +Allocate in `new()`: +```rust +let prev_aux_dir_prob = stream + .alloc_zeros::(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 # 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 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 ` → 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. diff --git a/docs/plans/2026-05-12-sp22-h6-next-session-prompt.md b/docs/plans/2026-05-12-sp22-h6-next-session-prompt.md new file mode 100644 index 000000000..03928d962 --- /dev/null +++ b/docs/plans/2026-05-12-sp22-h6-next-session-prompt.md @@ -0,0 +1,35 @@ +# Next-Session Follow-Up Prompt for SP22 H6 + +Paste this into the next session to pick up cleanly: + +--- + +Hi — picking up SP22 H6 implementation from yesterday's session. + +**Context**: We closed SP21 T2.2 (cascade wiring complete + eval pipeline fixed via 4-layer peel: 8.3+9, 8.5, 8.7, 8.4-fix). v10 hypothesis test confirmed WR is structurally pinned at 50.1-50.2% across all training durations and cascade configurations. + +H1 (label horizon) was falsified: aux head learns 78% directional acc at H=200 within-fold, but the policy is walled off via trunk separation (`pearl_separate_aux_trunk_when_shared_starves`). + +**Current hypothesis**: H6 — wire aux head's directional probability into policy STATE so the policy can use the proven directional signal. Detailed implementation runbook at: + + `docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md` + +That doc has all 10 implementation steps with file paths, kernel signatures, and the verification protocol. + +**State of the world**: +- Branch: `sp20-aux-h-fixed`, HEAD = `f50a974fb` (docs commit; H1 reverted at `e8814079d`) +- All training-side fixes (Phase 8.2/8.4/8.4-fix/8.6) are live +- Eval pipeline works end-to-end (compute-sanitizer 0 errors locally) +- No active workflows or monitors +- Local GPU (RTX 3050 Ti) + compute-sanitizer available for fast iteration + +**Today's goal**: Execute Steps 1-10 of the H6 runbook. Phase 1 (training-side + eval A3 fallback) only. If smoke shows WR > 50.5%, follow up with A2 (eval-side aux integration). If WR pinned, hypothesis falsified — pivot to H3 (reward density) or audit V/A. + +**Verification gates** before smoke dispatch: +1. `cargo check -p ml --features cuda` clean +2. `gpu_backtest_validation` tests pass +3. `compute-sanitizer --tool=memcheck` 0 errors + +**Then**: dispatch standard baseline smoke, watch cycle 1 WR. ~2 hours wall-clock to verdict. + +Start by reading the runbook, then begin Step 1 (state layout constants).