# Phase 2 v2 — Proper Dueling Distributional C51 Design

**Date:** 2026-05-29
**Status:** Spec (informed by 2026-05-29 post-mortem + RL literature research)
**Supersedes:** `2026-05-28-state-conditional-q-synthesis.md` (architecturally flawed)
**Baseline:** Plan A v2 (commit `fd3174262`, branch `ml-alpha-plan-a-v2`)

---

## 0. Why this spec exists

The original Phase 2 spec proposed three sub-phases that interacted incorrectly:

- **Phase 2.0** (V envelope clamp): Caused PPO advantage corruption when V was clamped tighter than actual return range. **DROPPED in Plan A v2** — confirmed harmful.
- **Phase 2.1** (Dueling C51): Implementation used scalar V as Bellman bootstrap, making Q-update action-INDEPENDENT (target = projection of `r + γ V(s') - V(s)` is the same for all actions in a given state). Provably collapses to no action preference at convergence. **REVERTED in Plan A v2**.
- **Phase 2.2** (Per-action heterogeneous atoms): Architecturally flawed per RL literature — categorical C51 with per-action atom supports requires a "support-alignment operator" that has no clean mathematical formulation. **NEVER ADOPTED in Plan A v2**.

**RL literature consensus (Perplexity research, 2026-05-29):**
> For C51 distributional Q, use a **single unified canonical support** across all actions. Per-action heterogeneous atoms are "last resort" — naturally handled by **quantile methods (IQN/QR-DQN)**, NOT categorical C51.

This spec proposes a properly-grounded Phase 2 v2 design.

## 1. Goal

Match the SURFER baseline (`dd049d9a4`, wr=0.56 + $10.6M @ 18k) AND beat the TREND baseline (`6d33f18b7`, $40M peak) by addressing the TWO architectural shortcomings the original Phase 2 spec correctly identified:

1. **Cross-regime value mediation** — FLAT-state Q values are small (per-step rewards), OPEN-state Q values are large (cumulative position value). Without a unified baseline, Q-learning can't transition between regimes coherently.
2. **Action-specific resolution** — Long/Short/Hold benefit from fine [-3, +1] resolution; FullFlat/HalfFlat benefit from wider range to capture trade-realization variance.

The **architecturally correct** way to achieve both is **proper dueling distributional Q** with:
- **V_head**: scalar value baseline (action-independent, shared across all actions)
- **A_dist**: per-action advantage distribution on a SHARED atom support with mean-centering identifiability constraint
- **Joint loss**: composed `Q = V + A − mean_a A` regressed against TD target

This is exactly what Phase 2.1 intended but implemented incorrectly.

## 2. Architectural Decisions

### 2.A Dueling decomposition (proper)

```
Q(s, a) = V(s) + A(s, a) − mean_{a'} A(s, a')
```

Where:
- `V(s)` = scalar state value, action-independent
- `A(s, a)` = advantage distribution (atoms) per action
- `mean_{a'} A(s, a')` = action-mean subtraction, enforces identifiability

**Identifiability**: Without mean subtraction, V and A are unidentifiable (any V+K and A−K gives same Q). The subtraction zero-centers A, making the decomposition unique.

### 2.B Shared atom support for A_dist

All actions share the same atom support: `[-A_MAX, +A_MAX]` with Q_N_ATOMS=21 atoms.

