# ISV-Unified Q-Value Support Range (per-branch)

> **Problem statement:** Three Q-support consumers (atom grid, quantile warm-start, loss projection) hold independent opinions about the valid Q range for each of the 4 action branches. When they disagree, atoms become a bad quantizer: probability mass collapses onto 1-2 boundary atoms, the policy can't differentiate actions, and reward signal doesn't flow into useful Q-gradient information.

## Observation driving this design

- `train-fpxnw` (current run, Q-clamp fix applied): atom utilization **11-15%** across 57 fold-0 epochs. Probability concentrates on 1-2 of 51 atoms.
- Peak fold-0 Sharpe trending down as correctness fixes land: **79 → 68 → 48**. Each fix removed an accidental exploration crutch (wide Q range = wide atom grid = noisy Q-differentiation = effective exploration).
- Hypothesis under test: the Sharpe regression is not caused by our fixes being wrong — it's caused by *remaining* range inconsistencies between consumers. Atoms sit in one range, loss projects onto another, quantile-init fills a third. Each fix exposed one layer of the mismatch but left the others.

## Design principles

1. **One signal, one source of truth** for Q-support range per branch.
2. **Adaptive, temporal, signal-driven** — per the `adaptive_not_tuned` rule. EMA-tracked from observed Q-stats.
3. **Safety-clamped** to static `config.v_{min,max}` — hard rail, not tuning dial.
4. **Per-branch** — direction, magnitude, order, urgency have genuinely different Q scales; forcing a shared range recreates the exact mismatch we're fixing, just moved one level deeper.
5. **ISV bus broadcast** — established pattern for cross-kernel signals; consumers read from the bus they already receive.

## Current state (before this spec)

| Component | Q-range source | Type |
|-----------|----------------|------|
| `atom_positions_buf[4, num_atoms]` | `adaptive_atom_positions` kernel with `config.v_{min,max}` | **Static, per-config** |
| `per_sample_support_buf[N, 3]` | `eval_v_range_pinned`, tiled per epoch | **Adaptive, single-EMA, shared across branches** |
| `warm_start_atom_positions` quantile clamp | `config.v_{min,max}` (commit `d38a8cf99`) | **Static, per-config** |
| `q_stats_reduce` kernel output | Global min/max/mean/var across all branches | **Not per-branch** |

The four components disagree along two axes:
- **Static vs adaptive**: atoms + quantile = static; projection = adaptive.
- **Per-branch vs pooled**: nothing is per-branch today; magnitude's Q=±5 and direction's Q=±50 share atoms and support bounds.

## Proposed state

| Component | New Q-range source | Type |
|-----------|-------------------|------|
| `atom_positions_buf[4, num_atoms]` | ISV per-branch center/half-width | **Adaptive, per-branch** |
| `per_sample_support_buf` — deprecated | ISV per-branch center/half-width (read per-branch inside loss kernel) | **Adaptive, per-branch** |
| `warm_start_atom_positions` quantile clamp | ISV per-branch center/half-width (read per-branch on host) | **Adaptive, per-branch** |
| `q_stats_reduce` kernel output | Per-branch `q_mean, q_std, q_gap` → 4 sets of stats | **Per-branch** |

One source of truth: the ISV bus. Broadcast written by one path, read by three consumers.

## ISV slot allocation

Current ISV layout: 23 slots (slot 22 = `sharpe_ema`). Extend to 31:

| Slot | Name | Semantics |
|------|------|-----------|
| 23 | `v_center_dir` | EMA Q-center, direction branch |
| 24 | `v_half_dir` | EMA Q-half-width, direction branch |
| 25 | `v_center_mag` | EMA Q-center, magnitude branch |
| 26 | `v_half_mag` | EMA Q-half-width, magnitude branch |
| 27 | `v_center_ord` | EMA Q-center, order branch |
| 28 | `v_half_ord` | EMA Q-half-width, order branch |
| 29 | `v_center_urg` | EMA Q-center, urgency branch |
| 30 | `v_half_urg` | EMA Q-half-width, urgency branch |

**Why center + half-width, not min + max:** the adaptive EMA computes these two quantities naturally. Center is mean-Q EMA; half-width is `max(k·std, k·gap, min_floor)`. Consumers compute `v_min = center - half_width`, `v_max = center + half_width` themselves. Storing (center, half-width) keeps the evolution independent — center can shift while width stays constant.

