feat: HER Future/Final strategies via GPU episode boundary tracking
Add GPU-native episode boundary tracking to enable HER Future and Final strategies, which previously could not be used (only Random worked). - Create her_episode_kernel.cu with 3 kernels: - fill_episode_ids_kernel: fills episode_ids[i] = i/L, zero CPU - her_sample_future_donors: samples random donor LATER in same episode via LCG - her_find_episode_end: finds last transition index of same episode - Add episode_ids CudaSlice<i32> field to GpuExperienceBatch; filled by fill_episode_ids_gpu() on every collect_experiences_gpu() call - Wire fill_episode_ids_kernel compilation into GpuExperienceCollector::new() - Add future_donors_func, episode_end_func, and rng_states fields to GpuHer - Add relabel_batch_with_strategy() dispatching to GPU episode kernels based on HerGpuStrategy (Future or Final); errors on Random (use existing path) - All episode ID computation uses direct integer arithmetic (i/L), no scan Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -299,6 +299,12 @@ pub struct GpuExperienceBatch {
|
||||
pub rewards: CudaSlice<f32>,
|
||||
/// Done flags `[total]` on GPU (i32, 0 or 1)
|
||||
pub dones: CudaSlice<f32>,
|
||||
/// 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<i32>,
|
||||
/// 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<f32>, // [alloc_episodes * alloc_timesteps]
|
||||
done_out: CudaSlice<f32>, // [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<f32>, // [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<CudaSlice<i32>>,
|
||||
/// 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<CudaSlice<f32>>,
|
||||
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::<i32>(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<CudaStream>,
|
||||
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<CudaStream>,
|
||||
) -> 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<CudaStream>,
|
||||
) -> Result<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}")))?;
|
||||
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<i32>` of length `total` with episode IDs.
|
||||
fn fill_episode_ids_gpu(
|
||||
stream: &Arc<CudaStream>,
|
||||
kernel: &CudaFunction,
|
||||
total: usize,
|
||||
timesteps_per_episode: usize,
|
||||
) -> Result<CudaSlice<i32>, MLError> {
|
||||
let mut episode_ids = stream.alloc_zeros::<i32>(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<CudaStream>,
|
||||
states: &CudaSlice<f32>,
|
||||
n_episodes: usize,
|
||||
timesteps: usize,
|
||||
sd: usize,
|
||||
shift: usize,
|
||||
) -> Result<CudaSlice<f32>, 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::<f32>(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)
|
||||
|
||||
@@ -139,8 +139,16 @@ pub struct GpuHer {
|
||||
source_indices: CudaSlice<i32>,
|
||||
donor_indices: CudaSlice<i32>,
|
||||
|
||||
// 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<u32>,
|
||||
|
||||
// Stream for kernel launches
|
||||
stream: Arc<CudaStream>,
|
||||
@@ -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<u32> = (0..her_batch as u32)
|
||||
.map(|i| i.wrapping_mul(2654435761u32).wrapping_add(1013904223u32))
|
||||
.collect();
|
||||
let mut rng_states = stream.alloc_zeros::<u32>(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<i32>,
|
||||
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<CudaStream>,
|
||||
) -> 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<CudaStream>,
|
||||
|
||||
112
crates/ml/src/cuda_pipeline/her_episode_kernel.cu
Normal file
112
crates/ml/src/cuda_pipeline/her_episode_kernel.cu
Normal file
@@ -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;
|
||||
}
|
||||
Reference in New Issue
Block a user