# 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`

> **Note: revised 2026-04-29** to match production C51 adaptive-support
> architecture. The prior wording (post-softmax `c51_probs_dir`, single
> linear `atom_values`, rollout-time `iqn_quantiles_dir`) didn't match
> production: the collector materialises only RAW `b_logits`, the C51
> support is per-sample per-direction adaptive (`per_sample_support
> [N, 4, 3]` + optional `atom_positions [4, num_atoms]`), and IQN
> inference is training-only — there is no rollout-time IQN quantile
> buffer to consume. Direction-branch Thompson is now
> single-distribution C51 with the production adaptive support; the
> principled posterior-sample fix to UCB asymmetry is preserved (sourced
> from the existing C51 rollout distribution). See the spec footer
> Status entry "2026-04-29 — Plan C T2 amendment" for the full triage.

- [ ] **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 the AMENDED set of buffers — RAW direction-branch C51 logits, the per-sample per-direction adaptive support tile, and (optionally) the non-linear atom positions. IQN buffers are NOT threaded; rollout-time IQN inference does not exist.

```cuda
extern "C" __global__ void experience_action_select(
    // ... existing args up to and including out_magnitude_conviction ...
    const float* __restrict__ b_logits_dir,       // [N, b0_size, n_atoms]   NEW (raw logits; collector's exp_b_logits slice)
    const float* __restrict__ per_sample_support, // [N, b0_size, 3] str-12  NEW (adaptive [v_min, v_max, delta_z])
    const float* __restrict__ atom_positions,     // [b0_size, n_atoms]      NEW (NULL = linear support)
    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 the single-distribution C51 Thompson loop. The amendment introduces two new device-inline helpers (`softmax_c51_inline`, `compute_atom_values_inline`) — installed alongside the four Plan C Task 1 primitives, in the same `experience_kernels.cu` Thompson math section.

```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).
     *
     * Single-distribution C51 with production's adaptive per-direction
     * support. Math reference: thompson_test_kernel.cu (Phase 0 standalone)
     * — the v2 entry points (`direction_thompson_v2_test`, `argmax_eq_v2_test`)
     * mirror this loop for Test 2.B bit-equivalence verification.
     *
     * E[Q]-derived `e_dir` must be reused for the conviction calculation
     * below (per spec §"Conviction stays E[Q]-based"). Computed once here.
     */
    #ifndef THOMPSON_MAX_ATOMS
    #define THOMPSON_MAX_ATOMS 128  /* matches c51_loss_kernel.cu MAX_ATOMS */
    #endif
    float probs_d[THOMPSON_MAX_ATOMS];
    float atom_vals_d[THOMPSON_MAX_ATOMS];
    const float* logits_i = b_logits_dir + (long long)i * b0_size * n_atoms;

    float e_dir[4];  // b0_size = 4 in current 4-direction action space
    for (int d = 0; d < b0_size; d++) {
        softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
        compute_atom_values_inline(per_sample_support, atom_positions,
                                   i, d, b0_size, n_atoms, atom_vals_d);
        e_dir[d] = compute_e_c51_inline(probs_d, atom_vals_d, n_atoms);
    }

    if (eval_mode) {
        /* EVAL: argmax of 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: per-direction Thompson sample, argmax over directions.
         * Recompute probs/atoms for each direction — caching across
         * directions would cost b0_size * 2 * MAX_ATOMS stack floats and
         * risk local-memory spill. */
        float best_sample = -CUDART_INF_F;
        int best_d = 0;
        for (int d = 0; d < b0_size; d++) {
            softmax_c51_inline(logits_i + d * n_atoms, n_atoms, probs_d);
            compute_atom_values_inline(per_sample_support, atom_positions,
                                       i, d, b0_size, n_atoms, atom_vals_d);
            float u = philox_uniform(i, timestep, rng_ctr++);
            float q_sample = sample_c51_inverse_cdf(probs_d, atom_vals_d, n_atoms, u);
            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`.

---

## Status

**2026-04-29 — Plan C T2 amendment.** The plan-as-authored Task 2 kernel
signature (`c51_probs_dir`, single linear `atom_values [n_atoms]`,
`iqn_quantiles_dir`) was incompatible with production:

1. **No post-softmax `c51_probs_dir` on the collector.** Only raw
   `exp_b_logits` is materialised; softmax happens inline in consumers.
2. **C51 support is per-sample per-direction adaptive.**
   `c51_loss_kernel.cu` (lines 615-744) consumes `per_sample_support
   [N, 4, 3]` (stride-12, per-direction `[v_min, v_max, delta_z]`) plus
   optional non-linear `atom_positions [4, num_atoms]`. A single
   global `atom_values [n_atoms]` does not exist.
3. **No rollout-time IQN inference.** `gpu_iqn_head.rs` is training-only;
   the collector does not produce per-direction IQN quantile buffers.

The amended Task 2 kernel signature passes raw `b_logits_dir`,
`per_sample_support`, optional `atom_positions`, and `n_atoms`.
Direction-branch Thompson is single-distribution C51 with adaptive
support — the principled posterior-sample fix to the UCB selector/target
asymmetry is preserved (the C51 distribution alone provides per-direction
posterior samples), without requiring infrastructure that doesn't exist.

Tasks 3+4 (buffer wiring) are unblocked — collector and evaluator already
own `per_sample_support_buf` and `exp_b_logits` in production.

Phase 0 standalone test kernel (`thompson_test_kernel.cu`) gained two new
entry points (`direction_thompson_v2_test`, `argmax_eq_v2_test`) that
mirror the amended production API for Test 2.B bit-equivalence parity.
The original Plan A entry points (`thompson_direction_test_batched`,
`argmax_eq_test`) are retained — they verify the underlying inverse-CDF
+ uniform-tau math primitives unchanged by the amendment, and Plan A
tests 0.B-0.F continue to consume them.
