diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index ac4c8dafe..71e33c357 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -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; diff --git a/crates/ml/src/cuda_pipeline/prefetch.rs b/crates/ml/src/cuda_pipeline/prefetch.rs deleted file mode 100644 index f27ae83a6..000000000 --- a/crates/ml/src/cuda_pipeline/prefetch.rs +++ /dev/null @@ -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)>; - -/// 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>, - #[allow(dead_code)] - handle: Option>, - completed: Arc, -} - -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(load_fn: F) -> Self - where - F: FnOnce() -> Result + 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 { - 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> { - 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); - } -} diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 58c692beb..5ffaebeb1 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -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>) -> Result { - // 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) -> Result { + // 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>) -> Result<(f32, f32), MLError> { + pub fn train_step(&mut self, batch: Option) -> 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>, + batch: Option, ) -> Result { // Skip during warmup if self.total_steps < self.config.warmup_steps as u64 { diff --git a/crates/ml/src/dqn/ensemble.rs b/crates/ml/src/dqn/ensemble.rs index 468f3d47d..11dccd049 100644 --- a/crates/ml/src/dqn/ensemble.rs +++ b/crates/ml/src/dqn/ensemble.rs @@ -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 diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml/src/dqn/regime_conditional.rs index 52922ee8b..4267a5067 100644 --- a/crates/ml/src/dqn/regime_conditional.rs +++ b/crates/ml/src/dqn/regime_conditional.rs @@ -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>) -> Result<(f32, f32), MLError> { + pub fn train_step(&mut self, batch: Option) -> 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>, + batch: Option, ) -> Result { 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; diff --git a/crates/ml/src/dqn/replay_buffer_type.rs b/crates/ml/src/dqn/replay_buffer_type.rs index 1709a9a86..4e71573ae 100644 --- a/crates/ml/src/dqn/replay_buffer_type.rs +++ b/crates/ml/src/dqn/replay_buffer_type.rs @@ -19,7 +19,20 @@ pub struct BatchSample { pub indices: Vec, // 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) -> 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>), diff --git a/crates/ml/src/dqn/trainable_adapter.rs b/crates/ml/src/dqn/trainable_adapter.rs index e01c3e95d..2f86b2a3b 100644 --- a/crates/ml/src/dqn/trainable_adapter.rs +++ b/crates/ml/src/dqn/trainable_adapter.rs @@ -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) -> Result { - 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) diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index c27319d5b..9ea798740 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -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 diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 6a1c80ae1..0880d1785 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -54,7 +54,7 @@ impl DQNAgentType { } /// Training step - returns (loss, grad_norm) - pub fn train_step(&mut self, batch: Option>) -> Result<(f32, f32), MLError> { + pub fn train_step(&mut self, batch: Option) -> 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>, + batch: Option, ) -> Result { 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) diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index bb36575a4..6d7248251 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -213,9 +213,6 @@ pub struct DQNTrainer { #[cfg(feature = "cuda")] gpu_experience_collector: Option, - /// Background prefetcher for overlapping disk I/O with GPU training - prefetcher: Option, - /// Reusable GPU staging buffers for zero-alloc fold transitions buffer_pool: Option, @@ -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> { - 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> = { + 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.