# Adaptive Training Dynamics Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Amplify reward signal (SNR 0.001→0.5) and add Q-gap momentum to sustain the val_Sharpe=21+ trajectory proven by train-w6qfd through 200 epochs.

**Architecture:** Two CUDA kernel changes: (1) a new reward shaping kernel that rank-normalizes rewards before replay buffer write, (2) a modification to the existing spread gradient in c51_grad_kernel.cu that modulates intensity based on Q-gap velocity. Both are additive — they build on the existing 32-fix foundation without changing any working code paths.

**Tech Stack:** CUDA 12.4, Rust (cudarc), nvcc compilation via build.rs, H100 GPU target

**Baseline:** train-w6qfd (commit 7eccfc53c) — val_Sharpe=21-25 sustained, Q-gap=0.13-0.14 growing, blind spread gradient working.

---

### Task 1: Reward Rank Normalization Kernel

**Files:**
- Create: `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu`
- Modify: `crates/ml/build.rs` (add to kernels_with_common list)
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`

- [ ] **Step 1: Create the CUDA kernel**

Create `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu`:

```cuda
/**
 * Rank-preserving signed reward standardization.
 *
 * r_shaped[i] = sign(r[i]) * rank(|r[i]|) * std_ema
 *
 * rank(|r[i]|) = count(|r[j]| <= |r[i]| for all j) / N
 * Counting-sort rank: O(N) per element, parallelizable.
 *
 * Grid: ceil(N/256), Block: 256. One thread per reward.
 */
extern "C" __global__ void reward_rank_normalize(
    float* __restrict__ rewards,     /* [N] in-place */
    int N,
    float std_ema                    /* running std of raw rewards */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    float r_i = rewards[i];
    float abs_i = fabsf(r_i);

    /* Counting-sort rank: count how many |rewards| are <= |r_i| */
    int count = 0;
    for (int j = 0; j < N; j++) {
        if (fabsf(rewards[j]) <= abs_i) count++;
    }

    /* Rank in [0, 1], sign preserved, scaled by std_ema */
    float rank = (float)count / (float)N;
    float sign = (r_i > 0.0f) ? 1.0f : (r_i < 0.0f) ? -1.0f : 0.0f;
    rewards[i] = sign * rank * fmaxf(std_ema, 1e-6f);
}
```

- [ ] **Step 2: Register kernel in build.rs**

In `crates/ml/build.rs`, add to the `kernels_with_common` array (after "nstep_kernel.cu"):

```rust
"reward_shaping_kernel.cu",
```

- [ ] **Step 3: Verify CUDA compilation**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

Expected: `Finished` with no nvcc errors.

- [ ] **Step 4: Load kernel in gpu_experience_collector.rs**

In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`, add to struct fields:

```rust
reward_rank_kernel: CudaFunction,
```

In the constructor, after loading other kernels from `exp_module_extra`:

```rust
let reward_module = context.load_cubin(
    include_bytes!(concat!(env!("OUT_DIR"), "/reward_shaping_kernel.cubin")).to_vec()
).map_err(|e| MLError::ModelError(format!("reward_shaping module: {e}")))?;
let reward_rank_kernel = reward_module.load_function("reward_rank_normalize")
    .map_err(|e| MLError::ModelError(format!("reward_rank_normalize load: {e}")))?;
```

Add to struct initialization.

- [ ] **Step 5: Wire kernel call after experience collection**

In `gpu_experience_collector.rs`, add a public method:

```rust
pub fn shape_rewards(&self, rewards: &mut CudaSlice<f32>, n: usize, reward_std: f32) -> Result<(), MLError> {
    let n_i32 = n as i32;
    let blocks = ((n + 255) / 256) as u32;
    unsafe {
        self.stream
            .launch_builder(&self.reward_rank_kernel)
            .arg(rewards)
            .arg(&n_i32)
            .arg(&reward_std)
            .launch(LaunchConfig {
                grid_dim: (blocks, 1, 1),
                block_dim: (256, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| MLError::ModelError(format!("reward_rank_normalize: {e}")))?;
    }
    Ok(())
}
```

- [ ] **Step 6: Call from training loop after experience collection**

In `crates/ml/src/trainers/dqn/trainer/training_loop.rs`, after `collect_experiences_gpu` returns the `GpuExperienceBatch`, call:

```rust
let reward_std = self.observed_reward_std;
collector.shape_rewards(
    &mut gpu_batch.rewards, gpu_batch.n_episodes * gpu_batch.timesteps * 2, reward_std as f32,
)?;
```

Note: the `* 2` accounts for counterfactual doubling.

- [ ] **Step 7: Build and verify**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -3
```

Expected: clean build.

- [ ] **Step 8: Commit**

```bash
git add crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu crates/ml/build.rs \
       crates/ml/src/cuda_pipeline/gpu_experience_collector.rs \
       crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: rank-preserving signed reward standardization — SNR 0.001→0.5

r_shaped = sign(r) * rank(|r|) * std_ema. Counting-sort rank within
batch gives outlier-resistant, scale-free normalization. Sign preserved
for absolute direction. std_ema rescales to original magnitude range.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
```

---

### Task 2: Q-Gap Momentum for Spread Gradient

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`

