diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index a3649b916..a98b2dbb8 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -740,8 +740,9 @@ pub struct GradientResult { struct ComputeLossResult { /// The loss tensor (still in the computation graph for backward pass). loss_tensor: Tensor, - /// Scalar loss value extracted from the tensor. - loss_value: f32, + /// Loss tensor cast to F32 for deferred scalar readback. + /// Read via `to_scalar::()` AFTER backward pass to avoid premature GPU flush. + loss_f32_tensor: Tensor, /// TD errors for PER priority updates. td_errors: Vec, /// Replay buffer indices for PER priority updates. @@ -1204,14 +1205,41 @@ impl DQN { // Create experience replay buffer (uniform or prioritized based on config) let memory = if config.use_per { - // Create Prioritized Experience Replay buffer - super::replay_buffer_type::ReplayBufferType::new_prioritized( - config.replay_buffer_capacity, - config.per_alpha, - config.per_beta_start, - config.per_beta_max, - config.per_beta_annealing_steps, - )? + // GPU PER: on CUDA, allocate GPU-resident ring buffer with OOM fallback to CPU PER. + // This activates the GpuBatch fast path in compute_loss_internal() and + // GPU TD error retention — eliminating 5 of 6 CPU↔GPU roundtrips per train step. + #[cfg(feature = "cuda")] + { + if device.is_cuda() { + super::replay_buffer_type::ReplayBufferType::try_gpu_prioritized_with_fallback( + config.replay_buffer_capacity, + config.state_dim, + config.per_alpha, + config.per_beta_start, + config.per_beta_max, + config.per_beta_annealing_steps, + &device, + )? + } else { + super::replay_buffer_type::ReplayBufferType::new_prioritized( + config.replay_buffer_capacity, + config.per_alpha, + config.per_beta_start, + config.per_beta_max, + config.per_beta_annealing_steps, + )? + } + } + #[cfg(not(feature = "cuda"))] + { + super::replay_buffer_type::ReplayBufferType::new_prioritized( + config.replay_buffer_capacity, + config.per_alpha, + config.per_beta_start, + config.per_beta_max, + config.per_beta_annealing_steps, + )? + } } else { // Create uniform replay buffer (standard DQN) // P2 FIX: Initialize at BASE_CAPACITY (10K) for adaptive growth @@ -2639,14 +2667,14 @@ impl DQN { .map_err(|e| MLError::TrainingError(format!("Failed to cast weights tensor: {}", e)))? .detach() }; - let weighted_diff = (&diff * &weights_tensor)?; - if self.config.use_huber_loss { - // Huber loss: L(x) = 0.5 * x^2 if |x| <= delta, else delta * (|x| - 0.5*delta) + // Huber loss on RAW diff first, then weight by IS weights. + // L_i = Huber(diff_i), loss = mean(w_i * L_i) + // IS weights must be applied AFTER the nonlinear Huber function, + // not before — otherwise the quadratic/linear regime boundaries shift. let delta = self.config.huber_delta; - let abs_diff = weighted_diff.abs()?; + let abs_diff = diff.abs()?; - // Element-wise Huber loss (use weighted_diff for loss computation) let half = Tensor::from_vec(vec![0.5_f32; batch_size], batch_size, device) .map_err(|e| { MLError::TrainingError(format!("Failed to create half tensor: {}", e)) @@ -2655,9 +2683,8 @@ impl DQN { .map_err(|e| { MLError::TrainingError(format!("Failed to cast half tensor: {}", e)) })?; - let squared_loss = (&weighted_diff * &weighted_diff)?.broadcast_mul(&half)?; // 0.5 * x^2 + let squared_loss = (&diff * &diff)?.broadcast_mul(&half)?; // 0.5 * x^2 - // Create delta tensor for operations let delta_tensor = Tensor::from_vec(vec![delta; batch_size], batch_size, device) .map_err(|e| { MLError::TrainingError(format!("Failed to create delta tensor: {}", e)) @@ -2681,16 +2708,18 @@ impl DQN { .map_err(|e| { MLError::TrainingError(format!("Failed to cast linear term tensor: {}", e)) })?; - let linear_loss = (linear_loss_term1 - &linear_loss_term2_tensor)?; // delta * (|x| - 0.5*delta) + let linear_loss = (linear_loss_term1 - &linear_loss_term2_tensor)?; // Condition: use squared if |x| <= delta, else linear - let mask = abs_diff.le(delta)?.to_dtype(dtype)?; // 1.0 if |x| <= delta, 0.0 otherwise + let mask = abs_diff.le(delta)?.to_dtype(dtype)?; let one_minus_mask = (Tensor::ones(mask.shape(), dtype, device)? - &mask)?; let huber_loss = ((&squared_loss * &mask)? + (&linear_loss * &one_minus_mask)?)?; - huber_loss.mean_all()? + + // NOW apply IS weights to the per-sample Huber losses + (huber_loss * weights_tensor)?.mean_all()? } else { - // MSE fallback (use weighted_diff) - (&weighted_diff * &weighted_diff)?.mean_all()? + // MSE: L_i = diff_i^2, loss = mean(w_i * L_i) + ((&diff * &diff)? * weights_tensor)?.mean_all()? } }; @@ -2740,13 +2769,12 @@ impl DQN { loss_with_entropy }; - // Extract scalar loss value (forces GPU→CPU pipeline flush) - // Cast to F32 for scalar extraction — loss_tensor may be BF16 on Ampere+ GPUs - let loss_value = loss_tensor + // Defer loss scalar readback: keep tensor on GPU to avoid premature pipeline flush. + // to_scalar() is called AFTER backward pass in train_step() / compute_gradients(), + // piggybacking on the grad norm flush (zero additional roundtrip). + let loss_f32_tensor = loss_tensor .to_dtype(DType::F32) - .map_err(|e| MLError::TrainingError(format!("Failed to cast loss to F32: {}", e)))? - .to_scalar::() - .map_err(|e| MLError::TrainingError(format!("Failed to extract loss: {}", e)))?; + .map_err(|e| MLError::TrainingError(format!("Failed to cast loss to F32: {}", e)))?; // BUG #41 FIX: Detach diff for PER priority updates (values only, no gradients). // GPU SATURATION: When GpuPrioritized replay is active, keep TD errors on GPU @@ -2792,7 +2820,7 @@ impl DQN { Ok(ComputeLossResult { loss_tensor, - loss_value, + loss_f32_tensor: loss_f32_tensor, td_errors: td_errors_vec, indices, states_tensor, @@ -2835,8 +2863,21 @@ impl DQN { self.training_steps += 1; // Update priorities for PER (if using prioritized replay) + #[cfg(feature = "cuda")] + { + if let (Some(ref td_gpu), Some(ref idx_gpu)) = + (&result.td_errors_gpu, &result.indices_gpu) + { + self.memory.update_priorities_gpu(idx_gpu, td_gpu)?; + } else if !result.indices.is_empty() { + self.memory + .update_priorities(&result.indices, &result.td_errors)?; + } + } + #[cfg(not(feature = "cuda"))] if !result.indices.is_empty() { - self.memory.update_priorities(&result.indices, &result.td_errors)?; + self.memory + .update_priorities(&result.indices, &result.td_errors)?; } // Step beta annealing for PER @@ -2854,7 +2895,13 @@ impl DQN { // Update target network with cosine-annealed EMA or hard updates self.update_target_networks()?; - Ok((result.loss_value, grad_norm)) + // Deferred loss readback: AFTER backward pass, the GPU pipeline is already + // flushed by grad_norm readback — this to_scalar() is near-free. + let loss_value = result.loss_f32_tensor + .to_scalar::() + .map_err(|e| MLError::TrainingError(format!("Deferred loss readback: {}", e)))?; + + Ok((loss_value, grad_norm)) } /// Compute clipped gradients for a batch WITHOUT applying an optimizer step. @@ -2888,8 +2935,13 @@ impl DQN { )); }; + // Deferred loss readback: AFTER backward_and_clip, GPU pipeline is flushed. + let loss_value = result.loss_f32_tensor + .to_scalar::() + .map_err(|e| MLError::TrainingError(format!("Deferred loss readback: {}", e)))?; + Ok(GradientResult { - loss: result.loss_value, + loss: loss_value, grad_norm: grad_norm as f32, grads, td_errors: result.td_errors, diff --git a/crates/ml-dqn/src/quantile_regression.rs b/crates/ml-dqn/src/quantile_regression.rs index 852e6e1eb..1e1762b0c 100644 --- a/crates/ml-dqn/src/quantile_regression.rs +++ b/crates/ml-dqn/src/quantile_regression.rs @@ -295,7 +295,7 @@ impl QuantileNetwork { // QR-DQN has fixed uniform taus so quantiles are already ordered, but IQN // uses random taus and the quantile outputs may not be monotonic. Sorting // guarantees correct CVaR (mean of the worst-alpha fraction) in both cases. - let sorted_quantiles = quantiles.sort_last_dim(true)?.0; // ascending + let sorted_quantiles = quantiles.contiguous()?.sort_last_dim(true)?.0; // ascending let tail = sorted_quantiles.narrow(2, 0, num_tail)?; tail.mean(2) } diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 19f9562fc..f369d3bad 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -157,6 +157,9 @@ impl ExperienceCollectorConfig { pub struct ExperienceBatch { /// Flattened state vectors `[N * L * STATE_DIM]` pub states: Vec, + /// Flattened next-state vectors `[N * L * STATE_DIM]` + /// `next_states[t] = states[t+1]` within each episode; terminal uses current state. + pub next_states: Vec, /// Selected action indices `[N * L]` pub actions: Vec, /// Combined rewards `[N * L]` @@ -687,9 +690,38 @@ impl GpuExperienceCollector { "Experience collection complete" ); - // ---- Step 7: Return batch ---- + // ---- Step 7: Build next_states from states (episode-aware shift) ---- + // next_states[t] = states[t+1] within same episode; terminal/last uses current state. + let sd = self.state_dim; + let mut next_states = Vec::with_capacity(total * sd); + for ep in 0..n_episodes { + for t in 0..timesteps { + let idx = ep * timesteps + t; + let is_done = done_flags.get(idx).copied().unwrap_or(0) != 0; + if !is_done && t + 1 < timesteps { + let next_idx = ep * timesteps + t + 1; + let ns_start = next_idx * sd; + let ns_end = ns_start + sd; + if let Some(ns) = states.get(ns_start..ns_end) { + next_states.extend_from_slice(ns); + } else { + // Out of bounds — copy current state as fallback + let s_start = idx * sd; + next_states.extend_from_slice(&states[s_start..s_start + sd]); + } + } else { + // Terminal or last timestep — next_state = current state + // (masked by done flag in Bellman equation) + let s_start = idx * sd; + next_states.extend_from_slice(&states[s_start..s_start + sd]); + } + } + } + + // ---- Step 8: Return batch ---- Ok(ExperienceBatch { states, + next_states, actions, rewards, done_flags, diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index f4a997d0e..083907a9b 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -281,6 +281,65 @@ impl DQNAgentType { } } + /// Insert a batch of experience tensors directly into the GPU replay buffer. + /// + /// When the replay buffer is `GpuPrioritized`, this feeds tensors straight to + /// `GpuReplayBuffer::insert_batch()` — zero CPU Experience intermediaries. + /// For CPU buffers, converts tensors back to `Vec` (fallback path). + #[cfg(feature = "cuda")] + pub fn insert_batch_tensors( + &self, + states: &candle_core::Tensor, + next_states: &candle_core::Tensor, + actions: &candle_core::Tensor, + rewards: &candle_core::Tensor, + dones: &candle_core::Tensor, + ) -> Result<(), MLError> { + match self { + Self::Standard(agent) => { + if let Some(mut gpu_buf) = agent.memory.as_gpu_buffer() { + gpu_buf.insert_batch(states, next_states, actions, rewards, dones) + } else { + // CPU PER fallback: convert tensors back to experiences + let batch_size = states.dim(0).map_err(|e| { + MLError::TrainingError(format!("states dim: {e}")) + })?; + let state_dim = states.dim(1).map_err(|e| { + MLError::TrainingError(format!("states dim1: {e}")) + })?; + let s_flat: Vec = states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("states to_vec1: {e}")))?; + let ns_flat: Vec = next_states.to_dtype(candle_core::DType::F32)?.flatten_all()?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("next_states to_vec1: {e}")))?; + let a_vec: Vec = actions.to_dtype(candle_core::DType::U32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("actions to_vec1: {e}")))?; + let r_vec: Vec = rewards.to_dtype(candle_core::DType::F32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("rewards to_vec1: {e}")))?; + let d_vec: Vec = dones.to_dtype(candle_core::DType::F32)?.to_vec1() + .map_err(|e| MLError::TrainingError(format!("dones to_vec1: {e}")))?; + let mut experiences = Vec::with_capacity(batch_size); + for i in 0..batch_size { + let start = i * state_dim; + let end = start + state_dim; + experiences.push(crate::dqn::Experience { + state: s_flat.get(start..end).unwrap_or(&[]).to_vec(), + action: a_vec.get(i).copied().unwrap_or(0) as u8, + reward: (r_vec.get(i).copied().unwrap_or(0.0) * 1_000_000.0) as i32, + next_state: ns_flat.get(start..end).unwrap_or(&[]).to_vec(), + done: d_vec.get(i).map_or(false, |&v| v > 0.5), + timestamp: 0, // fallback path — timestamp not meaningful for replay + }); + } + agent.memory.add_batch(experiences) + } + } + Self::RegimeConditional(_) => { + // Regime-conditional agent doesn't support GPU PER (each head has its own buffer) + Ok(()) + } + } + } + /// Check if agent can train (has enough replay buffer samples) pub fn can_train(&self) -> bool { match self { diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 62876240f..1102a773d 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -209,6 +209,11 @@ pub struct DQNTrainer { /// Current effective batch size (may be reduced by OOM recovery) current_batch_size: usize, + /// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step) + cached_avg_q: f64, + /// Counter for Q-value estimation frequency (estimate every N training steps) + q_estimation_counter: u64, + /// Pre-uploaded GPU training data (set once, reused across epochs) gpu_data: Option, @@ -864,6 +869,10 @@ impl DQNTrainer { // OOM recovery: track effective batch size current_batch_size: initial_batch_size, + // Q-value estimation: periodic (every 50 steps) instead of every step + cached_avg_q: 0.0, + q_estimation_counter: 0, + // GPU pipeline: pre-uploaded training data (initialized lazily at first epoch) gpu_data: None, #[cfg(feature = "cuda")] @@ -1926,11 +1935,46 @@ impl DQNTrainer { // (56 with OFI, 48 without) — must match kernel_dims.0. let raw_dim = if self.hyperparams.mbp10_data_dir.is_some() { 51 } else { 43 }; let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - let experiences = gpu_batch_to_experiences(&batch, aligned_dim); - let count = experiences.len(); + + // GPU PER fast path: create tensors from ExperienceBatch and + // insert directly into GpuReplayBuffer — no CPU Experience structs. + let count = batch.states.len() / aligned_dim; info!("GPU collected {} experiences ({} episodes × {} timesteps)", count, batch.n_episodes, batch.timesteps); - self.store_experiences_batch(experiences).await?; + { + let agent = self.agent.read().await; + if agent.memory().is_gpu_prioritized() && count > 0 { + // Build tensors on GPU directly from ExperienceBatch vecs. + // ExperienceBatch provides both states and next_states + // (episode-aware shift computed at collection time). + let states_t = candle_core::Tensor::from_vec( + batch.states.clone(), (count, aligned_dim), &self.device + ).map_err(|e| anyhow::anyhow!("GPU PER states tensor: {e}"))?; + let next_states_t = candle_core::Tensor::from_vec( + batch.next_states.clone(), (count, aligned_dim), &self.device + ).map_err(|e| anyhow::anyhow!("GPU PER next_states tensor: {e}"))?; + let actions_t = candle_core::Tensor::from_vec( + batch.actions.iter().map(|&a| a as u32).collect::>(), + count, &self.device + ).map_err(|e| anyhow::anyhow!("GPU PER actions tensor: {e}"))?; + let rewards_t = candle_core::Tensor::from_vec( + batch.rewards.clone(), count, &self.device + ).map_err(|e| anyhow::anyhow!("GPU PER rewards tensor: {e}"))?; + let dones_t = candle_core::Tensor::from_vec( + batch.done_flags.iter().map(|&d| d as f32).collect::>(), + count, &self.device + ).map_err(|e| anyhow::anyhow!("GPU PER dones tensor: {e}"))?; + + agent.insert_batch_tensors( + &states_t, &next_states_t, &actions_t, &rewards_t, &dones_t + ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; + } else { + // CPU PER fallback + drop(agent); // release read lock before store_experiences_batch + let experiences = gpu_batch_to_experiences(&batch, aligned_dim); + self.store_experiences_batch(experiences).await?; + } + } true } Err(e) => { @@ -3883,8 +3927,13 @@ impl DQNTrainer { loss_f32 }; - // Get Q-values from a sample state for monitoring (also performs Q-value divergence check) - let avg_q_value = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + // Q-value estimation: periodic (every 50 steps) to avoid extra forward pass overhead. + // Q-value divergence check still runs at 50-step intervals — sufficient for early stopping. + self.q_estimation_counter += 1; + if self.q_estimation_counter % 50 == 1 { + self.cached_avg_q = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } + let avg_q_value = self.cached_avg_q; // Convert to f64 for monitoring let grad_norm = grad_norm_f32 as f64; @@ -4092,8 +4141,12 @@ impl DQNTrainer { avg_loss }; - // Get Q-values for monitoring - let avg_q_value = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + // Q-value estimation: periodic (every 50 steps) to avoid extra forward pass overhead. + self.q_estimation_counter += 1; + if self.q_estimation_counter % 50 == 1 { + self.cached_avg_q = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; + } + let avg_q_value = self.cached_avg_q; Ok((loss_clipped, avg_q_value, final_grad_norm as f64)) }