plan: Adaptive Training Dynamics Layers 1+2 implementation plan

Task 1: Reward rank normalization kernel (SNR amplification) Task 2: Q-gap momentum for spread gradient (plateau modulation) Task 3: H100 integration test (50 epochs, compare to train-w6qfd) Layers 3+4 (dir→mag conditioning, variance sizing) deferred to separate plan after L1+L2 validation. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-14 21:42:57 +02:00
parent a5945db538
commit e9e4bd54ac
1 changed files with 310 additions and 0 deletions
--- a/docs/superpowers/plans/2026-04-14-adaptive-training-dynamics.md
+++ b/docs/superpowers/plans/2026-04-14-adaptive-training-dynamics.md
@@ -0,0 +1,310 @@
+# 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.