# Adaptive Training Dynamics v2 — Design Spec

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task.

**Goal:** Break the val_Sharpe plateau and achieve sustained OOS-positive trading by amplifying reward signal, preventing convergence traps, conditioning branch decisions, and using distributional risk for position sizing.

**Architecture:** 4-layer additive design on top of the existing 32-fix foundation (stochastic Expected SARSA, blind spread gradient, Boltzmann eval, scale-invariant temperature). Layers 1+2+4 are kernel-level changes. Layer 3 is a lightweight architecture change (one GEMM resize).

**Baseline:** train-w6qfd (commit 7eccfc53c) — val_Sharpe=21-25 sustained for 10+ epochs, Q-gap=0.138 and growing. This is the foundation to build on.

---

## Layer 1: Rank-Preserving Signed Standardization

### Problem
Reward SNR = 0.001 (mean ≈ 0, std ≈ 6.5). The C51 Bellman target `T_z = r + γ*z` is dominated by the bootstrapped z term. The model can barely distinguish good from bad actions.

### Solution
Transform raw rewards before replay buffer write:
```
r_shaped = sign(r) * rank(|r|) * std_ema
```

- `rank(|r|)`: rank the absolute reward within the batch [0.0, 1.0]. Parallel GPU-friendly via counting sort (count elements ≤ |r_i|, divide by B).
- `sign(r)`: preserved from raw reward. Ensures absolute direction (+gain/-loss) is not lost to relative ranking.
- `std_ema`: running EMA of raw reward std (already tracked as `observed_reward_std`). Rescales ranked values back to the original magnitude range.