- [ ] **Step 1: Add Q-gap EMA state to GpuDqnTrainer**

In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`, add fields to the struct (near `eval_q_mean_ema`):

```rust
/// Q-gap EMA for momentum-modulated spread gradient.
q_gap_ema: f32,
/// Spread velocity scalar — written by CPU, read by GPU via pinned memory.
spread_velocity_pinned: *mut f32,
spread_velocity_dev_ptr: u64,
```

Initialize in the constructor:

```rust
q_gap_ema: 0.0,
spread_velocity_pinned: /* allocate 1 float pinned device-mapped, same pattern as adaptive_clip_pinned */,
spread_velocity_dev_ptr: /* cuMemHostGetDevicePointer */,
```

Initialize the pinned value to `1.0` (full spread initially).

- [ ] **Step 2: Update Q-gap EMA and compute velocity**

In `reduce_current_q_stats`, after computing the Q-stats result, add:

```rust
// Update Q-gap EMA and write spread modulator to pinned memory
let q_gap = (result.avg_max_q as f32 - result.q_mean).max(0.0);
if self.q_gap_ema < 1e-12 {
    self.q_gap_ema = q_gap;
} else {
    let err = (q_gap - self.q_gap_ema).abs();
    let alpha = (err / (err + self.q_gap_ema.max(0.001))).clamp(0.01, 0.3);
    self.q_gap_ema = (1.0 - alpha) * self.q_gap_ema + alpha * q_gap;
}
let velocity = q_gap - self.q_gap_ema;
let delta_z_approx = self.eval_q_std_ema.max(0.01); // proxy for delta_z
let modulator = (1.0 - velocity / delta_z_approx).clamp(0.1, 2.0);
unsafe { *self.spread_velocity_pinned = modulator; }
```

- [ ] **Step 3: Pass velocity to c51_grad_kernel**

In `launch_c51_grad`, add the velocity device pointer as an argument:

```rust
.arg(&self.spread_velocity_dev_ptr)
```

- [ ] **Step 4: Modify c51_grad_kernel to use velocity**

In `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu`, add parameter:

```cuda
const float* __restrict__ spread_velocity)  /* [1] pinned device-mapped modulator */
```

Replace the current spread_scale line:

```cuda
float spread_scale = inv_batch * delta_z;
```

With:

```cuda
float velocity_mod = spread_velocity[0];
float spread_scale = inv_batch * delta_z * velocity_mod;
```

- [ ] **Step 5: Build and run smoke test**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -3
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "avg_grad\|test result" | tail -5
```

Expected: clean build, 19 tests pass, avg_grad non-zero.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
       crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
git commit -m "feat: Q-gap momentum modulates spread gradient — backs off during growth

spread_scale = inv_batch * delta_z * velocity_modulator
velocity_mod = clamp(1 - (Q_gap - Q_gap_ema) / delta_z, 0.1, 2.0)

Growing Q-gap → mod=0.1 (10% spread, let C51 learn).
Plateau → mod=1.0 (full spread, push Q-values apart).
Shrinking Q-gap → mod=2.0 (emergency spread, prevent collapse).
10% floor proven by train-w6qfd (val_Sharpe=21+ sustained).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"
```

---

### Task 3: Integration Test on H100

**Files:**
- No code changes — deployment test

- [ ] **Step 1: Push and launch**

```bash
git push origin main
./scripts/argo-train.sh dqn --epochs 50 --baseline --gpu-pool ci-training-h100
```

- [ ] **Step 2: Monitor first 10 epochs**

Check:
- `avg_grad` non-zero (~0.8-1.0)
- `val_Sharpe` positive and sustained (target: >15)
- `Q-gap` growing (target: >0.05 by epoch 5)
- No CUDA errors, no graph capture failures
- Reward values should now be in [-std_ema, +std_ema] range instead of [-20, +20]

```bash
argo logs <workflow-name> -n foxhunt -c main | grep -E "val_Sharpe|Q-gap|avg_grad|EPOCH_DIAG"
```

- [ ] **Step 3: Monitor epoch 20-25 (freeze zone)**

Check:
- `Q-gap` still growing (not frozen at 0.108 or 0.158)
- `val_Sharpe` not locked at a single value for 3+ epochs
- Spread velocity modulator should show variation (not stuck at 1.0)

- [ ] **Step 4: Document results**

Record: peak val_Sharpe, sustained val_Sharpe range, Q-gap trajectory, any anomalies.
Compare to train-w6qfd baseline (val_Sharpe=21-25, Q-gap=0.13-0.14).

---

### Future Tasks (Layers 3+4 — separate plan)

These require architecture changes and should be planned after Layers 1+2 are validated:

- **Layer 3: Direction-Conditioned Magnitude** — magnitude head input changes from `[h_s2]` to `[h_s2; Q_dir]`. Forward/backward GEMM resize. Separate spec + plan.
- **Layer 4: Distributional Variance Position Sizing** — extract Var[Q] from C51 atoms, scale position in portfolio kernel. Small kernel change but impacts reward distribution. Separate validation.
