# IntegratedTrainer Checkpoints + Eval Diag — 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:** Add full save/load surface to `IntegratedTrainer` (lite + full checkpoint variants), wire CLI for save/resume in `alpha_rl_train`, and make eval phase emit per-step diag to `eval_diag.jsonl`.

**Architecture:** Two checkpoint kinds (lite = weights only for inference; full = lite + Adam moments + ISV bus + scalars + optional replay). Bincode flat-binary, schema versioned. Diag emission extracted to a typed `DiagRecord` struct producible by trainer; train and eval loops each write to their own JSONL.

**Tech Stack:** Rust, cudarc, bincode, serde, AdamW (existing).

**Spec:** `docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md`

**Branch:** `ml-alpha-checkpoints-eval-diag` (off `ml-alpha-regime-observer`)

---

## File structure

| Path | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/src/trainer/checkpoint.rs` | Create | All `*Checkpoint` structs, `IntegratedCheckpoint` top-level, version constant |
| `crates/ml-alpha/src/trainer/diag_record.rs` | Create | Typed `DiagRecord` struct; serde::Serialize → one JSONL line |
| `crates/ml-alpha/src/trainer/optim.rs` | Modify | Add `AdamW::save_state` and `AdamW::load_state` |
| `crates/ml-alpha/src/trainer/integrated.rs` | Modify | Add `save_inference_checkpoint`, `save_full_checkpoint`, `load_checkpoint`, `snapshot_diag_record` methods |
| `crates/ml-alpha/src/trainer/mod.rs` | Modify | `pub mod checkpoint; pub mod diag_record;` |
| `crates/ml-alpha/examples/alpha_rl_train.rs` | Modify | New CLI flags; resume entry; save cadence; eval-diag writer |
| `crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs` | Create | Phase 1 unit test |
| `crates/ml-alpha/tests/checkpoint_full_roundtrip.rs` | Create | Phase 2 unit test |
| `crates/ml-alpha/tests/checkpoint_replay_roundtrip.rs` | Create | Phase 3 unit test |
| `crates/ml-alpha/tests/eval_diag_emission.rs` | Create | Phase 4 smoke |
| `crates/ml-alpha/tests/resume_continuation.rs` | Create | Phase 5 smoke |

---

## Phase 1: Schema + lite save/load

### Task 1.1: Define `IntegratedCheckpoint` schema

**Files:**
- Create: `crates/ml-alpha/src/trainer/checkpoint.rs`
- Modify: `crates/ml-alpha/src/trainer/mod.rs`

- [ ] **Step 1: Add `pub mod checkpoint;` to trainer/mod.rs**

```rust
// crates/ml-alpha/src/trainer/mod.rs — append at end of pub mod list
pub mod checkpoint;
```

- [ ] **Step 2: Create `checkpoint.rs` with all schema structs**

```rust
// crates/ml-alpha/src/trainer/checkpoint.rs
//! Bincode-serialized checkpoint format for IntegratedTrainer.
//!
//! Two variants:
//!   - "lite": weights only (encoder + heads + targets). For inference.
//!   - "full": lite + Adam moments + ISV bus + trainer scalars + optional replay.
//!
//! See docs/superpowers/specs/2026-05-31-checkpoints-and-eval-diag-design.md.

use serde::{Deserialize, Serialize};

/// Format version — bump on any layout-breaking change.
pub const CHECKPOINT_VERSION: u32 = 1;

#[derive(Serialize, Deserialize)]
pub struct IntegratedCheckpoint {
    pub version: u32,
    pub sha: String,
    pub step: u64,
    pub saved_at: String,
    pub kind: String,                                  // "lite" or "full"
    pub encoder: CfcTrunkCheckpointWire,
    pub heads: HeadWeightsCheckpoint,
    pub adam: Option<AdamStateCheckpoint>,
    pub isv: Option<Vec<f32>>,
    pub scalars: Option<TrainerScalarsCheckpoint>,
    pub replay: Option<ReplayBufferCheckpoint>,
}

/// Mirrors crate::cfc::trunk::Checkpoint (already bincode-serialized).
/// We embed by re-serializing into bytes so a struct-level change in
/// trunk::Checkpoint doesn't force an IntegratedCheckpoint version bump
/// — the trunk's own loader validates its bytes. Trade-off: one extra
/// serialize/deserialize per save/load (negligible cost).
#[derive(Serialize, Deserialize)]
pub struct CfcTrunkCheckpointWire {
    pub bytes: Vec<u8>,
}

#[derive(Serialize, Deserialize, Default)]
pub struct HeadWeightsCheckpoint {
    // DqnHead
    pub dqn_w: Vec<f32>,
    pub dqn_b: Vec<f32>,
    pub dqn_w_target: Vec<f32>,
    pub dqn_b_target: Vec<f32>,
    // IqnHead
    pub iqn_w_embed: Vec<f32>,
    pub iqn_b_embed: Vec<f32>,
    pub iqn_w_out: Vec<f32>,
    pub iqn_b_out: Vec<f32>,
    pub iqn_w_embed_target: Vec<f32>,
    pub iqn_b_embed_target: Vec<f32>,
    pub iqn_w_out_target: Vec<f32>,
    pub iqn_b_out_target: Vec<f32>,
    // DuelingQHead
    pub dueling_w_v: Vec<f32>,
    pub dueling_b_v: Vec<f32>,
    pub dueling_w_a: Vec<f32>,
    pub dueling_b_a: Vec<f32>,
    pub dueling_w_v_target: Vec<f32>,
    pub dueling_b_v_target: Vec<f32>,
    pub dueling_w_a_target: Vec<f32>,
    pub dueling_b_a_target: Vec<f32>,
    // PolicyHead
    pub policy_w: Vec<f32>,
    pub policy_b: Vec<f32>,
    // ValueHead
    pub value_w: Vec<f32>,
    pub value_b: Vec<f32>,
}

#[derive(Serialize, Deserialize, Default)]
pub struct AdamStateCheckpoint {
    /// One entry per optimizer slot the trainer holds; key names mirror
    /// trainer fields (e.g. "dqn_w", "iqn_w_embed", ...). Each value is
    /// (m_d, v_d, step_count).
    pub states: Vec<NamedAdamState>,
}

#[derive(Serialize, Deserialize)]
pub struct NamedAdamState {
    pub name: String,
    pub m: Vec<f32>,
    pub v: Vec<f32>,
    pub step: i32,
}

#[derive(Serialize, Deserialize, Default)]
pub struct TrainerScalarsCheckpoint {
    pub last_q_loss: f32,
    pub last_pi_loss: f32,
    pub last_v_loss: f32,
    pub train_step: u64,
}

#[derive(Serialize, Deserialize, Default)]
pub struct ReplayBufferCheckpoint {
    /// PER priority sumtree (length 2 × capacity − 1).
    pub priority_tree: Vec<f32>,
    /// Per-sample feature vectors flattened (capacity × feature_dim).
    pub features: Vec<f32>,
    /// Per-sample next-state features.
    pub next_features: Vec<f32>,
    /// Per-sample taken action indices.
    pub actions: Vec<i32>,
    /// Per-sample rewards.
    pub rewards: Vec<f32>,
    /// Per-sample done flags (0.0 or 1.0 stored as f32 for buffer
    /// uniformity).
    pub dones: Vec<f32>,
    /// Current write index in the circular buffer.
    pub write_head: u64,
    /// Number of samples actually populated (≤ capacity).
    pub n_populated: u64,
}
```

- [ ] **Step 3: Run `SQLX_OFFLINE=true cargo check -p ml-alpha` to verify schema compiles**

Expected: `Finished dev [unoptimized] target(s) in <time>`. No errors.

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/src/trainer/checkpoint.rs crates/ml-alpha/src/trainer/mod.rs
git commit -m "feat(trainer): IntegratedCheckpoint schema (1.1)"
```

### Task 1.2: Add `AdamW::save_state` and `load_state`

