# Adaptive C51 z-support via Device-Buffer v_range

> **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:** Replace fixed scalar v_min/v_max with a GPU-resident device buffer that all C51/MSE/CQL/expected_q kernels read via pointer. Initialize tight (±1.0) for cold-start atom resolution, monotonically expand based on observed Q-range. Also fix PopArt warmup (100→1) and variance floor (0.0001→1e-8).

**Architecture:** `v_range_buf: CudaSlice<f32>` [2] holds `[v_min, v_max]` on GPU. Host-side `v_range_host: [f32; 2]` for async HtoD — same proven pattern as `tau_host`/`tau_buf` for EMA. All 5 kernel files replace `float v_min, float v_max` scalar args with `const float* v_range_buf` pointer. CUDA graph captures the pointer (stable address), reads current values at replay time. Adaptive expansion at every q_stats readback (every 50 steps).

**Tech Stack:** CUDA 12.4 (nvcc cubin), Rust 1.85, cudarc

---

## File Map

| File | Action | Changes |
|------|--------|---------|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | Add `v_range_buf`/`v_range_host` fields, replace 10 `.arg(&v_min)/.arg(&v_max)` with `.arg(&self.v_range_buf)`, add `update_v_range()` + `expand_v_range()` methods |
| `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` | Modify | Replace `float v_min, float v_max` params with `const float* v_range_buf` |
| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | Modify | Same |
| `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` | Modify | Same |
| `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu` | Modify | Same |
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | `compute_expected_q`: same |
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | Replace `v_min`/`v_max` fields with `v_range_ptr: u64` pointing to trainer's buffer |
| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Modify | PopArt warmup 100→1, variance floor already 1e-8 |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Modify | Call `expand_v_range()` at q_stats readback |

---

### Task 1: Add v_range_buf to GpuDqnTrainer and initialize tight

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

- [ ] **Step 1: Add fields to GpuDqnTrainer struct**

After `tau_buf` / `tau_host` fields (around line 616):

```rust
/// Adaptive C51 z-support [v_min, v_max] on device.
/// Kernels read via pointer (graph-safe). Updated before graph replay.
v_range_buf: CudaSlice<f32>,
/// Host-side v_range for async HtoD (same pattern as tau_host).
v_range_host: [f32; 2],
```

- [ ] **Step 2: Allocate and initialize in constructor**

In `GpuDqnTrainer::new()`, after `tau_buf` allocation (around line 2400):

```rust
// Adaptive v_range: start tight for cold-start atom resolution.
// dz = 2.0/51 = 0.039 → normalized reward of 1.0 shifts 25.6 atoms.
let v_range_buf = stream.alloc_zeros::<f32>(2)
    .map_err(|e| MLError::ModelError(format!("alloc v_range_buf: {e}")))?;
let v_range_host = [-1.0_f32, 1.0_f32]; // tight cold-start
unsafe {
    cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
        v_range_buf.raw_ptr(),
        v_range_host.as_ptr().cast(),
        2 * std::mem::size_of::<f32>(),
        stream.cu_stream(),
    );
}
```

Add to `Self { ... }`:
```rust
v_range_buf,
v_range_host,
```

- [ ] **Step 3: Add update and expand methods**

After the `target_ema_update` method block:

```rust
/// Async HtoD update of v_range_buf from host values.
/// Called before graph replay when v_range_host changes.
fn sync_v_range(&self) {
    unsafe {
        cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
            self.v_range_buf.raw_ptr(),
            self.v_range_host.as_ptr().cast(),
            2 * std::mem::size_of::<f32>(),
            self.stream.cu_stream(),
        );
    }
}

/// Expand v_range based on observed Q-value range. Never shrinks.
/// Returns true if v_range changed (caller may want to log).
pub fn expand_v_range(&mut self, q_min: f32, q_max: f32) -> bool {
    let q_range = (q_max - q_min).abs();
    let margin = q_range.max(1.0) * 0.5;
    let new_v_min = (q_min - margin).min(self.v_range_host[0]);
    let new_v_max = (q_max + margin).max(self.v_range_host[1]);
    if new_v_min < self.v_range_host[0] || new_v_max > self.v_range_host[1] {
        self.v_range_host = [new_v_min, new_v_max];
        self.sync_v_range();
        true
    } else {
        false
    }
}

/// Current v_range for logging/diagnostics.
pub fn v_range(&self) -> [f32; 2] {
    self.v_range_host
}

/// Raw device pointer to v_range_buf (for experience collector).
pub fn v_range_buf_ptr(&self) -> u64 {
    self.v_range_buf.raw_ptr()
}
```