### Properties
- Outlier-resistant (rank is bounded [0,1])
- Scale-free (rank doesn't depend on absolute reward magnitude)
- Direction-preserving (sign maintained)
- No feedback loops (std_ema is slow, rank is per-batch)
- Zero hardcoded constants

### Files
- **Create:** `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu` (~40 lines)
  - `reward_rank_normalize` kernel: counting-sort rank + sign + scale
- **Modify:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
  - Call reward shaping kernel after experience collection, before replay buffer write
- **Modify:** `crates/ml/src/cuda_pipeline/experience_kernels.cu`
  - Wire reward_std_ema to the new kernel

### Testing
- Unit test: rank([3, -1, 5, -2]) with std_ema=2.0 should produce correctly signed ranked values
- Integration: avg_loss should converge faster (better signal → faster Bellman convergence)

---

## Layer 2: Temporal Q-Gap Momentum

### Problem
The blind spread gradient (from commit 7eccfc53c) is always active at `inv_batch * delta_z`. This works well (train-w6qfd val_Sharpe=21-25) but:
- During healthy Q-gap growth: the spread fights the C51 gradient slightly (wastes gradient budget)
- At plateaus: the spread is at the same intensity as during growth (could push harder)

### Solution
Modulate spread intensity by Q-gap velocity:
```
velocity = Q_gap - Q_gap_ema
spread_scale = inv_batch * delta_z * max(0.1, 1.0 - velocity / delta_z)
```

- Q-gap growing faster than delta_z/step: spread reduces to 10% floor
- Q-gap plateauing (velocity ≈ 0): spread at 100%
- Q-gap shrinking: spread exceeds 100% (emergency push)
- 10% floor: preserves minimum differentiation pressure (proven by train-w6qfd)

### State
- `q_gap_ema`: f32 field on GpuDqnTrainer, updated with adaptive alpha (same pattern as eval_q_mean_ema)
- Passed to c51_grad_kernel via per_sample_support buffer (or separate pinned scalar)

### Files
- **Modify:** `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`
  - Add `q_gap_velocity` parameter, modulate spread_scale
- **Modify:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
  - Add `q_gap_ema` field, update in `reduce_current_q_stats`, pass to kernel

### Properties
- Self-calibrating: delta_z is the natural velocity threshold (one atom of growth per step)
- 10% floor proven by train-w6qfd data
- Zero hardcoded constants (delta_z from per-sample IQL support)

---

## Layer 3: Direction-Conditioned Magnitude

### Problem
The 4-branch architecture treats branches as independent. But magnitude SHOULD depend on direction:
- `Long+Full` = aggressive trend trade (high conviction, high risk)
- `Long+Small` = cautious probe (low conviction, low risk)
- `Flat+Full` = nonsensical (no direction, full size)

The magnitude branch sees the same h_s2 regardless of which direction was selected.

### Solution
Feed direction Q-values as extra input to the magnitude head:
```
Current:  mag_logits = W_mag @ h_s2 + b_mag
Proposed: mag_logits = W_mag @ [h_s2; Q_dir] + b_mag
```

Where `Q_dir = [Q_short, Q_flat, Q_long]` (3 scalars, already computed by the direction head).

### Architecture Change
- Magnitude head FC layer input dimension: `h_s2_dim` → `h_s2_dim + 3`
- Weight matrix `W_mag`: `[adv_h, h_s2_dim]` → `[adv_h, h_s2_dim + 3]`
- Bias unchanged

### Implementation
- **Forward pass** (`batched_forward.rs`):
  - After direction head forward: extract E[Q_dir] from direction logits (3 scalars per sample)
  - Concat [h_s2, Q_dir] into a scratch buffer [B, h_s2_dim + 3]
  - Magnitude head uses the concatenated buffer as input
- **Backward pass** (`batched_backward.rs`):
  - Magnitude backward produces d_concat [B, h_s2_dim + 3]
  - Split: d_h_s2 += d_concat[:, :h_s2_dim], d_Q_dir = d_concat[:, h_s2_dim:]
  - d_Q_dir flows back through the direction head (additional backward through direction logits)
- **Weight init**: New 3 columns of W_mag initialized to zero (no initial conditioning bias)
- **CUDA Graph**: Requires recapture after shape change (one-time cost at epoch 0)

### Files
- **Modify:** `crates/ml/src/cuda_pipeline/batched_forward.rs`
  - Add concat buffer, modified magnitude forward
- **Modify:** `crates/ml/src/cuda_pipeline/batched_backward.rs`
  - Split d_concat, additional direction backward
- **Modify:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
  - Allocate concat scratch buffer, wire pointers
- **Modify:** `crates/ml/src/trainers/dqn/config.rs`
  - Update param_sizes for magnitude head

### Risk
- Breaks CUDA graph capture (shape change) — requires recapture
- Adds 3 * adv_h parameters (~768 for adv_h=256) — negligible
- The d_Q_dir backward path adds complexity to the backward pass

---

## Layer 4: Distributional Variance Position Sizing

### Problem
The C51 distribution encodes the FULL return distribution per action, but only the MEAN (E[Q]) is used. The VARIANCE (risk) is computed but discarded. Two actions with equal expected return but different risk should produce different position sizes.

### Solution
Extract distributional variance alongside expected Q:
```
E[Q_a] = Σ p_j * z_j          (already computed)
Var[Q_a] = Σ p_j * z_j² - E[Q_a]²   (5 extra lines)
```

Use variance for position sizing in the portfolio simulation:
```
position_scale = 1.0 / (1.0 + sqrt(Var[Q_taken]))
target_position = direction * magnitude * position_scale
```

- High variance → smaller position → less risk exposure
- Low variance → full position → capture the edge with confidence
- This is the Kelly criterion derived from the distributional atoms

### Files
- **Modify:** `crates/ml/src/cuda_pipeline/experience_kernels.cu`
  - In `compute_expected_q`: add `sum_z_sq` accumulation, compute variance, write to output buffer
  - In portfolio simulation kernel: read variance, apply position_scale
- **Modify:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
  - Allocate variance output buffer [B, total_actions]
  - Wire to compute_expected_q and portfolio kernel

### Properties
- Uses information already in the C51 distribution (zero extra computation for the atoms)
- 5 extra lines in compute_expected_q (one accumulator + one subtraction)
- 3 extra lines in portfolio kernel (read + sqrt + scale)
- Zero hardcoded constants (variance is a mathematical property)
- Naturally bounded: position_scale ∈ (0, 1]

---

## Implementation Order

1. **Layer 1** (Reward Shaping) — new kernel, ~40 lines. Foundation for all other layers.
2. **Layer 2** (Q-Gap Momentum) — modify existing spread gradient, ~15 lines. Builds on Layer 1 (better reward → faster Q-gap growth → momentum detects it).
3. **Layer 4** (Variance Sizing) — modify existing kernels, ~10 lines. Independent of Layer 3, quick win.
4. **Layer 3** (Dir→Mag Conditioning) — architecture change, ~200 lines across forward/backward. Most complex, highest risk, highest potential reward.

Each layer should be tested independently before adding the next.

---

## Success Criteria

| Metric | Current (train-w6qfd) | Target | Measurement |
|--------|----------------------|--------|-------------|
| val_Sharpe sustained | 21-25 (epochs 8-17) | >15 through epoch 50+ | val_Sharpe trajectory |
| Q-gap growth | 0.138 at epoch 17 | >0.2 by epoch 50 | Q-gap trajectory |
| OOS Sharpe | ~0 (oscillating) | >0 sustained | OOS Sharpe from trade stats |
| No freeze | Proven through epoch 17 | Through 200 epochs | Q-gap + val_Sharpe not locked |
| Position sizing | Fixed 100% | Variance-scaled | Max drawdown reduction |

---

## Risks and Mitigations

| Risk | Mitigation |
|------|-----------|
| Rank normalization changes C51 atom placement | Adaptive eval v_range (already implemented) tracks new reward scale |
| Q-gap momentum over-suppresses spread | 10% floor proven by train-w6qfd baseline |
| Dir→mag conditioning breaks graph capture | One-time recapture cost, well-understood pattern |
| Variance sizing reduces returns | Net Sharpe should improve (lower vol outweighs lower return) |
| Layer interaction effects | Implement in order, test between layers |