**Files:**
- Modify: `crates/ml-alpha/src/trainer/optim.rs`

- [ ] **Step 1: Add save/load methods to AdamW**

```rust
// crates/ml-alpha/src/trainer/optim.rs — after impl AdamW { ... }

impl AdamW {
    /// Download m, v moment buffers + step count for checkpointing.
    pub fn save_state(&self) -> Result<(Vec<f32>, Vec<f32>, i32)> {
        let mut m_host = vec![0.0_f32; self.m_d.len()];
        let mut v_host = vec![0.0_f32; self.v_d.len()];
        self.stream.memcpy_dtoh(&self.m_d, m_host.as_mut_slice())?;
        self.stream.memcpy_dtoh(&self.v_d, v_host.as_mut_slice())?;
        Ok((m_host, v_host, self.t))
    }

    /// Upload m, v moment buffers + restore step count. Validates length.
    pub fn load_state(&mut self, m_host: &[f32], v_host: &[f32], step: i32) -> Result<()> {
        anyhow::ensure!(
            m_host.len() == self.m_d.len(),
            "AdamW::load_state: m length mismatch (got {}, expected {})",
            m_host.len(), self.m_d.len()
        );
        anyhow::ensure!(
            v_host.len() == self.v_d.len(),
            "AdamW::load_state: v length mismatch (got {}, expected {})",
            v_host.len(), self.v_d.len()
        );
        self.stream.memcpy_htod(m_host, &mut self.m_d)?;
        self.stream.memcpy_htod(v_host, &mut self.v_d)?;
        self.t = step;
        Ok(())
    }
}
```

- [ ] **Step 2: Build to verify compiles**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: clean.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/src/trainer/optim.rs
git commit -m "feat(trainer): AdamW save_state/load_state (1.2)"
```

### Task 1.3: Implement `save_inference_checkpoint`

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Add helper at top of integrated.rs imports**

```rust
// At the top of integrated.rs (within existing use { ... } block or as additional)
use crate::trainer::checkpoint::{
    IntegratedCheckpoint, CfcTrunkCheckpointWire, HeadWeightsCheckpoint,
    CHECKPOINT_VERSION,
};
```

- [ ] **Step 2: Add helper for downloading f32 CudaSlices**

```rust
// crates/ml-alpha/src/trainer/integrated.rs — in impl IntegratedTrainer { ... }

/// Download a CudaSlice<f32> into a host Vec. Uses self.stream
/// (same stream as the head's owning module, so no cross-stream sync
/// is needed for read-only state-dump operations).
fn download_f32(&self, slice: &CudaSlice<f32>) -> Result<Vec<f32>> {
    let mut v = vec![0.0_f32; slice.len()];
    self.stream.memcpy_dtoh(slice, v.as_mut_slice())?;
    Ok(v)
}
```

- [ ] **Step 3: Implement HeadWeightsCheckpoint construction**

```rust
// crates/ml-alpha/src/trainer/integrated.rs — in impl IntegratedTrainer

fn collect_head_weights(&self) -> Result<HeadWeightsCheckpoint> {
    let dqn = &self.dqn_head;
    let iqn = &self.iqn_head;
    let dueling = &self.dueling_q_head;
    let policy = &self.policy_head;
    let value = &self.value_head;
    Ok(HeadWeightsCheckpoint {
        dqn_w:               self.download_f32(&dqn.w_d)?,
        dqn_b:               self.download_f32(&dqn.b_d)?,
        dqn_w_target:        self.download_f32(&dqn.w_target_d)?,
        dqn_b_target:        self.download_f32(&dqn.b_target_d)?,
        iqn_w_embed:         self.download_f32(&iqn.w_embed_d)?,
        iqn_b_embed:         self.download_f32(&iqn.b_embed_d)?,
        iqn_w_out:           self.download_f32(&iqn.w_out_d)?,
        iqn_b_out:           self.download_f32(&iqn.b_out_d)?,
        iqn_w_embed_target:  self.download_f32(&iqn.w_embed_target_d)?,
        iqn_b_embed_target:  self.download_f32(&iqn.b_embed_target_d)?,
        iqn_w_out_target:    self.download_f32(&iqn.w_out_target_d)?,
        iqn_b_out_target:    self.download_f32(&iqn.b_out_target_d)?,
        dueling_w_v:         self.download_f32(&dueling.w_v_d)?,
        dueling_b_v:         self.download_f32(&dueling.b_v_d)?,
        dueling_w_a:         self.download_f32(&dueling.w_a_d)?,
        dueling_b_a:         self.download_f32(&dueling.b_a_d)?,
        // DuelingQHead may not have target nets — verify when writing;
        // if absent, leave the *_target fields as Vec::new() (Default).
        // VERIFY: read dueling_q.rs structure before this task and
        // either fill the target accessors here or drop the four
        // *_target_d fields from HeadWeightsCheckpoint to match.
        dueling_w_v_target:  Vec::new(),  // PLACEHOLDER — confirm exists
        dueling_b_v_target:  Vec::new(),
        dueling_w_a_target:  Vec::new(),
        dueling_b_a_target:  Vec::new(),
        policy_w:            self.download_f32(&policy.w_d)?,
        policy_b:            self.download_f32(&policy.b_d)?,
        value_w:             self.download_f32(&value.w_d)?,
        value_b:             self.download_f32(&value.b_d)?,
    })
}
```

> **Implementer note**: `dueling_q_head` target-net fields existence
> isn't confirmed in the spec exploration. Before this task, read
> `crates/ml-alpha/src/rl/dueling_q.rs` once; if target buffers exist
> (`w_v_target_d` etc.), fill those four fields and remove the
> `PLACEHOLDER` comments. If not, delete the four `dueling_*_target`
> fields from `HeadWeightsCheckpoint` in Task 1.1's schema.

- [ ] **Step 4: Implement `save_inference_checkpoint`**

```rust
// crates/ml-alpha/src/trainer/integrated.rs — in impl IntegratedTrainer

/// Save weights-only checkpoint for inference deployment. ~55 MB.
pub fn save_inference_checkpoint(&self, path: &std::path::Path) -> Result<()> {
    use std::io::Write;
    let trunk_bytes = self.perception.trunk_checkpoint_bytes()
        .context("serialize encoder trunk")?;
    let heads = self.collect_head_weights()
        .context("collect head weights")?;
    // Synthesize a UTC timestamp WITHOUT calling Date.now()/SystemTime
    // directly (the trainer is meant to be deterministic-ish; timestamps
    // are bookkeeping only). Use SystemTime::now() — this is the host
    // path, not the GPU graph, so it's fine here.
    let saved_at = chrono::Utc::now().to_rfc3339();
    let ckpt = IntegratedCheckpoint {
        version: CHECKPOINT_VERSION,
        sha: env!("GIT_SHA").to_string(),  // set by build.rs (see Task 1.3-bis)
        step: 0,
        saved_at,
        kind: "lite".to_string(),
        encoder: CfcTrunkCheckpointWire { bytes: trunk_bytes },
        heads,
        adam: None,
        isv: None,
        scalars: None,
        replay: None,
    };
    let bytes = bincode::serialize(&ckpt)
        .context("bincode serialize IntegratedCheckpoint")?;
    if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; }
    let mut f = std::fs::File::create(path)
        .with_context(|| format!("create {}", path.display()))?;
    f.write_all(&bytes).context("write checkpoint bytes")?;
    Ok(())
}
```

- [ ] **Step 5: Add `PerceptionTrainer::trunk_checkpoint_bytes`**

```rust
// crates/ml-alpha/src/trainer/perception.rs — add helper

