plan(dqn): Thompson sampling rollout — Plans A+B+C+D (4 sequential phases)

Implementation plans for the distributional-RL aggregation spec
(docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md).

Plan A — Phase 0: TDD hypothesis verification (8 tasks)
  - Standalone GPU test kernel: sample_c51_inverse_cdf,
    sample_iqn_quantile_interp, compute_e_c51, compute_e_iqn,
    thompson_direction_test, argmax_eq_test
  - 6 unit tests (5 GPU + 1 CPU synthetic edge)
  - GPU integration on converged checkpoint (Test 0.F)

Plan B — Phase 1: existing-lever audit (8 tasks)
  - 6 unit tests covering B.2/CF/PopArt/Q-target audit fixture
  - Exit gate: all 6 PASS = no reward-shaping bug; Phase 2 unblocked

Plan C — Phase 2: Thompson sampling integration (11 tasks)
  - Direction branch of experience_action_select rewritten:
    eps-greedy + Boltzmann → Thompson (training) + argmax E[Q] (eval)
  - C51 + IQN buffers wired to gpu_dqn_trainer + gpu_backtest_evaluator
  - train_active_frac HEALTH_DIAG instrumentation
  - 4 GPU-direct tests against production kernel
  - Aggregation Contract table + memory pearl

Plan D — Phase 3: long verification + dead-code cleanup (8 tasks)
  - Test 3.A: regression anchor against original C51 Flat bias
  - L4: 1 seed × 6 folds × 30 epoch (~1 hour)
  - L5: 5 seed × 6 fold matrix per Plan 5 Task 5 (Tier 1+2+3 PASS)
  - Direction-only eps_dir adaptive boost + Boltzmann tau-floor cleanup
    (per feedback_no_partial_refactor.md)
  - nsys regression check; final architecture/spec footer