**Bootstrap value**: `A_MAX = 3.0` (matches Plan A v2's atom span scale).

**Why shared not per-action**: Per RL literature, advantages have similar magnitudes across actions (centered around 0 by mean-subtraction). The differential is in the magnitudes, NOT the typical range. Shared atoms with reward_scale normalization captures both regimes:
- Long/Short: typical A in [-1, +0.3] (small per-step deviation)
- FullFlat/HalfFlat: typical A in [-2, +2] (realization variance)

Both fit comfortably in [-3, +3] with adequate resolution.

### 2.C Bellman target — joint composed Q loss

The KEY architectural difference from broken Phase 2.1:

**Phase 2.1 (WRONG)**: A trained against scalar `y − V(s) = r + γ V(s') − V(s)` projected onto A atoms. Target is action-INDEPENDENT.

**Phase 2 v2 (CORRECT)**: Composed Q distribution `Q_dist(s, a_t) = V(s) + A_dist(s, a_t) − mean_{a'} A_dist(s, a')` is regressed against TD target distribution. Target is the C51 projection of `r + γ × Q_target_dist(s', a*)` onto shared atom support, where `a* = argmax_a E[Q_target(s', a)]` (Double-DQN).

**This is action-DEPENDENT through a*** → A naturally learns action-specific structure.

Gradient flows through composed Q to BOTH V and A:
- ∂L/∂V: full TD error magnitude
- ∂L/∂A(s, a_t): TD error per-atom contribution
- Mean-subtraction term: gradient distributes negatively to all A(s, a') for a' ≠ a_t

### 2.D V head training — auxiliary regression

V head trains independently via MSE against returns (n-step):
```
loss_V = (V(s) − R_n(s))²
```
where `R_n = r + γ R_{n-1} ... + γ^n V_target(s_{t+n})`.

**This is in ADDITION to V's gradient from composed Q loss.** Dual training:
1. From composed Q loss (gradient via dueling decomposition)
2. From auxiliary return regression (PPO advantage signal)

Both push V toward the true value function. The auxiliary loss accelerates V convergence and provides PPO with stable advantages.

### 2.E State-conditional action mask (deferred)

Phase 2.3 (state mask) is intentionally DEFERRED. The mask is independently orthogonal to dueling — can be added later as a wrapper on top of either Phase 2 v1 or v2 architecture.

## 3. Implementation Sub-phases

### Sub-phase 2v2.0 — Plan A v2 baseline merge (already done)

Branch: `ml-alpha-plan-a-v2` (commit `fd3174262`).

- ✅ Distributional C51 Q (dd049d9a4 architecture)
- ✅ V_head for PPO advantage (dd049d9a4 architecture)
- ✅ Bug fixes: VSN stride, h_mag multi-warp, compute_advantage_return done-branch
- ✅ Cluster validated: wr=0.57 + $3.5M @ step 7000 (run `alpha-rl-8ksnj`)

**This is the foundation.** All sub-phases build on this.

### Sub-phase 2v2.1 — V joint training (preparation, 1 day)

**Goal**: V gets gradient from BOTH PPO advantage signal AND the upcoming composed Q loss.

**Files to modify:**
- `crates/ml-alpha/cuda/v_head_fwd_bwd.cu` — V head already done as scalar regression. Keep.
- `crates/ml-alpha/src/trainer/integrated.rs` — ensure V_head.backward fires on EVERY step (already does for PPO).
- New: `crates/ml-alpha/cuda/composed_q_grad_to_v.cu` — kernel that adds ∂L_Q/∂V into V's grad_h_t accumulator.

**Specific changes:**
1. After distributional Q loss kernel fires, add a kernel that:
   - Reads per-batch CE-loss gradient w.r.t. composed Q
   - Sums contribution to V (constant across atoms): `dL/dV = Σ_z (dL/dQ_atom[z]) × 1` (V adds the same to every atom)
   - Adds to `grad_h_t` (V's encoder gradient accumulator)

**Smoke test**: local b=128 × 1000 steps. Pass if `l_v` decreases monotonically and `V_pred` converges to observed return mean within 2σ.

### Sub-phase 2v2.2 — Dueling composed Q kernel (the core change, 2-3 days)

**Goal**: Implement `Q = V + A − mean_a A` composition in forward + backward kernels.

**Files to modify:**
- `crates/ml-alpha/cuda/dqn_distributional_q.cu` — reinterpret Q-head output as A_dist (no architectural change in shape, just semantics).
- New: `crates/ml-alpha/cuda/composed_q_fwd.cu` — kernel computing composed Q distribution per batch per action from V + A_dist.
- New: `crates/ml-alpha/cuda/composed_q_bwd.cu` — gradient through composition.
- `crates/ml-alpha/cuda/bellman_target_projection.cu` — UNCHANGED (still projects scalar bootstrap onto shared atoms).
- `crates/ml-alpha/src/rl/dqn.rs` — wire new kernels; `forward` returns composed Q; `backward` routes gradient.

**Specific changes:**

1. **Forward composition** (per batch, per action, per atom):
   ```
   composed_logits[b, a, z] = V(s_b) + A_logits[b, a, z] − mean_{a'} A_logits[b, a', z]
   ```
   Note: subtraction happens PER ATOM (z) for proper distributional decomposition.

2. **Action selection** (rollout):
   ```
   Q_value(s, a) = softmax(composed_logits)[z] · atom_z
   action = argmax_a E[Q(s, a)]  // or Thompson sample
   ```
   At rollout time, computed from composed Q.

3. **Bellman target** (TD bootstrap):
   ```
   a* = argmax_a E[Q_target(s', a)]   // Double-DQN
   y = r + γ × atom_z_values    // bootstrap per atom (same atoms shared)
   target_dist = C51 projection of (y, Q_target_softmax(s', a*)) onto shared atoms
   ```

4. **Loss**: CE between composed_Q_dist(s_t, a_t) and target_dist.

5. **Backward**:
   - `∂L/∂V[s_b]`: sum gradient over atoms → V's grad
   - `∂L/∂A[b, a_t, z]`: (1 − 1/N_ACTIONS) × ∂L/∂composed_logits[b, a_t, z]
   - `∂L/∂A[b, a, z]` for a ≠ a_t: −(1/N_ACTIONS) × ∂L/∂composed_logits[b, a_t, z]

The mean-subtraction Jacobian distributes negative gradient to all OTHER actions, which is the identifiability constraint.

**Smoke test**: local b=128 × 2000 steps. Pass criteria:
- `l_q` (composed Q loss) stable, declining trend after warmup
- `l_v` (V loss) declining
- `q_pi_agree` (action-mean A agreement) > 0.3 sustained (vs Phase 2.1's 0.01 collapse)
- Action distribution diverse (entropy > 1.5)

### Sub-phase 2v2.3 — Cluster validation (1 day)

**Submission:**
```
./scripts/alpha-rl-run.sh --branch ml-alpha-phase2-v2-dueling --steps 20000 --batch 1024 --gpu l40s
```

**Pass criteria:**
- Hold% at step 100 ≥ 50% (sub-90% is acceptable; dueling shifts initial Mamba2 RNG chain consumption)
- wr trajectory ≥ Plan A v2 baseline by step 5000
- pnl_cum profitable (> $0) by step 10000
- entropy stable in [1.5, 2.2] band
- No PPO l_pi explosion (l_pi < 1M by step 20k — Plan A v2 ranges 0-millions, that's normal)

**Failure handling:**
- Hold% crashes to 95%+ early → mean-subtraction term has wrong sign → revisit backward kernel signs
- l_q diverges (> 100 sustained) → composed Q computation has bug → revisit forward kernel
- wr stuck at 0.50 (no improvement) → V is being trained but A isn't getting useful gradient → check V's contribution dominates A's

## 4. Future Phase 2.3 (state mask) — design only, defer implementation

Per the original spec's §2.C. The mask is independent of dueling — it operates on action SELECTION (force illegal actions to -∞ before argmax/Thompson). Can be added as a thin layer over the dueling Q outputs.

**Defer until v2 dueling is validated at cluster.**

## 5. Future Phase 2.4 (per-action atom support) — abandoned

Per RL literature: this is the wrong direction for categorical C51. If we ever need per-action ranges (we may not — dueling A with shared support should suffice), the proper path is IQN, not per-action C51 atoms.

The existing `IqnHead` infrastructure (already in codebase) can be promoted from secondary head to primary Q-learning head if cluster experiments show shared-support C51 insufficient. **Defer until v2 dueling is validated and we have evidence for IQN need.**

## 6. Risks + Mitigations

| Risk | Impact | Mitigation |
|---|---|---|
| Mean-subtraction backward gradient sign wrong | Q diverges; opposite of dueling | Unit test backward against finite-difference oracle BEFORE cluster |
| V's auxiliary regression overwhelms composed Q signal | A doesn't learn action-specific structure | Balance λ_v and λ_q in loss balance controller; both already exist |
| Composed Q forward kernel race condition on `mean A` | NaN; non-determinism | Compute mean in single thread-0 reduce per batch (no atomicAdd, follows `feedback_no_atomicadd` pattern) |
| Cluster Mamba2 RNG init shift makes start state different | Hold% at step 100 won't match dd049d9a4's 90% | Document expected delta; rely on long-horizon convergence to wr=0.56+ |
| Adam optimizer not loading new kernel weights correctly | Silent training failure | Smoke test before cluster: log `w_d.shape` and `grad_w` per Adam step |

## 7. What this spec does NOT do

- No state-conditional action mask (defer to 2.3 if v2 succeeds)
- No per-action heterogeneous atoms (abandoned per research)
- No V envelope clamp (Phase 2.0 — confirmed harmful in Plan A v1)
- No new ISV controllers (P1-P4 from May 29 controller work was based on broken Phase 2.1 — also abandoned)

## 8. Success definition

Phase 2 v2 succeeds if it achieves:
- wr ≥ 0.57 at step 10k cluster (matches dd049d9a4)
- pnl_cum ≥ $5M at step 20k cluster (matches mid-trajectory dd049d9a4)
- `q_pi_agree_ema` ≥ 0.3 sustained (proves Q-π coupling — was 0.01 in broken Phase 2.1)
- No NaN spikes (V envelope clamp not needed because V regresses to bounded scaled targets naturally)

If it surpasses ($10M+, beats trend's $40M plateau): pursued architectural intent achieved.

## 9. Implementation order

```
2v2.0 ✅ Plan A v2 baseline merge (DONE)
   ↓ (smoke gate: wr=0.57 confirmed @ cluster)
2v2.1 V joint training preparation
   ↓ (smoke gate: V converges with composed Q gradient)
2v2.2 Dueling composed Q kernel (forward + backward)
   ↓ (smoke gate: q_pi_agree > 0.3 sustained at local b=128)
2v2.3 Cluster validation
   ↓ (cluster gate: wr=0.57+, pnl profitable by step 10k)
Future: 2.3 state mask if needed
Future: IQN promotion if per-action resolution proves needed
```

Each sub-phase has a standalone smoke test. Failures halt progression and trigger root-cause analysis.

## 10. Lessons captured

This redesign is informed by an exhausting 2026-05-29 day of cluster failures (alpha-rl-77c54, fcw67, z76v4, rkzkg, tdtqv, zss69, 4sjzw — 7 failed cluster runs). Pearls written:

- `pearl_state_mask_starves_q_for_closing_actions` (D test)
- `pearl_violation_rate_blind_to_hold_attractor` (H controller postmortem)
- `pearl_atomicadd_masks_v_instability` (atomicAdd removal)
- `pearl_dd049d9a4_surfer_baseline_verified` (dd049d9a4 reproduction)

The meta-lesson: **architectural intent ≠ implementation correctness**. The original Phase 2 spec captured real RL design intuitions (V mediation for cross-regime, per-action resolution for varying scales) but implemented them incorrectly. The proper implementation requires careful attention to:
1. Identifiability (mean-subtraction in dueling)
2. Cross-action mathematical coherence (shared atom support)
3. Joint training (composed Q loss, not separate Q and V losses)

Plan A v2 is the proven foundation. Phase 2 v2 is the proper architectural completion of the original intent.