impl PerceptionTrainer {
    /// Serialize the trunk's existing Checkpoint into bytes — used by
    /// IntegratedTrainer's save_*_checkpoint to embed.
    pub fn trunk_checkpoint_bytes(&self) -> Result<Vec<u8>> {
        // We need the trunk's internal Checkpoint serialized. Cheapest:
        // call its existing save_checkpoint into a tempfile, read back.
        // BUT that touches disk. Cleaner: factor trunk's save_checkpoint
        // body into a `to_checkpoint() -> Checkpoint` then serialize
        // here. For Phase 1 just add a `to_bytes` method on CfcTrunk
        // that mirrors save_checkpoint's body sans file-write.
        self.trunk.checkpoint_bytes()
    }
}
```

- [ ] **Step 6: Refactor `CfcTrunk::save_checkpoint` to use new `checkpoint_bytes()`**

```rust
// crates/ml-alpha/src/cfc/trunk.rs — refactor

impl CfcTrunk {
    /// Build the internal Checkpoint struct + bincode-serialize.
    /// Used both by save_checkpoint (writes to file) and by the
    /// IntegratedTrainer's embedding path.
    pub fn checkpoint_bytes(&self) -> Result<Vec<u8>> {
        // ... move existing body of save_checkpoint here up to and
        // including the `let bytes = bincode::serialize(&ckpt)...` line.
        // Return bytes instead of writing.
        Ok(bytes)
    }

    pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> {
        use std::io::Write;
        let bytes = self.checkpoint_bytes()?;
        let mut f = std::fs::File::create(path)
            .with_context(|| format!("create {}", path.display()))?;
        f.write_all(&bytes).context("write checkpoint bytes")?;
        Ok(())
    }
}
```

- [ ] **Step 7: Wire `GIT_SHA` build env var**

```rust
// crates/ml-alpha/build.rs — add near top of main()
let sha = std::process::Command::new("git")
    .args(["rev-parse", "--short", "HEAD"])
    .output()
    .ok()
    .and_then(|o| String::from_utf8(o.stdout).ok())
    .map(|s| s.trim().to_string())
    .unwrap_or_else(|| "unknown".to_string());
println!("cargo:rustc-env=GIT_SHA={}", sha);
println!("cargo:rerun-if-changed=.git/HEAD");
```

- [ ] **Step 8: Add `chrono` dep if missing**

Check `crates/ml-alpha/Cargo.toml` for `chrono`. If absent:

```toml
# crates/ml-alpha/Cargo.toml — under [dependencies]
chrono = { workspace = true, features = ["serde"] }
```

(If `chrono` is not in the workspace, add it to workspace Cargo.toml too.)

- [ ] **Step 9: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: clean.

- [ ] **Step 10: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs \
        crates/ml-alpha/src/trainer/perception.rs \
        crates/ml-alpha/src/cfc/trunk.rs \
        crates/ml-alpha/build.rs \
        crates/ml-alpha/Cargo.toml
git commit -m "feat(trainer): save_inference_checkpoint (1.3)"
```

### Task 1.4: Implement `load_checkpoint` (lite branch)

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Add helpers for uploading + load entrypoint**

```rust
// crates/ml-alpha/src/trainer/integrated.rs — in impl IntegratedTrainer

fn upload_f32(stream: &CudaStream, host: &[f32], dst: &mut CudaSlice<f32>) -> Result<()> {
    anyhow::ensure!(host.len() == dst.len(),
        "upload_f32: length mismatch (got {}, expected {})", host.len(), dst.len());
    stream.memcpy_htod(host, dst)?;
    Ok(())
}

fn apply_head_weights(&mut self, h: &HeadWeightsCheckpoint) -> Result<()> {
    let s = &self.stream;
    Self::upload_f32(s, &h.dqn_w,              &mut self.dqn_head.w_d)?;
    Self::upload_f32(s, &h.dqn_b,              &mut self.dqn_head.b_d)?;
    Self::upload_f32(s, &h.dqn_w_target,       &mut self.dqn_head.w_target_d)?;
    Self::upload_f32(s, &h.dqn_b_target,       &mut self.dqn_head.b_target_d)?;
    Self::upload_f32(s, &h.iqn_w_embed,        &mut self.iqn_head.w_embed_d)?;
    Self::upload_f32(s, &h.iqn_b_embed,        &mut self.iqn_head.b_embed_d)?;
    Self::upload_f32(s, &h.iqn_w_out,          &mut self.iqn_head.w_out_d)?;
    Self::upload_f32(s, &h.iqn_b_out,          &mut self.iqn_head.b_out_d)?;
    Self::upload_f32(s, &h.iqn_w_embed_target, &mut self.iqn_head.w_embed_target_d)?;
    Self::upload_f32(s, &h.iqn_b_embed_target, &mut self.iqn_head.b_embed_target_d)?;
    Self::upload_f32(s, &h.iqn_w_out_target,   &mut self.iqn_head.w_out_target_d)?;
    Self::upload_f32(s, &h.iqn_b_out_target,   &mut self.iqn_head.b_out_target_d)?;
    Self::upload_f32(s, &h.dueling_w_v,        &mut self.dueling_q_head.w_v_d)?;
    Self::upload_f32(s, &h.dueling_b_v,        &mut self.dueling_q_head.b_v_d)?;
    Self::upload_f32(s, &h.dueling_w_a,        &mut self.dueling_q_head.w_a_d)?;
    Self::upload_f32(s, &h.dueling_b_a,        &mut self.dueling_q_head.b_a_d)?;
    // Target nets — confirm dueling_q.rs surface in Task 1.3.
    Self::upload_f32(s, &h.policy_w, &mut self.policy_head.w_d)?;
    Self::upload_f32(s, &h.policy_b, &mut self.policy_head.b_d)?;
    Self::upload_f32(s, &h.value_w,  &mut self.value_head.w_d)?;
    Self::upload_f32(s, &h.value_b,  &mut self.value_head.b_d)?;
    Ok(())
}

/// Load a checkpoint of either kind. Returns the saved train `step`
/// for the caller to use as resume cursor.
pub fn load_checkpoint(&mut self, path: &std::path::Path) -> Result<u64> {
    let bytes = std::fs::read(path)
        .with_context(|| format!("read {}", path.display()))?;
    let ckpt: IntegratedCheckpoint = bincode::deserialize(&bytes)
        .context("bincode deserialize IntegratedCheckpoint")?;
    anyhow::ensure!(ckpt.version == CHECKPOINT_VERSION,
        "checkpoint version mismatch: got {}, expected {}",
        ckpt.version, CHECKPOINT_VERSION);

    if ckpt.sha != env!("GIT_SHA") {
        eprintln!("checkpoint SHA mismatch (saved {} vs current {}) — attempting load anyway",
                  ckpt.sha, env!("GIT_SHA"));
    }

    // Encoder
    self.perception.load_trunk_from_bytes(&ckpt.encoder.bytes)
        .context("load encoder trunk")?;

    // Heads (always present)
    self.apply_head_weights(&ckpt.heads)?;

    // Full-only fields
    if ckpt.kind == "full" {
        if let Some(adam) = ckpt.adam.as_ref() {
            self.apply_adam_state(adam)?;
        }
        if let Some(isv) = ckpt.isv.as_ref() {
            anyhow::ensure!(isv.len() == crate::rl::isv_slots::RL_SLOTS_END,
                "ISV length mismatch: got {} expected {}",
                isv.len(), crate::rl::isv_slots::RL_SLOTS_END);
            for (i, v) in isv.iter().enumerate() {
                self.isv_mapped.write_record(i, *v);
            }
        }
        if let Some(s) = ckpt.scalars.as_ref() {
            self.last_q_loss  = s.last_q_loss;
            self.last_pi_loss = s.last_pi_loss;
            self.last_v_loss  = s.last_v_loss;
        }
        if let Some(r) = ckpt.replay.as_ref() {
            self.apply_replay_buffer(r)?;
        }
    }

    self.stream.synchronize()?;
    Ok(ckpt.step)
}
```

- [ ] **Step 2: Add `PerceptionTrainer::load_trunk_from_bytes`**

