diff --git a/docs/superpowers/specs/2026-05-03-sp7-loss-balance-controller-design.md b/docs/superpowers/specs/2026-05-03-sp7-loss-balance-controller-design.md new file mode 100644 index 000000000..944ac9edf --- /dev/null +++ b/docs/superpowers/specs/2026-05-03-sp7-loss-balance-controller-design.md @@ -0,0 +1,250 @@ +# SP7: Loss-Balance Controller — Design + +**Status**: approved (2026-05-03) +**Branch**: `sp5-magnitude-differentiation` (extending SP5 closure) +**Predecessor**: SP5 close-out (commits `ab378327e` … `04df20bdc`) + +## Goal + +Replace the dormant per-branch CQL budget output (hardcoded to 0 by Pearl 2, floored to 0.02 in the consumer) and the floor-pinned per-branch C51 budget output with a single **outcome-driven controller** that adapts both heads' per-branch budgets so each loss head's gradient contribution stays in a target ratio relative to IQN. + +This unblocks magnitude differentiation: at production scale (50-epoch baseline at HEAD `2fb7d7f57`, smoke `train-multi-seed-n4qv2` ep4-7) we observed `grad_split_bwd cql=4.13` vs `iqn=0.07` (≈60×), with magnitude Q-values within 0.003 of each other, producing `eval_dist eq=1.000` (full eval-collapse to Quarter via argmax). + +## Non-goals + +- Not redesigning Pearl 2. Pearl 2 keeps writing `flatness[b]` to `FLATNESS_BASE` (consumed by NoisyNet σ — `pearl_per_branch_noisy_sigma`) and `budget_iqn[b]` to `BUDGET_IQN_BASE`. IQN is the reference; needs to stay stable. +- Not adding new diagnostic instrumentation. The grad-norm signals the controller consumes already exist (per-loss `grad_decomp_launch_*` pinned slots, populated every step). +- Not adapting trail/distill/rec/pred budgets. Out of scope. Same controller pattern can be applied later. + +## Architecture + +``` + Existing per-step pipeline (every training step) + ──────────────────────────────────────────────── + Forward pass → Backward (IQN, CQL, C51, Ens, distill, …) + │ + Each loss component runs `grad_decomp_launch_*` + which writes a 3-float pinned slot: + [mag_norm, dir_norm, trunk_norm] + │ + ▼ + ┌────────────────────────────────────────────────────────┐ + │ loss_balance_controller_kernel (NEW — single block, 8 threads)│ + │ │ + │ Reads: │ + │ iqn_decomp_pinned[3] ← grad_decomp_launch_iqn │ + │ cql_decomp_pinned[3] ← grad_decomp_launch_cql_sx │ + │ c51_decomp_pinned[3] ← grad_decomp_launch_c51 │ + │ isv[BUDGET_CQL_BASE+0..4] ← previous step │ + │ isv[BUDGET_C51_BASE+0..4] ← previous step │ + │ │ + │ Computes (per head h ∈ {CQL, C51}, per slice s): │ + │ actual_ratio = h_norm[s] / max(iqn_norm[s], EPS) │ + │ correction = target_ratio[h] / max(ratio, EPS) │ + │ new_budget = (1-α)·old + α·(old × correction) │ + │ new_budget = clamp(new_budget, EPS, MAX_BUDGET) │ + │ │ + │ Writes (raw, pre-smoothing): │ + │ scratch[BUDGET_CQL_BASE+0..4] │ + │ scratch[BUDGET_C51_BASE+0..4] │ + └────────────────────────────────────────────────────────┘ + │ + ▼ + apply_pearls_ad_kernel (existing) — Pearl A sentinel-bootstrap + + Pearl D smoothing → ISV[BUDGET_CQL_BASE+0..4], + ISV[BUDGET_C51_BASE+0..4] + │ + ▼ + Next step's compute_adaptive_budgets() reads these slots + → drives apply_cql_saxpy / apply_c51_budget_scale. +``` + +## Inputs (existing — no new producers needed) + +All three are pinned 3-float slots `[mag_norm, dir_norm, trunk_norm]`, computed every step by the corresponding `grad_decomp_launch_*` after each loss component's backward pass. The pinned device pointers are already accessible as kernel args. + +| Field | Source | Wired? | +|---|---|---| +| `iqn_norm[mag, dir, trunk]` | `grad_decomp_launch_iqn` (per-step) | yes | +| `cql_norm[mag, dir, trunk]` | `grad_decomp_launch_cql_sx` (per-step) | yes | +| `c51_norm[mag, dir, trunk]` | `grad_decomp_launch_c51` (per-step) | yes | + +The `cql_sx` decomp captures the SAXPY-applied delta (post-budget); `iqn` and `c51` capture the pre-budget delta. We use `cql_sx` for parity with what actually lands in `grad_buf`. + +## Outputs (existing dormant slots) + +| Slot | Current state | New state | +|---|---|---| +| `BUDGET_CQL_BASE+0..4` (ISV 198..202) | Pearl 2 writes 0 → floors to 0.02 in `compute_adaptive_budgets` | Controller writes adaptive per-branch budget | +| `BUDGET_C51_BASE+0..4` (ISV 190..194) | Pearl 2 writes `flatness × …` → in production drops below floor 0.05 | Controller writes adaptive per-branch budget | + +`BUDGET_ENS_BASE+0..4` and `BUDGET_IQN_BASE+0..4` remain Pearl 2's responsibility for now (Ens is residual; IQN is reference). + +## Math + +For each loss head `h ∈ {CQL, C51}` and each slice `s ∈ {dir, mag, trunk}`: + +``` +target_ratio[CQL] = TARGET_CQL_OVER_IQN // Invariant 1: 2.0 +target_ratio[C51] = TARGET_C51_OVER_IQN // Invariant 1: 1.0 + +iqn_n = max(iqn_norm[s], EPS_DIV) +h_n = h_norm[s] +actual_ratio = h_n / iqn_n +correction = target_ratio[h] / max(actual_ratio, EPS_DIV) + +old_budget = isv[BUDGET_h_BASE + b] // for the branch this slice maps to +candidate = old_budget * correction +new_budget = old_budget + ALPHA_META * (candidate - old_budget) +new_budget = clamp(new_budget, EPS_DIV, MAX_BUDGET) + +scratch[BUDGET_h_BASE + b] = new_budget +``` + +**Slice→branch mapping**: + +| ISV slot | Branch | Slice used | +|---|---|---| +| `BUDGET_CQL_BASE+0` | direction | `dir` | +| `BUDGET_CQL_BASE+1` | magnitude | `mag` | +| `BUDGET_CQL_BASE+2` | order | `trunk` | +| `BUDGET_CQL_BASE+3` | urgency | `trunk` | +| `BUDGET_C51_BASE+0..3` | same pattern | same pattern | + +Order and urgency reuse the trunk grad-norm ratio because `grad_decomp_kernel` only emits trunk/dir/mag slices. This is acceptable because: +- Order (3 actions) and urgency (3 actions) get the same gradient pressure as trunk receives; +- The TRUE per-branch grad-norm extension would require modifying `grad_decomp_kernel` to track 4 branches, which is out of scope and not load-bearing for the immediate problem (CQL dominance affects magnitude head, which we DO have a dedicated slice for). + +## Cold start (Pearl A — sentinel bootstrap) + +ISV slots start at sentinel 0 (per `pearl_first_observation_bootstrap`). The +key insight: `cql_norm[s]` reflects the *floor-applied* budget (the consumer +`compute_adaptive_budgets` already floors a 0-valued ISV read to +`COLD_START_FLOOR=0.02` before passing to `apply_cql_saxpy`). So the +"current effective budget" on a sentinel step is the floor, not zero. + +First controller invocation seeds: + +``` +if cql_norm[s] < EPS_DIV || iqn_norm[s] < EPS_DIV: + // Either side is uninstrumented or warmup hasn't started. + seed = COLD_START_FLOOR // 0.02 — Invariant 1 anchor +else if old_budget < EPS_DIV: // sentinel 0 → first observation + // The norm we observed was produced under cql_budget=COLD_START_FLOOR + // (consumer-side floor). Multiplicative correction with that floor as + // the basis seeds at the value that achieves target_ratio next step. + actual_ratio = cql_norm[s] / iqn_norm[s] + correction = target_ratio[h] / max(actual_ratio, EPS_DIV) + seed = COLD_START_FLOOR * correction + seed = clamp(seed, EPS_DIV, MAX_BUDGET) +else: + // Standard EMA update path (above). +``` + +Same logic for C51. + +**Contract**: the cold-start basis (`COLD_START_FLOOR=0.02`) must equal the +floor used by `compute_adaptive_budgets` in `fused_training.rs:3386`. If +that floor changes, this constant moves with it. Same anchor, same value, +defined once in `loss_balance_controller_kernel.cu` and referenced from a +header comment in `fused_training.rs`. + +## Fold-boundary reset (Pearl A) + +The slots `BUDGET_CQL_BASE+0..4` and `BUDGET_C51_BASE+0..4` are already registered in `StateResetRegistry` (`state_reset_registry.rs:585-595`) with sentinel 0. **No new registration needed.** On fold boundary the controller automatically re-runs cold-start logic. + +## Smoothing (Pearl D) + +The kernel writes raw new-budget values to a scratch buffer. The existing `apply_pearls_ad_kernel` then runs: + +- Pearl A: detects sentinel 0 in source ISV, replaces directly (skips smoothing on first observation). +- Pearl D: Wiener-optimal EMA based on `(diff_var / (diff_var + sample_var))`. + +This is the standard SP4 pattern (`pearl_first_observation_bootstrap` + `pearl_wiener_optimal_adaptive_alpha`). The new controller integrates with no new smoothing kernel. + +## Pearl 2 changes + +`pearl_2_budget_kernel.cu` is modified to **stop writing** to `scratch[BUDGET_CQL_BASE..]`, `scratch[BUDGET_C51_BASE..]`, `scratch[BUDGET_ENS_BASE..]`. It still writes: +- `BUDGET_IQN_BASE+0..4` — IQN budget (reference, needs to stay stable) +- `FLATNESS_BASE+0..4` — consumed by NoisyNet σ pearl + +Pearl 2 launch order in `training_loop.rs:3704-3710` is unchanged. The new controller launches **after** all `grad_decomp_launch_*` calls (so iqn/cql/c51 norms are populated) and **before** `apply_pearls_ad_kernel` (so smoothing applies to the new outputs). + +The `compute_adaptive_budgets` function in `fused_training.rs:3376` is unchanged — it reads ISV slots regardless of who wrote them. + +## Stale doc cleanup (per `feedback_trust_code_not_docs`) + +The docstring at `fused_training.rs:3409` claims: +```rust +/// B4/G5: Last adaptive CQL gradient budget (0.10×(1−regime)×health). +``` + +This formula is not implemented anywhere — the actual code at `fused_training.rs:3386` just floors to 0.02. The "regime_stability allocator" referenced in `state_reset_registry.rs:588` does not exist. **Both docstrings get rewritten** to describe the new controller (or deleted if redundant). + +## Files touched + +| File | Change | +|---|---| +| **NEW** `crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu` | ~140 LOC kernel (1×8 grid, 2 heads × 4 branches) | +| `crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu` | drop CQL/C51/ENS writes; keep IQN + flatness | +| `crates/ml/src/trainers/dqn/fused_training.rs` | add launch site after grad_decomp; update docstring at line 3409 | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | kernel slot in trainer struct + load + capture path; new `launch_loss_balance_controller` fn | +| `crates/ml/build.rs` | cubin manifest entry for new kernel | +| `docs/dqn-wire-up-audit.md` | Fix 31 entry — what landed, what was deleted, what's new | +| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_loss_balance_controller.md` | new memory pearl | +| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md` | index entry for the new pearl | + +## Invariant-1 anchors (per `feedback_isv_for_adaptive_bounds`) + +These are constants because they encode structural design choices, not tuned values that should adapt: + +| Constant | Value | Justification | +|---|---|---| +| `TARGET_CQL_OVER_IQN` | 2.0 | CQL is regularizer; carries 67% of (CQL+IQN) signal weight by design | +| `TARGET_C51_OVER_IQN` | 1.0 | C51 and IQN are both distributional; equal weight | +| `ALPHA_META` | 0.01 | Slow EMA; matches `pearl_adaptive_moe_lambda` baseline | +| `EPS_DIV` | 1e-8 | Numerical-stability anchor; same as Pearl 2 | +| `COLD_START_FLOOR` | 0.02 | Cold-start budget when grad signal unavailable; matches existing fused_training floor | +| `MAX_BUDGET` | 1.0 | Upper bound; budget is a multiplier in [0, 1] | + +## Acceptance criteria + +1. **Code lands clean**: `SQLX_OFFLINE=true cargo check -p ml --offline` and `SQLX_OFFLINE=true cargo build -p ml --release --offline --features cuda` both green at branch HEAD. +2. **Local RTX 3050 Ti smoke**: `multi_fold_convergence` passes with non-trivial `cql_budget_per_branch` drift away from 0.02. +3. **L40S 5-epoch smoke**: `grad_split_bwd cql/iqn` ratio at HEALTH_DIAG[4] is within `[1.0, 4.0]` (target 2.0; allow 2× tolerance during transient). Same for `c51/iqn` within `[0.5, 2.0]` (target 1.0). +4. **L40S 50-epoch full validation**: magnitude Q-spread `q_full − q_quarter > 0.01` by ep20 (currently <0.003 at ep7). `eval_dist` shows non-trivial Half/Full mass (currently `eq=1.0 eh=0 ef=0`). +5. **No regressions**: smoke `magnitude_distribution` still passes. `label_scale` stays in [0.03, 0.6] range (Bug-1 fix preserved). + +## Risk + rollback + +**Risk 1: feedback loop instability**. If `α=0.01` is too aggressive, budget oscillation could cascade into Q-value oscillation. Mitigation: clamp range `[EPS_DIV, MAX_BUDGET=1.0]` prevents catastrophic drift. Pearl D smoothing further dampens. + +**Risk 2: target ratio is wrong**. If CQL=2× IQN turns out to suppress IQN learning, magnitude differentiation may fail in different ways. Mitigation: target ratios are Invariant-1 anchors but the framework supports trivially changing them; smoke metrics will surface the issue at ep5. + +**Risk 3: cold-start interaction with fold boundary**. The `cql_norm < EPS_DIV` check seeds with `COLD_START_FLOOR=0.02` — same as current behavior. If fold boundary lands a step where the gradient hasn't been computed yet (chunk boundaries), seed is benign. Mitigation: `EPS_DIV` is `1e-8`, well below any real gradient signal. + +**Rollback**: revert the commit. Pearl 2's CQL/C51/ENS writes restored automatically (preserved as removed-then-re-added in single revert). No data migration concerns. + +## Out-of-scope but related + +- **C51 budget below floor mystery**: the `c51_budget_per_branch [dir=0.0500 …]` flat-at-floor observation suggests Pearl 2's flatness logic produces values below the floor. This is independent of SP7 and resolved automatically by SP7 since Pearl 2 stops writing to `BUDGET_C51_BASE`. (Pearl 2's flatness signal is preserved at `FLATNESS_BASE` for NoisyNet σ.) +- **Per-branch grad_decomp**: extending `grad_component_delta_norm` from 3-slice (trunk/dir/mag) to 4-branch (trunk/dir/mag/ord/urg) would let SP7 use dedicated grad-norms for ord and urg branches instead of trunk fallback. Worth doing later as a generalization, but the current trunk fallback is the right MVP. +- **Wiener-optimal α for the controller itself**: per `pearl_wiener_optimal_adaptive_alpha`, α could be derived from `diff_var / (diff_var + sample_var)`. Skipping for MVP — fixed `α=0.01` is consistent with `pearl_adaptive_moe_lambda` baseline. Phase 2 if the smoke shows oscillation. + +## Memory pearls created/updated + +**NEW** `pearl_loss_balance_controller.md`: +> When two loss heads contribute to the same parameter buffer at structurally +> different magnitudes (e.g., CQL's softmax-of-logsumexp emits gradient on +> all actions × atoms while IQN only emits on the action taken), a flatness- +> driven budget controller cannot balance them — it operates on Q-value +> distribution shape, not on observed gradient magnitudes. The fix is an +> outcome-driven controller that targets a fixed gradient-norm RATIO between +> the heads and adapts the multiplier to achieve it. Reference head (IQN) +> stays under existing flatness control; managed heads (CQL, C51) get +> ratio-driven budgets. Generalizes `pearl_adaptive_moe_lambda` to +> multi-head loss balancing. + +**Why**: future-me dispatching a similar adaptive-coefficient task should reach for ratio-driven (not signal-driven) controllers when the loss heads have known magnitude asymmetry. + +**How to apply**: when a new loss head is added to the training stack, decide whether it should be a reference (own its budget) or managed (driven by ratio to a reference). Add the controller's launch site after the head's `grad_decomp_launch_*` so the norm is populated.