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>
This commit is contained in:
jgrusewski
2026-03-30 10:08:54 +02:00
parent 09b6bd6eaa
commit 3056f08535

View File

@@ -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")?;