```rust
// crates/ml-alpha/src/trainer/perception.rs

impl PerceptionTrainer {
    pub fn load_trunk_from_bytes(&mut self, bytes: &[u8]) -> Result<()> {
        // Wrap CfcTrunk::load_checkpoint to accept bytes instead of a
        // path. Equivalent: write a `load_from_bytes(&mut self, ...)`
        // method on CfcTrunk that takes (dev, cfg, bytes) and uploads
        // into existing buffers. The current `load_checkpoint` returns
        // a new CfcTrunk by value; refactor to populate &mut self
        // instead (mirroring AdamW::load_state's in-place upload).
        self.trunk.load_from_bytes(bytes)
    }
}
```

> **Implementer note**: This is a non-trivial refactor of `CfcTrunk::load_checkpoint`. The current `load_checkpoint` returns `Self` (a new trunk). To support in-place restore, add `CfcTrunk::load_from_bytes(&mut self, bytes: &[u8]) -> Result<()>` that uploads weights to the EXISTING buffers (matches the `upload_f32` pattern in `apply_head_weights` above). Keep the old `load_checkpoint` (returns new) for the existing tests that use it. Read `crates/ml-alpha/src/cfc/trunk.rs:499` for the current load body — it's a long but mechanical refactor.

- [ ] **Step 3: Stub `apply_adam_state` and `apply_replay_buffer` (filled in Phase 2/3)**

```rust
// crates/ml-alpha/src/trainer/integrated.rs

fn apply_adam_state(&mut self, _a: &AdamStateCheckpoint) -> Result<()> {
    anyhow::bail!("apply_adam_state: not yet implemented (Phase 2)")
}

fn apply_replay_buffer(&mut self, _r: &ReplayBufferCheckpoint) -> Result<()> {
    anyhow::bail!("apply_replay_buffer: not yet implemented (Phase 3)")
}
```

- [ ] **Step 4: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: clean.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs crates/ml-alpha/src/trainer/perception.rs crates/ml-alpha/src/cfc/trunk.rs
git commit -m "feat(trainer): load_checkpoint (lite branch) (1.4)"
```

### Task 1.5: Phase 1 round-trip unit test

**Files:**
- Create: `crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs`

- [ ] **Step 1: Write the test**

```rust
// crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs

use anyhow::Result;
use ml_alpha::trainer::integrated::IntegratedTrainer;
use ml_core::device::MlDevice;
use tempfile::tempdir;

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn lite_checkpoint_save_load_preserves_weights() -> Result<()> {
    let Ok(dev) = MlDevice::cuda(0) else {
        eprintln!("skipping: no CUDA device");
        return Ok(());
    };

    // Build trainer A (random init via seed).
    let mut trainer_a = IntegratedTrainer::with_controllers_bootstrapped(
        &dev,
        /* b_size */ 16,
        /* seed */ 42,
    )?;

    // Save lite checkpoint.
    let tmp = tempdir()?;
    let ckpt_path = tmp.path().join("test.lite.ckpt");
    trainer_a.save_inference_checkpoint(&ckpt_path)?;

    // Build trainer B (different seed → different random init).
    let mut trainer_b = IntegratedTrainer::with_controllers_bootstrapped(
        &dev,
        /* b_size */ 16,
        /* seed */ 999,
    )?;

    // Verify A and B have DIFFERENT weights pre-load.
    let a_dqn_w_pre = trainer_a.download_f32_pub(&trainer_a.dqn_head.w_d)?;
    let b_dqn_w_pre = trainer_b.download_f32_pub(&trainer_b.dqn_head.w_d)?;
    assert_ne!(a_dqn_w_pre, b_dqn_w_pre, "different seeds → different weights pre-load");

    // Load A's checkpoint into B.
    let loaded_step = trainer_b.load_checkpoint(&ckpt_path)?;
    assert_eq!(loaded_step, 0, "lite checkpoint stores step=0");

    // Verify weights now identical.
    let b_dqn_w_post = trainer_b.download_f32_pub(&trainer_b.dqn_head.w_d)?;
    let max_diff = a_dqn_w_pre.iter().zip(b_dqn_w_post.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0.0_f32, f32::max);
    assert!(max_diff < 1e-6,
        "post-load weights match (max_diff={})", max_diff);

    eprintln!("Phase 1 OK — lite checkpoint round-trips weights with max_diff={}", max_diff);
    Ok(())
}
```

> **Implementer note**: `download_f32_pub` is a public test-only wrapper to be added on `IntegratedTrainer`:
> ```rust
> #[cfg(test)]
> pub fn download_f32_pub(&self, slice: &CudaSlice<f32>) -> Result<Vec<f32>> {
>     self.download_f32(slice)
> }
> ```
> Or just make `download_f32` `pub(crate)` for test access.

- [ ] **Step 2: Run the test**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test --release -p ml-alpha --test checkpoint_lite_roundtrip -- --ignored --nocapture
```

Expected: `test lite_checkpoint_save_load_preserves_weights ... ok`. Print: `Phase 1 OK — lite checkpoint round-trips weights with max_diff=0`.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/tests/checkpoint_lite_roundtrip.rs crates/ml-alpha/src/trainer/integrated.rs
git commit -m "test(trainer): Phase 1 lite checkpoint round-trip (1.5)"
```

---

## Phase 2: Full save/load (Adam + ISV + scalars)

### Task 2.1: Implement `collect_adam_state` + `apply_adam_state`

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Catalogue every Adam in IntegratedTrainer**

Read `crates/ml-alpha/src/trainer/integrated.rs` and grep for `AdamW`:

```bash
grep -nE "AdamW|_adam:" crates/ml-alpha/src/trainer/integrated.rs | head -30
```

> **Implementer note**: Capture every Adam instance the trainer owns (dqn_w_adam, dqn_b_adam, iqn_*_adam, policy_*_adam, value_*_adam, dueling_*_adam, outcome_*_adam, perception_*_adam, etc.). The exact name list isn't pre-enumerable; produce a complete list during this task from the actual struct.

- [ ] **Step 2: Implement collection method**

```rust
// crates/ml-alpha/src/trainer/integrated.rs

fn collect_adam_state(&self) -> Result<AdamStateCheckpoint> {
    let mut states = Vec::new();
    // Macro-like helper to reduce repetition:
    macro_rules! add {
        ($name:expr, $field:expr) => {{
            let (m, v, t) = $field.save_state()?;
            states.push(NamedAdamState { name: $name.to_string(), m, v, step: t });
        }};
    }
    add!("dqn_w",     self.dqn_w_adam);
    add!("dqn_b",     self.dqn_b_adam);
    add!("iqn_w_embed", self.iqn_w_embed_adam);
    add!("iqn_b_embed", self.iqn_b_embed_adam);
    add!("iqn_w_out",   self.iqn_w_out_adam);
    add!("iqn_b_out",   self.iqn_b_out_adam);
    add!("dueling_w_v", self.dueling_w_v_adam);
    add!("dueling_b_v", self.dueling_b_v_adam);
    add!("dueling_w_a", self.dueling_w_a_adam);
    add!("dueling_b_a", self.dueling_b_a_adam);
    add!("policy_w",  self.policy_w_adam);
    add!("policy_b",  self.policy_b_adam);
    add!("value_w",   self.value_w_adam);
    add!("value_b",   self.value_b_adam);
    add!("outcome_w", self.outcome_w_adam);
    add!("outcome_b", self.outcome_b_adam);
    // ... add any other Adams discovered in step 1.
    // Perception trainer has its own Adams (BCE / aux); skip them
    // here UNLESS the trainer struct directly owns them. If owned by
    // self.perception, add a perception.collect_adam_states() method
    // that returns Vec<NamedAdamState> and extend here.
    states.extend(self.perception.collect_adam_states()?);
    Ok(AdamStateCheckpoint { states })
}

