From 494125c10471f25f0f3db6b1d8bf2beabc60ca68 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 26 Apr 2026 23:24:25 +0200 Subject: [PATCH] =?UTF-8?q?spec(dqn):=20distributional=20RL=20aggregation?= =?UTF-8?q?=20=E2=80=94=20UCB=20on=20C51+IQN=20=CF=83=20for=20action=20sel?= =?UTF-8?q?ection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Designs the structural fix for the C51 expected-Q Hold/Flat bias exposed by the Kelly val-Flat-collapse fix. The bias is a manifestation of a deeper project-wide pearl: "Distributional RL aggregation discards uncertainty; action selection must restore it." Any value head representing Q as a distribution (atoms, quantiles, ensembles) MUST expose both E[Q] and σ(Q) to action selection. Boltzmann on E[Q] alone produces structural bias toward low-variance actions regardless of expected payoff — the C51+IQN Flat-attractor is one instance of this lost-information pattern. Fix: extract σ(Q) from BOTH C51 atoms (closed form) and IQN quantiles (IQR/1.349), blend by loss-time weight, feed Q_eff = E[Q] + κ·σ (κ=1.0 structural identity) to direction-branch Boltzmann ONLY. 4-phase implementation: Phase 0 — TDD hypothesis verification (Rust mirror functions + 5 unit tests including GPU integration on real checkpoint) Phase 1 — Audit existing reward levers (B.2, CF, PopArt, Q-target) via 6 unit tests; fix any bugs found Phase 2 — UCB integration: new compute_q_with_uncertainty kernel, modified action_select, Rust orchestration, project-wide aggregation contract in dqn-wire-up-audit.md Phase 3 — Verification per Plan 5 Task 5 multi-seed × multi-fold 5-layer verification gate; 8 risks with mitigations; 4 stop conditions that halt execution and force redesign. Existing band-aid fixes (Kelly cap, eps-floor, tau-floor) stay — they address symptoms at different layers. UCB adds the missing aggregation step that was the common root across all the symptoms. Direction-branch only — magnitude/order/urgency don't have the Flat-attractor (atom-mass collapse asymmetry). 5-7 days active work across 4 sub-plans; each gets its own writing-plans cycle. Co-Authored-By: Claude Opus 4.7 (1M context) --- ...26-distributional-rl-aggregation-design.md | 399 ++++++++++++++++++ 1 file changed, 399 insertions(+) create mode 100644 docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md diff --git a/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md b/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md new file mode 100644 index 000000000..294b84ded --- /dev/null +++ b/docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md @@ -0,0 +1,399 @@ +# Distributional RL Aggregation: Restoring Uncertainty at Action Selection + +**Date**: 2026-04-26 +**Status**: Design approved, awaiting implementation plans (4 sub-plans) +**Origin**: Task #178 — Design proper C51 expected-Q Hold/Flat bias mitigation +**Predecessor fixes (kept as real fixes for symptoms)**: Kelly cap warm-branch (`0c9d1ee39`), val-Flat-collapse (`Kelly` chain), eps-floor adaptive boost (`d54b49efc`), ISV-adaptive Boltzmann tau (`7a3d88646`) + +--- + +## The Pearl (project-wide principle) + +> **Distributional RL aggregation discards uncertainty; action selection must restore it.** +> +> Any value head that internally represents Q as a distribution (atoms, quantiles, ensembles, etc.) MUST expose both `E[Q]` AND `σ(Q)` to action selection. Boltzmann/argmax on `E[Q]` alone produces structural bias toward low-variance actions regardless of expected payoff — because aggregation collapses `(mean=0, σ=0)` and `(mean=−ε, σ=large)` to neighbouring scalars where the deterministic action wins. + +This pearl is upstream of every value-baseline asymmetric bug found in the DQN system to date (Kelly bootstrap deadlock, val-Flat-collapse, C51 Hold/Flat bias, magnitude saturation). They are all manifestations of the same lost-information pattern. + +The pearl extends to any future distributional value head (TFT, Mamba2, ensemble) — encoded as a project-wide aggregation contract enforced via `docs/dqn-wire-up-audit.md`. + +--- + +## The Problem + +### Statement + +C51 represents per-direction Q as a discrete distribution over atoms `[v_min, v_max]`. IQN represents the same Q as quantile values at fixed `τ ∈ {0.05, 0.25, 0.5, 0.75, 0.95}`. + +Action selection currently uses scalar `E[Q_d]`: +``` +P(d) ∝ exp(E[Q_d] / τ_boltzmann) +``` + +For each direction `d`: + +``` +E[Q_d] = Σᵢ p_d_i · v_i (C51) + = mean({Q_τ(d)}) (IQN) +``` + +### The asymmetry the model rationally learns + +- **Flat**: trajectory return ≡ 0 (no position → no PnL change). After training, `C51_flat → δ(v=0)`, so `E[Q_flat] = 0` AND `σ(Q_flat) ≈ 0`. +- **Long/Short**: stochastic trajectory return minus tx_cost. Without discovered edge, `E[Q_directional] ≈ −tx_cost < 0`, AND `σ(Q_directional) > 0` (real return variance). + +The model is statistically correct under the data it has seen. The issue is that aggregation to scalar `E[Q]` discards the σ information. Boltzmann sees only `(0, −ε, −ε, 0)` for `(short, hold, long, flat)` and rationally prefers Flat — because there is no signal in `E[Q]` alone that `Long` has option value the deterministic Flat does not. + +### Empirical confirmation + +Three runs (`train-dhmp6`, `train-bscl2` post-Kelly-fix) show consistent post-training behaviour: + +``` +Epoch 1 (random init): val_dir_dist [S=.236 H=.181 L=.411 F=.173] active=64.6% +Epoch 2 (one epoch): val_dir_dist [S=.135 H=.358 L=.142 F=.364] active=27.7% +``` + +The 35→72% Hold+Flat shift in one epoch is the C51 attractor in action. Two attempts at the symptom layer (ISV-adaptive Boltzmann tau floor, adaptive eps_dir floor) produced no measurable effect — confirming the issue is in the aggregation step, not in sampling sharpness or exploration rate. + +--- + +## The Fix Design + +### Architecture overview + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Per-direction Q-distribution sources │ +│ ┌──────────────┐ ┌──────────────┐ │ +│ │ C51 atoms │ │ IQN quantiles │ │ +│ │ p_d_i, v_i │ │ Q_τ(d) │ │ +│ └──────┬───────┘ └──────┬────────┘ │ +│ │ │ │ +│ ▼ ▼ │ +│ E[Q_C51], σ_C51 E[Q_IQN], σ_IQN │ +│ │ │ │ +│ └────────┬───────────────┘ │ +│ ▼ │ +│ Aggregation Contract (NEW): │ +│ E[Q_d] = c51_alpha · E_C51 + (1−c51_alpha) · E_IQN │ +│ σ²(Q_d) = c51_alpha · σ²_C51 + (1−c51_alpha) · σ²_IQN │ +│ │ │ +│ ▼ │ +│ Action selection (modified): │ +│ Q_eff(d) = E[Q_d] + κ · σ(Q_d), κ = 1.0 │ +│ Boltzmann samples on Q_eff with adaptive tau │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### σ computation + +**From C51 atoms (closed-form)**: +``` +σ_C51(d) = sqrt(Σᵢ p_d_i · (v_i − E[Q_C51_d])²) +``` + +**From IQN quantiles (IQR/normal const)**: +``` +σ_IQN(d) ≈ (Q_τ=0.75(d) − Q_τ=0.25(d)) / 1.349 +``` + +**Combined via c51_alpha blend** (same weight that drives the loss): +``` +σ²_combined(d) = c51_alpha · σ²_C51(d) + (1−c51_alpha) · σ²_IQN(d) +σ_combined(d) = sqrt(σ²_combined(d)) +``` + +NaN/Inf guards: `var = max(var, 0)` before sqrt; clamp output `σ ≤ 0.5 · (v_max − v_min)` (the maximum σ a two-point distribution at the support extremes could produce — anything larger is numerical artefact). + +### Direction-branch only + +UCB applies ONLY to the direction softmax in `experience_action_select`. Magnitude/order/urgency branches stay vanilla. + +**Rationale**: The bias is structurally direction-specific because Flat is the only action with deterministic 0-PnL distribution. Magnitude (Quarter/Half/Full) all generate information regardless of which is chosen — no atom-mass collapse asymmetry. UCB on magnitude would systematically prefer Full over Quarter (larger σ from larger position) even when Quarter is more rewarding. That's a different bias, not a fix. + +### κ is a structural identity, not a tuned constant + +`κ = 1.0` — use σ at face value. Standard deviation and expected value share the same units; combining them at unit weight is the structurally meaningful choice. Any other κ would be a magnitude tradeoff requiring justification. + +If empirically too aggressive, future ISV-driven adaptation (e.g., `κ × isv[learning_health]`) is a follow-up requiring its own design. + +### Conviction stays Q-based, not Q_eff-based + +Conviction (`q_range / max(isv[21], q_range)`) is the per-sample policy confidence used by the Kelly cap warmup floor. It must reflect the model's actual Q-spread, NOT the inflated UCB Q_eff. We compute conviction on raw `q_b0[a]`, not on `q_b0[a] + κ·σ`. + +--- + +## Phased Implementation + +### Phase 0 — Test-Driven Hypothesis Verification (no behavioural changes) + +**Goal**: Prove (or disprove) the bias hypothesis with controlled inputs BEFORE changing any kernel behaviour. + +**Deliverables**: +- Rust mirror functions in `crates/ml/src/trainers/dqn/distributional_q.rs`: + - `expected_value(atoms: &[f32], probs: &[f32]) -> f32` + - `std_dev_c51(atoms: &[f32], probs: &[f32], mean: f32) -> f32` + - `std_dev_iqn(quantiles: &[f32]) -> f32` + - `blend_uncertainty(sigma_c51: f32, sigma_iqn: f32, c51_alpha: f32) -> f32` + - `boltzmann(q: &[f32], tau: f32) -> Vec` +- 5 unit tests: + - **0.A**: Bias mechanism reproduces under controlled C51 atoms (vanilla Boltzmann picks Flat; UCB reverses) + - **0.B**: σ_C51 matches numpy reference within 1e-5 + - **0.C**: σ_IQN matches IQR/1.349 reference within 10% on Gaussian samples + - **0.D**: Aggregation contract preserves uncertainty (monotonicity in c51_alpha) + - **0.E** (`#[ignore]`, GPU): Real checkpoint extraction confirms hypothesis on trained model + +**Exit gate**: All 5 tests pass on a real trained checkpoint from `train-bscl2` (the run that exhibited collapse). If 0.E fails, the hypothesis was wrong — stop and redesign. + +**Time budget**: 1 day. + +### Phase 1 — Audit & Repair Existing Reward Levers (TDD, conditional) + +**Goal**: Verify B.2 novelty bonus, counterfactual reward, PopArt, and Q-target propagation work as designed. **If a bug is found, fix it before Phase 2.** If no bugs, Phase 1 is a no-op gate. + +This phase is "fix any real bugs in existing levers" — NOT "scale magnitudes." Magnitude tuning is a band-aid for a missing UCB; we don't paper over Phase 2's job. + +**Deliverables**: 6 unit tests: +- **1.A**: B.2 novelty bonus producer formula matches kernel math +- **1.B**: B.2 consumer wires to correct (i,t) reward slot +- **1.C**: Counterfactual reward sign + symmetry +- **1.D**: Reward composition matches HEALTH_DIAG reward_split +- **1.E**: Q-target propagates realized rewards (atom mass shift toward observed return) +- **1.F**: PopArt preserves relative reward ordering between directions + +**Exit gate**: All 6 pass; OR a bug is found, fixed, and tests re-run green. + +**Time budget**: 1 day (longer if bug found). + +### Phase 2 — UCB Integration + +**Goal**: Land the structural fix. + +**Deliverables**: + +#### Component 4.1: `compute_q_with_uncertainty` kernel + +New file `crates/ml/src/cuda_pipeline/q_with_uncertainty_kernel.cu`. Replaces the direction-branch path of the existing `compute_expected_q` (other branches stay). + +Inputs: C51 atom probabilities `[B, b0_size, n_atoms]`, atom values `[n_atoms]`, IQN quantiles `[B, b0_size, 5]`, ISV signals (for `c51_alpha`). +Outputs: `q_dir_out [B, b0_size]` (E[Q]) AND `sigma_dir_out [B, b0_size]` (σ). + +Math per (sample i, direction d): +``` +E_C51 = Σ p_i_d_a · v_a +var_C51 = Σ p_i_d_a · (v_a − E_C51)² +σ_C51 = sqrt(max(var_C51, 0)) +σ_C51 = min(σ_C51, 0.5 · (v_max − v_min)) // numerical clamp + +E_IQN = mean of quantiles +σ_IQN = (Q_τ=0.75 − Q_τ=0.25) / 1.349 + +α = isv_signals[c51_alpha_idx] +E[Q] = α · E_C51 + (1 − α) · E_IQN +var = α · σ²_C51 + (1 − α) · σ²_IQN +σ(Q) = sqrt(max(var, 0)) +``` + +Single block-per-sample reduction. Estimated <100 µs per batch. + +#### Component 4.2: `experience_action_select` modification + +`crates/ml/src/cuda_pipeline/experience_kernels.cu`: + +Add `const float* sigma_dir` input. Modify direction-branch Boltzmann ONLY: + +```c +// Existing: +float qv = q_sign * q_b0[a]; + +// New: +float qv = q_sign * (q_b0[a] + 1.0f * sigma_dir[a]); +// ^^^^ structural identity κ +``` + +Magnitude/order/urgency: unchanged. +Conviction calc downstream: uses raw `q_b0[a]`, NOT Q_eff (per design). + +#### Component 4.3: Rust orchestration + +`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` and `gpu_backtest_evaluator.rs`: + +- Allocate `q_sigma_buf: CudaSlice` parallel to direction slice of `q_values_buf` +- Launch `compute_q_with_uncertainty` after forward pass, before action_select +- Pass `q_sigma_buf` as new arg to action_select + +#### Component 4.4: Phase 2 unit tests + +- **2.A**: End-to-end UCB flips picks on the failure mode (CPU mirror) +- **2.B**: σ buffer correctly populated by GPU kernel (matches CPU mirror within 1e-5) +- **2.C**: Magnitude/order/urgency unaffected (numerical equality after RNG seeding) +- **2.D**: Real-batch e2e test on a trained checkpoint reproduces the expected behaviour + +#### Component 4.5: Documentation + audit + +- Update `docs/dqn-wire-up-audit.md` with Distributional Q-head Aggregation Contract table: + +| Distributional Q-head | Produces E[Q]? | Produces σ(Q)? | Wired to action_select? | +|---|---|---|---| +| C51 atom kernel | ✅ | ✅ (Phase 2) | ✅ direction only | +| IQN quantile kernel | ✅ | ✅ (Phase 2) | ✅ direction only via blend | +| (future) TFT Q-head | required | required | required | + +- Pre-commit hook (existing Invariant 7) extended to require σ column for any new distributional kernel +- Pearl entry to memory: `pearl_distributional_rl_aggregation.md` + +**Exit gate**: All Phase 0 + Phase 1 + Phase 2 tests pass; L3 smoke verification primary gates pass. + +**Time budget**: 2-3 days. + +### Phase 3 — Verification & Close-out (per Plan 5 Task 5) + +**Deliverables**: +- L4 multi-fold deploy (1 seed × 6 folds × 10 epoch L40S): verify secondary gates +- L5 full validation matrix (per Plan 5 Task 5): Tier 1/2/3 PASS +- nsys regression check (Phase H continues independently; UCB overhead measured) +- Final architecture-doc updates + +**Time budget**: per Plan 5 existing budget. + +--- + +## Verification Protocol (5-Layer Gate) + +Each layer must pass before promoting to the next. No layer-skipping. + +| Layer | Scope | Pass criterion | Runtime | +|---|---|---|---| +| **L1**: CPU unit tests | Phase 0+1+2 mirror functions | All asserts pass | seconds | +| **L2**: GPU integration test | Tests 0.E + 2.B against real checkpoint | σ matches CPU mirror within 1e-5 | seconds | +| **L3**: Smoke deploy (3-fold × 5-epoch L40S) | Primary gates | active_frac > 40% across all epochs; val_dir_dist faithful | ~12 min | +| **L4**: Multi-fold deploy (1 seed × 6 folds × 10 epoch) | Secondary gates | val_sharpe non-negative trend; val_win_rate ≥ 50% by ep 5+ | ~2 hours | +| **L5**: Full validation matrix (Plan 5 Task 5) | Tier 1/2/3 exit criteria | per Plan 5 spec | ~6 hours | + +### Success criteria (the gates that define "fixed") + +**Primary (gate A — REQUIRED)**: +- `val_dir_dist` stays within ±20% of `val_picked_dir_dist` for all directions across all epochs (no Hold/Flat collapse > 70%) +- `active_frac` (Long+Short) stays above ~40% — the C51 attractor does not lock in + +**Secondary (gate B — REQUIRED for full close-out)**: +- `val_sharpe` trend across training is non-negative (model improves, not degrades) +- `val_win_rate` ≥ 50% by epoch 5-10 on training data with edge + +Phase 1 success = A. Phase 2 success = A + B. + +--- + +## Risk Register + +### R1: σ collapses across all directions late in training + +**Scenario**: Model trains long enough that ALL C51 distributions tighten to near-δ. σ_combined → 0 for every direction. UCB bonus → 0. + +**Assessment**: Correct behaviour — model is confident. Not a risk. + +### R2: σ inflated by numerical instability in C51 atoms + +**Scenario**: Bad atom updates concentrate mass at v_max or v_min. σ becomes artificially huge. UCB gives Long/Short an unrealistic bonus. + +**Mitigation**: Clamp σ ≤ 0.5 × (v_max − v_min). The 0.5 ratio is structural — a two-point distribution at extremes has σ = (v_max − v_min) / 2. Anything larger is numerical artefact. + +### R3: κ = 1.0 too aggressive in some training regimes + +**Scenario**: Late in training when Q-magnitudes are large, σ at face value might still be a meaningful bonus, pushing exploration too hard. + +**Mitigation strategy** (deferred unless seen): if empirically problematic, link κ to ISV[learning_health]. But default κ=1.0 ships first; any change requires its own justification + tests. + +### R4: IQN σ approximation underestimates true variance + +**Scenario**: `(Q_0.75 − Q_0.25) / 1.349` assumes Gaussian. Real return distributions are leptokurtic (fat tails). Underestimates σ. + +**Assessment**: Conservative bias — better to underestimate σ than overestimate. Acceptable as-is for v1. + +**Future improvement**: full 5-quantile range with kurtosis-aware weighting. Track as follow-up. + +### R5: c51_alpha drift breaks σ_combined + +**Mitigation**: Use the same clamped read (ISV[21]-style outlier guard) as `c51_loss_kernel.cu:215`. NaN-guard final σ_combined. + +### R6: train/eval inconsistency + +**Assessment**: UCB applies in BOTH training and eval (same kernel, same code path). No mismatch. Test 2.D verifies. + +### R7: Performance regression from extra kernel launch + +**Mitigation**: Single block-per-sample reduction (~100 µs / batch). Profile with nsys after landing; fold into Plan 5 Task 3 nsys harness. Estimated overhead < 1% of forward pass time. + +### R8: Hypothesis is wrong — UCB doesn't fix the bias + +**Diagnostic ladder** (apply in order if L3 smoke fails): +1. Re-read HEALTH_DIAG: are σ values populating? +2. Is κ=1.0 actually applied? (One-shot debug printf on first batch.) +3. Are σ magnitudes reasonable? σ_long ≈ σ_flat → both distributions collapsed. +4. If σ correct but bias persists: value-baseline fix isn't sufficient. Bug elsewhere (Q-target propagation, reward composition). + +**Rollback plan**: Phase 2 is encapsulated in 1-2 commits. Revert cleanly. Phase 0+1 tests stay (no harm). Investigation continues from Phase 0 audit data. + +### Stop conditions + +If ANY of these occur, halt and reassess: + +1. Phase 0 Test 0.E fails (real checkpoint doesn't exhibit hypothesised σ structure) +2. Phase 1 reveals a structural bug in reward shaping (priority pivot to that fix) +3. L3 smoke shows Hold/Flat collapse persists after Phase 2 (hypothesis was wrong) +4. L4 reveals new bias not predicted (something else in the picture) + +Each stop condition produces an investigation task analogous to #178. + +--- + +## What Is Deliberately NOT Touched + +- Loss kernels (C51 KL-loss, IQN Huber loss): gradient flow unchanged +- Reward shaping (B.2 novelty bonus, counterfactual, PopArt, etc.): unchanged unless Phase 1 finds a bug +- Magnitude/order/urgency action selection: unchanged +- Q-target propagation: uses E[Q] only, no σ +- ISV slots: no new slots added +- Existing band-aid fixes (Kelly cap, eps-floor, tau-floor): kept — they address symptoms at different layers; UCB adds the missing aggregation step + +--- + +## Cumulative Effort + +5-7 days of active work across 4 sub-plans, gated by 5 verification layers. + +| Plan | Description | Time | +|---|---|---| +| Plan A | Phase 0: TDD hypothesis verification | 1 day | +| Plan B | Phase 1: existing-lever audit | 1 day (longer if bug) | +| Plan C | Phase 2: UCB integration | 2-3 days | +| Plan D | Phase 3: verification + close-out | per Plan 5 budget | + +Each plan gets its own writing-plans cycle. + +--- + +## References + +- Predecessor design: `docs/superpowers/specs/2026-04-24-dqn-v2-unified-design.md` +- Plan 5 (validation): `docs/superpowers/plans/2026-04-24-dqn-v2-plan-5-validation.md` +- Pearl precedents: + - `memory/pearl_blend_formulas_must_have_permanent_floor.md` (Kelly bootstrap deadlock) + - `memory/pearl_one_unbounded_signal_per_reward.md` (reward-shaping discipline) + - `memory/pearl_cold_path_no_exception_to_gpu_drives.md` (compute placement) +- Wire-up audit: `docs/dqn-wire-up-audit.md` +- Empirical evidence runs: `train-dhmp6` (pre-Kelly-fix), `train-pp9d8`, `train-4r6p8` (post-Kelly-fix), `train-bscl2` (post-eps-floor) + +--- + +## Aggregation Contract (project-wide pearl) + +For any future distributional value head added to the system: + +1. The kernel/module MUST produce a `σ(Q)` buffer of the same shape as its `E[Q]` output +2. The σ buffer MUST be wired to `action_select` (or whatever value-aware policy mechanism is in use) +3. Boltzmann/argmax/etc. MUST consume `Q + κ·σ`, not bare `E[Q]` +4. The `dqn-wire-up-audit.md` Distributional Q-head Aggregation Contract table MUST be updated to list the new head +5. The pre-commit hook will refuse merges of distributional heads missing σ output + +This contract turns the principle into an engineering invariant. New methods can join the system without re-discovering the bias class.