diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 026e7e416..5a1514bcd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -596,6 +596,113 @@ impl GpuIqnHead { pub fn per_sample_loss(&self) -> &CudaSlice { &self.per_sample_loss } + + /// Compute CVaR-based position scaling from IQN quantiles. + /// + /// Runs the IQN forward-only kernel on `h_s2` (trunk activation), + /// then computes CVaR at `alpha` (e.g. 0.05 = 5th percentile) for + /// each sample's selected exposure action. + /// + /// Returns a GPU buffer `[batch_size]` with scaling factors in [0.25, 1.0]: + /// - CVaR ≥ 0 (expected profit even in worst case) → scale = 1.0 + /// - CVaR < 0 (expected loss in worst case) → scale = max(0.25, 1.0 + CVaR * 5.0) + /// + /// The scale is applied to target_position in the env_step kernel. + pub fn compute_cvar_scales( + &mut self, + h_s2: &CudaSlice, + actions: &CudaSlice, + batch_size: usize, + alpha: f32, + ) -> Result, MLError> { + let b = batch_size; + let n = self.config.num_quantiles; + let tba = self.config.total_branch_actions(); + + // 1. Sample τ values for inference + let total_taus = (b * n) as i32; + let tau_blocks = (b * n + 255) / 256; + let tau_config = LaunchConfig { + grid_dim: (tau_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let rng_step = self.rng_step as u32; + self.rng_step += 1; + unsafe { + self.stream + .launch_builder(&self.sample_taus_kernel) + .arg(&mut self.online_taus) + .arg(&total_taus) + .arg(&rng_step) + .launch(tau_config) + .map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?; + } + + // 2. Run IQN forward-only kernel + let fwd_config = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + let batch_i32 = b as i32; + unsafe { + self.stream + .launch_builder(&self.forward_kernel) + .arg(h_s2) + .arg(&self.online_taus) + .arg(&self.online_params) + .arg(&mut self.save_q_online) + .arg(&batch_i32) + .launch(fwd_config) + .map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?; + } + + // 3. Compute CVaR on CPU (small: B scalars from B*N*TBA quantiles) + // Read quantile Q-values and actions to host + let q_size = b * n * tba; + let q_view = self.save_q_online.slice(..q_size); + let mut q_host = vec![0.0_f32; q_size]; + self.stream.memcpy_dtoh(&q_view, &mut q_host) // gpu-exit: CVaR readback (~4KB for B=32) + .map_err(|e| MLError::ModelError(format!("IQN CVaR q readback: {e}")))?; + + let mut act_host = vec![0_i32; b]; + let act_view = actions.slice(..b); + self.stream.memcpy_dtoh(&act_view, &mut act_host) // gpu-exit: action readback + .map_err(|e| MLError::ModelError(format!("IQN CVaR action readback: {e}")))?; + + // Compute CVaR per sample + let alpha_count = ((alpha * n as f32) as usize).max(1); + let mut cvar_scales = vec![1.0_f32; b]; + for i in 0..b { + // Extract quantile values for this sample's exposure action + let exposure_idx = (act_host[i] / 9) as usize; // factored: exposure * 9 + order * 3 + urgency + let exposure_idx = exposure_idx.min(self.config.branch_0_size - 1); + + let mut quantiles: Vec = (0..n) + .map(|t| q_host[i * n * tba + t * tba + exposure_idx]) + .collect(); + quantiles.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + // CVaR = mean of lowest alpha_count quantiles + let cvar: f32 = quantiles[..alpha_count].iter().sum::() / alpha_count as f32; + + // Scale: [0.25, 1.0]. CVaR ≥ 0 → full size. CVaR < 0 → reduce. + cvar_scales[i] = if cvar >= 0.0 { + 1.0 + } else { + (1.0 + cvar * 5.0).clamp(0.25, 1.0) + }; + } + + // Upload scales to GPU + let mut scales_buf = self.stream.alloc_zeros::(b) + .map_err(|e| MLError::ModelError(format!("IQN CVaR alloc scales: {e}")))?; + self.stream.memcpy_htod(&cvar_scales, &mut scales_buf) + .map_err(|e| MLError::ModelError(format!("IQN CVaR upload scales: {e}")))?; + + Ok(scales_buf) + } } // ---------------------------------------------------------------------------