fn apply_adam_state(&mut self, a: &AdamStateCheckpoint) -> Result<()> {
    for ns in &a.states {
        match ns.name.as_str() {
            "dqn_w"          => self.dqn_w_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "dqn_b"          => self.dqn_b_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "iqn_w_embed"    => self.iqn_w_embed_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "iqn_b_embed"    => self.iqn_b_embed_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "iqn_w_out"      => self.iqn_w_out_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "iqn_b_out"      => self.iqn_b_out_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "dueling_w_v"    => self.dueling_w_v_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "dueling_b_v"    => self.dueling_b_v_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "dueling_w_a"    => self.dueling_w_a_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "dueling_b_a"    => self.dueling_b_a_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "policy_w"       => self.policy_w_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "policy_b"       => self.policy_b_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "value_w"        => self.value_w_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "value_b"        => self.value_b_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "outcome_w"      => self.outcome_w_adam.load_state(&ns.m, &ns.v, ns.step)?,
            "outcome_b"      => self.outcome_b_adam.load_state(&ns.m, &ns.v, ns.step)?,
            // Perception-owned states routed to perception.
            n if n.starts_with("perception_") => {
                self.perception.apply_adam_state(n, &ns.m, &ns.v, ns.step)?;
            }
            other => anyhow::bail!("apply_adam_state: unknown name '{}'", other),
        }
    }
    Ok(())
}
```

- [ ] **Step 3: Add `PerceptionTrainer::collect_adam_states` + `apply_adam_state`**

```rust
// crates/ml-alpha/src/trainer/perception.rs

impl PerceptionTrainer {
    pub fn collect_adam_states(&self) -> Result<Vec<NamedAdamState>> {
        let mut out = Vec::new();
        macro_rules! add {
            ($name:expr, $field:expr) => {{
                let (m, v, t) = $field.save_state()?;
                out.push(NamedAdamState { name: $name.to_string(), m, v, step: t });
            }};
        }
        // Enumerate every Adam in PerceptionTrainer here. Read
        // perception.rs to discover them.
        // Example placeholders — REPLACE with actual fields:
        add!("perception_bce_w", self.bce_w_adam);
        add!("perception_aux_w", self.aux_w_adam);
        // ... etc.
        Ok(out)
    }

    pub fn apply_adam_state(&mut self, name: &str, m: &[f32], v: &[f32], step: i32) -> Result<()> {
        match name {
            "perception_bce_w" => self.bce_w_adam.load_state(m, v, step)?,
            "perception_aux_w" => self.aux_w_adam.load_state(m, v, step)?,
            // ... etc.
            other => anyhow::bail!("PerceptionTrainer::apply_adam_state: unknown '{}'", other),
        }
        Ok(())
    }
}
```

- [ ] **Step 4: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: clean.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs crates/ml-alpha/src/trainer/perception.rs
git commit -m "feat(trainer): Adam state save/load (2.1)"
```

### Task 2.2: Implement `save_full_checkpoint` (no replay)

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Add the save method**

```rust
// crates/ml-alpha/src/trainer/integrated.rs

/// Save full checkpoint: weights + Adam + ISV + scalars. ~65 MB without
/// replay; +~8 GB if `include_replay=true`.
pub fn save_full_checkpoint(
    &self,
    path: &std::path::Path,
    train_step: u64,
    include_replay: bool,
) -> Result<()> {
    use std::io::Write;
    let trunk_bytes = self.perception.trunk_checkpoint_bytes()?;
    let heads = self.collect_head_weights()?;
    let adam  = Some(self.collect_adam_state()?);
    // ISV snapshot — read all RL_SLOTS_END entries via mapped-pinned host_ptr.
    // mapped-pinned host_ptr is the canonical CPU<->GPU path; per
    // feedback_no_htod_htoh_only_mapped_pinned this avoids HtoH copies
    // through any non-mapped staging buffer.
    let n_slots = crate::rl::isv_slots::RL_SLOTS_END;
    let mut isv_vec = vec![0.0_f32; n_slots];
    for i in 0..n_slots {
        isv_vec[i] = self.isv_mapped.read_record(i);
    }
    let isv = Some(isv_vec);
    let scalars = Some(TrainerScalarsCheckpoint {
        last_q_loss:  self.last_q_loss,
        last_pi_loss: self.last_pi_loss,
        last_v_loss:  self.last_v_loss,
        train_step,
    });
    let replay = if include_replay { Some(self.collect_replay_buffer()?) } else { None };
    let saved_at = chrono::Utc::now().to_rfc3339();
    let ckpt = IntegratedCheckpoint {
        version: CHECKPOINT_VERSION,
        sha: env!("GIT_SHA").to_string(),
        step: train_step,
        saved_at,
        kind: "full".to_string(),
        encoder: CfcTrunkCheckpointWire { bytes: trunk_bytes },
        heads, adam, isv, scalars, replay,
    };
    let bytes = bincode::serialize(&ckpt)?;
    if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; }
    let mut f = std::fs::File::create(path)?;
    f.write_all(&bytes)?;
    Ok(())
}

fn collect_replay_buffer(&self) -> Result<ReplayBufferCheckpoint> {
    anyhow::bail!("collect_replay_buffer: not yet implemented (Phase 3)")
}
```

- [ ] **Step 2: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: clean.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs
git commit -m "feat(trainer): save_full_checkpoint (no replay) (2.2)"
```

### Task 2.3: Phase 2 round-trip unit test

**Files:**
- Create: `crates/ml-alpha/tests/checkpoint_full_roundtrip.rs`

- [ ] **Step 1: Write the test**

```rust
// crates/ml-alpha/tests/checkpoint_full_roundtrip.rs

use anyhow::Result;
use ml_alpha::trainer::integrated::IntegratedTrainer;
use ml_core::device::MlDevice;
use tempfile::tempdir;

#[test]
#[ignore = "requires CUDA"]
fn full_checkpoint_save_load_preserves_isv_and_adam() -> Result<()> {
    let Ok(dev) = MlDevice::cuda(0) else { return Ok(()); };
    let mut trainer = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 42)?;

    // Mutate ISV to a known non-bootstrap state.
    trainer.isv_mapped.write_record(0, 0.12345);
    trainer.isv_mapped.write_record(123, -0.98765);
    trainer.last_pi_loss = 7.7;
    trainer.last_q_loss  = 11.0;
    trainer.last_v_loss  = 0.001;

    let tmp = tempdir()?;
    let ckpt_path = tmp.path().join("test.full.ckpt");
    trainer.save_full_checkpoint(&ckpt_path, /*train_step*/ 42, /*include_replay*/ false)?;

    // Reset trainer to a fresh state, then load.
    let mut trainer_b = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 999)?;
    let loaded_step = trainer_b.load_checkpoint(&ckpt_path)?;
    assert_eq!(loaded_step, 42);

    // ISV restored.
    assert!((trainer_b.isv_mapped.read_record(0) - 0.12345).abs() < 1e-6);
    assert!((trainer_b.isv_mapped.read_record(123) - (-0.98765)).abs() < 1e-6);

    // Scalars restored.
    assert!((trainer_b.last_pi_loss - 7.7).abs() < 1e-6);
    assert!((trainer_b.last_q_loss  - 11.0).abs() < 1e-6);

    eprintln!("Phase 2 OK — full checkpoint round-trips ISV + scalars");
    Ok(())
}
```

- [ ] **Step 2: Run the test**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test --release -p ml-alpha --test checkpoint_full_roundtrip -- --ignored --nocapture
```

Expected: `Phase 2 OK`.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/tests/checkpoint_full_roundtrip.rs
git commit -m "test(trainer): Phase 2 full checkpoint round-trip (2.3)"
```

---

## Phase 3: Replay buffer save/load

### Task 3.1: Add PER snapshot methods

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (or wherever the PER buffer lives — verify in step 1)

- [ ] **Step 1: Locate the PER buffer**

```bash
grep -nE "PerBuffer|per_buffer|priority_tree|capacity" crates/ml-alpha/src/trainer/integrated.rs | head -10
```

Read the buffer's struct to enumerate its CudaSlice fields. Update `ReplayBufferCheckpoint` in Task 1.1's schema if field layout differs.

- [ ] **Step 2: Implement `collect_replay_buffer`**

```rust
// crates/ml-alpha/src/trainer/integrated.rs — replace stub

