docs(sp22): H6 implementation runbook + next-session prompt
Two new plan docs to enable clean session handoff: 1. docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md Detailed 10-step implementation runbook for H6 (aux→policy state bridge). Each step has file paths, kernel signatures, expected diff, and risk callouts. Estimated ~5hr for Phase 1 (training-side + A3 eval fallback). A2 (eval-side aux integration) is +1 day follow-up if H6 confirmed. 2. docs/plans/2026-05-12-sp22-h6-next-session-prompt.md Concise context-loading prompt to paste into next session. Includes state of the world, runbook pointer, verification gates, and start-here pointer. Why staged: H6 implementation crosses kernel-level state-layout contract (every consumer of STATE_DIM must migrate atomically per feedback_no_partial_refactor). Doing it RIGHT requires careful incremental verification with compute-sanitizer between steps — not safely batched into a tail-end session. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
261
docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md
Normal file
261
docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md
Normal file
@@ -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<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.
|
||||
35
docs/plans/2026-05-12-sp22-h6-next-session-prompt.md
Normal file
35
docs/plans/2026-05-12-sp22-h6-next-session-prompt.md
Normal file
@@ -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).
|
||||
Reference in New Issue
Block a user