From 3056f08535ccdceca5aa66c4f355b5761cd596a7 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 30 Mar 2026 10:08:54 +0200
Subject: [PATCH] feat(generalization): #32 vaccine batch sampling from replay
 buffer

Wire the gradient vaccine's validation batch through the training loop:
- Pre-sample a SECOND batch from PER alongside each training batch
- Pass via FusedTrainingCtx::pending_vaccine_batch (consumed per step)
- Separate PER indices ensure train and val batches are independent

The vaccine's non-graph forward+backward runs on this held-out batch
to produce g_val. Gradient projection (dot + SAXPY) ensures g_train
only moves in directions where both batches agree.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../src/trainers/dqn/trainer/training_loop.rs | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs
index bbf97667b..e852321b6 100644
--- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs
+++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs
@@ -1356,7 +1356,7 @@ impl DQNTrainer {
             let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);
 
             let sample_start = std::time::Instant::now();
-            let batches = {
+            let (batches, vaccine_batches) = {
                 let agent = self.agent.read().await;
                 let buffer = agent.memory();
                 let can = buffer.can_sample(self.current_batch_size);
@@ -1369,14 +1369,24 @@ impl DQNTrainer {
                     "Pre-sampling replay buffer"
                 );
                 let mut b = Vec::with_capacity(chunk_end - chunk_start);
+                // #32 Gradient Vaccine: sample a second batch for validation
+                let mut vaccine_batches = Vec::with_capacity(chunk_end - chunk_start);
+                let vaccine_enabled = self.hyperparams.enable_gradient_vaccine;
                 for _ in chunk_start..chunk_end {
                     b.push(buffer.can_sample(self.current_batch_size).then(|| {
                         buffer
                             .sample(self.current_batch_size)
                             .map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
                     }).transpose()?);
+                    // Sample separate vaccine batch (different indices from PER)
+                    if vaccine_enabled && buffer.can_sample(self.current_batch_size) {
+                        vaccine_batches.push(buffer.sample(self.current_batch_size)
+                            .map_err(|e| anyhow::anyhow!("Vaccine sample: {e}")).ok());
+                    } else {
+                        vaccine_batches.push(None);
+                    }
                 }
-                b
+                (b, vaccine_batches)
             };
             sample_total_us += sample_start.elapsed().as_micros() as u64;
 
@@ -1386,12 +1396,17 @@ impl DQNTrainer {
                 let accum_steps = self.hyperparams.gradient_accumulation_steps;
 
                 if accum_steps <= 1 {
+                    let mut vaccine_iter = vaccine_batches.into_iter();
                     for explicit_batch in batches {
                         let fused_start = std::time::Instant::now();
                         let _gpu_result = if let Some(ref mut fused) = self.fused_ctx {
                             let batch_data = explicit_batch.as_ref().ok_or_else(|| {
                                 anyhow::anyhow!("No batch data for fused training step")
                             })?;
+                            // #32 Set vaccine batch for this step (consumed by run_full_step)
+                            if let Some(Some(vb)) = vaccine_iter.next() {
+                                fused.pending_vaccine_batch = vb.gpu_batch;
+                            }
                             let result = fused.run_full_step(batch_data, &mut *agent, &self.device)
                                 .map_err(|e| { eprintln!("!!! FUSED STEP ERROR: {:#}", e); e })
                                 .context("Fused CUDA training step failed")?;