**Bootstrap values** (pre any Q observations, set at construction + fold-boundary reset):
- `v_center_* = 0.0`
- `v_half_* = (config.v_max - config.v_min) / 2.0`

At epoch 1 these produce `[config.v_min, config.v_max]` for every branch — byte-identical to current behavior.

**Safety clamp** on every ISV write:
- `center = center_raw.clamp(config.v_min + min_half, config.v_max - min_half)`
- `half = half_raw.clamp(min_half_floor, (config.v_max - config.v_min) / 2.0)`

where `min_half_floor = 0.1 * (config.v_max - config.v_min)` prevents atom-grid collapse. Guarantees atoms always span a usable fraction of the config range.

## Prerequisite: per-branch Q-stats kernel

The ISV write path needs per-branch `q_mean`, `q_std`, `q_gap` statistics. Currently `q_stats_reduce` produces one global tuple across all 108 actions. This must be reworked first.

**Option A (extend existing kernel):** modify `q_stats_reduce` to emit `[4 × 7]` instead of `[7]`. Output format per branch: `[avg_max_q, q_min, q_max, q_mean, q_var, atom_entropy, atom_utilization]`.

**Option B (separate per-branch kernel):** add `q_stats_per_branch_reduce` alongside the existing one. Keeps the legacy kernel for any global consumers, adds a parallel path.

**Choice: Option A.** The global stats are already computed implicitly in the per-branch reduction (sum-across-branches of per-branch sums). Callers that need the global tuple can aggregate from the per-branch output. Avoids kernel duplication.

Branch offsets into `q_out_buf[N, total_actions]`:
- direction: `[0, b0)`
- magnitude: `[b0, b0+b1)`
- order: `[b0+b1, b0+b1+b2)`
- urgency: `[b0+b1+b2, b0+b1+b2+b3)`

Pass `b0, b1, b2, b3` to the kernel; iterate per sample, per branch, reduce.

**Readback shape change:** host side `QValueStatsResult` becomes `[QValueStatsResult; 4]` (or a new `PerBranchQValueStats`). Propagate the type change through `reduce_current_q_stats`, `update_eval_v_range`, consumers in `training_loop.rs`.

## Write path

Only `GpuDqnTrainer::update_eval_v_range` writes the ISV slots. It already runs at epoch boundary after `compute_q_stats`. New code flow:

```rust
pub fn update_eval_v_range(&mut self, per_branch_stats: &[QValueStatsResult; 4]) {
    for (branch_idx, stats) in per_branch_stats.iter().enumerate() {
        // Adaptive-rate EMA per branch
        let baseline = self.eval_q_std_ema[branch_idx].max(0.001);
        let mean_err = (stats.q_mean - self.eval_q_mean_ema[branch_idx]).abs();
        let std_err = (stats.q_std() - self.eval_q_std_ema[branch_idx]).abs();
        let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.3);
        let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.3);
        self.eval_q_mean_ema[branch_idx] =
            (1.0 - alpha_mean) * self.eval_q_mean_ema[branch_idx] + alpha_mean * stats.q_mean;
        self.eval_q_std_ema[branch_idx] =
            (1.0 - alpha_std) * self.eval_q_std_ema[branch_idx] + alpha_std * stats.q_std();

        // Derived center + half-width
        let gap_width = (10.0 * stats.q_gap()).max(3.0 * self.eval_q_std_ema[branch_idx]).max(min_half_floor);
        let abs_half = (self.config.v_max - self.config.v_min) as f32 * 0.5;
        let half = gap_width.min(abs_half).max(min_half_floor);
        let center = self.eval_q_mean_ema[branch_idx].clamp(
            (self.config.v_min as f32) + half,
            (self.config.v_max as f32) - half,
        );

        // Pinned device-mapped: writes visible to next kernel replay
        unsafe {
            *self.isv_signals_pinned.add(23 + 2 * branch_idx) = center;
            *self.isv_signals_pinned.add(24 + 2 * branch_idx) = half;
        }
    }
}
```

Per-branch EMA state replaces the current single `eval_q_mean_ema, eval_q_std_ema` scalars — they become `[f32; 4]`.

