gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ml): ungate GPU experience collector + async PER pre-sampling

Browse Source 
Enable curiosity module by default (weight 0.0→0.1) to satisfy the
three-way gate (dueling + target + curiosity) that was blocking the
GPU experience collector CUDA kernel. This eliminates the 30-40% CPU
experience collection phase that was the main GPU idle bottleneck.

Additional changes:
- BatchSample API: train_step/compute_gradients now accept
  Option<BatchSample> instead of Option<Vec<Experience>>, preserving
  PER importance-sampling weights and indices through pre-sampling.
- Async PER pre-sampling: train_step_single_batch and
  train_step_with_accumulation now pre-sample from the replay buffer
  using a READ lock before acquiring the WRITE lock for GPU training.
  This separates CPU sampling (~250μs) from GPU forward/backward (~3ms).
- Delete dead EpochPrefetcher: the binary (train_baseline_rl.rs) already
  implements fold prefetching with background thread + mpsc channel +
  GPU double-buffering, making the trainer's EpochPrefetcher redundant.
- Hyperopt DQN bounds: curiosity_weight min 0.0→0.01 so PSO can never
  fully disable curiosity (which would re-gate the GPU collector).

10 files changed, -195 net lines. 2497 tests pass, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-02 21:14:14 +01:00
parent 0c227bb7d0
commit 148d4cbb43
10 changed files with 85 additions and 280 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -10,7 +10,6 @@ use crate::MLError;
pub mod double_buffer;
pub mod multi_gpu;
pub mod prefetch;
#[cfg(feature = "cuda")]
pub mod gpu_portfolio;

183
crates/ml/src/cuda_pipeline/prefetch.rs
View File

