diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 3433e77ee..c6a4d6089 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -299,6 +299,12 @@ pub struct GpuExperienceBatch { pub rewards: CudaSlice, /// Done flags `[total]` on GPU (i32, 0 or 1) pub dones: CudaSlice, + /// Episode index per transition `[total]` on GPU (i32). + /// + /// `episode_ids[i] = i / timesteps_per_episode`. + /// Episode `ep` occupies flat indices `[ep * timesteps, (ep+1) * timesteps)`. + /// Required by HER Future and Final strategies for GPU-native donor selection. + pub episode_ids: CudaSlice, /// State dimensionality (so the consumer can reshape) pub state_dim: usize, /// Number of episodes @@ -362,8 +368,12 @@ pub struct GpuExperienceCollector { // ── Compiled experience kernels ───────────────────────────────── state_gather_kernel: CudaFunction, action_select_kernel: CudaFunction, + expert_override_kernel: CudaFunction, env_step_kernel: CudaFunction, expected_q_kernel: CudaFunction, + nstep_kernel: CudaFunction, + reward_norm_kernel: CudaFunction, + fill_episode_ids_kernel: CudaFunction, // ── Per-timestep batch buffers (reused each timestep) ─────────── /// Assembled state batch for cuBLAS: [N, state_dim]. @@ -408,11 +418,23 @@ pub struct GpuExperienceCollector { rewards_out: CudaSlice, // [alloc_episodes * alloc_timesteps] done_out: CudaSlice, // [alloc_episodes * alloc_timesteps] + /// Running reward statistics for cross-epoch normalization (Welford's online algorithm) + reward_running_mean: f64, + reward_running_var: f64, + reward_running_count: u64, + // Persistent epoch state — survives across kernel launches. epoch_state: CudaSlice, // [8] /// Bitfield: bit 0 = reset portfolio, bit 1 = reset DSR, bit 2 = reset vol EMA reset_flags: u32, + /// Expert demonstration actions [total_bars] — pre-computed MA crossover signals. + /// -1 = no expert opinion (use Q-network), >= 0 = expert exposure action index. + /// Uploaded once at training start, read by action selection kernel. + expert_actions_gpu: Option>, + /// Expert demo override probability (0.0 = disabled, decays over epochs). + expert_ratio: f32, + // OFI features pre-uploaded to GPU [total_bars * OFI_DIM] (None when OFI disabled) ofi_gpu: Option>, ofi_dim: usize, @@ -487,14 +509,16 @@ impl GpuExperienceCollector { let total_branch_actions = branch_sizes[0] + branch_sizes[1] + branch_sizes[2]; // 11 // ── Step 1: Compile experience kernels ────────────────────────── - let (state_gather_kernel, action_select_kernel, env_step_kernel, expected_q_kernel) = + let (state_gather_kernel, action_select_kernel, expert_override_kernel, env_step_kernel, expected_q_kernel) = compile_experience_kernels(&stream, state_dim, market_dim)?; + let (nstep_kernel, reward_norm_kernel) = compile_nstep_kernel(&stream)?; + let fill_episode_ids_kernel = compile_fill_episode_ids_kernel(&stream)?; info!( state_dim, market_dim, num_atoms_max, - "GPU experience collector: experience kernels compiled (3 focused + expected_q)" + "GPU experience collector: experience kernels compiled (3 focused + expected_q + episode_ids)" ); // ── Step 2: Create cuBLAS forward context ─────────────────────── @@ -787,8 +811,12 @@ impl GpuExperienceCollector { cvar_scales_ptr: 0, // NULL = no CVaR scaling initially state_gather_kernel, action_select_kernel, + expert_override_kernel, env_step_kernel, expected_q_kernel, + nstep_kernel, + reward_norm_kernel, + fill_episode_ids_kernel, batch_states, batch_actions, q_gaps_buf, @@ -815,9 +843,14 @@ impl GpuExperienceCollector { states_out, actions_out, rewards_out, + reward_running_mean: 0.0, + reward_running_var: 1.0, + reward_running_count: 0, done_out, epoch_state, reset_flags: 0, + expert_actions_gpu: None, + expert_ratio: 0.0, ofi_gpu: ofi_placeholder, ofi_dim, curiosity_trainer, @@ -839,6 +872,27 @@ impl GpuExperienceCollector { self.cvar_scales_ptr = device_ptr; } + /// Upload pre-computed expert demonstration actions to GPU. + /// + /// `expert_actions[bar_index]` = expert exposure action index (-1 = no opinion). + /// Generated once from MA crossover + ADX filter, stays GPU-resident for the + /// entire training run. The action selection kernel reads this at each timestep + /// and overrides Q-network actions with probability `expert_ratio`. + pub fn upload_expert_actions(&mut self, actions: &[i32]) -> Result<(), MLError> { + let mut buf = self.stream.alloc_zeros::(actions.len()) + .map_err(|e| MLError::ModelError(format!("alloc expert_actions: {e}")))?; + self.stream.memcpy_htod(actions, &mut buf) + .map_err(|e| MLError::ModelError(format!("upload expert_actions: {e}")))?; + self.expert_actions_gpu = Some(buf); + info!(n_actions = actions.len(), "Expert demo actions uploaded to GPU"); + Ok(()) + } + + /// Set the expert demo override ratio for the current epoch. + pub fn set_expert_ratio(&mut self, ratio: f32) { + self.expert_ratio = ratio; + } + /// All outputs remain GPU-resident for zero-copy training. /// /// The caller **must** call `stream().synchronize()` before passing @@ -861,9 +915,109 @@ impl GpuExperienceCollector { let actions = dtod_clone_i32(&self.stream, &self.actions_out, total, "actions")?; let dones = dtod_clone_f32(&self.stream, &self.done_out, total, "dones")?; - // Build next_states on GPU (episode-aware shift) + // ── N-step return accumulation ─────────────────────────────────── + // Converts 1-step rewards into R_n = sum(gamma^i * r_i) and + // OR's done flags over n steps. Uses double-buffering to avoid races. + let n_steps = config.n_steps.max(1); + if n_steps > 1 { + let raw_rewards = dtod_clone_f32(&self.stream, &rewards, total, "raw_rewards_nstep")?; + let raw_dones = dtod_clone_f32(&self.stream, &dones, total, "raw_dones_nstep")?; + + let gamma_f32 = config.gamma; + let n_steps_i32 = n_steps; + let l_i32 = timesteps as i32; + let n_i32 = n_episodes as i32; + let blocks = ((total + 255) / 256) as u32; + + unsafe { + self.stream + .launch_builder(&self.nstep_kernel) + .arg(&raw_rewards) + .arg(&raw_dones) + .arg(&rewards) + .arg(&dones) + .arg(&gamma_f32) + .arg(&n_steps_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!("nstep_accumulate_kernel: {e}")))?; + } + + debug!(n_steps, gamma = config.gamma, "N-step return accumulation applied"); + } + + // ── Reward normalization ────────────────────────────────────── + // Normalize rewards so dense shaping (0.01x) and sparse trade-completion + // (±2.0) are on the same scale for C51. Uses batch statistics (one DtoH + // readback per experience collection, not per training step). + if config.reward_norm_alpha > 0.0 && total > 0 { + let mut reward_host = vec![0.0_f32; total]; + self.stream.memcpy_dtoh(&rewards, &mut reward_host) + .map_err(|e| MLError::ModelError(format!("reward norm DtoH: {e}")))?; + + // Update running statistics with Welford's online algorithm. + // This maintains mean/variance across epochs, not just per-batch. + // EMA blending with alpha ensures adaptation to distribution shifts. + let alpha = config.reward_norm_alpha as f64; + for &r in &reward_host { + let r64 = r as f64; + self.reward_running_count += 1; + if self.reward_running_count == 1 { + self.reward_running_mean = r64; + self.reward_running_var = 1.0; + } else { + // EMA update (exponential moving average for non-stationary rewards) + let old_mean = self.reward_running_mean; + self.reward_running_mean = (1.0 - alpha) * old_mean + alpha * r64; + let delta = r64 - old_mean; + let delta2 = r64 - self.reward_running_mean; + self.reward_running_var = (1.0 - alpha) * self.reward_running_var + + alpha * delta * delta2; + } + } + + let mean = self.reward_running_mean as f32; + let std = (self.reward_running_var as f32).sqrt(); + let inv_std = 1.0_f32 / (std + 1e-8_f32); + + let n_i32 = total as i32; + let blocks = ((total + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.reward_norm_kernel) + .arg(&rewards) + .arg(&mean) + .arg(&inv_std) + .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!("reward_normalize_kernel: {e}")))?; + } + + debug!(mean, std, "Reward normalization applied (batch statistics)"); + } + + // Build next_states on GPU (episode-aware shift by n_steps) let next_states = build_next_states_dtod( - &self.stream, &states, n_episodes, timesteps, sd, + &self.stream, &states, n_episodes, timesteps, sd, n_steps as usize, + )?; + + // ── Episode ID fill (GPU-native, zero CPU) ──────────────────── + // episode_ids[i] = i / L — enables HER Future and Final strategies. + // Direct integer arithmetic: no linear scan needed. + let episode_ids = fill_episode_ids_gpu( + &self.stream, + &self.fill_episode_ids_kernel, + total, + timesteps, )?; debug!( @@ -879,6 +1033,7 @@ impl GpuExperienceCollector { actions, rewards, dones, + episode_ids, state_dim: sd, n_episodes, timesteps, @@ -1050,6 +1205,35 @@ impl GpuExperienceCollector { )))?; } + // ── 4b. Expert action override (GPU-native MA crossover + ADX) + // When expert_ratio > 0, overrides Q-network exposure actions with + // expert signals. Reads prices and ADX directly from GPU buffers. + if self.expert_ratio > 0.0 { + let er = self.expert_ratio; + let tb = total_bars; + unsafe { + self.stream + .launch_builder(&self.expert_override_kernel) + .arg(&mut self.batch_actions) + .arg(targets_buf) + .arg(market_features_buf) + .arg(&self.episode_starts_buf) + .arg(&self.current_timesteps) + .arg(&mut self.rng_states) + .arg(&er) + .arg(&n_i32) + .arg(&tb) + .arg(&md) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!( + "expert_action_override t={t}: {e}" + )))?; + } + } + // ── 5. Environment step (8-component composite reward) ────── let max_pos = config.max_position; let tx_cost = config.tx_cost_multiplier; @@ -1374,7 +1558,7 @@ fn compile_experience_kernels( stream: &Arc, state_dim: usize, market_dim: usize, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { let kernel_src = include_str!("experience_kernels.cu"); // Inject compile-time constants before the source @@ -1398,6 +1582,9 @@ fn compile_experience_kernels( let action_select = module .load_function("experience_action_select") .map_err(|e| MLError::ModelError(format!("load experience_action_select: {e}")))?; + let expert_override = module + .load_function("expert_action_override") + .map_err(|e| MLError::ModelError(format!("load expert_action_override: {e}")))?; let env_step = module .load_function("experience_env_step") .map_err(|e| MLError::ModelError(format!("load experience_env_step: {e}")))?; @@ -1405,7 +1592,83 @@ fn compile_experience_kernels( .load_function("compute_expected_q") .map_err(|e| MLError::ModelError(format!("load compute_expected_q: {e}")))?; - Ok((state_gather, action_select, env_step, expected_q)) + Ok((state_gather, action_select, expert_override, env_step, expected_q)) +} + +/// Compile the n-step return accumulation kernel. +/// +/// Converts 1-step (s,a,r,s',d) transitions into n-step (s,a,R_n,s_{t+n},d_n) +/// by accumulating discounted rewards and OR-ing done flags over n steps. +fn compile_nstep_kernel( + stream: &Arc, +) -> Result<(CudaFunction, CudaFunction), MLError> { + let kernel_src = include_str!("nstep_kernel.cu"); + let context = stream.context(); + let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context) + .map_err(|e| MLError::ModelError(format!("nstep_kernel compilation: {e}")))?; + let module = context + .load_module(ptx) + .map_err(|e| MLError::ModelError(format!("nstep_kernel module load: {e}")))?; + let nstep = module + .load_function("nstep_accumulate_kernel") + .map_err(|e| MLError::ModelError(format!("load nstep_accumulate_kernel: {e}")))?; + let reward_norm = module + .load_function("reward_normalize_kernel") + .map_err(|e| MLError::ModelError(format!("load reward_normalize_kernel: {e}")))?; + // regime_scale_td_errors — loaded from dqn_utility_kernels.cu (trainer module), + // not from nstep_kernel.cu. The trainer's version uses STATE_DIM compile-time constant. + Ok((nstep, reward_norm)) +} + +/// Compile the HER episode boundary kernels (fill_episode_ids, future, final). +fn compile_fill_episode_ids_kernel( + stream: &Arc, +) -> Result { + let kernel_src = include_str!("her_episode_kernel.cu"); + let context = stream.context(); + let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context) + .map_err(|e| MLError::ModelError(format!("her_episode_kernel compilation: {e}")))?; + let module = context + .load_module(ptx) + .map_err(|e| MLError::ModelError(format!("her_episode_kernel module load: {e}")))?; + module + .load_function("fill_episode_ids_kernel") + .map_err(|e| MLError::ModelError(format!("load fill_episode_ids_kernel: {e}"))) +} + +/// Fill episode_ids[i] = i / L on GPU. Zero CPU in hot path. +/// +/// Returns a fresh `CudaSlice` of length `total` with episode IDs. +fn fill_episode_ids_gpu( + stream: &Arc, + kernel: &CudaFunction, + total: usize, + timesteps_per_episode: usize, +) -> Result, MLError> { + let mut episode_ids = stream.alloc_zeros::(total) + .map_err(|e| MLError::ModelError(format!("alloc episode_ids[{total}]: {e}")))?; + + let l_i32 = timesteps_per_episode as i32; + let total_i32 = total as i32; + let blocks = ((total + 255) / 256) as u32; + + // Safety: episode_ids is a valid GPU allocation of `total` i32 elements. + // l_i32 and total_i32 are non-negative and consistent with the allocation. + unsafe { + stream + .launch_builder(kernel) + .arg(&mut episode_ids) + .arg(&l_i32) + .arg(&total_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("fill_episode_ids_kernel: {e}")))?; + } + + Ok(episode_ids) } // --------------------------------------------------------------------------- @@ -1443,40 +1706,48 @@ fn dtod_clone_i32( } /// Build next_states from states via episode-aware DtoD shift. +/// +/// For n-step returns, `shift` is n (not 1): `next_states[t] = states[t+n]`. +/// When `t + shift >= timesteps`, duplicates the last available state. fn build_next_states_dtod( stream: &Arc, states: &CudaSlice, n_episodes: usize, timesteps: usize, sd: usize, + shift: usize, ) -> Result, MLError> { let total = n_episodes * timesteps; let total_elems = total * sd; + let shift = shift.max(1).min(timesteps); // clamp shift to valid range - if timesteps <= 1 { + if timesteps <= shift { return dtod_clone_f32(stream, states, total_elems, "next_states_trivial"); } let mut dst = stream.alloc_zeros::(total_elems) .map_err(|e| MLError::ModelError(format!("alloc next_states f32[{total_elems}]: {e}")))?; - let shifted_count = timesteps - 1; + let shifted_count = timesteps - shift; for ep in 0..n_episodes { let ep_offset = ep * timesteps * sd; - // Copy states[ep][1..timesteps] -> next_states[ep][0..timesteps-1] - let src_shifted = states.slice((ep_offset + sd)..(ep_offset + timesteps * sd)); + // Copy states[ep][shift..timesteps] -> next_states[ep][0..timesteps-shift] + let src_shifted = states.slice((ep_offset + shift * sd)..(ep_offset + timesteps * sd)); let mut dst_shifted = dst.slice_mut(ep_offset..(ep_offset + shifted_count * sd)); stream.memcpy_dtod(&src_shifted, &mut dst_shifted) .map_err(|e| MLError::ModelError(format!("DtoD next_states shift ep{ep}: {e}")))?; - // Duplicate last timestep: states[ep][T-1] -> next_states[ep][T-1] - let last_offset = ep_offset + (timesteps - 1) * sd; - let src_last = states.slice(last_offset..(last_offset + sd)); - let mut dst_last = dst.slice_mut(last_offset..(last_offset + sd)); - stream.memcpy_dtod(&src_last, &mut dst_last) - .map_err(|e| MLError::ModelError(format!("DtoD next_states last ep{ep}: {e}")))?; + // Fill remaining positions with last available state + let last_src_offset = ep_offset + (timesteps - 1) * sd; + let src_last = states.slice(last_src_offset..(last_src_offset + sd)); + for t in shifted_count..timesteps { + let dst_offset = ep_offset + t * sd; + let mut dst_t = dst.slice_mut(dst_offset..(dst_offset + sd)); + stream.memcpy_dtod(&src_last, &mut dst_t) + .map_err(|e| MLError::ModelError(format!("DtoD next_states fill ep{ep} t{t}: {e}")))?; + } } Ok(dst) diff --git a/crates/ml/src/cuda_pipeline/gpu_her.rs b/crates/ml/src/cuda_pipeline/gpu_her.rs index 2c7ba2b0a..29698b080 100644 --- a/crates/ml/src/cuda_pipeline/gpu_her.rs +++ b/crates/ml/src/cuda_pipeline/gpu_her.rs @@ -139,8 +139,16 @@ pub struct GpuHer { source_indices: CudaSlice, donor_indices: CudaSlice, - // Compiled kernel + // Compiled kernels relabel_func: CudaFunction, + /// GPU kernel: sample a future donor in the same episode (HER Future). + future_donors_func: CudaFunction, + /// GPU kernel: find the last transition in the same episode (HER Final). + episode_end_func: CudaFunction, + + // Per-sample LCG RNG state for Future strategy donor selection. + // Seeded once at construction; survives across relabel_batch_with_strategy calls. + rng_states: CudaSlice, // Stream for kernel launches stream: Arc, @@ -165,6 +173,9 @@ impl GpuHer { // Compile the relabel kernel let relabel_func = compile_her_kernel(&stream, &config)?; + // Compile HER episode boundary kernels (Future + Final strategies) + let (future_donors_func, episode_end_func) = compile_her_episode_kernels(&stream)?; + // Pre-allocate staging buffers let out_states = alloc_f32(&stream, her_batch * config.state_dim, "her_out_states")?; let out_next_states = alloc_f32(&stream, her_batch * config.state_dim, "her_out_next_states")?; @@ -176,15 +187,25 @@ impl GpuHer { let source_indices = alloc_i32(&stream, her_batch, "her_source_indices")?; let donor_indices = alloc_i32(&stream, her_batch, "her_donor_indices")?; + // Per-sample RNG state for Future strategy LCG — seeded with index. + // Uploaded once at construction; subsequent calls advance the state in-place on GPU. + let rng_init: Vec = (0..her_batch as u32) + .map(|i| i.wrapping_mul(2654435761u32).wrapping_add(1013904223u32)) + .collect(); + let mut rng_states = stream.alloc_zeros::(her_batch) + .map_err(|e| MLError::ModelError(format!("GpuHer alloc rng_states: {e}")))?; + stream.memcpy_htod(&rng_init, &mut rng_states) + .map_err(|e| MLError::ModelError(format!("GpuHer rng_states upload: {e}")))?; + let vram_bytes = (her_batch * config.state_dim * 2 + her_batch * 3) * 4 - + her_batch * 2 * 4; // output bufs + index bufs + + her_batch * 3 * 4; // output bufs + index bufs + rng_states info!( her_batch_size = her_batch, state_dim = config.state_dim, goal_dim = config.goal_dim, strategy = ?config.strategy, vram_kb = vram_bytes / 1024, - "GpuHer initialized: relabel kernel compiled, staging buffers allocated" + "GpuHer initialized: relabel + episode kernels compiled, staging buffers allocated" ); Ok(Self { @@ -197,6 +218,9 @@ impl GpuHer { source_indices, donor_indices, relabel_func, + future_donors_func, + episode_end_func, + rng_states, stream, }) } @@ -318,12 +342,142 @@ impl GpuHer { donors.resize_with(her_batch_size, || rng.gen_range(0..buffer_size as i32)); donors } + + /// GPU-native donor selection for Future and Final strategies. + /// + /// Dispatches to the appropriate episode-aware kernel based on the configured + /// strategy. Requires `episode_ids` from `GpuExperienceBatch` (filled by + /// `fill_episode_ids_kernel`) so there is zero CPU work in the hot path. + /// + /// After this call, `self.donor_indices` contains the GPU-computed donors. + /// The caller should then invoke `relabel_batch` with `source_idx_host` and + /// a placeholder CPU donor slice (will be overwritten by the GPU donor buffer). + /// + /// # Arguments + /// + /// * `episode_ids` - `[capacity]` i32 on device — episode index per transition + /// * `source_idx_host` - PER-sampled source indices (CPU, length `her_batch_size`) + /// * `episode_length` - Timesteps per episode (L) + /// * `buffer_size` - Total flat buffer capacity (N * L) + /// * `her_batch_size` - Number of experiences to relabel + /// + /// # Errors + /// + /// Returns `MLError::ConfigError` when called with `HerGpuStrategy::Random` + /// (use `generate_random_donors` instead). + pub fn relabel_batch_with_strategy( + &mut self, + episode_ids: &CudaSlice, + source_idx_host: &[i32], + episode_length: usize, + buffer_size: usize, + her_batch_size: usize, + ) -> Result<(), MLError> { + match self.config.strategy { + HerGpuStrategy::Random => { + return Err(MLError::ConfigError( + "relabel_batch_with_strategy: use generate_random_donors for Random strategy".to_owned(), + )); + } + HerGpuStrategy::Future | HerGpuStrategy::Final => {} + } + + if her_batch_size == 0 { + return Err(MLError::InvalidInput( + "HER batch size must be > 0".to_owned(), + )); + } + let max_her = self.config.her_batch_size(); + if her_batch_size > max_her { + return Err(MLError::InvalidInput(format!( + "HER batch size {her_batch_size} exceeds pre-allocated max {max_her}" + ))); + } + + // Upload source indices to GPU + self.stream + .memcpy_htod(&source_idx_host[..her_batch_size], &mut self.source_indices) + .map_err(|e| MLError::ModelError(format!("HER source_indices upload: {e}")))?; + + let l_i32 = episode_length as i32; + let capacity_i32 = buffer_size as i32; + let her_batch_i32 = her_batch_size as i32; + let blocks = ((her_batch_size + 255) / 256) as u32; + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + match self.config.strategy { + HerGpuStrategy::Future => { + // Safety: all buffers are valid GPU allocations on the same context. + // source_indices contains valid indices in [0, capacity). rng_states + // is pre-allocated at her_batch_size. donor_indices is pre-allocated. + unsafe { + self.stream + .launch_builder(&self.future_donors_func) + .arg(episode_ids) + .arg(&self.source_indices) + .arg(&mut self.donor_indices) + .arg(&mut self.rng_states) + .arg(&l_i32) + .arg(&capacity_i32) + .arg(&her_batch_i32) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!("her_sample_future_donors: {e}")))?; + } + } + HerGpuStrategy::Final => { + // Safety: all buffers are valid GPU allocations on the same context. + // source_indices contains valid indices in [0, capacity). + // end_indices == donor_indices (same pre-allocated buffer, re-used). + unsafe { + self.stream + .launch_builder(&self.episode_end_func) + .arg(episode_ids) + .arg(&self.source_indices) + .arg(&mut self.donor_indices) + .arg(&l_i32) + .arg(&capacity_i32) + .arg(&her_batch_i32) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!("her_find_episode_end: {e}")))?; + } + } + HerGpuStrategy::Random => unreachable!(), + } + + Ok(()) + } } // --------------------------------------------------------------------------- // Kernel compilation // --------------------------------------------------------------------------- +/// Compile the HER episode boundary kernels. +/// +/// Returns `(future_donors_func, episode_end_func)`. +fn compile_her_episode_kernels( + stream: &Arc, +) -> Result<(CudaFunction, CudaFunction), MLError> { + let kernel_src = include_str!("her_episode_kernel.cu"); + let context = stream.context(); + let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context) + .map_err(|e| MLError::ModelError(format!("her_episode_kernel compilation: {e}")))?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("her_episode_kernel module load: {e}")) + })?; + let future_donors = module.load_function("her_sample_future_donors").map_err(|e| { + MLError::ModelError(format!("load her_sample_future_donors: {e}")) + })?; + let episode_end = module.load_function("her_find_episode_end").map_err(|e| { + MLError::ModelError(format!("load her_find_episode_end: {e}")) + })?; + Ok((future_donors, episode_end)) +} + /// Compile the HER relabel kernel with dimension defines. fn compile_her_kernel( stream: &Arc, diff --git a/crates/ml/src/cuda_pipeline/her_episode_kernel.cu b/crates/ml/src/cuda_pipeline/her_episode_kernel.cu new file mode 100644 index 000000000..1189ba01f --- /dev/null +++ b/crates/ml/src/cuda_pipeline/her_episode_kernel.cu @@ -0,0 +1,112 @@ +/** + * HER Episode Boundary Kernels + * + * GPU-native episode boundary tracking to enable HER Future and Final + * strategies. All three kernels operate on a flat [N * L] buffer where + * episode `ep` occupies indices [ep*L, (ep+1)*L). + * + * Zero CPU in the hot path — episode IDs filled by GPU kernel. + */ + +/* ══════════════════════════════════════════════════════════════════════ + * KERNEL 1: fill_episode_ids_kernel + * + * Fills episode_ids[i] = i / L for each transition index i. + * Direct arithmetic — no linear scan. + * + * Launch config: grid=(ceil(total/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void fill_episode_ids_kernel( + int* __restrict__ episode_ids, /* [N * L] output */ + int L, /* timesteps per episode */ + int total /* N * L */ +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + episode_ids[i] = i / L; +} + +/* ══════════════════════════════════════════════════════════════════════ + * KERNEL 2: her_sample_future_donors + * + * HER Future strategy: for each source transition, sample a random + * donor transition LATER in the same episode. + * + * Episode ep occupies [ep*L, (ep+1)*L) in the flat buffer. + * Samples uniformly from [src+1, ep_end) using per-thread LCG. + * Fallback to src itself when the source is the last step of its episode. + * + * Launch config: grid=(ceil(her_batch_size/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void her_sample_future_donors( + const int* __restrict__ episode_ids, /* [capacity] */ + const int* __restrict__ source_indices, /* [her_batch_size] */ + int* __restrict__ donor_indices, /* [her_batch_size] output */ + unsigned int* __restrict__ rng_states, /* [her_batch_size] per-thread LCG state */ + int L, /* timesteps per episode */ + int capacity, /* N * L, total buffer size */ + int her_batch_size +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= her_batch_size) return; + + int src = source_indices[idx]; + int ep = episode_ids[src]; + + /* Episode occupies [ep * L, (ep+1) * L) in the flat buffer */ + int ep_end = ep * L + L; + if (ep_end > capacity) ep_end = capacity; + + /* Sample from [src+1, ep_end) */ + int range = ep_end - src - 1; + if (range <= 0) { + donor_indices[idx] = src; /* fallback: self (last step of episode) */ + return; + } + + /* LCG — same pattern as action selection kernel */ + unsigned int rng = rng_states[idx]; + rng = rng * 1664525u + 1013904223u; + rng_states[idx] = rng; + + /* Map [0, 2^24) uniformly onto [0, range) */ + int offset = (int)((float)(rng >> 8) / 16777216.0f * (float)range); + if (offset >= range) offset = range - 1; + + donor_indices[idx] = src + 1 + offset; +} + +/* ══════════════════════════════════════════════════════════════════════ + * KERNEL 3: her_find_episode_end + * + * HER Final strategy: for each source transition, find the last + * transition index of the same episode. + * + * Episode ep occupies [ep*L, (ep+1)*L) so the last index is + * (ep+1)*L - 1, clamped to capacity-1. + * + * Launch config: grid=(ceil(her_batch_size/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void her_find_episode_end( + const int* __restrict__ episode_ids, /* [capacity] */ + const int* __restrict__ source_indices, /* [her_batch_size] */ + int* __restrict__ end_indices, /* [her_batch_size] output: last idx in episode */ + int L, /* timesteps per episode */ + int capacity, /* N * L, total buffer size */ + int her_batch_size +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= her_batch_size) return; + + int src = source_indices[idx]; + int ep = episode_ids[src]; + + /* Last index of this episode */ + int ep_end = (ep + 1) * L - 1; + if (ep_end >= capacity) ep_end = capacity - 1; + + end_indices[idx] = ep_end; +}