fn collect_replay_buffer(&self) -> Result<ReplayBufferCheckpoint> {
    let per = &self.per_buffer;  // exact path TBD per step 1
    Ok(ReplayBufferCheckpoint {
        priority_tree:  self.download_f32(&per.priority_tree_d)?,
        features:       self.download_f32(&per.features_d)?,
        next_features:  self.download_f32(&per.next_features_d)?,
        actions:        {
            let mut v = vec![0i32; per.actions_d.len()];
            self.stream.memcpy_dtoh(&per.actions_d, v.as_mut_slice())?;
            v
        },
        rewards:        self.download_f32(&per.rewards_d)?,
        dones:          self.download_f32(&per.dones_d)?,
        write_head:     per.write_head as u64,
        n_populated:    per.n_populated as u64,
    })
}

fn apply_replay_buffer(&mut self, r: &ReplayBufferCheckpoint) -> Result<()> {
    let s = &self.stream;
    let per = &mut self.per_buffer;
    Self::upload_f32(s, &r.priority_tree, &mut per.priority_tree_d)?;
    Self::upload_f32(s, &r.features,       &mut per.features_d)?;
    Self::upload_f32(s, &r.next_features,  &mut per.next_features_d)?;
    anyhow::ensure!(r.actions.len() == per.actions_d.len());
    s.memcpy_htod(&r.actions, &mut per.actions_d)?;
    Self::upload_f32(s, &r.rewards, &mut per.rewards_d)?;
    Self::upload_f32(s, &r.dones,   &mut per.dones_d)?;
    per.write_head  = r.write_head  as usize;
    per.n_populated = r.n_populated as usize;
    Ok(())
}
```

- [ ] **Step 3: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs
git commit -m "feat(trainer): PER replay save/load (3.1)"
```

### Task 3.2: Phase 3 unit test

**Files:**
- Create: `crates/ml-alpha/tests/checkpoint_replay_roundtrip.rs`

- [ ] **Step 1: Write the test**

```rust
// crates/ml-alpha/tests/checkpoint_replay_roundtrip.rs

#[test]
#[ignore = "requires CUDA"]
fn full_checkpoint_with_replay_preserves_sample_order() -> Result<()> {
    let Ok(dev) = MlDevice::cuda(0) else { return Ok(()); };
    let mut a = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 42)?;

    // Push some synthetic transitions so the buffer has content.
    for _ in 0..100 {
        a.push_synthetic_transition_for_test()?;  // helper to add — see below
    }

    let tmp = tempdir()?;
    let ckpt = tmp.path().join("test.replay.ckpt");
    a.save_full_checkpoint(&ckpt, 100, /*include_replay*/ true)?;

    let mut b = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 999)?;
    b.load_checkpoint(&ckpt)?;

    // Sample identical batches from A and B; expect identical indices.
    let idx_a = a.per_buffer.sample_indices_for_test(16, /*seed*/ 7)?;
    let idx_b = b.per_buffer.sample_indices_for_test(16, /*seed*/ 7)?;
    assert_eq!(idx_a, idx_b, "PER deterministic sample after replay restore");

    eprintln!("Phase 3 OK — replay buffer round-trips");
    Ok(())
}
```

> **Implementer note**: This test needs two helpers — `push_synthetic_transition_for_test` and `sample_indices_for_test`. Both should be `#[cfg(any(test, feature = "test-utils"))]`-gated to avoid leaking test surface into release builds. Implement them as thin wrappers over whatever internal push/sample methods already exist.

- [ ] **Step 2: Run + Commit (same pattern as Phase 2)**

---

## Phase 4: Eval-phase diag emission

### Task 4.1: Extract typed `DiagRecord`

**Files:**
- Create: `crates/ml-alpha/src/trainer/diag_record.rs`
- Modify: `crates/ml-alpha/src/trainer/mod.rs`, `crates/ml-alpha/src/trainer/diag_staging.rs`

- [ ] **Step 1: Read current diag emission**

```bash
grep -nE "json!|write_record|DiagStaging" crates/ml-alpha/src/trainer/diag_staging.rs | head -20
grep -nE "json!" crates/ml-alpha/examples/alpha_rl_train.rs | head -5
```

Identify the current inline JSON structure (the keys we already saw in cluster JSONL: `step`, `loss.*`, `popart.*`, `risk_stack.*`, etc.).

- [ ] **Step 2: Create typed `DiagRecord`**

```rust
// crates/ml-alpha/src/trainer/diag_record.rs

use serde::Serialize;

/// One per-step diagnostic record. Schema is the same for train and eval.
/// Mirrors the existing inline `json!{ ... }` block in alpha_rl_train.rs;
/// serializing this struct produces identical JSONL output (verified by
/// the round-trip test in Task 4.3).
#[derive(Serialize)]
pub struct DiagRecord {
    pub step: u64,
    pub elapsed_s: f32,
    pub loss: LossBlock,
    pub popart: PopartBlock,
    pub risk_stack: RiskStackBlock,
    // ... mirror every nested block currently emitted.
}

#[derive(Serialize)]
pub struct LossBlock {
    pub q: f32,
    pub pi: f32,
    pub v: f32,
    pub bce: f32,
    pub frd: f32,
    pub aux: f32,
    pub total: f32,
}

#[derive(Serialize)]
pub struct PopartBlock {
    pub mean: f32,
    pub sigma: f32,
    pub var: f32,
}

#[derive(Serialize)]
pub struct RiskStackBlock {
    pub cmdp: CmdpBlock,
    pub kelly: KellyBlock,
    pub regime: RegimeBlock,
    pub atom_calibration: AtomCalBlock,
}

// ... full type enumeration (mechanical, reading the current JSONL schema).
// The exact nested layout matches what's currently emitted in
// alpha_rl_train.rs's json!{} block.
```

> **Implementer note**: This is a substantial but mechanical task. Read every key in `crates/ml-alpha/examples/alpha_rl_train.rs`'s `json!{}` block (in the train loop) and add a typed field for it. Match nesting exactly. Field order does NOT need to match — serde uses struct order but JSON consumers don't care about order. If any field can be null, use `Option<f32>`.

- [ ] **Step 3: Update `mod.rs`**

```rust
// crates/ml-alpha/src/trainer/mod.rs
pub mod diag_record;
```

- [ ] **Step 4: Build to verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/src/trainer/diag_record.rs crates/ml-alpha/src/trainer/mod.rs
git commit -m "feat(trainer): typed DiagRecord (4.1)"
```

### Task 4.2: Add `IntegratedTrainer::snapshot_diag_record`

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Implement the snapshot method**

```rust
// crates/ml-alpha/src/trainer/integrated.rs

use crate::trainer::diag_record::{
    DiagRecord, LossBlock, PopartBlock, RiskStackBlock, /* ... */
};

impl IntegratedTrainer {
    /// Snapshot the trainer state into a typed DiagRecord. Used by both
    /// train loop (replacing the existing inline json!{} construction)
    /// and eval loop (new — Phase 5 wires this).
    pub fn snapshot_diag_record(&self, step: u64, elapsed_s: f32) -> Result<DiagRecord> {
        // Read from ISV + last_*_loss + DiagStaging's last-readable
        // buffer. The exact reads mirror the existing
        // `step_with_lobsim_reward_and_train` JSONL build path —
        // duplicate them here. The trainer already does these reads
        // every step for the train JSONL; centralising them keeps the
        // SAME values flowing into both files.
        let loss = LossBlock {
            q:     self.last_q_loss,
            pi:    self.last_pi_loss,
            v:     self.last_v_loss,
            bce:   self.last_bce_loss,
            frd:   self.last_frd_loss,
            aux:   self.last_aux_loss,
            total: self.last_total_loss,
        };
        let popart = PopartBlock {
            mean:  self.isv_mapped.read_record(POPART_MEAN_INDEX),
            sigma: self.isv_mapped.read_record(POPART_SIGMA_INDEX),
            var:   self.isv_mapped.read_record(POPART_VAR_INDEX),
        };
        let risk_stack = RiskStackBlock { /* ... assemble */ };
        Ok(DiagRecord { step, elapsed_s, loss, popart, risk_stack })
    }
}
```

- [ ] **Step 2: Replace inline json!{} in alpha_rl_train.rs train loop**

Find the existing `json!{}` block in the train loop:

```bash
grep -nE "json!" crates/ml-alpha/examples/alpha_rl_train.rs
```

Replace the inline construction with:

```rust
let record = trainer.snapshot_diag_record(step as u64, t_start.elapsed().as_secs_f32())?;
diag_staging.write_record_typed(&record)?;
```

And update `DiagStaging::write_record_typed` to take `&DiagRecord` (serde_json::to_string).

- [ ] **Step 3: Smoke test — verify train JSONL schema unchanged**

```bash
# 100-step local smoke
SQLX_OFFLINE=true target/release/examples/alpha_rl_train --n-steps 100 \
  --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \
  --predecoded-dir /tmp/rl-smoke-lpi-diag/data \
  --out /tmp/diag-schema-test \
  --instrument-mode front-month --n-backtests 16 --log-every 50 --seed 42
