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