- [ ] **Step 4: Verify compilation (will fail — kernels still expect scalar args)**

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

Expected: compiles (Rust side added fields, kernels not changed yet).

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat: add v_range_buf device buffer for adaptive C51 z-support"
```

---

### Task 2: Modify all 5 CUDA kernels to read v_range from device pointer

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu`
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu`

- [ ] **Step 1: c51_loss_kernel.cu — replace scalar v_min/v_max with pointer**

In the kernel signature (around line 205-206), replace:
```c
    float v_min,
    float v_max,
```
with:
```c
    const float* __restrict__ v_range_buf, /* [2]: v_range_buf[0]=v_min, v_range_buf[1]=v_max */
```

At the start of the kernel body (after variable declarations), add:
```c
    float v_min = v_range_buf[0];
    float v_max = v_range_buf[1];
```

Also update the inner `bellman_project_branch` helper function (around line 118) — if it takes `float v_min, float v_max` as params, those come from the caller which now reads from the buffer. Check if `bellman_project_branch` is `__device__` (inline) — if so, its callers pass the local `v_min`/`v_max` variables which are already correct.

- [ ] **Step 2: mse_loss_kernel.cu — same change**

Replace `float v_min, float v_max` (lines 159-160) with `const float* __restrict__ v_range_buf`. Add reads at kernel body start.

- [ ] **Step 3: mse_grad_kernel.cu — same change**

Replace `float v_min, float v_max` (line 24) with `const float* __restrict__ v_range_buf`. Add reads at kernel body start.

- [ ] **Step 4: cql_grad_kernel.cu — same change**

Replace `float v_min, float v_max` (line 26) with `const float* __restrict__ v_range_buf`. Add reads at kernel body start.

- [ ] **Step 5: experience_kernels.cu compute_expected_q — same change**

Replace `float v_min, float v_max` (around line 2108-2109) with `const float* __restrict__ v_range_buf`. Add reads at kernel body start.

- [ ] **Step 6: Force rebuild and verify all kernels compile**

```bash
touch crates/ml/build.rs && SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "error|FAILED|Precompiled"
```

Expected: all cubins compile. Rust will fail (launch sites still pass scalars).

- [ ] **Step 7: Commit**

```bash
git add crates/ml/src/cuda_pipeline/c51_loss_kernel.cu crates/ml/src/cuda_pipeline/mse_loss_kernel.cu crates/ml/src/cuda_pipeline/mse_grad_kernel.cu crates/ml/src/cuda_pipeline/cql_grad_kernel.cu crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "feat: all C51/MSE/CQL/expected_q kernels read v_range from device pointer"
```

---

### Task 3: Update all Rust launch sites to pass v_range_buf pointer

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

- [ ] **Step 1: Replace all 10 `.arg(&v_min)/.arg(&v_max)` pairs in gpu_dqn_trainer.rs**

At each of the 5 launch sites (lines 1745-1746, 4178-4179, 5263-5264, 5485-5486, 5572-5573), replace the TWO lines:
```rust
                .arg(&v_min)
                .arg(&v_max)
```
or:
```rust
                .arg(&v_min_f)
                .arg(&v_max_f)
```
with ONE line:
```rust
                .arg(&self.v_range_buf)
```

Also remove the local `let v_min = self.config.v_min;` / `let v_max = self.config.v_max;` / `let v_min_f = ...` / `let v_max_f = ...` declarations that are no longer needed at each site.

- [ ] **Step 2: Update experience collector**

In `gpu_experience_collector.rs`, replace the `v_min`/`v_max` fields with a pointer to the trainer's v_range_buf. Add `v_range_ptr: u64` field (set via `set_v_range_ptr()` alongside `set_trainer_params_ptr()`).

Replace the launch site (lines 1935-1936):
```rust
                        .arg(&v_min_f)
                        .arg(&v_max_f)