# Compare last line key set vs pre-refactor sample
head -1 /tmp/diag-schema-test/diag.jsonl | jq -r 'keys[]' | sort > /tmp/new_keys.txt
head -1 /tmp/rl-smoke-lpi-diag/diag.jsonl | jq -r 'keys[]' | sort > /tmp/old_keys.txt
diff /tmp/old_keys.txt /tmp/new_keys.txt
```

Expected: empty diff (same top-level keys). Drill down further with `jq -r 'paths|join(".")'` to compare nested.

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs crates/ml-alpha/src/trainer/diag_staging.rs crates/ml-alpha/examples/alpha_rl_train.rs
git commit -m "refactor(trainer): replace inline json! with typed DiagRecord (4.2)"
```

### Task 4.3: Wire eval-phase diag emission

**Files:**
- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs`

- [ ] **Step 1: Add eval-diag writer setup**

```rust
// crates/ml-alpha/examples/alpha_rl_train.rs — around eval phase entry (line ~1339)

let eval_diag_path = cli.eval_diag_jsonl.clone()
    .unwrap_or_else(|| cli.out.join("eval_diag.jsonl"));
let mut eval_diag_writer = std::io::BufWriter::new(
    std::fs::File::create(&eval_diag_path)
        .with_context(|| format!("create {}", eval_diag_path.display()))?
);
eprintln!("eval per-step diag JSONL: {}", eval_diag_path.display());
```

- [ ] **Step 2: Emit one line per eval step**

```rust
// Inside the eval loop, after `step_with_lobsim_gpu`:
let record = trainer.snapshot_diag_record(eval_step as u64, t_start.elapsed().as_secs_f32())?;
let line = serde_json::to_string(&record)?;
use std::io::Write;
writeln!(&mut eval_diag_writer, "{}", line)?;
// Flush every log_every steps to keep tail -f working.
if eval_step % cli.log_every == 0 {
    eval_diag_writer.flush()?;
}
```

- [ ] **Step 3: Flush at end of eval phase**

```rust
// After the eval loop:
eval_diag_writer.flush()?;
drop(eval_diag_writer);
```

- [ ] **Step 4: Add CLI flag (placeholder — fleshed out in Phase 5)**

```rust
// crates/ml-alpha/examples/alpha_rl_train.rs — in Cli struct

/// Eval-phase per-step diag JSONL path. Default: `<out>/eval_diag.jsonl`.
#[arg(long)]
eval_diag_jsonl: Option<PathBuf>,
```

- [ ] **Step 5: Smoke test — verify eval_diag.jsonl emits 5000 lines**

```bash
# Local 1k train + 100 eval
SQLX_OFFLINE=true target/release/examples/alpha_rl_train \
  --n-steps 1000 --n-eval-steps 100 --fold-idx 1 --n-folds 3 \
  --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \
  --predecoded-dir /tmp/rl-smoke-lpi-diag/data \
  --out /tmp/eval-diag-test \
  --instrument-mode front-month --n-backtests 16 --log-every 50 --seed 42
wc -l /tmp/eval-diag-test/diag.jsonl /tmp/eval-diag-test/eval_diag.jsonl
```

Expected: `diag.jsonl` = 1000 lines, `eval_diag.jsonl` = 100 lines.

- [ ] **Step 6: Commit**

```bash
git add crates/ml-alpha/examples/alpha_rl_train.rs
git commit -m "feat(rl): emit eval-phase per-step diag to eval_diag.jsonl (4.3)"
```

---

## Phase 5: Resume integration

### Task 5.1: Add CLI flags + resume entry

**Files:**
- Modify: `crates/ml-alpha/examples/alpha_rl_train.rs`

- [ ] **Step 1: Add CLI flags**

```rust
// crates/ml-alpha/examples/alpha_rl_train.rs — in Cli struct

/// Save a full (resumable) checkpoint every N train steps. 0 disables.
#[arg(long, default_value_t = 5000)]
checkpoint_every: usize,

/// Rolling-window count of full checkpoints to keep.
#[arg(long, default_value_t = 3)]
checkpoint_keep: usize,

/// Include the PER replay buffer in full checkpoints (~8 GB/ckpt).
#[arg(long)]
checkpoint_replay: bool,

/// Save a lite (inference) checkpoint every N train steps. 0 disables.
#[arg(long, default_value_t = 5000)]
inference_ckpt_every: usize,

/// Resume training from this checkpoint at startup. Trainer state is
/// fully restored; training continues from saved step + 1.
#[arg(long)]
resume_from: Option<PathBuf>,
```

- [ ] **Step 2: Wire resume**

```rust
// Inside main, after `let mut trainer = ...` construction:

let start_step: usize = if let Some(path) = &cli.resume_from {
    eprintln!("resuming from {}", path.display());
    let loaded_step = trainer.load_checkpoint(path)? as usize;
    eprintln!("resumed at step {} — continuing from step {}", loaded_step, loaded_step + 1);
    // TODO: advance the data loader by `loaded_step` sequences so the
    // next data window matches where we left off. If MultiHorizonLoader
    // doesn't have a seek method, add one OR document the known
    // limitation that resume re-uses the same data starting point
    // (acceptable for short-horizon resume).
    loaded_step + 1
} else { 0 };
```

- [ ] **Step 3: Wire periodic saves in train loop**

```rust
// Inside the train loop body, after each step:

if cli.checkpoint_every > 0 && step > 0 && step % cli.checkpoint_every == 0 {
    let ckpt_dir = cli.out.join("checkpoints");
    let path = ckpt_dir.join(format!("train_step_{}.full.ckpt", step));
    trainer.save_full_checkpoint(&path, step as u64, cli.checkpoint_replay)?;
    eprintln!("saved full checkpoint {} ({} MB)", path.display(),
              std::fs::metadata(&path)?.len() / (1024 * 1024));
    rolling_cleanup(&ckpt_dir, "train_step_", ".full.ckpt", cli.checkpoint_keep)?;
}
if cli.inference_ckpt_every > 0 && step > 0 && step % cli.inference_ckpt_every == 0 {
    let ckpt_dir = cli.out.join("checkpoints");
    let path = ckpt_dir.join(format!("inference_step_{}.lite.ckpt", step));
    trainer.save_inference_checkpoint(&path)?;
    rolling_cleanup(&ckpt_dir, "inference_step_", ".lite.ckpt", cli.checkpoint_keep)?;
}
```

- [ ] **Step 4: Add rolling cleanup helper**

```rust
// crates/ml-alpha/examples/alpha_rl_train.rs — top-level fn

