diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 993de248a..084c61060 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -297,6 +297,14 @@ pub struct ExperienceCollectorConfig { pub eps_start: f32, /// v8: Epsilon schedule end value. pub eps_end: f32, + /// v8: TD(λ) trace parameter. 0.0 = 1-step TD (n-step only), >0.01 = TD(λ) overwrites n-step. + pub td_lambda: f32, + /// v8: Maximum trace length for truncated TD(λ). + pub max_trace_length: i32, + /// v8: Fraction of replay batch relabeled with hindsight optimal exit (0.0 = disabled). + pub hindsight_fraction: f32, + /// v8: Lookahead bars for hindsight optimal exit. + pub hindsight_lookahead: i32, } impl Default for ExperienceCollectorConfig { @@ -368,6 +376,10 @@ impl Default for ExperienceCollectorConfig { total_epochs: 20, eps_start: 0.3, eps_end: 0.02, + td_lambda: 0.0, + max_trace_length: 7, + hindsight_fraction: 0.0, + hindsight_lookahead: 10, } } } @@ -1425,6 +1437,48 @@ impl GpuExperienceCollector { debug!(n_steps, gamma = config.gamma, "N-step return accumulation applied"); } + // v8: TD(λ) — overwrites n-step results with geometrically-weighted multi-horizon returns + let td_lambda = config.td_lambda; + let max_trace = config.max_trace_length; + if td_lambda > 0.01 && max_trace > 1 { + // Clone 1-step rewards/dones as raw inputs (before n-step overwrote them, use + // the current state — TD(λ) independently accumulates from the same base signals) + let raw_rewards_td = dtod_clone_f32_native(&self.stream, &self.rewards_out, total, "raw_rewards_td")?; + let raw_dones_td = dtod_clone_f32_native(&self.stream, &self.done_out, total, "raw_dones_td")?; + // Self-bootstrap: use rewards as Q(s') approximation (improves as model learns) + let q_next = dtod_clone_f32_native(&self.stream, &self.rewards_out, total, "q_next_bootstrap")?; + + let gamma_f32 = config.gamma; + let lambda_f32 = td_lambda; + let max_trace_i32 = max_trace; + let l_i32 = timesteps as i32; + let n_i32 = (n_episodes * 2) as i32; // includes counterfactual + let blocks = ((total + 255) / 256) as u32; + + unsafe { + self.stream + .launch_builder(&self.td_lambda_kernel) + .arg(&raw_rewards_td) // raw 1-step rewards + .arg(&raw_dones_td) // raw 1-step dones + .arg(&q_next) // bootstrapped Q(s') + .arg(&rewards) // output: overwrite with G^λ + .arg(&dones) // output: overwrite with done_n + .arg(&gamma_f32) + .arg(&lambda_f32) + .arg(&max_trace_i32) + .arg(&l_i32) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("td_lambda_kernel: {e}")))?; + } + + debug!(td_lambda, max_trace, "TD(λ) return accumulation applied (replaces n-step)"); + } + // Reward normalization: DISABLED for reward v4. // Per-bar percentage returns (~0.001) are naturally scaled. // Normalizing would scramble the C51 atom distribution. @@ -1445,6 +1499,64 @@ impl GpuExperienceCollector { timesteps, )?; + // v8: Hindsight relabeling — replace fraction of rewards with optimal exit PnL + if config.hindsight_fraction > 0.0 { + // Build bar_indices: bar_idx[ep * L + t] = episode_starts[ep] + t + let mut bar_indices_cpu = vec![0_i32; total]; + for ep in 0..n_episodes { + for t in 0..timesteps { + let bar = episode_starts[ep.min(episode_starts.len() - 1)] + t as i32; + // Base batch + bar_indices_cpu[ep * timesteps + t] = bar; + // Counterfactual batch (offset by base_total) + bar_indices_cpu[base_total + ep * timesteps + t] = bar; + } + } + let mut bar_indices_gpu = self.stream.alloc_zeros::(total) + .map_err(|e| MLError::ModelError(format!("alloc bar_indices: {e}")))?; + self.stream.memcpy_htod(&bar_indices_cpu, &mut bar_indices_gpu) + .map_err(|e| MLError::ModelError(format!("upload bar_indices: {e}")))?; + + let b0 = self.branch_sizes[0] as i32; + let b1 = self.branch_sizes[1] as i32; + let b2 = self.branch_sizes[2] as i32; + let max_pos = config.max_position; + let hl = config.hindsight_lookahead.max(1); + let tb = config.total_bars; + let hf = config.hindsight_fraction; + let bs = total as i32; + let blocks = ((total + 255) / 256) as u32; + + unsafe { + self.stream + .launch_builder(&self.hindsight_relabel_kernel) + .arg(targets_buf) // const bfloat16* targets + .arg(&actions) // const int* actions + .arg(&bar_indices_gpu) // const int* bar_indices + .arg(&rewards) // float* rewards + .arg(&b0) // int b0_size + .arg(&b1) // int b1_size + .arg(&b2) // int b2_size + .arg(&max_pos) // float max_position + .arg(&hl) // int lookahead + .arg(&tb) // int total_bars + .arg(&hf) // float hindsight_fraction + .arg(&bs) // int batch_size + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("hindsight_relabel_kernel: {e}")))?; + } + + debug!( + hindsight_fraction = config.hindsight_fraction, + lookahead = config.hindsight_lookahead, + "Hindsight relabeling applied" + ); + } + debug!( n_episodes, timesteps, @@ -2159,6 +2271,11 @@ impl GpuExperienceCollector { &self.best_exposure_out } + /// Mutable access to the episode starts GPU buffer (for curriculum learning). + pub fn episode_starts_buf_mut(&mut self) -> &mut CudaSlice { + &mut self.episode_starts_buf + } + /// Train curiosity forward model directly on GPU using experience data. pub fn train_curiosity_gpu( &mut self, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 05f52fcd7..dd3de65b1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -761,15 +761,9 @@ impl FusedTrainingCtx { } } - // ── Step 2b: v8 stubs (hindsight + curriculum — disabled by default) ── - // v8: Hindsight relabeling (wired but requires bar_indices in PER batch — future work) - if self.hindsight_fraction > 0.0 { - tracing::debug!("Hindsight relabeling: kernel loaded, awaiting bar_indices in PER batch"); - } - // v8: Curriculum learning (wired but requires sorted episode starts — future work) - if self.curriculum_enabled { - tracing::debug!("Curriculum learning: kernel loaded, awaiting sorted episode starts"); - } + // ── Step 2b: v8 hindsight + curriculum are now wired in the experience collector ── + // Hindsight relabeling runs in collect_experiences_gpu (before PER insert). + // Curriculum learning runs in the training loop (restricts episode starts per epoch). // ── Step 2c: HER donor computation (outside graph_aux) ─────────── // Donor indices vary per step (random/future/final). The computation diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index e470983cf..a56cd4ecf 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -243,6 +243,34 @@ impl DQNTrainer { } } + // v8: Curriculum learning — compute difficulty scores and sort episode starts + let curriculum_sorted_starts: Option> = if self.hyperparams.curriculum_enabled { + if let (Some(ref collector), Some(ref targets_cuda)) = (&self.gpu_experience_collector, &self.targets_raw_cuda) { + let total_bars = training_data.len(); + match collector.compute_difficulty_scores(targets_cuda, total_bars) { + Ok(scores_gpu) => { + let mut scores_cpu = vec![0.0_f32; total_bars]; + if let Err(e) = collector.stream().memcpy_dtoh(&scores_gpu, &mut scores_cpu) { + warn!("Curriculum DtoH failed: {e}"); + None + } else { + let mut indices: Vec = (0..total_bars as i32).collect(); + indices.sort_by(|&a, &b| { + scores_cpu[a as usize].partial_cmp(&scores_cpu[b as usize]) + .unwrap_or(std::cmp::Ordering::Equal) + }); + info!("v8: Curriculum learning initialized — {} bars sorted by difficulty", total_bars); + Some(indices) + } + } + Err(e) => { + warn!("Curriculum scoring failed (using default order): {e}"); + None + } + } + } else { None } + } else { None }; + // Training loop for epoch in 0..self.hyperparams.epochs { self.current_epoch = epoch; @@ -295,6 +323,39 @@ impl DQNTrainer { self.init_gpu_experience_collector().await?; let phase1_ms = phase1_start.elapsed().as_secs_f64() * 1000.0; + // v8: Curriculum — restrict episode starts to easy bars early, expand over training + if let Some(ref sorted) = curriculum_sorted_starts { + let warmup_frac = self.hyperparams.curriculum_warmup_fraction as f32; + let total_epochs = self.hyperparams.epochs.max(1) as f32; + let epoch_frac = (epoch as f32 / total_epochs).min(1.0); + let usable_frac = if warmup_frac > 0.0 { + (warmup_frac + (1.0 - warmup_frac) * (epoch_frac / warmup_frac).min(1.0)).min(1.0) + } else { + 1.0 + }; + let usable_count = ((usable_frac * sorted.len() as f32) as usize).max(100); + let restricted = &sorted[..usable_count.min(sorted.len())]; + + if let Some(ref mut collector) = self.gpu_experience_collector { + // Subsample restricted sorted indices as episode starts + let alloc_cap = collector.alloc_episodes(); + let n_episodes = alloc_cap.min(usable_count); + let stride = (usable_count / n_episodes).max(1); + let episode_starts: Vec = (0..n_episodes) + .map(|i| restricted[(i * stride).min(restricted.len() - 1)]) + .collect(); + if let Some(ref stream) = self.cuda_stream { + let _ = stream.memcpy_htod(&episode_starts, collector.episode_starts_buf_mut()); + } + if epoch % 5 == 0 { + debug!( + epoch, usable_frac, usable_count, + n_episodes, "Curriculum: restricting episode starts" + ); + } + } + } + // Causal feature masking { let frac = self.hyperparams.feature_mask_fraction; @@ -357,13 +418,8 @@ impl DQNTrainer { } } - // v8: TD(lambda) returns — kernel loaded, awaiting value estimates in collector - if self.hyperparams.td_lambda > 0.01 { - tracing::debug!( - td_lambda = self.hyperparams.td_lambda, - "TD(lambda): kernel loaded, awaiting V(s) estimates for full wiring" - ); - } + // v8: TD(λ) is now wired in collect_experiences_gpu — self-bootstrap with + // rewards as Q(s') approximation, overwrites n-step when td_lambda > 0.01. // Periodic shrink-and-perturb let sp_interval = self.hyperparams.shrink_perturb_interval; @@ -1105,6 +1161,10 @@ impl DQNTrainer { total_epochs: self.hyperparams.epochs as i32, eps_start: self.hyperparams.epsilon_start as f32, eps_end: self.hyperparams.epsilon_end as f32, + td_lambda: self.hyperparams.td_lambda as f32, + max_trace_length: self.hyperparams.max_trace_length as i32, + hindsight_fraction: self.hyperparams.hindsight_fraction as f32, + hindsight_lookahead: self.hyperparams.hindsight_lookahead as i32, ..Default::default() };