@@ -1,183 +0,0 @@
//! Background data prefetcher for overlapping disk I/O with GPU training.
//!
//! Spawns a `std::thread` to load and extract features for the next
//! walk-forward fold while the current fold trains on GPU. The result
//! can be collected via `try_take()` (non-blocking) or `take()` (blocking).
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use tracing::{info, warn};
use crate::MLError;
/// Result of a prefetch operation: training data ready for GPU upload.
pub type PrefetchResult = Vec<([f64; 51], Vec<f64>)>;
/// Background data prefetcher.
///
/// Spawns a thread to run a data-loading closure. The caller collects
/// the result when the current training fold finishes.
///
/// # Example
///
/// ```ignore
/// let prefetcher = EpochPrefetcher::spawn(move || load_next_fold(&path));
/// // ... train current fold ...
/// let next_data = prefetcher.take()?;
/// double_buffer.upload_to_staging(&next_data)?;
/// double_buffer.swap()?;
/// ```
pub struct EpochPrefetcher {
receiver: mpsc::Receiver<Result<PrefetchResult, MLError>>,
#[allow(dead_code)]
handle: Option<JoinHandle<()>>,
completed: Arc<AtomicBool>,
}
impl std::fmt::Debug for EpochPrefetcher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EpochPrefetcher")
.field("handle", &self.handle.as_ref().map(|_| "..."))
.field("completed", &self.completed.load(Ordering::Acquire))
.finish()
}
}
impl EpochPrefetcher {
/// Spawn a background thread that runs the given data-loading closure.
///
/// The closure should load training data from disk, extract features,
/// and return the result. It runs on a dedicated OS thread to avoid
/// blocking the async runtime.
pub fn spawn<F>(load_fn: F) -> Self
where
F: FnOnce() -> Result<PrefetchResult, MLError> + Send + 'static,
{
let (tx, rx) = mpsc::channel();
let completed = Arc::new(AtomicBool::new(false));
let completed_clone = Arc::clone(&completed);
let handle = thread::Builder::new()
.name("epoch-prefetch".into())
.spawn(move || {
info!("EpochPrefetcher: background data load started");
let result = load_fn();
match &result {
Ok(data) => info!(
"EpochPrefetcher: loaded {} bars",
data.len(),
),
Err(e) => warn!("EpochPrefetcher: load failed: {e}"),
}
// Signal completion before sending so is_ready() becomes
// true even if the receiver has been dropped.
completed_clone.store(true, Ordering::Release);
// Receiver may have been dropped if the trainer finished early
drop(tx.send(result));
})
.ok();
Self {
receiver: rx,
handle,
completed,
}
}
/// Block until the prefetch completes and return the result.
pub fn take(self) -> Result<PrefetchResult, MLError> {
let result = self.receiver.recv().map_err(|_| {
MLError::ModelError("EpochPrefetcher: background thread dropped sender".into())
})?;
if let Some(handle) = self.handle {
drop(handle.join());
}
result
}
/// Non-blocking check: returns `Some(result)` if ready, `None` if still loading.
pub fn try_take(&self) -> Option<Result<PrefetchResult, MLError>> {
self.receiver.try_recv().ok()
}
/// Whether the prefetch has completed (result is ready to take).
///
/// Non-destructive: uses an atomic flag set by the background thread
/// after `load_fn()` returns, so this can be called repeatedly without
/// consuming the result.
pub fn is_ready(&self) -> bool {
self.completed.load(Ordering::Acquire)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prefetch_success() {
let prefetcher = EpochPrefetcher::spawn(|| {
let data: PrefetchResult = (0..100)
.map(|i| {
let features = [i as f64 * 0.01; 51];
let targets = vec![100.0, 101.0, 100.5, 101.5];
(features, targets)
})
.collect();
Ok(data)
});
let result = prefetcher.take().unwrap();
assert_eq!(result.len(), 100);
}
#[test]
fn test_prefetch_error() {
let prefetcher = EpochPrefetcher::spawn(|| {
Err(MLError::ModelError("disk read failed".into()))
});
let result = prefetcher.take();
assert!(result.is_err());
}
#[test]
fn test_try_take_not_ready_initially() {
let prefetcher = EpochPrefetcher::spawn(|| {
std::thread::sleep(std::time::Duration::from_millis(100));
Ok(vec![])
});
// Immediately after spawn, result likely not ready yet
// (but this is timing-dependent, so we just verify the API works)
let _maybe = prefetcher.try_take();
}
#[test]
fn test_is_ready_returns_true_when_complete() {
let prefetcher = EpochPrefetcher::spawn(|| Ok(vec![]));
std::thread::sleep(std::time::Duration::from_millis(50));
assert!(
prefetcher.is_ready(),
"is_ready should return true after task completes"
);
}
#[test]
fn test_prefetch_large_data() {
let prefetcher = EpochPrefetcher::spawn(|| {
let data: PrefetchResult = (0..50_000)
.map(|i| {
let features = [i as f64 * 0.001; 51];
let targets = vec![100.0, 101.0, 100.5, 101.5];
(features, targets)
})
.collect();
Ok(data)
});
let result = prefetcher.take().unwrap();
assert_eq!(result.len(), 50_000);
}
}

14
crates/ml/src/dqn/dqn.rs
View File

@@ -1351,12 +1351,10 @@ impl DQN {
/// # Arguments
///
/// * `batch` - Optional externally-provided batch. If `None`, samples from replay buffer.
fn compute_loss_internal(&mut self, batch: Option<Vec<Experience>>) -> Result<ComputeLossResult, MLError> {
// Get batch of experiences
let (experiences, weights, indices) = if let Some(batch) = batch {
// External batch provided - create uniform weights
let len = batch.len();
(batch, vec![1.0; len], vec![])
fn compute_loss_internal(&mut self, batch: Option<super::replay_buffer_type::BatchSample>) -> Result<ComputeLossResult, MLError> {
// Get batch of experiences with PER IS-weights preserved when pre-sampled
let (experiences, weights, indices) = if let Some(batch_sample) = batch {
(batch_sample.experiences, batch_sample.weights, batch_sample.indices)
} else {
// Sample from replay buffer (uniform or prioritized)
if !self.memory.can_sample(self.config.min_replay_size) {
@@ -1977,7 +1975,7 @@ impl DQN {
/// Training step with experience batch
///
/// Returns (loss, gradient_norm) tuple
pub fn train_step(&mut self, batch: Option<Vec<Experience>>) -> Result<(f32, f32), MLError> {
pub fn train_step(&mut self, batch: Option<super::replay_buffer_type::BatchSample>) -> Result<(f32, f32), MLError> {
// Skip gradient updates during warmup period
if self.total_steps < self.config.warmup_steps as u64 {
return Ok((0.0, 0.0)); // Return dummy values, no training during warmup
@@ -2040,7 +2038,7 @@ impl DQN {
/// or if the forward/backward pass fails.
pub fn compute_gradients(
&mut self,
batch: Option<Vec<Experience>>,
batch: Option<super::replay_buffer_type::BatchSample>,
) -> Result<GradientResult, MLError> {
// Skip during warmup
if self.total_steps < self.config.warmup_steps as u64 {

2
crates/ml/src/dqn/ensemble.rs
View File

@@ -560,7 +560,7 @@ impl DQNEnsemble {
let batch = buffer.sample(cfg.batch_size)?;
drop(buffer); // Release lock before training
let (loss, _grad_norm) = agent.train_step(Some(batch))?;
let (loss, _grad_norm) = agent.train_step(Some(super::replay_buffer_type::BatchSample::uniform(batch)))?;
losses.push(loss);
} else {
losses.push(0.0); // Not enough data yet

24
crates/ml/src/dqn/regime_conditional.rs
View File

@@ -324,10 +324,10 @@ impl RegimeConditionalDQN {
/// # Returns
///
/// Tuple of (average_loss, average_grad_norm) across all heads
pub fn train_step(&mut self, batch: Option<Vec<Experience>>) -> Result<(f32, f32), MLError> {
pub fn train_step(&mut self, batch: Option<super::replay_buffer_type::BatchSample>) -> Result<(f32, f32), MLError> {
// Get batch of experiences from trending head's buffer (all heads have same experiences)
let experiences = if let Some(batch) = batch {
batch
let experiences = if let Some(batch_sample) = batch {
batch_sample.experiences
} else {
// Use constants for min_replay_size and batch_size
if !self.trending_head.memory.can_sample(100) {
@@ -359,7 +359,7 @@ impl RegimeConditionalDQN {
let mut num_heads_trained = 0;
if !trending_batch.is_empty() {
let (loss, grad_norm) = self.trending_head.train_step(Some(trending_batch))?;
let (loss, grad_norm) = self.trending_head.train_step(Some(super::replay_buffer_type::BatchSample::uniform(trending_batch)))?;
total_loss += loss;
total_grad_norm += grad_norm;
num_heads_trained += 1;
@@ -375,7 +375,7 @@ impl RegimeConditionalDQN {
}
if !ranging_batch.is_empty() {
let (loss, grad_norm) = self.ranging_head.train_step(Some(ranging_batch))?;
let (loss, grad_norm) = self.ranging_head.train_step(Some(super::replay_buffer_type::BatchSample::uniform(ranging_batch)))?;
total_loss += loss;
total_grad_norm += grad_norm;
num_heads_trained += 1;
@@ -391,7 +391,7 @@ impl RegimeConditionalDQN {
}
if !volatile_batch.is_empty() {
let (loss, grad_norm) = self.volatile_head.train_step(Some(volatile_batch))?;
let (loss, grad_norm) = self.volatile_head.train_step(Some(super::replay_buffer_type::BatchSample::uniform(volatile_batch)))?;
total_loss += loss;
total_grad_norm += grad_norm;
num_heads_trained += 1;
@@ -455,12 +455,12 @@ impl RegimeConditionalDQN {
/// contains no key collisions.
pub fn compute_gradients(
&mut self,
batch: Option<Vec<Experience>>,
batch: Option<super::replay_buffer_type::BatchSample>,
) -> Result<GradientResult, MLError> {
use candle_core::backprop::GradStore;
let experiences = if let Some(b) = batch {
b
let experiences = if let Some(batch_sample) = batch {
batch_sample.experiences
} else {
if !self.trending_head.memory.can_sample(100) {
return Err(MLError::TrainingError(
@@ -519,7 +519,7 @@ impl RegimeConditionalDQN {
}
if !trending_batch.is_empty() {
let result = self.trending_head.compute_gradients(Some(trending_batch))?;
let result = self.trending_head.compute_gradients(Some(super::replay_buffer_type::BatchSample::uniform(trending_batch)))?;
let vars = self.trending_head.optimizer_vars()?.to_vec();
merge_grads(&mut merged_grads, result.grads, &vars)?;
total_loss += result.loss;
@@ -530,7 +530,7 @@ impl RegimeConditionalDQN {
}
if !ranging_batch.is_empty() {
let result = self.ranging_head.compute_gradients(Some(ranging_batch))?;
let result = self.ranging_head.compute_gradients(Some(super::replay_buffer_type::BatchSample::uniform(ranging_batch)))?;
let vars = self.ranging_head.optimizer_vars()?.to_vec();
merge_grads(&mut merged_grads, result.grads, &vars)?;
total_loss += result.loss;
@@ -541,7 +541,7 @@ impl RegimeConditionalDQN {
}
if !volatile_batch.is_empty() {
let result = self.volatile_head.compute_gradients(Some(volatile_batch))?;
let result = self.volatile_head.compute_gradients(Some(super::replay_buffer_type::BatchSample::uniform(volatile_batch)))?;
let vars = self.volatile_head.optimizer_vars()?.to_vec();
merge_grads(&mut merged_grads, result.grads, &vars)?;
total_loss += result.loss;

15
crates/ml/src/dqn/replay_buffer_type.rs
View File

@@ -19,7 +19,20 @@ pub struct BatchSample {
pub indices: Vec<usize>, // Indices for priority updates (empty for uniform)
}
/// Enum wrapper for runtime buffer type switching
impl BatchSample {
/// Create a batch with uniform weights (no importance sampling correction).
/// Use this when experiences were not sampled via PER.
pub fn uniform(experiences: Vec<Experience>) -> Self {
let len = experiences.len();
Self { experiences, weights: vec![1.0; len], indices: vec![] }
}
}
/// Enum wrapper for runtime buffer type switching.
///
/// Clone is cheap (Arc inner types) and enables sharing the buffer between
/// the training thread and a background pre-sample thread for async PER.
#[derive(Clone)]
pub enum ReplayBufferType {
/// Uniform random sampling (standard DQN)
Uniform(Arc<Mutex<ExperienceReplayBuffer>>),

2
crates/ml/src/dqn/trainable_adapter.rs
View File

@@ -80,7 +80,7 @@ impl DQNTrainableAdapter {
///
/// This is a convenience method that combines forward, backward, and optimizer_step
pub fn train_batch(&mut self, experiences: Vec<Experience>) -> Result<f64, MLError> {
let (loss, _grad_norm) = self.dqn.train_step(Some(experiences))?;
let (loss, _grad_norm) = self.dqn.train_step(Some(super::replay_buffer_type::BatchSample::uniform(experiences)))?;
self.current_step += 1;
self.loss_history.push(loss as f64);
Ok(loss as f64)

12
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -403,7 +403,7 @@ impl Default for DQNParams {
kelly_min_trades: 20, // WAVE 19: Default 20 trades minimum
volatility_window: 20, // WAVE 19: Default 20-bar volatility window
warmup_ratio: 0.0, // WAVE 26 P1.5: Default no warmup (production stable)
curiosity_weight: 0.0, // WAVE 26 P1.8: Default no curiosity (disabled)
curiosity_weight: 0.1, // WAVE 26 P1.8: Default curiosity enabled (ungates GPU experience collector)
// WAVE 26 P1.4: Ensemble Uncertainty (default: DISABLED)
use_ensemble_uncertainty: false, // Default: disabled (use noisy networks)
ensemble_size: 5.0, // Default: 5 heads
@@ -495,7 +495,7 @@ impl ParameterSpace for DQNParams {
(0.0, 0.2), // 28: warmup_ratio (0-20% warmup)
// WAVE 26 P1.8: Curiosity-driven exploration (29D → 30D)
(0.0, 0.5), // 29: curiosity_weight (intrinsic reward scaling)
(0.01, 0.5), // 29: curiosity_weight (intrinsic reward scaling, >0 required for GPU experience collector)
// WAVE 26 P1.12: Polyak soft update coefficient (30D → 31D)
(0.0001_f64.ln(), 0.01_f64.ln()), // 30: tau (log scale, 0.0001-0.01, Rainbow default: 0.001)
@@ -579,7 +579,7 @@ impl ParameterSpace for DQNParams {
let warmup_ratio = x[28].clamp(0.0, 0.2);
// WAVE 26 P1.8: Extract curiosity weight
let curiosity_weight = x[29].clamp(0.0, 0.5);
let curiosity_weight = x[29].clamp(0.01, 0.5);
// WAVE 26 P1.12: Extract tau (Polyak soft update coefficient)
let tau = x[30].exp().clamp(0.0001, 0.01); // Log scale: 0.0001-0.01, default: 0.001
@@ -3337,7 +3337,7 @@ mod tests {
ensemble_size: 5.0,
beta_variance: 0.5, beta_disagreement: 0.5, beta_entropy: 0.1,
warmup_ratio: 0.1, // WAVE 26 P1.5
curiosity_weight: 0.0, // WAVE 26 P1.8
curiosity_weight: 0.1, // WAVE 26 P1.8: enabled (GPU experience collector)
// WAVE 26 P0: TD Error and Batch Diversity
td_error_clamp_max: 10.0,
batch_diversity_cooldown: 50.0,
@@ -3476,7 +3476,7 @@ mod tests {
// WAVE 26 P1.5: Warmup ratio
0.1, // warmup_ratio (10% warmup)
// WAVE 26 P1.8: Curiosity weight
0.0, // curiosity_weight
0.1, // curiosity_weight (>0 required for GPU experience collector)
// WAVE 26 P1.12: Tau (Polyak soft update)
0.005_f64.ln(), // tau (about -5.3)
// WAVE 26 P0: TD Error and Batch Diversity
@@ -3514,7 +3514,7 @@ mod tests {
// WAVE 26 P1.5: Warmup ratio min
0.0, // warmup_ratio min
// WAVE 26 P1.8: Curiosity weight min
0.0, // curiosity_weight min
0.01, // curiosity_weight min (>0 required for GPU experience collector)
// WAVE 26 P1.12: Tau min
0.0001_f64.ln(), // tau min (about -9.21)
// WAVE 26 P0: TD Error and Batch Diversity min

6
crates/ml/src/trainers/dqn/config.rs
View File

@@ -54,7 +54,7 @@ impl DQNAgentType {
}
/// Training step - returns (loss, grad_norm)
pub fn train_step(&mut self, batch: Option<Vec<Experience>>) -> Result<(f32, f32), MLError> {
pub fn train_step(&mut self, batch: Option<crate::dqn::replay_buffer_type::BatchSample>) -> Result<(f32, f32), MLError> {
match self {
Self::Standard(agent) => agent.train_step(batch),
Self::RegimeConditional(agent) => agent.train_step(batch),
@@ -300,7 +300,7 @@ impl DQNAgentType {
/// Only supported for Standard DQN. RegimeConditional returns an error.
pub fn compute_gradients(
&mut self,
batch: Option<Vec<Experience>>,
batch: Option<crate::dqn::replay_buffer_type::BatchSample>,
) -> Result<crate::dqn::dqn::GradientResult, MLError> {
match self {
Self::Standard(agent) => agent.compute_gradients(batch),
@@ -778,7 +778,7 @@ impl DQNHyperparameters {
variance_cap: 1.0, // Default: cap at 1.0 to prevent over-penalization
// WAVE 26 P1.8: Curiosity-Driven Exploration
curiosity_weight: 0.0, // Default: disabled (hyperopt will tune 0.0-0.5)
curiosity_weight: 0.1, // Default: enabled (ungates GPU experience collector; hyperopt tunes 0.01-0.5)
// WAVE 26 P2.2: Gradient Accumulation
gradient_accumulation_steps: 1, // Default: no accumulation (standard training)

106
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -213,9 +213,6 @@ pub struct DQNTrainer {
#[cfg(feature = "cuda")]
gpu_experience_collector: Option<crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector>,
/// Background prefetcher for overlapping disk I/O with GPU training
prefetcher: Option<crate::cuda_pipeline::prefetch::EpochPrefetcher>,
/// Reusable GPU staging buffers for zero-alloc fold transitions
buffer_pool: Option<crate::cuda_pipeline::GpuBufferPool>,
@@ -845,9 +842,6 @@ impl DQNTrainer {
#[cfg(feature = "cuda")]
gpu_experience_collector: None,
// GPU pipeline: background prefetcher (activated via set_prefetcher())
prefetcher: None,
// GPU pipeline: staging buffer pool (auto-initialized on CUDA devices)
buffer_pool,
@@ -859,22 +853,6 @@ impl DQNTrainer {
})
}
/// Set a background prefetcher for overlapping disk I/O with GPU training.
///
/// The prefetcher loads the next walk-forward fold's data on a background
/// thread while the current fold trains on GPU.
pub fn set_prefetcher(&mut self, prefetcher: crate::cuda_pipeline::prefetch::EpochPrefetcher) {
self.prefetcher = Some(prefetcher);
}
/// Take the current prefetcher result (blocking), if one is active.
///
/// Returns `None` if no prefetcher was set, otherwise blocks until
/// the background load completes and returns the prefetched data.
pub fn take_prefetched_data(&mut self) -> Option<Result<crate::cuda_pipeline::prefetch::PrefetchResult>> {
self.prefetcher.take().map(|p| p.take().map_err(Into::into))
}
/// Get a reference to the double-buffered loader, if GPU is active.
pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
self.double_buffer.as_ref()
@@ -3284,26 +3262,26 @@ impl DQNTrainer {
}
/// Standard single-batch training step (no gradient accumulation)
///
/// Pre-samples from PER buffer using READ lock before acquiring WRITE lock
/// for GPU training. This preserves PER IS-weights and indices for proper
/// importance sampling correction and priority updates.
async fn train_step_single_batch(&mut self) -> Result<(f64, f64, f64)> {
let mut agent = self.agent.write().await;
// OOM recovery fix: when the batch size has been halved by the OOM retry
// loop, we must sample from the replay buffer ourselves at the reduced
// size and pass the explicit batch. Otherwise `agent.train_step(None)`
// would sample at the agent's original (un-halved) `config.batch_size`.
let explicit_batch = if self.current_batch_size < self.hyperparams.batch_size {
// Pre-sample batch OUTSIDE the write lock using read-only access to the buffer.
// PER IS-weights and indices are preserved for correct importance sampling.
// The write lock is only held during GPU forward/backward + optimizer step.
let explicit_batch = {
let agent = self.agent.read().await;
let buffer = agent.memory();
if buffer.can_sample(self.current_batch_size) {
let batch_sample = buffer
.sample(self.current_batch_size)
.map_err(|e| anyhow::anyhow!("Failed to sample reduced batch: {}", e))?;
Some(batch_sample.experiences)
} else {
None // Fall back to agent's default sampling
}
} else {
None // Use agent's default sampling (original batch size)
};
let sample_size = self.current_batch_size;
buffer.can_sample(sample_size).then(|| {
buffer
.sample(sample_size)
.map_err(|e| anyhow::anyhow!("Failed to pre-sample batch: {}", e))
}).transpose()?
}; // READ lock released here
let mut agent = self.agent.write().await;
// Call the agent's train_step which implements real Q-learning
// Now returns (loss, grad_norm) tuple with actual gradient norm from optimizer
@@ -3375,6 +3353,28 @@ impl DQNTrainer {
self.current_batch_size * accumulation_steps
);
// Pre-sample ALL mini-batches using READ lock (no GPU contention).
let pre_sampled: Vec<Option<crate::dqn::replay_buffer_type::BatchSample>> = {
let agent = self.agent.read().await;
let buffer = agent.memory();
let sample_size = self.current_batch_size;
let mut batches = Vec::with_capacity(accumulation_steps);
for step_idx in 0..accumulation_steps {
batches.push(
buffer.can_sample(sample_size).then(|| {
buffer.sample(sample_size).map_err(|e| {
anyhow::anyhow!(
"Failed to pre-sample batch (accum step {}): {}",
step_idx,
e
)
})
}).transpose()?,
);
}
batches
}; // READ lock released
let mut agent = self.agent.write().await;
// === Phase 1: Accumulate gradients across N mini-batches ===
@@ -3384,32 +3384,10 @@ impl DQNTrainer {
let mut all_indices = Vec::new();
let mut final_grad_norm = 0.0_f32;
for step in 0..accumulation_steps {
// OOM recovery: sample at the reduced batch size when OOM
// recovery has halved `current_batch_size`.
let explicit_batch = if self.current_batch_size < self.hyperparams.batch_size {
let buffer = agent.memory();
if buffer.can_sample(self.current_batch_size) {
let batch_sample = buffer
.sample(self.current_batch_size)
.map_err(|e| {
anyhow::anyhow!(
"Failed to sample reduced batch (accum step {}): {}",
step,
e
)
})?;
Some(batch_sample.experiences)
} else {
None
}
} else {
None
};
for (step, batch) in pre_sampled.into_iter().enumerate() {
// Compute forward pass + backward WITHOUT optimizer step
let result = agent
.compute_gradients(explicit_batch)
.compute_gradients(batch)
.map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?;
// Get vars for accumulation. Var is an Arc wrapper so cloning is cheap.