fn rolling_cleanup(dir: &Path, prefix: &str, suffix: &str, keep: usize) -> Result<()> {
    let mut entries: Vec<_> = std::fs::read_dir(dir)?
        .filter_map(|e| e.ok())
        .filter(|e| {
            let name = e.file_name().to_string_lossy().into_owned();
            name.starts_with(prefix) && name.ends_with(suffix)
        })
        .collect();
    entries.sort_by_key(|e| e.metadata().and_then(|m| m.modified()).ok());
    while entries.len() > keep {
        let oldest = entries.remove(0);
        let _ = std::fs::remove_file(oldest.path());
    }
    Ok(())
}
```

- [ ] **Step 5: Add end-of-phase saves**

```rust
// After train loop completes (before eval phase):
let final_path = cli.out.join("checkpoints/inference_final.lite.ckpt");
trainer.save_inference_checkpoint(&final_path)?;

// After eval phase completes:
let post_eval_path = cli.out.join("checkpoints/inference_post_eval.lite.ckpt");
trainer.save_inference_checkpoint(&post_eval_path)?;
```

- [ ] **Step 6: Smoke test — save + resume**

```bash
# Run 200 steps, expecting a checkpoint at step 100
SQLX_OFFLINE=true target/release/examples/alpha_rl_train \
  --n-steps 200 --checkpoint-every 100 --checkpoint-keep 2 \
  --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \
  --predecoded-dir /tmp/rl-smoke-lpi-diag/data \
  --out /tmp/resume-smoke --instrument-mode front-month \
  --n-backtests 16 --log-every 50 --seed 42

# Resume from step 100, run 100 more steps
SQLX_OFFLINE=true target/release/examples/alpha_rl_train \
  --n-steps 200 --resume-from /tmp/resume-smoke/checkpoints/train_step_100.full.ckpt \
  --mbp10-data-dir /tmp/rl-smoke-lpi-diag/data \
  --predecoded-dir /tmp/rl-smoke-lpi-diag/data \
  --out /tmp/resume-smoke-cont --instrument-mode front-month \
  --n-backtests 16 --log-every 50 --seed 42
```

Expected: second run prints `resumed at step 100 — continuing from step 101`. Final losses in `/tmp/resume-smoke-cont/diag.jsonl` should be close to (but not identical to) those of an unbroken 200-step run, modulo data-loader-position drift.

- [ ] **Step 7: Commit**

```bash
git add crates/ml-alpha/examples/alpha_rl_train.rs
git commit -m "feat(rl): CLI flags + resume + rolling checkpoints (5.1)"
```

### Task 5.2: Resume continuation smoke test

**Files:**
- Create: `crates/ml-alpha/tests/resume_continuation.rs`

- [ ] **Step 1: Write the test**

```rust
// crates/ml-alpha/tests/resume_continuation.rs

#[test]
#[ignore = "requires CUDA + test_data"]
fn resume_from_checkpoint_continues_trajectory() -> Result<()> {
    // 1. Run trainer A: 100 steps, save full checkpoint.
    // 2. Construct trainer B, load A's checkpoint.
    // 3. Run B for 1 more step.
    // 4. Run trainer A for 1 more step (101 total).
    // 5. Compare A.last_pi_loss to B.last_pi_loss; expect within ε=1e-3
    //    (modulo non-determinism in PER sampling).
    // Detailed implementation matches the harness used in checkpoint_full_roundtrip
    // but with actual train iterations rather than synthetic mutations.
    Ok(())
}
```

- [ ] **Step 2: Run + commit**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test --release -p ml-alpha --test resume_continuation -- --ignored --nocapture
git add crates/ml-alpha/tests/resume_continuation.rs
git commit -m "test(rl): resume continuation smoke (5.2)"
```

---

## Phase 6: Cluster validation

### Task 6.1: Inference-deployment smoke (separate process)

**Files:**
- Create: `crates/ml-alpha/examples/alpha_rl_inference_smoke.rs`

- [ ] **Step 1: Write the smoke binary**

```rust
// crates/ml-alpha/examples/alpha_rl_inference_smoke.rs

//! Loads a lite checkpoint and runs N inference steps on held-out data.
//! Confirms the checkpoint is usable for deployment.

use anyhow::Result;
use clap::Parser;
use ml_alpha::trainer::integrated::IntegratedTrainer;
// ... full imports

#[derive(Parser)]
struct Cli {
    #[arg(long)] checkpoint: PathBuf,
    #[arg(long)] data_dir:   PathBuf,
    #[arg(long, default_value_t = 1000)] n_steps: usize,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    let dev = MlDevice::cuda(0)?;
    let mut trainer = IntegratedTrainer::with_controllers_bootstrapped(&dev, 16, 0)?;
    let _step = trainer.load_checkpoint(&cli.checkpoint)?;
    // ... build loader, run N step_with_lobsim_gpu calls, print summary.
    eprintln!("inference smoke OK: ran {} steps with checkpoint loaded", cli.n_steps);
    Ok(())
}
```

- [ ] **Step 2: Build + commit**

```bash
SQLX_OFFLINE=true cargo build --release --example alpha_rl_inference_smoke -p ml-alpha
git add crates/ml-alpha/examples/alpha_rl_inference_smoke.rs
git commit -m "feat(rl): inference deployment smoke binary (6.1)"
```

### Task 6.2: Push + cluster validation

- [ ] **Step 1: Push**

```bash
git push
```

- [ ] **Step 2: Submit cluster run**

```bash
scripts/argo-alpha-rl.sh \
  --sha $(git rev-parse --short HEAD) \
  --branch ml-alpha-checkpoints-eval-diag \
  --gpu-pool ci-training-l40s \
  --n-steps 20000 --n-backtests 1024 --seed 16962 \
  --instrument-mode front-month \
  --fold-idx 1 --n-folds 3 --n-eval-steps 5000
```

Note: cluster script may need a flag update to pass the new checkpoint args through to the in-pod alpha_rl_train invocation. If `argo-alpha-rl.sh` doesn't accept `--checkpoint-every`, add it as a pass-through.

- [ ] **Step 3: Verify on cluster**

After run completes:
1. `kubectl exec` into the pod or a fresh investigator pod (PVC unlocked post-completion).
2. `ls /feature-cache/alpha-rl-runs/<sha>/fold1/checkpoints/` → expect ~5 files (3× lite, 3× full, inference_final, inference_post_eval).
3. `wc -l /feature-cache/alpha-rl-runs/<sha>/fold1/eval_diag.jsonl` → expect 5000.
4. Copy `inference_post_eval.lite.ckpt` locally, run `alpha_rl_inference_smoke` against it on 1k held-out steps.

- [ ] **Step 4: Commit pearl + update memory**

```bash
# Pearl describing the design
cat > /home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_checkpoint_two_variant_lite_full.md <<'EOF'
---
name: pearl-checkpoint-two-variant-lite-full
description: IntegratedTrainer checkpoints come in two variants — lite (55 MB, weights only, for inference) and full (65 MB, +Adam +ISV +scalars, for resume). PER replay is opt-in (~8 GB).
metadata:
  type: pearl
---
[body]
EOF
```

- [ ] **Step 5: Update MEMORY.md index**

Add the pearl link under the "Patterns / Pearls" section.

---

## Self-review checklist (before declaring done)

- [ ] Every checkpoint round-trip test passes on local RTX 3050
- [ ] `diag.jsonl` schema unchanged after Task 4.2 refactor (jq diff = empty)
- [ ] `eval_diag.jsonl` emits the same schema as `diag.jsonl`
- [ ] Local resume smoke shows continuation within 1e-3 loss tolerance
- [ ] Inference smoke binary loads a lite checkpoint and runs without errors
- [ ] Cluster run produces all expected files on PVC
- [ ] Pearl + MEMORY.md updated
- [ ] No `unwrap()` calls added in any new code (per pre-commit hooks)
- [ ] No `atomicAdd` added in any kernel (we don't touch kernels in this plan, but defensive)

## Done means

- All Phase 1–5 tests pass
- One cluster run completes producing the full checkpoint set
- `alpha_rl_inference_smoke` runs successfully against the cluster's `inference_post_eval.lite.ckpt`
- Pearl saved + memory index updated
- Branch ready to merge into `ml-alpha-regime-observer`