Plans gated sequentially: each phase exit gate must pass before next.
This commit is contained in:
jgrusewski
2026-04-27 08:10:42 +02:00
parent 021bb0ef73
commit 3c61ee52ce
4 changed files with 3006 additions and 0 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,484 @@
# Plan B — Phase 1: Audit & Repair Existing Reward Levers
> **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:** Verify B.2 novelty bonus, counterfactual reward, PopArt normalization, and Q-target propagation work as designed in the existing codebase. **If a bug is found, fix it before Plan C / Phase 2.** If no bugs, Phase 1 is a no-op gate that locks the audit results in tests.
**Architecture:** GPU-direct tests against the EXISTING `experience_env_step` and related kernels — no new code paths created. Each test launches the kernel with controlled state/action inputs and asserts properties of the output (reward components, Q-target shifts).
**Tech Stack:** Same as Plan A (Rust 1.85, CUDA 12.4, cudarc, `#[ignore]` GPU tests).
**Prerequisite:** Plan A complete and exit gate passed. If Plan A failed, do not start Plan B.
**Spec reference:** Phase 1 of `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`.
---
## File Structure
| Path | Status | Responsibility |
|---|---|---|
| `crates/ml/src/trainers/dqn/reward_audit_tests.rs` | Create | Phase 1 audit unit tests against existing reward shaping kernels |
| `crates/ml/src/trainers/dqn/mod.rs` | Modify | Add `pub mod reward_audit_tests;` (cfg(test) only) |
No new kernels. All 6 tests launch existing `experience_env_step` and related kernels with controlled inputs.
---
### Task 1: Create reward_audit_tests module skeleton
**Files:**
- Create: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- Modify: `crates/ml/src/trainers/dqn/mod.rs`
- [ ] **Step 1: Add module declaration**
In `crates/ml/src/trainers/dqn/mod.rs`:
```rust
#[cfg(test)]
pub mod reward_audit_tests;
```
- [ ] **Step 2: Write skeleton with helpers**
Create `crates/ml/src/trainers/dqn/reward_audit_tests.rs`:
```rust
//! Phase 1 audit tests for existing reward shaping mechanisms.
//!
//! Each test launches the existing experience_env_step kernel (or a
//! focused subset) with controlled inputs and asserts properties of
//! reward_components_per_sample output. If any test fails, fix the
//! bug before Plan C / Phase 2.
//!
//! Run: cargo test -p ml --lib reward_audit_tests -- --ignored --nocapture
use std::sync::Arc;
use cudarc::driver::{CudaContext, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
const B0_SIZE: usize = 4;
const N_REWARD_COMPONENTS: usize = 6;
// reward_components_per_sample slot indices (per c51_loss_kernel.cu / experience_kernels.cu):
// [0] popart final reward (input to PopArt normalizer)
// [1] cf counterfactual reward
// [2] trail trail-stop bonus
// [3] micro micro-reward (small per-bar PnL signal)
// [4] opp_cost opportunity cost penalty
// [5] bonus B.2 novelty bonus + C.4/D.4* timing/persistence
fn make_test_stream() -> Arc<CudaStream> {
let ctx = CudaContext::new(0).expect("CUDA context");
ctx.default_stream()
}
// Tests follow.
```
- [ ] **Step 3: Verify compile**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib --tests`
Expected: clean.
- [ ] **Step 4: Commit**
```bash
git add crates/ml/src/trainers/dqn/reward_audit_tests.rs crates/ml/src/trainers/dqn/mod.rs
git commit -m "test(dqn): Phase 1 audit module skeleton"
```
---
### Task 2: Test 1.A — B.2 novelty bonus producer formula
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Read the existing B.2 producer**
Open `crates/ml/src/cuda_pipeline/experience_kernels.cu` lines ~2487-2493. The formula per project memory is:
```
novelty = max(0, 1 - ISV[71]/max(ISV[72], 1e-4))
bonus = conviction × vol_proxy × novelty
```
Verify the kernel matches this formula by reading source. Compare line-by-line against expected.
- [ ] **Step 2: Write the test**
Append to `reward_audit_tests.rs`:
```rust
/// Phase 1 Test 1.A — B.2 novelty bonus producer formula matches expected math.
///
/// Per spec §"Existing levers", B.2 fires on Flat→Positioned transition.
/// Setup synthetic state where:
/// - Pre-trade position = 0 (flat)
/// - Action = Long (induces Flat→Positioned)
/// - ISV[71] (TRADE_ATTEMPT_RATE_EMA) = 0.0001
/// - ISV[72] (TRADE_TARGET_RATE) = 0.001
/// - vol_proxy ≈ 0.005 (from features)
/// - conviction = 0.5 (from action_select)
///
/// Expected:
/// novelty = max(0, 1 - 0.0001/0.001) = 0.9
/// bonus = 0.5 × 0.005 × 0.9 = 0.00225
/// reward_components[5] = 0.00225 (within 1e-7)
#[test]
#[ignore]
fn test_1a_b2_novelty_bonus_formula() {
let stream = make_test_stream();
// The full integration test requires invoking experience_env_step with
// a complete state buffer + ISV signals. The audit approach: extract
// the formula into a Rust mirror and verify the kernel's output for
// the same inputs matches.
//
// For Phase 1 audit, we test the FORMULA correctness via a focused
// micro-integration:
// 1. Load existing experience_env_step kernel
// 2. Set up minimal portfolio_state, action, isv_signals to trigger
// Flat→Positioned with controlled conviction + vol_proxy
// 3. Run one timestep
// 4. Read reward_components_per_sample[i*L*6 + t*6 + 5]
// 5. Assert bonus matches conviction × vol_proxy × novelty
//
// Implementation requires reusing GpuExperienceCollector setup. The
// simplest path: spawn the existing test infra at
// `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs::tests` —
// most of those tests are ignored by default and exercise the same
// kernel.
// For this audit test, we use the existing collector test fixture.
// If the fixture API surface differs, adapt to it. The key assertion
// is on reward_components[i,t,5].
// Concrete test (pseudo-code; concrete bindings depend on existing API):
//
// let mut collector = GpuExperienceCollector::new_for_test(&stream)?;
// collector.set_isv(71, 0.0001);
// collector.set_isv(72, 0.001);
// collector.set_position(0, 0.0); // start flat
// collector.run_one_step(action=DIR_LONG, conviction=0.5, vol_proxy=0.005)?;
// let bonus = collector.read_reward_component(0, 0, 5)?;
// assert!((bonus - 0.00225).abs() < 1e-7,
// "B.2 bonus expected 0.00225, got {}", bonus);
// Until we wire the collector test fixture, this test is a SCAFFOLD —
// mark as panic with TODO so it fails until properly implemented.
panic!("test_1a wiring incomplete — implement using GpuExperienceCollector test fixture");
}
```
- [ ] **Step 3: Run test (expected to fail with panic)**
Run: `cargo test -p ml --lib reward_audit_tests::test_1a_b2_novelty_bonus_formula -- --ignored --nocapture`
Expected: `panic!` with the TODO message — confirming the test is registered.
- [ ] **Step 4: Wire test fixture against `GpuExperienceCollector`**
The existing kernel takes many inputs. Three approaches:
(a) **Use existing test fixture** if `GpuExperienceCollector` already has a `new_for_test()` or similar fixture with controllable ISVs and minimal state. Check `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` for `#[cfg(test)]` blocks.
(b) **Add a `new_for_phase1_audit()` factory** in `gpu_experience_collector.rs` that takes minimal inputs (1 episode, 1 timestep, controllable position + action + ISVs) and exposes the reward_components buffer for read-back.
(c) **Source-only audit**: rather than launching the kernel, statically verify the formula in the kernel source matches expected. Less rigorous but cheap.
Recommended: (b) for tests 1.A-1.D (formula correctness); (c) is acceptable for 1.E (Q-target propagation), 1.F (PopArt) which require full training step.
Implementation of (b):
```rust
// In gpu_experience_collector.rs, add:
#[cfg(test)]
pub fn new_for_phase1_audit(stream: Arc<CudaStream>) -> Result<Self, MLError> {
// Minimal config: 1 episode, 1 timestep, deterministic
Self::new_inner(/* minimal config */ ...)
}
#[cfg(test)]
pub fn run_one_step_for_audit(
&mut self,
action: i32,
conviction: f32,
vol_proxy: f32,
) -> Result<(), MLError> { ... }
#[cfg(test)]
pub fn read_reward_component(&self, i: usize, t: usize, c: usize) -> Result<f32, MLError> { ... }
```
Implement these three methods; they're simple wrappers over existing infrastructure.
- [ ] **Step 5: Replace panic with real assertion**
Replace the `panic!` in test_1a with the actual assertion using the new fixture. Run again; should now pass (if formula is correct) or fail (if formula has a bug — investigate).
- [ ] **Step 6: Commit**
```bash
git add crates/ml/src/trainers/dqn/reward_audit_tests.rs crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
git commit -m "test(dqn): Phase 1 Test 1.A — B.2 novelty bonus formula audit
Verifies kernel-side B.2 producer matches the documented formula
(conviction × vol_proxy × novelty). Uses new GpuExperienceCollector
audit fixture for isolated kernel exercise."
```
---
### Task 3: Test 1.B — B.2 consumer wires to correct (i,t) slot
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 1 Test 1.B — B.2 consumer wires to correct reward_components slot.
///
/// Setup [N=2 episodes, L=3 timesteps]. Trigger Flat→Positioned at (i=1, t=2).
/// Verify reward_components[i*L*6 + t*6 + 5] is the only non-zero bonus slot.
#[test]
#[ignore]
fn test_1b_b2_consumer_correct_slot() {
let stream = make_test_stream();
// Use audit fixture with N=2, L=3
// ... setup positions: ep 0 stays flat; ep 1 enters Long at t=2 ...
// Run; read reward_components_per_sample full buffer
// Assert: only index [1*3*6 + 2*6 + 5] is non-zero in bonus column
panic!("test_1b — implement using GpuExperienceCollector audit fixture with N=2, L=3");
}
```
- [ ] **Step 2: Implement using the audit fixture from Task 2 Step 4**
(Extend the fixture if needed for multi-episode multi-timestep setups.)
- [ ] **Step 3: Run + verify pass**
Run: `cargo test -p ml --lib reward_audit_tests::test_1b_b2_consumer_correct_slot -- --ignored --nocapture`
Expected: passes if existing kernel correctly indexes; fails if there's an indexing bug.
- [ ] **Step 4: Commit**
```bash
git add crates/ml/src/trainers/dqn/reward_audit_tests.rs
git commit -m "test(dqn): Phase 1 Test 1.B — B.2 wires correct reward slot"
```
---
### Task 4: Test 1.C — Counterfactual reward sign + symmetry
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 1 Test 1.C — Counterfactual reward signs + symmetry.
///
/// Scenario: model picks Flat at (i,t). Hindsight calculation says:
/// true_long_return = +0.005
/// true_short_return = -0.005
///
/// Expected:
/// reward_components[i,t,1] when picked=Flat (cf reward going to Long bin) = +0.005 × cf_ratio
/// Symmetric scenario (model picked Long, hindsight says Long was correct):
/// cf reward should NOT go to Flat (cf is for the un-taken alternative)
#[test]
#[ignore]
fn test_1c_cf_reward_sign_symmetry() {
// Setup: launch experience_env_step with two scenarios:
// Scenario 1: action=Flat, hindsight long_return = +0.005, cf_ratio=0.5
// Scenario 2: action=Long, hindsight long_return = +0.005, cf_ratio=0.5
//
// For Scenario 1: cf component should be +0.0025 (cf for un-taken Long)
// For Scenario 2: cf component should be ~0 (we took the action; no cf needed)
panic!("test_1c — implement using GpuExperienceCollector audit fixture with cf_ratio control");
}
```
- [ ] **Step 2-4: Implement, run, commit**
Same pattern as Tasks 2-3. The audit fixture needs to expose `cf_ratio` and `hindsight_returns`.
```bash
git commit -m "test(dqn): Phase 1 Test 1.C — counterfactual reward sign + symmetry audit"
```
---
### Task 5: Test 1.D — Reward composition matches reward_split
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 1 Test 1.D — Sum of reward components matches HEALTH_DIAG reward_split totals.
///
/// HEALTH_DIAG reports per-component contribution as fraction of |reward| variance.
/// This test verifies sum(reward_components[i,t,k] for k in 0..6) ≈ total_reward_per_sample[i,t]
/// (within PopArt normalization scaling).
#[test]
#[ignore]
fn test_1d_reward_composition_consistency() {
// Run a multi-step rollout, read total_reward_per_sample AND reward_components_per_sample.
// Assert: sum across components ≈ total reward (within numerical tolerance).
panic!("test_1d — implement using audit fixture with multi-step rollout");
}
```
- [ ] **Step 2-4: Implement, run, commit**
```bash
git commit -m "test(dqn): Phase 1 Test 1.D — reward component sum vs total"
```
---
### Task 6: Test 1.E — Q-target propagates realized rewards
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 1 Test 1.E — Q-target shifts toward observed return after one training step.
///
/// Captures pre-update Q distribution for action Long; runs one Q-learning gradient
/// step on a synthetic transition with reward=+0.01; verifies post-update Q distribution
/// has shifted atom mass toward the +0.01 atom.
///
/// This is a critical gate: if Q-target propagation is broken, no exploration
/// strategy (Thompson included) will help.
#[test]
#[ignore]
fn test_1e_q_target_propagates_rewards() {
// 1. Construct a fixed state. Run forward to get C51 atoms p_pre.
// 2. Run one C51 KL-loss gradient step with target = +0.01 (force atom mass toward this).
// 3. Run forward again to get p_post.
// 4. Assert: atom mass at index nearest +0.01 has INCREASED in p_post vs p_pre.
panic!("test_1e — requires single-step training fixture; implement against existing trainer test infra");
}
```
- [ ] **Step 2: Implement using existing DQN trainer single-step fixture**
If `DQNTrainer::run_single_step_for_test()` doesn't exist, add it (similar pattern to Plan B Task 2 Step 4 fixture).
- [ ] **Step 3: Run + verify**
```bash
git commit -m "test(dqn): Phase 1 Test 1.E — Q-target propagation audit"
```
---
### Task 7: Test 1.F — PopArt preserves directional reward ordering
**Files:**
- Modify: `crates/ml/src/trainers/dqn/reward_audit_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 1 Test 1.F — PopArt normalization preserves relative reward magnitudes.
///
/// Construct synthetic raw rewards across directions:
/// raw_long = +0.001
/// raw_short = -0.001
/// raw_flat = 0.0
/// raw_hold = 0.0
///
/// After PopArt normalization (mean μ, std σ):
/// long_norm > flat_norm > short_norm
///
/// PopArt should preserve ordering. If it doesn't, the bias may be aggravated
/// by mis-normalized rewards reaching Q-target.
#[test]
#[ignore]
fn test_1f_popart_preserves_ordering() {
// Read existing PopArt state (mean, std). Apply to synthetic raw rewards.
// Assert: normalized rewards preserve raw rewards' relative order.
panic!("test_1f — implement using PopArt state read-back from existing kernel");
}
```
- [ ] **Step 2-4: Implement, run, commit**
```bash
git commit -m "test(dqn): Phase 1 Test 1.F — PopArt ordering preservation audit"
```
---
### Task 8: Phase 1 exit gate
**Files:** none.
- [ ] **Step 1: Run all 6 audit tests**
Run: `cargo test -p ml --lib reward_audit_tests -- --ignored --nocapture`
Expected: all 6 pass.
- [ ] **Step 2: If any test fails, STOP**
If 1.A-1.D fails → bug in reward shaping kernel. Open a separate fix task; resolve before Plan C.
If 1.E fails → Q-target propagation is broken. Investigate gradient flow + projection.
If 1.F fails → PopArt is corrupting magnitudes. Investigate.
Each failure is a separate investigation; do not proceed to Plan C until all 6 pass.
- [ ] **Step 3: Commit Phase 1 exit confirmation**
If all pass:
```bash
git commit --allow-empty -m "phase1(dqn): all 6 audit tests pass — existing reward levers verified
No bugs found in B.2 producer/consumer wiring, counterfactual reward
sign/symmetry, reward composition consistency, Q-target propagation, or
PopArt directional ordering. Plan C / Phase 2 unblocked."
```
---
## Self-Review
**Spec coverage:**
| Spec requirement | Implemented in |
|---|---|
| 1.A B.2 producer formula | Task 2 |
| 1.B B.2 wires correct slot | Task 3 |
| 1.C CF sign + symmetry | Task 4 |
| 1.D Reward composition | Task 5 |
| 1.E Q-target propagation | Task 6 |
| 1.F PopArt ordering | Task 7 |
| Exit gate (all pass before Plan C) | Task 8 |
All 6 tests covered. ✓
**Placeholder scan:** Tests 1.A-1.F have `panic!("...implement using...")` markers. These are scaffold panics, NOT spec placeholders — they'll be replaced with real implementations in each test's Step 4. The tests EXIST as registered tests from creation; the panic is a forced fail until implementation completes.
**Type consistency:** `B0_SIZE = 4`, `N_REWARD_COMPONENTS = 6` consistent across all tests. Reward component slot indices documented in module-level comment.
**Time:** 1 day if all 6 pass green; 2-5 days if bugs found per spec.
---
## Plan complete and saved to `docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-B-phase-1.md`.

View File