## Read-path changes

**1. `adaptive_atom_positions` kernel (CUDA)**

Current signature:
```cpp
extern "C" __global__ void adaptive_atom_positions(
    const float* __restrict__ spacing_raw,   // [NA]
    float* __restrict__ positions_out,       // [NA]
    int num_atoms,
    float v_min,  // per-branch in call site loop, from config
    float v_max
);
```

New signature:
```cpp
extern "C" __global__ void adaptive_atom_positions(
    const float* __restrict__ spacing_raw,       // [NA]
    float* __restrict__ positions_out,           // [NA]
    int num_atoms,
    int branch_idx,                              // 0..4
    const float* __restrict__ isv_signals        // read slots 23+2*branch_idx, 24+2*branch_idx
);
```

Kernel body:
```cpp
float center = isv_signals[23 + 2 * branch_idx];
float half   = isv_signals[24 + 2 * branch_idx];
float v_min  = center - half;
float v_max  = center + half;
// (rest unchanged — softmax + cumsum scaled by v_min, v_max)
```

Rust caller `recompute_atom_positions` changes: instead of passing config.v_min/v_max as f32, pass `branch_idx` and `isv_signals_dev_ptr`. The f64→f32 ABI fix from commit `768cc7d82` becomes moot (no more scalar f32 arg for bounds).

**2. `warm_start_atom_positions` (Rust host side)**

Current: reads reward quantiles, clamps with `hyperparams.v_{min,max}`, writes to atom_positions_buf.

New: reads per-branch bounds from `isv_signals_pinned` (pinned, zero-cost host read), clamps each of the 4 branch quantile arrays independently:

```rust
for branch_idx in 0..4 {
    let center = unsafe { *isv_pinned.add(23 + 2 * branch_idx) };
    let half   = unsafe { *isv_pinned.add(24 + 2 * branch_idx) };
    let v_min  = center - half;
    let v_max  = center + half;
    for (atom_idx, q) in branch_quantiles[branch_idx].iter_mut().enumerate() {
        *q = q.clamp(v_min, v_max);
    }
}
```

Currently `compute_reward_quantiles` returns a single `Vec<f32>` of size num_atoms shared across branches. Either:
- Keep shared quantiles (the reward distribution is the same regardless of branch), just clamp per-branch.
- Extract per-branch quantiles (requires per-branch reward attribution). **Out of scope** — shared quantiles with per-branch clamp is sufficient for the fix.

**3. C51 loss + MSE loss kernels**

Currently read `per_sample_support_buf[N, 3]` which is tiled from the single `eval_v_range_pinned`. For per-branch atom positions, the projection support in the loss kernel must also be per-branch.

Two options:
- **Keep `per_sample_support_buf` but make it `[N, 4, 3]`** (per-sample, per-branch). Tiling cost: 4× but still one HtoD per epoch. Minimal kernel change — index `per_sample_support[sample, branch_idx, {min,max,dz}]`.
- **Remove `per_sample_support_buf`, read from ISV directly inside the loss kernel.** Smaller kernel arg surface, one fewer buffer. But every Bellman-projection iteration does an extra ISV read per atom.

**Choice: option 1** (per-sample × per-branch support tile). Simpler kernel diff, caches well, matches the atom_positions layout. Tiling happens on epoch boundary from ISV slots.

**4. `update_per_sample_support` (Rust) — rename & restructure**

`update_per_sample_support(v_min, v_max)` becomes `update_per_sample_support_per_branch()` — reads all 8 ISV slots, tiles into `[N, 4, 3]`.

## Cross-fold semantics

`reset_eval_v_range_state()` extends to:
- Per-branch EMA state reset to 0.0 (center) and config-default (std).
- ISV slots 23..30 reset to bootstrap values.

Fold N+1 starts with atoms spanning full config range per branch. After epoch 1's per-branch Q-stats, adapts to new fold's distribution. No cross-fold bleed.

## Feedback-loop safety

Same safety rails as before, now per-branch:
1. **Hard clamp** at `config.v_{min,max}` — per-branch half-width cannot exceed `(v_max - v_min) / 2`.
2. **Minimum half-width floor** — prevents atom-grid collapse when Q-distribution is narrow.
3. **Anti-LR controller**, **Q-divergence EMA**, **gradient clipping** — unchanged, all still apply.