```
with:
```rust
                        .arg(&self.v_range_ptr)
```

Where `self.v_range_ptr` is a raw u64 device pointer to the trainer's v_range_buf.

Add `set_v_range_ptr(&mut self, ptr: u64)` method alongside `set_trainer_params_ptr`.

Wire it in `training_loop.rs` where `set_trainer_params_ptr` is called:
```rust
collector.set_v_range_ptr(ctx.trainer_v_range_buf_ptr());
```

- [ ] **Step 3: Verify compilation and run tests**

```bash
SQLX_OFFLINE=true cargo check -p ml --lib
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget config monitoring
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: all launch sites pass v_range_buf pointer — graph-safe adaptive v_range"
```

---

### Task 4: Wire adaptive expansion + PopArt warmup fix

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`

- [ ] **Step 1: Expose expand_v_range through FusedTrainingCtx**

In `fused_training.rs`, add a method:
```rust
pub(crate) fn expand_v_range(&mut self, q_min: f32, q_max: f32) -> bool {
    self.trainer.expand_v_range(q_min, q_max)
}

pub(crate) fn v_range(&self) -> [f32; 2] {
    self.trainer.v_range()
}

pub(crate) fn trainer_v_range_buf_ptr(&self) -> u64 {
    self.trainer.v_range_buf_ptr()
}
```

- [ ] **Step 2: Call expand_v_range at q_stats readback**

In `training_loop.rs`, find where q_stats are read back (search for `compute_q_stats` or `q_stats` or `avg_max_q`). After the Q-stats readback, call:

```rust
if let Some(ref mut fused) = self.fused_ctx {
    if let Some(stats) = &q_diagnostics {
        let (q_min, q_max) = (stats.q_min as f32, stats.q_max as f32);
        if fused.expand_v_range(q_min, q_max) {
            let vr = fused.v_range();
            info!(v_min = vr[0], v_max = vr[1], "C51 v_range expanded");
        }
    }
}
```

- [ ] **Step 3: Fix PopArt warmup 100→1**

In `gpu_dqn_trainer.rs`, replace both instances of:
```rust
let warmup = 100_i32;
```
with:
```rust
let warmup = 1_i32; // Normalize from step 1 — delayed warmup starves C51 of gradient signal
```

- [ ] **Step 4: Verify compilation and run all tests**

```bash
touch crates/ml/build.rs && SQLX_OFFLINE=true cargo build -p ml --lib
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget config monitoring
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: adaptive v_range expansion + PopArt warmup 100→1"
```

---

### Task 5: Smoke test — verify Q-values grow with adaptive v_range

**Files:** (no changes — verification only)

- [ ] **Step 1: Run 50-epoch test and check Q-value trajectory**

```bash
FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::walk_forward::test_walk_forward_no_overfitting_50_epochs --ignored --test-threads=1 --nocapture 2>&1 | grep -E "Q-value range|v_range expanded"
```

Expected: Q-values should grow from O(0.001) to O(0.1-1.0) within 50 epochs. `v_range expanded` log entries should appear as Q-range grows.

- [ ] **Step 2: Run full smoke test suite**

```bash
FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests --ignored --test-threads=1
```

Expected: 19 passed, 0 failed.

- [ ] **Step 3: Run compute-sanitizer**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib --no-run 2>&1 | tail -1
FOXHUNT_TEST_DATA=test_data/futures-baseline compute-sanitizer --tool memcheck target/debug/deps/ml-*.test "smoke_tests::feature_coverage" --ignored --test-threads=1
```

Expected: `ERROR SUMMARY: 0 errors`

- [ ] **Step 4: Commit and push**

```bash
git add -A
git commit -m "fix: adaptive C51 v_range + PopArt warmup — Q-values grow from cold start"
git push origin main
```