@@ -0,0 +1,671 @@
# Plan C — Phase 2: Thompson Sampling Integration
> **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:** Integrate Thompson sampling on the direction branch into production `experience_action_select` kernel, replacing the existing eps-greedy + Boltzmann path for direction. Magnitude/order/urgency unchanged. Add `train_active_frac` HEALTH_DIAG instrumentation. Verify via 4 GPU-direct unit tests.
**Architecture:** Modify the direction branch of `experience_action_select`: remove eps-greedy + Boltzmann; add Thompson (training) + argmax E[Q] (eval). Thread C51 atom probabilities, atom values, and IQN quantiles into the kernel as new arguments. Reuse the math implemented in Phase 0's `thompson_test_kernel.cu` as the reference.
**Tech Stack:** Same as Plans A+B.
**Prerequisite:** Plan A and Plan B complete. All 12 tests (Phase 0 + Phase 1) passing.
**Spec reference:** Phase 2 of `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`.
---
## File Structure
| Path | Status | Responsibility |
|---|---|---|
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | Direction branch of `experience_action_select` rewritten for Thompson + argmax |
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | Thread C51 + IQN buffers to action_select call site |
| `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` | Modify | Same wiring for val path |
| `crates/ml/src/trainers/dqn/trainer/metrics.rs` | Modify | Add `train_active_frac` HEALTH_DIAG line |
| `crates/ml/src/trainers/dqn/distributional_q_tests.rs` | Modify | Add Phase 2 tests 2.A-2.D |
| `docs/dqn-wire-up-audit.md` | Modify | Aggregation Contract table |
| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_thompson_for_distributional_action_selection.md` | Create | Memory pearl entry |
---
### Task 1: Add Thompson math as device-inline functions in experience_kernels.cu
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
- [ ] **Step 1: Copy device-inline math from thompson_test_kernel.cu**
Open `crates/ml/src/cuda_pipeline/experience_kernels.cu`. Find a suitable location near the top of the file (after includes, before the first kernel definition). Add the same `__device__ __forceinline__` functions as in `thompson_test_kernel.cu`:
```cuda
#ifndef N_IQN_QUANTILES
#define N_IQN_QUANTILES 5
#endif
__device__ __forceinline__ float sample_c51_inverse_cdf(
const float* __restrict__ p,
const float* __restrict__ atom_values,
int n_atoms,
float u
) {
float cum = 0.0f;
for (int a = 0; a < n_atoms; a++) {
cum += p[a];
if (u < cum) return atom_values[a];
}
return atom_values[n_atoms - 1];
}
__device__ __forceinline__ float sample_iqn_quantile_interp(
const float* __restrict__ q,
float u
) {
const float taus[N_IQN_QUANTILES] = {0.05f, 0.25f, 0.50f, 0.75f, 0.95f};
if (u <= taus[0]) return q[0];
if (u >= taus[N_IQN_QUANTILES - 1]) return q[N_IQN_QUANTILES - 1];
for (int i = 0; i < N_IQN_QUANTILES - 1; i++) {
if (u <= taus[i + 1]) {
float w = (u - taus[i]) / (taus[i + 1] - taus[i]);
return q[i] * (1.0f - w) + q[i + 1] * w;
}
}
return q[N_IQN_QUANTILES - 1];
}
__device__ __forceinline__ float compute_e_c51_inline(
const float* __restrict__ p,
const float* __restrict__ atom_values,
int n_atoms
) {
float e = 0.0f;
for (int a = 0; a < n_atoms; a++) e += p[a] * atom_values[a];
return e;
}
__device__ __forceinline__ float compute_e_iqn_inline(const float* __restrict__ q) {
float s = 0.0f;
for (int i = 0; i < N_IQN_QUANTILES; i++) s += q[i];
return s / (float)N_IQN_QUANTILES;
}
```
- [ ] **Step 2: Verify compile**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
Expected: clean. nvcc rebuilds with new device functions.
- [ ] **Step 3: Commit**
```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "feat(dqn): add Thompson + argmax-of-E[Q] device-inline functions
Mirrors the Phase 0 standalone test kernel's math into the production
kernel file. Used by Phase 2 direction-branch action selection."
```
---
### Task 2: Replace direction branch in experience_action_select
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
- [ ] **Step 1: Read the existing direction branch**
Open `experience_kernels.cu` and find `experience_action_select` (around line 756). Find the direction branch around lines 858-887 (the eps-greedy + Boltzmann block).
- [ ] **Step 2: Add new kernel parameters**
Modify the kernel signature to accept three new buffers + n_atoms:
```cuda
extern "C" __global__ void experience_action_select(
// ... existing args up to and including out_conviction ...
const float* __restrict__ c51_probs_dir, // [N, b0_size, n_atoms] NEW
const float* __restrict__ atom_values, // [n_atoms] NEW
const float* __restrict__ iqn_quantiles_dir, // [N, b0_size, N_IQN_QUANTILES] NEW
int n_atoms // NEW
) {
```
- [ ] **Step 3: Replace direction branch with Thompson + argmax**
Replace the existing direction-branch code (the if-else with eps-greedy + Boltzmann at ~lines 858-887) with:
```cuda
/* Direction branch: Thompson sampling at training, argmax E[Q] at eval.
* Replaces the prior eps-greedy + Boltzmann; ε floor and tau-floor for
* direction become unreachable (cleaned up in Phase 3 / Plan D).
*
* Math reference: thompson_test_kernel.cu (Phase 0 standalone, same math).
* Sampling at training: 0.5 × (sample_C51 + sample_IQN), argmax across dirs.
* Eval: argmax of (E_C51 + E_IQN) — /2 dropped (monotonic).
*
* E[Q]-derived q_b0 must be reused for the conviction calculation below
* (per spec §"Conviction stays E[Q]-based"). We compute e_dir[d] once
* here and reuse.
*/
const float* c51_probs_i = c51_probs_dir + (long long)i * b0_size * n_atoms;
const float* iqn_quantiles_i = iqn_quantiles_dir + (long long)i * b0_size * N_IQN_QUANTILES;
/* Compute per-direction E[Q] for: (a) eval-mode argmax;
* (b) conviction read-back below; (c) Thompson uses raw samples,
* not E[Q], but conviction needs E[Q]. */
float e_dir[4]; // b0_size = 4 in current 4-direction action space
for (int d = 0; d < b0_size; d++) {
float e_c51 = compute_e_c51_inline(c51_probs_i + d * n_atoms, atom_values, n_atoms);
float e_iqn = compute_e_iqn_inline(iqn_quantiles_i + d * N_IQN_QUANTILES);
e_dir[d] = e_c51 + e_iqn; // joint sum; /2 monotonic for argmax
}
if (eval_mode) {
/* EVAL: argmax of joint E[Q]. Pure exploitation. */
float best_q = -CUDART_INF_F;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
if (e_dir[d] > best_q) { best_q = e_dir[d]; best_d = d; }
}
dir_idx = best_d;
} else {
/* TRAINING: Thompson sample. */
float best_sample = -CUDART_INF_F;
int best_d = 0;
for (int d = 0; d < b0_size; d++) {
float u_c51 = philox_uniform(i, timestep, rng_ctr++);
float s_c51 = sample_c51_inverse_cdf(
c51_probs_i + d * n_atoms, atom_values, n_atoms, u_c51
);
float u_iqn = philox_uniform(i, timestep, rng_ctr++);
float s_iqn = sample_iqn_quantile_interp(
iqn_quantiles_i + d * N_IQN_QUANTILES, u_iqn
);
float q_sample = s_c51 + s_iqn; // /2 dropped
if (q_sample > best_sample) { best_sample = q_sample; best_d = d; }
}
dir_idx = best_d;
}
```
- [ ] **Step 4: Reuse e_dir for conviction**
Find the conviction calculation downstream (around line 1091 `out_conviction != NULL`). It currently reads `q_b0[a]` to compute q_max_raw / q_min_raw. Modify it to use `e_dir` (the array we already computed):
```cuda
if (out_conviction != NULL) {
float q_max_raw = e_dir[0];
float q_min_raw = e_dir[0];
for (int a = 1; a < b0_size; a++) {
q_max_raw = fmaxf(q_max_raw, e_dir[a]);
q_min_raw = fminf(q_min_raw, e_dir[a]);
}
float q_range = q_max_raw - q_min_raw;
// ... rest of conviction unchanged ...
}
```
This honors the spec's "Conviction stays E[Q]-based" requirement — using `e_dir` (joint E[Q]) rather than Thompson samples.
- [ ] **Step 5: Compile + grep for `eps_dir`**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
Expected: clean.
Then grep: `grep -n "eps_dir" crates/ml/src/cuda_pipeline/experience_kernels.cu`
If `eps_dir` still has consumers in non-direction branches (mag/ord/urg), it stays in the signature. If unused after the direction-branch removal, the kernel signature can drop `eps_exp_mult` for direction-only — verify by tracing.
- [ ] **Step 6: Commit**
```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "feat(dqn): direction-branch Thompson sampling in experience_action_select
Replaces eps-greedy + Boltzmann on direction with:
- Training: argmax of (sample_C51 + sample_IQN) per direction (Thompson)
- Eval: argmax of (E_C51 + E_IQN) per direction (no exploration)
Conviction calculation now reads e_dir[] (joint E[Q] computed inline)
preserving the spec's requirement that conviction uses raw E[Q] not
samples (avoids Kelly cap jitter).
Magnitude/order/urgency branches: unchanged."
```
---
### Task 3: Buffer wiring in gpu_dqn_trainer.rs
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- [ ] **Step 1: Find the action_select call site**
Run: `grep -n "experience_action_select" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
This file has a call to launch the action_select kernel (the existing one). Find the call.
- [ ] **Step 2: Identify existing C51 + IQN buffers**
The trainer maintains C51 atom probabilities and IQN quantiles after each forward pass. Find the buffers (likely named something like `c51_probs_buf`, `iqn_quantiles_buf`, `atom_values_buf`). Verify they exist on GPU at the moment action_select is launched.
If buffers don't exist or are stored in different layouts, three options:
(a) Compute them just-in-time before action_select via the existing C51/IQN forward kernels
(b) Persist them in the trainer struct as new fields populated by the forward pass
(c) Refactor existing forward to expose them
Choose based on existing code patterns.
- [ ] **Step 3: Pass new arguments to launch_builder**
Modify the `launch_builder` chain (the `.arg(...)` sequence) to add the four new args after the existing args:
```rust
.arg(&self.c51_probs_buf)
.arg(&self.atom_values_buf)
.arg(&self.iqn_quantiles_buf)
.arg(&(self.config.n_atoms as i32))
```
The order MUST match the order added to the kernel signature in Task 2 Step 2.
- [ ] **Step 4: Compile**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
Expected: clean.
- [ ] **Step 5: Commit**
```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat(dqn): wire C51+IQN buffers to action_select in trainer path
Threads c51_probs, atom_values, iqn_quantiles, and n_atoms through
the experience_action_select kernel launch in the training path."
```
---
### Task 4: Buffer wiring in gpu_backtest_evaluator.rs
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs`
- [ ] **Step 1: Find the action_select call site**
Run: `grep -n "experience_action_select\|action_select_kernel" crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs`
The evaluator has its own action_select call for val. Find it (likely in `submit_dqn_step_loop_cublas` per memory).
- [ ] **Step 2: Pass new arguments**
Same as Task 3 Step 3. Add four new args to the `launch_builder` chain:
```rust
.arg(ch_c51_probs) // chunked c51_probs buffer
.arg(&self.atom_values_buf)
.arg(ch_iqn_quantiles) // chunked iqn_quantiles buffer
.arg(&n_atoms_i32)
```
If chunked buffers don't exist for c51_probs / iqn_quantiles in the evaluator (only flat per-step buffers exist), allocate chunked versions matching the existing chunk pattern (`chunked_states_buf`, etc.). New CudaSlice<f32> fields in the evaluator struct, allocated in `ensure_action_select_ready`.
- [ ] **Step 3: Compile + commit**
```bash
git add crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
git commit -m "feat(dqn): wire C51+IQN buffers to action_select in val/backtest path"
```
---
### Task 5: Add `train_active_frac` HEALTH_DIAG instrumentation
**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/metrics.rs`
- [ ] **Step 1: Find the HEALTH_DIAG val_active_frac line**
Run: `grep -n "active_frac" crates/ml/src/trainers/dqn/trainer/metrics.rs`
Identify how val_active_frac is computed (likely from per-direction action counts in the val window).
- [ ] **Step 2: Add training-rollout counterpart**
The training rollout stores actions in the experience collector. Find where actions histogrammed (likely in `monitor.action_counts` per HEALTH_DIAG metrics). Add a `train_active_frac` calculation:
```rust
// Per-direction action counts during training rollout
let train_short = monitor.action_counts[0..3].iter().sum::<usize>() as f32; // d=0
let train_long = monitor.action_counts[6..9].iter().sum::<usize>() as f32; // d=2
let train_active = train_short + train_long;
let train_total: f32 = monitor.action_counts.iter().sum::<usize>() as f32;
let train_active_frac = if train_total > 0.0 { train_active / train_total } else { 0.0 };
tracing::info!(
"HEALTH_DIAG[{}]: train_active_frac={:.4} (Long+Short during training rollout)",
epoch, train_active_frac,
);
```
(The action_counts indexing follows the project's 12-bin layout per memory `MEMORY.md` — verify against existing code; if action_counts is structured differently, adapt.)
- [ ] **Step 3: Compile + run a quick smoke**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib` — clean.
- [ ] **Step 4: Commit**
```bash
git add crates/ml/src/trainers/dqn/trainer/metrics.rs
git commit -m "feat(dqn): add train_active_frac HEALTH_DIAG metric
Counterpart to existing val_active_frac. Reports the fraction of
TRAINING rollout actions that were Long or Short. Required for L3
verification gate of Plan D (Phase 3) which checks Thompson is
generating directional exploration during training."
```
---
### Task 6: Phase 2 Test 2.A — End-to-end Thompson + argmax modes
**Files:**
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs`
- [ ] **Step 1: Write the test**
Append to `distributional_q_tests.rs`:
```rust
/// Phase 2 Test 2.A — Production action_select kernel produces correct mode
/// behavior for training (stochastic Thompson) and eval (deterministic argmax).
#[test]
#[ignore]
fn test_2a_production_kernel_modes() {
let stream = make_test_stream();
// Reuse Phase 0 Test 0.D failure-mode setup
// Launch experience_action_select (production kernel) directly via
// GpuExperienceCollector audit fixture (extended in Plan B Task 2 Step 4).
//
// Assert:
// eval_mode=true: 10 calls produce identical dir_idx (deterministic)
// eval_mode=false: 10000 calls produce diverse dir_idx with
// (Long+Short) ≥ 30% on the failure-mode setup
panic!("test_2a — implement using GpuExperienceCollector audit fixture with eval_mode flag");
}
```
- [ ] **Step 2: Implement using the production kernel directly**
Use the audit fixture from Plan B (which already exposes the production kernel). Set up controlled C51+IQN inputs matching Phase 0 Test 0.D failure-mode. Toggle eval_mode and assert behavior.
- [ ] **Step 3: Run + commit**
```bash
git add crates/ml/src/trainers/dqn/distributional_q_tests.rs
git commit -m "test(dqn): Phase 2 Test 2.A — production kernel mode behavior"
```
---
### Task 7: Phase 2 Test 2.B — Production kernel matches Phase 0 standalone wrapper
**Files:**
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 2 Test 2.B — Production kernel output matches Phase 0 standalone wrapper.
///
/// Run BOTH `experience_action_select` (production, Phase 2) and
/// `thompson_direction_test` (Phase 0 standalone) with identical inputs and
/// the same Philox seed. Outputs must be bit-identical.
#[test]
#[ignore]
fn test_2b_production_matches_standalone() {
// Setup synthetic C51+IQN inputs.
// Launch standalone (using launch_thompson_direction from Plan A).
// Launch production (using audit fixture from Plan B with same seed).
// Assert: standalone_dir_idx == production_dir_idx for 100 different seeds.
panic!("test_2b — implement; assert byte-identical output between standalone and production");
}
```
- [ ] **Step 2-3: Implement, run, commit**
The production kernel uses Philox internally; standalone uses LCG. To make them bit-identical, EITHER:
- Use the same RNG in both (refactor standalone to use Philox)
- Or compare DISTRIBUTIONS rather than per-call output (looser but still meaningful)
Choose: refactor standalone to Philox. Update `thompson_test_kernel.cu` (Plan A Task 1) to use `philox_uniform` rather than LCG. Then 2.B becomes truly bit-identical comparison.
```bash
git commit -m "test(dqn): Phase 2 Test 2.B — production matches Phase 0 standalone (bit-identical)"
```
---
### Task 8: Phase 2 Test 2.C — Magnitude/order/urgency unaffected
**Files:**
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 2 Test 2.C — Magnitude/order/urgency picks unchanged after Phase 2.
///
/// Compare against snapshot of pre-Phase-2 kernel output: with the same
/// inputs and same Philox seed, mag/ord/urg picks must match the snapshot.
/// (Direction is allowed to differ — that's the whole point of Phase 2.)
#[test]
#[ignore]
fn test_2c_mag_ord_urg_unaffected() {
// Generate a snapshot file `test_data/phase2_pre_change_snapshot.json`
// BEFORE landing the Phase 2 kernel changes. Snapshot is per-direction
// mag/ord/urg picks for 100 seed values.
//
// After landing changes: rerun, assert mag/ord/urg picks match snapshot
// for identical seeds. (Direction differs — that's fine.)
panic!("test_2c — generate snapshot pre-change; assert post-change matches");
}
```
- [ ] **Step 2: Generate snapshot**
This test requires generating the snapshot BEFORE Task 2-3 (the kernel/wiring changes). Workflow:
(a) Before starting Plan C, run a small test program that launches the existing kernel with controlled inputs + 100 seeds, captures `(mag_idx, ord_idx, urg_idx)` per seed, writes JSON to `test_data/phase2_pre_change_snapshot.json`.
(b) Start Plan C (Tasks 1-7).
(c) After Plan C kernel changes, run the test which loads the snapshot and re-runs with same inputs+seeds, asserting mag/ord/urg match.
If the snapshot wasn't captured first, this test must be re-engineered — perhaps by reverting Plan C temporarily, capturing snapshot, reapplying. Plan ahead.
- [ ] **Step 3: Implement, run, commit**
```bash
git add crates/ml/src/trainers/dqn/distributional_q_tests.rs test_data/phase2_pre_change_snapshot.json
git commit -m "test(dqn): Phase 2 Test 2.C — mag/ord/urg picks unaffected by direction Thompson"
```
---
### Task 9: Phase 2 Test 2.D — Real-batch e2e on converged checkpoint
**Files:**
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs`
- [ ] **Step 1: Write the test**
Append:
```rust
/// Phase 2 Test 2.D — Real-batch e2e: load converged checkpoint, run full
/// action select on a representative val batch, verify:
/// - Training mode: active_frac (Long+Short) > 40% over the batch
/// - Eval mode: dir_idx matches argmax(E[Q]) for each (state) tuple
#[test]
#[ignore]
fn test_2d_real_batch_e2e() {
let stream = make_test_stream();
// Reuse the converged checkpoint from Phase 0 Test 0.F.
// Run full forward → action_select on a batch of representative val states.
// Eval mode: verify dir_idx = argmax(E[Q]) per state.
// Train mode: count active dir_idx (Long + Short) across batch; assert frac > 40%.
panic!("test_2d — implement using converged checkpoint loader from Phase 0");
}
```
- [ ] **Step 2-3: Implement, run, commit**
```bash
git commit -m "test(dqn): Phase 2 Test 2.D — real-batch e2e on converged checkpoint"
```
---
### Task 10: Aggregation Contract table + memory pearl
**Files:**
- Modify: `docs/dqn-wire-up-audit.md`
- Create: `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_thompson_for_distributional_action_selection.md`
- Modify: `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md`
- [ ] **Step 1: Add Aggregation Contract table to audit doc**
In `docs/dqn-wire-up-audit.md`, add a new section after the existing distributional Q row:
```markdown
## Distributional Q-head Aggregation Contract
For any module producing per-(state, action) Q estimates as a distribution
(atoms, quantiles, ensembles), the contract requires:
| Q-head | Produces E[Q]? | Sampleable? | Wired to training Thompson? | Wired to eval argmax? |
|---|---|---|---|---|
| C51 atom kernel | ✅ | ✅ via inverse-CDF | ✅ direction only (v1) | ✅ direction only (v1) |
| IQN quantile kernel | ✅ | ✅ via uniform-τ interpolation | ✅ direction only (v1) | ✅ direction only (v1) |
| Ensemble Q-head (existing — currently used only for `ens_agree` metric) | needs verification | ✅ pick member uniformly | v2.1 enhancement | v2.1 enhancement |
| Future Q-head additions | required | required | required | required |
**Note**: TFT (task #148) is a Variable Selection Network for trunk feature
processing — not a Q-head. The contract applies specifically to modules that
produce per-(state, action) Q estimates.
The pre-commit hook (Invariant 7) will refuse merges of new distributional or
ensemble Q-heads that fail this table.
```
- [ ] **Step 2: Write the memory pearl**
Create `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_thompson_for_distributional_action_selection.md`:
```markdown
---
name: Thompson sampling for distributional RL action selection
description: When using C51/IQN/ensemble distributional Q, action selection MUST sample from the distribution (Thompson) — never aggregate to scalar E[Q] then Boltzmann/argmax — because aggregation discards uncertainty and structurally biases low-variance actions
type: feedback
---
When a value head represents Q as a distribution (atoms, quantiles, ensembles), action selection must use a sampling-based selector (Thompson sampling: draw one value from each action's distribution, pick argmax of those draws). Aggregating to scalar `E[Q]` then doing Boltzmann or argmax produces structural bias toward low-variance actions because aggregation collapses `(mean=0, σ=0)` and `(mean=−ε, σ=large)` to neighbouring scalars where the deterministic action wins regardless of latent option value.
**Why:** This pearl was distilled from the val-Flat-collapse investigation (commits Kelly fix `0c9d1ee39` and follow-ups). After fixing the Kelly cap, val_dir_dist still collapsed to 70%+ Hold/Flat in 1 epoch. The model rationally learned that Long/Short directional Q's were slightly negative from tx_cost while Flat had Q=0 from its `δ(0)` distribution. Aggregating to scalar `E[Q]` then Boltzmann/argmax produced this collapse. Two band-aid attempts at the symptom layer (ISV-adaptive Boltzmann tau, adaptive eps_dir floor) had no measurable effect — confirming the issue was in the AGGREGATION step, not in sampling sharpness.
**How to apply:** Before any new value head is added to the system (TFT Q-head, Mamba2 Q-head, etc.), verify it exposes a sampleable interface. Add a row to the Distributional Q-head Aggregation Contract table in `docs/dqn-wire-up-audit.md`. The pre-commit hook (Invariant 7) will refuse merges that fail.
For action selection: training uses Thompson sampling (drawing one value per action's distribution, picking argmax). Eval uses raw `argmax(E[Q])` (no exploration overlay; honest report of model's mode). The two are distinct selectors — never use Thompson at eval (would inflate reported edge by injecting exploration noise into reported decisions).
Magnitude/order/urgency branches are EXEMPT — they don't have the structural σ asymmetry that Flat has on direction. Applying Thompson to magnitude would prefer Full over Quarter (larger σ from larger position) and worsen the documented magnitude saturation pattern.
Reference: spec `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`.
```
- [ ] **Step 3: Add MEMORY.md index entry**
In `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md`, add to the pearl list:
```markdown
- [pearl_thompson_for_distributional_action_selection.md](pearl_thompson_for_distributional_action_selection.md) — distributional Q-heads must use sampling-based action selection (Thompson at training, argmax at eval) because aggregation to E[Q] discards uncertainty and structurally biases low-variance actions. From val-Flat-collapse investigation.
```
- [ ] **Step 4: Commit**
```bash
git add docs/dqn-wire-up-audit.md
git commit -m "docs(dqn): Distributional Q-head Aggregation Contract table
Documents the project-wide invariant: any new distributional Q-head
must expose sampleable interface and wire into Thompson (training) +
argmax (eval) action selection. Pre-commit hook (Invariant 7) will
enforce."
```
(Memory is local to the user; not committed to repo.)
---
### Task 11: Phase 2 exit gate
- [ ] **Step 1: Run all 4 Phase 2 tests**
Run: `cargo test -p ml --lib distributional_q_tests::test_2 -- --ignored --nocapture`
Expected: 2.A, 2.B, 2.C, 2.D all pass.
- [ ] **Step 2: Run all Phase 0 tests still pass**
Run: `cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture`
Expected: all 6 Phase 0 tests + 4 Phase 2 tests pass.
- [ ] **Step 3: L3 short smoke verification**
Run: `./scripts/argo-train.sh dqn --baseline --epochs 5 --folds 3 --gpu-pool ci-training-l40s --tag plan-c-l3-smoke`
Wait ~12 min. Check HEALTH_DIAG output:
- `train_active_frac > 0.40` across all 5 epochs of all 3 folds → primary gate A passes
- val_dir_dist == val_picked_dir_dist (Kelly fix property) → already verified
If train_active_frac < 0.40 → Thompson is not exploring; debug PRNG quality or interpolation logic.
- [ ] **Step 4: Commit Phase 2 exit confirmation**
```bash
git commit --allow-empty -m "phase2(dqn): Plan C complete — Thompson integration verified
All 4 Phase 2 unit tests pass. L3 smoke shows train_active_frac > 40%
across all training epochs. Plan D / Phase 3 unblocked."
```
---
## Self-Review
**Spec coverage:**
| Spec requirement | Implemented in |
|---|---|
| Component 4.1: kernel modification | Tasks 1, 2 |
| Component 4.2: buffer wiring (trainer) | Task 3 |
| Component 4.2: buffer wiring (evaluator) | Task 4 |
| Component 4.3: conviction stays E[Q]-based | Task 2 Step 4 |
| Component 4.4: Phase 2 unit tests (2.A-2.D) | Tasks 6, 7, 8, 9 |
| Component 4.5: aggregation contract + pearl | Task 10 |
| `train_active_frac` instrumentation (NEW) | Task 5 |
| Phase 2 exit gate (L3 smoke) | Task 11 |
All requirements covered. ✓
**Placeholder scan:** Tests 2.A-2.D have `panic!("...implement...")` scaffolds — these are explicit implementation steps to follow in Step 2-3 of each task, not spec placeholders. Cleared as the steps are completed.
**Type consistency:** `B0_SIZE = 4`, `N_IQN_QUANTILES = 5`, n_atoms (variable) consistent across all kernels and tests. Direction indexing (Short=0, Hold=1, Long=2, Flat=3) consistent.
---
## Plan complete and saved to `docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md`.

View File

@@ -0,0 +1,753 @@
# Plan D — Phase 3: Long Verification + Direction-Branch Dead-Code Cleanup 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:** Close out the Thompson-sampling rollout. Run L4 (long smoke, 1 seed × 6 folds × 30 epoch) and L5 (full validation matrix per Plan 5 Task 5) to verify Thompson holds up over many epochs. Then delete unreachable direction-only `eps_dir` adaptive-boost and Boltzmann tau-floor code paths inside `experience_action_select`. Magnitude/order/urgency code paths stay 100 % intact.
**Architecture:** No new mechanisms. Two verification layers (L4, L5) plus a tightly-scoped cleanup pass on `experience_kernels.cu`. The cleanup is mandated by `feedback_no_partial_refactor.md`: Phase 2 changed the shared "action-select-direction" contract from `eps-greedy + Boltzmann` to `Thompson + argmax`; every consumer of the old contract must migrate in the same logical change. Leaving direction-only `eps_dir` clamps + tau floors compiled but unreachable is the partial-refactor anti-pattern.
**Tech Stack:** Same as Plans A+B+C. New tooling: `argo-train.sh --multi-seed`, `argo-train.sh --folds`, `scripts/validation/check_all_tiers.py`, nsys profile harness from Plan 5 Task 3.
**Prerequisite:** Plans A + B + C complete. All 16 unit tests (6 Phase 0 + 6 Phase 1 + 4 Phase 2) passing. Plan C's L3 short smoke pass with `train_active_frac > 0.40` across all epochs.
**Spec reference:** Phase 3 of `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md` (lines 445-460), 5-Layer Verification Protocol (lines 464-489), Risk Register R3+R8 (lines 507-553).
---
## File Structure
| Path | Status | Responsibility |
|---|---|---|
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | Delete direction-only `eps_dir` adaptive-boost + EPS_FLOOR clamp + tau-floor lines; remove `eps_dir` from kernel signature if unused elsewhere |
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | Drop the `eps_dir`-related kernel-arg if signature changed |
| `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` | Modify | Same as above |
| `crates/ml/src/trainers/dqn/distributional_q_tests.rs` | Modify | Add Phase 3 regression test 3.A: load converged checkpoint, replay val-Flat-collapse setup, assert Thompson keeps `train_active_frac > 0.40` |
| `docs/dqn-v2-final-validation.md` | Create | Final L5 validation report (one row per Tier; same template as Plan 5 Task 5 Step 5.6) |
| `docs/dqn-wire-up-audit.md` | Modify | Mark direction-branch eps/tau rows as REMOVED with commit SHA reference |
| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md` | Modify | Reference Plan D close-out commit SHA in pearl entry created in Plan C |
| Argo deploy logs | Read | L4 + L5 evidence |
---
### Task 1: Phase 3 Test 3.A — Regression test against the original bug
**Files:**
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs`
**Why this test:** All 16 prior unit tests verify the new Thompson mechanism. None replays the *original failure*. If a future refactor accidentally reverts to scalar-aggregate-then-Boltzmann selection, no existing test catches it. Test 3.A is the regression anchor: it replays the post-Kelly val-Flat-collapse on the converged checkpoint and asserts Thompson reverses it.
- [ ] **Step 1: Write the failing test**
Append to `crates/ml/src/trainers/dqn/distributional_q_tests.rs`:
```rust
/// Phase 3 Test 3.A — Regression anchor against C51 expected-Q Hold/Flat
/// bias. Setup mirrors the original val-Flat-collapse pathology:
/// - Synthetic distributions: Flat = δ(0); Long/Short = small negative
/// mean with non-trivial σ (≥ 0.05).
/// Acceptance:
/// - Argmax(E[Q]) picks Flat for ≥ 95 % of states (replays the bug).
/// - Thompson sampling on the SAME distributions picks Long+Short for
/// ≥ 40 % of states (verifies the fix).
/// - Run BOTH on the production `experience_action_select` kernel,
/// not the standalone test kernel — production is what would regress.
///
/// This test is the durable artefact of the design lesson; if it ever
/// fails after a future refactor, that refactor reintroduced the bias
/// class and must be reverted.
#[test]
#[ignore]
fn test_3a_regression_anchor_c51_flat_bias() {
let stream = make_test_stream();
let n_states = 512;
let n_atoms = 51;
let v_min = -0.5_f32;
let v_max = 0.5_f32;
let dz = (v_max - v_min) / (n_atoms as f32 - 1.0);
let atom_values: Vec<f32> = (0..n_atoms).map(|a| v_min + a as f32 * dz).collect();
// Synth C51: Flat is δ(0); Short/Long/Hold = Gaussian centred at -0.005,
// σ = 0.05 (one-tx-cost expected loss with realistic uncertainty).
let mut c51_probs = vec![0.0_f32; n_states * 4 * n_atoms];
let mut iqn_quantiles = vec![0.0_f32; n_states * 4 * 5];
let directional_atom_pdf = gaussian_atom_pdf(&atom_values, -0.005, 0.05);
let flat_atom_pdf = delta_zero_atom_pdf(&atom_values);
for s in 0..n_states {
for d in 0..4 {
let pdf = if d == 3 { &flat_atom_pdf } else { &directional_atom_pdf };
let off = (s * 4 + d) * n_atoms;
c51_probs[off..off + n_atoms].copy_from_slice(pdf);
// IQN: same shape, expressed as five equispaced quantiles.
let q_off = (s * 4 + d) * 5;
let q = if d == 3 {
[0.0_f32; 5]
} else {
[-0.082, -0.039, -0.005, 0.029, 0.072]
};
iqn_quantiles[q_off..q_off + 5].copy_from_slice(&q);
}
}
// Push to GPU and call the production kernel twice: once eval (argmax),
// once train (Thompson).
let argmax_picks = launch_production_action_select(&stream, &c51_probs, &iqn_quantiles, &atom_values, /*eval_mode=*/true);
let thompson_picks = launch_production_action_select(&stream, &c51_probs, &iqn_quantiles, &atom_values, /*eval_mode=*/false);
let argmax_flat_frac = (argmax_picks.iter().filter(|&&d| d == 3).count() as f32) / n_states as f32;
let thompson_active_frac = (thompson_picks.iter().filter(|&&d| d == 0 || d == 2).count() as f32) / n_states as f32;
assert!(argmax_flat_frac >= 0.95, "argmax should still demonstrate the bias on this synth setup, got {argmax_flat_frac}");
assert!(thompson_active_frac >= 0.40, "Thompson must reverse the bias, got active_frac={thompson_active_frac}");
}
fn gaussian_atom_pdf(atom_values: &[f32], mean: f32, sigma: f32) -> Vec<f32> {
let raw: Vec<f32> = atom_values.iter().map(|&v| {
let z = (v - mean) / sigma;
(-0.5 * z * z).exp()
}).collect();
let total: f32 = raw.iter().sum();
raw.into_iter().map(|p| p / total).collect()
}
fn delta_zero_atom_pdf(atom_values: &[f32]) -> Vec<f32> {
let mut p = vec![0.0_f32; atom_values.len()];
let zero_idx = atom_values.iter().position(|&v| v.abs() < 1e-6).expect("atom grid must contain zero");
p[zero_idx] = 1.0;
p
}
```
`launch_production_action_select` is a thin Rust wrapper that fills the same kernel inputs `experience_action_select` would receive in production, then reads back the resulting `actions` GPU buffer and decodes `dir_idx = action / 27`. Reuse the buffer-wiring pattern established in Plan C Tasks 3-4.
- [ ] **Step 2: Run test to verify it fails before checkpoint loaded**
Run: `cargo test -p ml --lib distributional_q_tests::test_3a_regression_anchor_c51_flat_bias -- --ignored --nocapture`
Expected before completing helpers: FAIL with "function not defined" or `launch_production_action_select` panic.
- [ ] **Step 3: Implement helpers + run test to PASS**
Implement `launch_production_action_select` (~ 60 LOC). Run again; expect PASS — argmax_flat_frac ≥ 0.95, thompson_active_frac ≥ 0.40.
- [ ] **Step 4: Commit**
```bash
git add crates/ml/src/trainers/dqn/distributional_q_tests.rs
git commit -m "test(dqn): Phase 3 Test 3.A — regression anchor against C51 expected-Q Flat bias
Replays the original val-Flat-collapse setup on the production
experience_action_select kernel: argmax(E[Q]) picks Flat ≥ 95% (bug
reproduces); Thompson picks Long+Short ≥ 40% (fix verified). Durable
guard against future refactors reintroducing the aggregation-bias
class."
```
---
### Task 2: L4 long smoke (1 seed × 6 folds × 30 epoch)
**Files:**
- Read: Argo logs at `argo-server.foxhunt.internal/workflows/foxhunt/...`
- Read: HEALTH_DIAG output
**Why L4:** L3 (Plan C exit gate) was 5 epochs, three folds, primary gate `train_active_frac > 0.40` only. L4 verifies Phase 3 Risk R8 (long-term edge discovery): does val_sharpe trend non-negative over 30 epochs? Does val_active_frac evolve as the model finds edge? If after 30 epochs val_sharpe stays flat AND val_active_frac stays low, the data has no edge — that is correct model behaviour, not a bug.
- [ ] **Step 1: Push to origin/main**
Per Invariants: always `git push origin main` before deploy.
```bash
git push origin main
```
- [ ] **Step 2: Launch L4**
```bash
./scripts/argo-train.sh dqn \
--baseline \
--epochs 30 \
--folds 6 \
--commit "$(git rev-parse HEAD)" \
--gpu-pool ci-training-l40s \
--tag plan-d-l4-long-smoke
```
Wait ~1 hour. Monitor via Argo UI or:
```bash
argo list -l foxhunt-tag=plan-d-l4-long-smoke
argo wait @latest -l foxhunt-tag=plan-d-l4-long-smoke --timeout 90m
```
- [ ] **Step 3: Gather HEALTH_DIAG metrics**
```bash
./scripts/gather-multi-seed-metrics.sh plan-d-l4-long-smoke > /tmp/l4-metrics.json
```
- [ ] **Step 4: Verify L4 acceptance criteria**
Pass criteria from spec line 473:
- `val_sharpe` trend across 30 epochs is non-negative (linreg slope ≥ 0)
- `val_active_frac` evolves: not stuck at 0 across all 30 epochs
- `train_active_frac > 0.40` across every epoch (regression check vs L3)
Use:
```bash
python3 scripts/validation/check_l4_long_smoke.py /tmp/l4-metrics.json
```
If `check_l4_long_smoke.py` does not exist, write it (~ 50 LOC) computing slope of last-fold val_sharpe via numpy.polyfit and asserting the three criteria. Add it under `scripts/validation/` next to existing `check_tier1.py` etc.
Exit 0 required. If exit non-zero:
- val_sharpe slope < 0 → escalate per Risk R3 / R8: open follow-up task to investigate; this is NOT cause for L4 retry. Per `feedback_stop_on_anomaly.md`: re-running hoping for a different outcome is anti-pattern.
- train_active_frac < 0.40 in any epoch → Phase 2 Thompson regression; revisit Plan C kernel.
- [ ] **Step 5: Commit L4 evidence**
```bash
git add scripts/validation/check_l4_long_smoke.py
git commit -m "validation(dqn): Plan D L4 long smoke — train_active_frac ≥ 40%, val_sharpe non-negative slope, val_active_frac evolves"
```
---
### Task 3: L5 full validation matrix (per Plan 5 Task 5)
**Files:**
- Read: `docs/superpowers/plans/2026-04-24-dqn-v2-plan-5-validation.md` Task 5
- Run: `scripts/validation/check_all_tiers.py`
- Create: `docs/dqn-v2-final-validation.md`
L5 is exactly Plan 5 Task 5 re-executed against the post-Thompson commit. The Plan 5 Task 5 protocol is the source of truth for procedure; this task is a thin wrapper that links into it and adds Thompson-specific evidence rows.
- [ ] **Step 1: Pre-flight checks**
```bash
git status --porcelain | test -z "$(cat)"
git fetch origin main
test "$(git rev-parse main)" = "$(git rev-parse origin/main)"
# All Phase 0+1+2+3 tests pass locally
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture 2>&1 | tee /tmp/dist-tests.log
grep -q "test result: ok" /tmp/dist-tests.log
# Smoke tests still pass
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tee /tmp/smoke.log
grep -q "test result: ok" /tmp/smoke.log
```
- [ ] **Step 2: Launch L5 = Plan 5 Task 5 multi-seed × multi-fold matrix**
```bash
git push origin main
./scripts/argo-train.sh \
--multi-seed 5 \
--folds 6 \
--commit "$(git rev-parse HEAD)" \
--tag dqnv2-thompson-final
```
Wait ~ 6 hours. Monitor:
```bash
argo list -l foxhunt-tag=dqnv2-thompson-final
argo wait @latest -l foxhunt-tag=dqnv2-thompson-final --timeout 8h
```
- [ ] **Step 3: Gather L5 metrics**
```bash
./scripts/gather-multi-seed-metrics.sh dqnv2-thompson-final > /tmp/dqnv2-thompson-final-metrics.json
./scripts/aggregate-norm-stats.py \
--input-dir /mnt/training-data/norm-stats/dqnv2-thompson-final \
--output /mnt/training-data/norm-stats/dqnv2-thompson-final-aggregate.json
```
- [ ] **Step 4: Run Tier 1 + Tier 2 + Tier 3 validation**
```bash
python3 scripts/validation/check_all_tiers.py /tmp/dqnv2-thompson-final-metrics.json --warmup-end 15
```
Exit 0 required on the FIRST run. If non-zero:
- DO NOT retry with different seeds or a different tag (per `feedback_stop_on_anomaly.md` and Plan 5 Step 5.4)
- Open the failing tier as a separate task for investigation
- Plan D does not land until Tier 1 + Tier 2 + Tier 3 all PASS on the first run
- [ ] **Step 5: nsys regression check (Phase H continuity)**
```bash
mc cp minio/foxhunt-training-artifacts/profiles/baseline-plan5/profile-seed0-fold0.nsys-rep /tmp/baseline.nsys-rep
mc cp minio/foxhunt-training-artifacts/profiles/$(git rev-parse --short HEAD)/profile-seed0-fold0.nsys-rep /tmp/current.nsys-rep
python3 scripts/compare-nsys-profiles.py /tmp/baseline.nsys-rep /tmp/current.nsys-rep
```
Exit 0 required. The Thompson kernel is ~ 50 µs per batch (per Risk R1 mitigation); regression must stay within Plan 5 Task 3 tolerance (typically ≤ 5 % step-time delta).
If nsys regression > tolerance, open as Plan H follow-up but DO NOT block Plan D — Phase 3 spec line 455 says "Phase H continues independently".
- [ ] **Step 6: Document final run**
Create `docs/dqn-v2-final-validation.md`:
```markdown
# DQN v2 — Final Validation Run (post-Thompson)
**Date:** <DATE>
**Commit SHA:** <SHA>
**Argo tag:** dqnv2-thompson-final
**Matrix:** 5 seeds × 6 folds = 30 training jobs
**GPU hours:** ~<H> hours L40S
**Thompson rollout:** Plans A+B+C+D landed at <SHA>
## Tier 1 — Convergence discipline
[paste check_tier1.py output]
## Tier 2 — Behavioural parity
[paste check_tier2.py output, including val_trades_per_bar, val_active_frac,
direction entropy, and the new train_active_frac line]
## Tier 3 — Profitability
[paste check_tier3.py output with val_sharpe, win_rate, profit_factor,
std across seeds]
## Thompson-specific evidence
| Metric | Value | Acceptance |
|---|---|---|
| train_active_frac mean across all (seed, fold, epoch) | <X> | ≥ 0.40 |
| val_active_frac final-epoch mean across seeds | <X> | reported honestly (no gate) |
| val_dir_dist == val_picked_dir_dist | ✓ / ✗ | Kelly-fix invariant retained |
| Phase 3 Test 3.A (regression anchor) | PASS / FAIL | PASS required |
## nsys regression
[paste compare-nsys-profiles.py output]
## HEALTH_DIAG sample lines (last 3 epochs, seed 0, fold 5)
[paste 3 lines showing reward_split, temporal_reward, attn, aux,
controller, train_active_frac, val_active_frac]
## Sign-off
Thompson sampling integrated into DQN v2 direction branch. C51 expected-Q
Hold/Flat aggregation bias resolved structurally — see
`docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`.
All 9 invariants preserved across every commit in Plans A-D.
```
- [ ] **Step 7: Commit L5 evidence**
```bash
git add docs/dqn-v2-final-validation.md
git commit -m "validation(dqn): Plan D L5 final matrix — Tier 1+2+3 PASS first run"
```
---
### Task 4: Direction-branch dead-code cleanup — `eps_dir` adaptive boost + EPS_FLOOR clamp
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (lines 806-865 region)
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (kernel-arg call site for `experience_action_select`)
- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` (same)
**Why this cleanup:** Phase 2 replaced direction-branch action selection wholesale (eps-greedy + Boltzmann → Thompson + argmax). The `eps_dir` clamp at line 816 and adaptive-boost block at lines 822-861 are now structurally unreachable for direction (the eps-greedy branch is gone). Per `feedback_no_partial_refactor.md`: every consumer of the changed contract must migrate in the same logical change. Dead-but-compiled code is the partial-refactor anti-pattern. Magnitude/order/urgency continue to use the EPS_FLOOR clamp + adaptive boost is direction-only so removing only the direction lines is a safe surgical edit.
- [ ] **Step 1: Read the current state**
```bash
sed -n '800,870p' crates/ml/src/cuda_pipeline/experience_kernels.cu
```
Confirm:
- Line ~806 declares `float eps_dir = ...`
- Line ~816 has `eps_dir = fmaxf(eps_dir, EPS_FLOOR);`
- Lines ~822-861 are an adaptive-boost block scaling `eps_dir` from ISV[71/72] passive_pressure
- Line ~861 ends with `eps_dir = fmaxf(eps_dir, adaptive_floor);`
- Line ~922 has `if (!eval_mode && philox_uniform(i, timestep, rng_ctr++) < eps_dir) {` (the eps-greedy direction-sample branch — replaced by Thompson in Plan C Task 2)
If any line numbers differ from the spec, that's expected — work from semantic location, not literal line numbers (Plan C will have shifted lines around the kernel).
- [ ] **Step 2: Identify safe deletion span**
After Plan C Task 2, the direction-branch action select reads:
```cuda
if (eval_mode) { ... argmax E[Q] ... } else { ... Thompson ... }
```
NEITHER path consumes `eps_dir`. Therefore:
- Delete the `float eps_dir = fminf(epsilon * eps_exp_mult, 1.0f);` declaration (line ~806).
- Delete `eps_dir = fmaxf(eps_dir, EPS_FLOOR);` (line ~816).
- Delete the entire adaptive-boost block (lines ~822-861) including its containing comment block.
- Delete the now-orphaned `if (!eval_mode && philox_uniform(...) < eps_dir) {...}` direction-eps-greedy branch (line ~922 region) — Plan C should have removed this already; verify.
KEEP unchanged:
- `eps_mag`, `eps_ord`, `eps_urg` declarations
- `eps_mag = fmaxf(eps_mag, EPS_FLOOR);` and the two sibling lines
- All uses of `eps_mag`, `eps_ord`, `eps_urg` in their respective magnitude/order/urgency branches
- [ ] **Step 3: Apply the edits**
Edit `crates/ml/src/cuda_pipeline/experience_kernels.cu` removing the direction-only lines listed in Step 2.
- [ ] **Step 4: Grep for residual `eps_dir` references**
```bash
grep -rn "eps_dir" crates/ml/
```
Expected: zero matches in `.cu` files. If any matches remain (e.g. in a comment, in a Rust `*.rs` file kernel-arg call, in `gpu_dqn_trainer.rs` or `gpu_backtest_evaluator.rs` where the C++ kernel signature is mirrored), follow up:
- If `experience_action_select` kernel signature still includes `eps_dir`/passive_pressure parameters, drop them. Update Rust call sites to match.
- If only ISV[71]/ISV[72] producer kernels remain (these write `passive_pressure` even though direction no longer reads it), DO NOT delete those producers — `feedback_isv_for_adaptive_bounds.md`: ISV signal bus values may be consumed by future modules. Adding a comment that no current consumer uses these is acceptable; deleting them is not.
- [ ] **Step 5: Build and run unit tests**
```bash
SQLX_OFFLINE=true cargo build --release -p ml 2>&1 | tee /tmp/build.log
grep -q "error\[" /tmp/build.log && { echo "BUILD FAILED"; exit 1; }
cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture
```
Expected: build succeeds. All 17 tests (Phases 0+1+2+3) PASS — including Test 3.A regression anchor on the cleaned-up kernel.
- [ ] **Step 6: Re-run L3 short smoke as a regression check**
```bash
git push origin HEAD:main
./scripts/argo-train.sh dqn --baseline --epochs 5 --folds 3 \
--commit "$(git rev-parse HEAD)" --gpu-pool ci-training-l40s \
--tag plan-d-cleanup-l3-regression
argo wait @latest -l foxhunt-tag=plan-d-cleanup-l3-regression --timeout 30m
./scripts/gather-multi-seed-metrics.sh plan-d-cleanup-l3-regression > /tmp/l3-regression.json
python3 -c "
import json
m = json.load(open('/tmp/l3-regression.json'))
mins = [e['train_active_frac'] for e in m['epochs']]
assert min(mins) >= 0.40, f'cleanup regression: min train_active_frac={min(mins)}'
print('OK: cleanup preserves train_active_frac >= 0.40')
"
```
If `train_active_frac` regresses below 0.40 after cleanup, REVERT the cleanup commit immediately and open as a follow-up task. The tau-floor / eps_dir code may have been a hidden Thompson dependency that needs deeper investigation.
- [ ] **Step 7: Commit cleanup**
```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu \
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
git commit -m "cleanup(dqn): Plan D — remove direction-only eps_dir adaptive boost + EPS_FLOOR clamp
Phase 2 replaced direction-branch action selection (eps-greedy +
Boltzmann → Thompson + argmax). The eps_dir adaptive boost (ISV[71/72]
passive_pressure → eps floor scaled to 0.5) and the EPS_FLOOR clamp
on eps_dir are now structurally unreachable for the direction branch.
Per feedback_no_partial_refactor.md: every consumer of a changed shared
contract must migrate in the same logical change.
Magnitude/order/urgency keep their EPS_FLOOR + eps_mag/ord/urg paths
intact — they were never reached by Thompson.
Verified: 17/17 distributional_q_tests pass, L3 regression smoke
maintains train_active_frac >= 0.40."
```
---
### Task 5: Direction-branch dead-code cleanup — Boltzmann tau-floor
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` (around lines 887-944)
**Why:** Same partial-refactor argument. The direction Boltzmann tau-floor (lines 935-944 region, where `q_dir_abs_ref` clamps `tau_d`) is unreachable now that direction uses Thompson (no Boltzmann). The mag/ord/urg branches keep their tau-floor logic verbatim.
- [ ] **Step 1: Read the current state**
```bash
sed -n '880,950p' crates/ml/src/cuda_pipeline/experience_kernels.cu
```
Confirm:
- Comment around line 887: "Branch 0: direction — Boltzmann softmax with ISV-adaptive tau floor."
- Around line 935-944: `q_dir_abs_ref` computation and `float tau_floor = fmaxf(q_dir_abs_ref, 0.01f); float tau_d = fmaxf(q_max_d - q_min_d, tau_floor);`
If Plan C Task 2 already removed these along with the direction Boltzmann path, this task may be a no-op. Verify by searching:
```bash
grep -n "q_dir_abs_ref\|tau_d\|tau_floor.*direction" crates/ml/src/cuda_pipeline/experience_kernels.cu
```
If grep returns no direction-only lines, skip to Step 5 ("commit no-op").
- [ ] **Step 2: Identify safe deletion span**
Delete:
- The "Branch 0: direction — Boltzmann softmax with ISV-adaptive tau floor" header comment.
- The `q_dir_abs_ref` computation (the EMA from ISV slot — but only if it's used ONLY for `tau_floor` for direction; if mag branch reads it too, KEEP the producer + mark it unused-by-direction in a one-line comment).
- The direction-only `tau_floor`, `tau_d`, and the direction Boltzmann softmax block (replaced by Thompson in Plan C Task 2 — this verifies Plan C didn't leave residue).
KEEP unchanged:
- All `tau_mag`, `tau_ord`, `tau_urg` declarations and their `fmaxf(_, 0.01f)` clamps.
- All `q_mag_abs_ref`, `q_ord_abs_ref`, `q_urg_abs_ref` ISV producers and consumers.
- [ ] **Step 3: Apply the edits**
Edit `experience_kernels.cu` removing the direction-only Boltzmann tau-floor lines.
- [ ] **Step 4: Build, run unit tests, regression smoke**
```bash
SQLX_OFFLINE=true cargo build --release -p ml 2>&1 | tee /tmp/build.log
grep -q "error\[" /tmp/build.log && { echo "BUILD FAILED"; exit 1; }
cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture
```
Expected: 17/17 pass.
If the deletion span was non-trivial (anything beyond a comment), re-run the L3 regression smoke as in Task 4 Step 6.
- [ ] **Step 5: Commit cleanup (or no-op)**
If real edits made:
```bash
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "cleanup(dqn): Plan D — remove direction-only Boltzmann tau-floor
Phase 2 Thompson replaces direction Boltzmann; the q_dir_abs_ref-driven
tau_floor for the direction branch is now structurally unreachable.
Magnitude/order/urgency keep tau_floor + Boltzmann intact.
Per feedback_no_partial_refactor.md: every consumer of the changed
shared contract must migrate in the same logical change."
```
If no-op (Plan C already cleaned this):
```bash
git commit --allow-empty -m "cleanup(dqn): Plan D Task 5 — verified Plan C removed direction tau-floor; no residual code"
```
---
### Task 6: Audit-doc + memory-pearl cross-reference
**Files:**
- Modify: `docs/dqn-wire-up-audit.md`
- Modify: `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md`
- [ ] **Step 1: Audit-doc update**
In `docs/dqn-wire-up-audit.md`, find the rows added in Plan C Task 10 (Distributional Q-head Aggregation Contract) and the rows for `eps_dir` adaptive boost / direction tau-floor (if present). Update:
- Append to the Aggregation Contract table description: "Phase 3 (Plan D, commit `<SHA>`) cleaned up direction-only `eps_dir` adaptive-boost and Boltzmann tau-floor code; magnitude/order/urgency paths intact."
- If there's a per-component "wire-up status" row for `eps_dir` adaptive boost, mark it: `REMOVED in <SHA> — direction-branch unreachable after Thompson migration; magnitude/order/urgency EPS_FLOOR + eps_mag/ord/urg paths preserved.`
- Same for direction Boltzmann tau-floor row.
- [ ] **Step 2: Memory-pearl cross-reference**
In `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md`, find the line added in Plan C Task 10 referencing `pearl_thompson_for_distributional_action_selection.md`. Append:
```
- [pearl_thompson_for_distributional_action_selection.md](pearl_thompson_for_distributional_action_selection.md) — distributional Q-heads must use sampling-based action selection. Plans A-D landed: Phase 0/1 verification, Phase 2 production integration (commit `<SHA>`), Phase 3 long verification + direction-branch dead-code cleanup (commit `<SHA>`).
```
(Replace existing line in-place; do not duplicate.)
- [ ] **Step 3: Commit**
```bash
git add docs/dqn-wire-up-audit.md
git commit -m "docs(dqn): Plan D close-out — Aggregation Contract notes cleanup; eps_dir + direction tau-floor rows REMOVED with Plan D SHA reference"
```
(Memory file is local; no commit needed.)
---
### Task 7: Final spec status footer
**Files:**
- Modify: `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`
- [ ] **Step 1: Append status footer**
Append at end of file:
```markdown
---
## Status (post-implementation)
**LANDED:** <DATE>
**Final commit SHA:** <SHA>
**Plans:** A (`<SHA>`), B (`<SHA>`), C (`<SHA>`), D (`<SHA>`)
**Verification matrix:**
- L1 (unit tests): 17/17 PASS — 6 Phase 0 + 6 Phase 1 + 4 Phase 2 + 1 Phase 3 regression
- L2 (GPU integration on converged checkpoint): PASS — Tests 0.F + 2.B + 2.D + 3.A
- L3 (3-fold × 5-epoch L40S): train_active_frac > 0.40 across all epochs
- L4 (1-seed × 6-fold × 30-epoch L40S): val_sharpe non-negative slope; val_active_frac evolves
- L5 (5-seed × 6-fold matrix): Tier 1 + Tier 2 + Tier 3 PASS first run; nsys within tolerance
**Cleanup**: direction-only `eps_dir` adaptive-boost (`d54b49efc`) and direction Boltzmann tau-floor (`7a3d88646`) deleted as unreachable code. Magnitude/order/urgency paths preserved. Per `feedback_no_partial_refactor.md`.
**v2 enhancements deferred:** triple-source Thompson, persistent Thompson, CVaR-eval, IDS, hierarchical Thompson, curiosity coupling — each a follow-up spec.
**Aggregation Contract enforced:** any future distributional Q-head must expose sampleable interface and wire into Thompson (training) + argmax (eval). Pre-commit hook (Invariant 7) refuses non-compliant merges.
```
- [ ] **Step 2: Commit**
```bash
git add docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md
git commit -m "spec(dqn): Plan D close-out — Thompson sampling LANDED; status footer added"
```
---
### Task 8: Phase 3 exit gate
- [ ] **Step 1: All unit tests pass**
```bash
cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture
```
Expected: 17/17 pass (Phases 0+1+2+3).
- [ ] **Step 2: All smoke tests pass**
```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
cargo test -p ml --lib -- smoke_tests --ignored --nocapture
```
Expected: all PASS.
- [ ] **Step 3: L4 evidence file present**
```bash
test -f /tmp/l4-metrics.json && echo "L4 metrics gathered"
```
- [ ] **Step 4: L5 final validation doc committed**
```bash
test -f docs/dqn-v2-final-validation.md
git log --format=%H -- docs/dqn-v2-final-validation.md | head -1
```
- [ ] **Step 5: Cleanup commits landed**
```bash
git log --grep="cleanup(dqn): Plan D" --oneline
```
Expected: at least one commit (Task 4); possibly two (Task 5 if non-no-op).
- [ ] **Step 6: nsys regression within tolerance**
Already verified in Task 3 Step 5. Confirm by re-reading `scripts/compare-nsys-profiles.py` exit-0 status from that step.
- [ ] **Step 7: Final close-out commit**
```bash
git push origin main
git commit --allow-empty -m "phase3(dqn): Plan D complete — Thompson sampling rollout LANDED
L1: 17/17 unit tests pass (Phases 0+1+2+3 incl. Test 3.A regression anchor)
L3: train_active_frac > 0.40 across all epochs (re-verified post-cleanup)
L4: 30-epoch val_sharpe non-negative slope; val_active_frac evolves
L5: Tier 1 + Tier 2 + Tier 3 PASS first run; nsys within tolerance
Cleanup: direction-only eps_dir adaptive-boost + Boltzmann tau-floor removed.
Spec docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md
status footer updated. Aggregation Contract pre-commit hook enforced.
C51 expected-Q Hold/Flat aggregation bias structurally resolved."
git push origin main
```
---
## Self-Review
**Spec coverage check (against spec lines 445-460 + 464-489):**
| Spec deliverable | Task |
|---|---|
| L4 multi-fold deploy (1 seed × 6 folds × 30 epoch) | Task 2 |
| L5 full validation matrix per Plan 5 Task 5 (Tier 1/2/3 PASS) | Task 3 |
| Direction-branch eps_dir adaptive-boost cleanup | Task 4 |
| Direction-branch Boltzmann tau-floor cleanup | Task 5 |
| nsys regression check | Task 3 Step 5 |
| Final architecture doc updates | Tasks 6 + 7 |
| L1 (unit tests) gate | Task 8 Step 1 |
| L2 (integration tests) gate | Test 3.A in Task 1 covers this for Phase 3 |
| Spec line 454 detail: `experience_kernels.cu` lines 814-865 cleanup | Task 4 explicit |
| KEEP Kelly cap fix (`0c9d1ee39`), KEEP train return + sharpe annualization fixes (non-tau parts of `7a3d88646`) | Implicit (no removal task touches these) |
| Aggregation Contract pre-commit hook continues to enforce | Task 6 references it; the hook itself ships in Plan C Task 10 |
All deliverables covered.
**Placeholder scan:**
- No "TBD"/"TODO"/"implement later" strings.
- Test 3.A includes complete code for the kernel-input synthesis and assertions; only the `launch_production_action_select` wrapper is described not coded — it reuses the buffer-wiring pattern landed in Plan C Tasks 3-4 (referenced explicitly), which is the appropriate cross-plan reference rather than copy-paste.
- L4 acceptance script `check_l4_long_smoke.py` is sized (~ 50 LOC) and pinned to numpy.polyfit; this is a complete instruction.
**Type/symbol consistency:**
- `train_active_frac` HEALTH_DIAG metric introduced in Plan C Task 5; consumed in Task 2 Step 4 and Task 4 Step 6 here. Same name throughout.
- `experience_action_select` kernel name consistent with Plan C.
- `EPS_FLOOR` constant referenced consistently with the kernel source (`experience_kernels.cu` line 815).
- `dqnv2-thompson-final` Argo tag matches the L5 metric-gather command and the validation doc.
- `pearl_thompson_for_distributional_action_selection.md` filename matches Plan C Task 10 Step 2.
- `docs/dqn-v2-final-validation.md` matches Plan 5 Task 5 Step 5.6 template.
**Sequencing check:**
- Test 3.A (Task 1) lands BEFORE the cleanup tasks (4, 5) — gives a regression anchor that survives the cleanup.
- L4 (Task 2) lands BEFORE L5 (Task 3) — short before long, per the L3→L4→L5 spec progression.
- Cleanup (Tasks 4, 5) lands AFTER L5 (Task 3) — verifies the production system is healthy before doing surgical removals.
- Wait — the spec says cleanup is "+ 0.5 day" on top of Plan 5 budget but doesn't strictly mandate sequencing. Re-reading spec line 460: cleanup is treated as part of Phase 3 alongside L4/L5. Acceptable to do cleanup either before or after L5; doing it AFTER L5 is safer (we know the production code passes L5 first; cleanup is then a low-risk surgical pass). Keep current order.
- Audit doc + spec footer (Tasks 6, 7) reference commit SHAs from earlier tasks — must come last. Correct.
- Final close-out (Task 8) is the empty-commit gate — must come last. Correct.
**Risk-coverage check (against spec Risk Register R1-R8):**
- R1 (Thompson cost): nsys regression in Task 3 Step 5 catches.
- R2 (PRNG quality): Test 3.A re-exercises Phase 0 Test 0.B's PRNG path on production kernel.
- R3 (Thompson at training, argmax at eval — wasted exploration): explicitly accepted as by-design tradeoff per spec; no code-level mitigation, but L5 Tier 3 catches if it becomes a profitability problem.
- R8 (slow edge discovery): L4 30 epochs targets this directly per spec line 543.
All risks covered or explicitly accepted.
**Self-review complete. No fixes needed.**
---
## Execution Handoff
After Plans A + B + C + D all written, the writing-plans skill workflow concludes by offering execution choice:
**1. Subagent-Driven (recommended)** — fresh subagent per task, two-stage review between tasks.
**2. Inline Execution** — execute tasks in this session using executing-plans, batch with checkpoints.
Plan D is the last plan; once it's accepted, execution can begin from Plan A.