A per-branch runaway (e.g., magnitude Q grows without bound) is observable via `isv[24+2k]` hitting the config-clamp ceiling. Diagnosable directly from HEALTH_DIAG, which already reads ISV slots.

## Test plan

1. **Unit test: per-branch stats extraction** — feed known `q_out_buf` into `q_stats_reduce`, assert output matches per-branch hand-computed stats.
2. **Smoke test: ISV write-read round-trip** — run `update_eval_v_range` with synthetic per-branch stats, read back ISV slots 23..30, assert matches expected center + half.
3. **Smoke test: atom range tracks ISV** — inject ISV slots with known values, run `recompute_atom_positions`, assert atom_positions span `[center - half, center + half]` for each branch.
4. **Smoke test: quantile clamp uses per-branch bounds** — inject ISV values that differ per branch, verify `warm_start_atom_positions` clamps each branch's atoms to that branch's range.
5. **Smoke test: fold-boundary reset** — call `reset_for_fold`, assert all 8 ISV slots at bootstrap values.
6. **Smoke test: atom utilization** — target metric. Expect ≥40% after fix (up from 11-15%).
7. **L40S 60-epoch run** — compare Sharpe trajectory, peak Sharpe, atom util vs `train-fpxnw` baseline.

## Implementation phases

**Phase 0 — prerequisite: per-branch Q-stats kernel (~1 day)**
- Modify `q_stats_reduce` to emit `[4 × 7]` stats.
- Propagate `PerBranchQValueStats` type through Rust side.
- Unit-test equivalence: aggregated per-branch stats == old global stats.
- Commit as a data-structure refactor, no semantic behavior change.

**Phase 1 — ISV plumbing (~half day, no functional change at epoch 1)**
- Grep-update `ISV_DIM` 23 → 31.
- Add per-branch EMA state `[f32; 4]` for mean + std.
- Bootstrap values set at construction.
- `update_eval_v_range` writes ISV slots 23..30.
- `reset_eval_v_range_state` resets slots.
- Consumers still reach for `config.v_{min,max}` — no semantic change.
- Commit: "ISV plumbing, zero behavior change."

**Phase 2 — consumer migration (~half day, breaks behavior)**
- `adaptive_atom_positions` kernel: new signature, reads ISV.
- `recompute_atom_positions` Rust caller: passes branch_idx + isv_dev_ptr.
- `warm_start_atom_positions`: per-branch clamp from ISV pinned.
- `update_per_sample_support` rename + restructure to `[N, 4, 3]` tiling.
- C51/MSE loss kernels: index per-branch into support tile.
- Commit: "v-range unified via ISV, per-branch atom grids."

**Phase 3 — validation + cleanup (~half day + 2.5h GPU)**
- Run smoke tests 1-6.
- L40S run, compare to `train-fpxnw`.
- Delete `eval_v_range_pinned` if redundant (probably is).

Total: ~2.5 days engineering + 1 validation run.

## Success criteria

All four must hold:
1. **Correctness:** ISV slots 23..30 written from one place, read by three consumers; no other path reaches for `config.v_{min,max}` as a runtime range except through ISV.
2. **Atom utilization improved:** ≥40% on fold 0 (up from 11-15% currently).
3. **Sharpe not regressed:** peak fold-0 Sharpe ≥ current +48.19.
4. **Fold-1 stability holds:** Q-range at fold 1 epoch 1 is within ISV-derived bounds (no ±333k or ±144k excursions as in pre-fix runs).

## Out of scope

- Removing `eval_v_range_pinned` (deferred to Phase 3 cleanup).
- Per-branch reward quantile extraction (shared quantile with per-branch clamp is sufficient).
- Cross-fold inherit (fold-reset is the proven default; revisit only if convergence speed becomes a bottleneck).
- Additional ISV signals beyond the 8 slots for v-range.

## Revision history

- **v1 (earlier today)**: proposed 2 ISV slots (global center + half-width). Revised after recognizing that direction-branch Q-scale (~±50) differs from magnitude-branch (~±5) — a shared range recreates the exact mismatch the fix targets.
- **v2 (this version)**: 8 slots (4 centers + 4 half-widths). Per-branch Q-stats kernel added as explicit prerequisite.
