diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 1659a7102..b1cf2843b 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -234,6 +234,11 @@ __device__ void block_bellman_project( t_z = fminf(fmaxf(t_z, V_MIN), V_MAX); float b_pos = (t_z - V_MIN) / DELTA_Z; + /* Keep b_pos away from exact upper boundary to prevent phantom + * upper==lower when ceilf(N-1) == floorf(N-1) = N-1. + * Without this, probability mass piles on boundary atoms. */ + b_pos = fminf(b_pos, (float)(NUM_ATOMS - 1) - 1e-6f); + b_pos = fmaxf(b_pos, 1e-6f); int lower = (int)floorf(b_pos); int upper = (int)ceilf(b_pos); lower = max(min(lower, NUM_ATOMS - 1), 0); diff --git a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu index 883c1180f..6d7a84f36 100644 --- a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu +++ b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu @@ -23,7 +23,8 @@ extern "C" __global__ void epsilon_greedy_select( unsigned int* actions_out, /* [batch_size] — output action indices */ const float epsilon, const int batch_size, - const int num_actions + const int num_actions, + const float q_gap_threshold /* min Q-gap for trade entry (0.0 = disabled) */ ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; @@ -47,6 +48,13 @@ extern "C" __global__ void epsilon_greedy_select( action = (unsigned int)a; } } + /* Q-gap conviction filter: force flat when conviction is low */ + if (q_gap_threshold > 0.0f && num_actions > 2) { + float q_flat = q[2]; /* index 2 = Flat */ + if (best_q - q_flat < q_gap_threshold) { + action = 2; + } + } } rng_states[idx] = rng; @@ -90,7 +98,8 @@ extern "C" __global__ void epsilon_greedy_routed( const float limit_fill_min, const float limit_fill_max, const float spread_cost_frac, - const float spread_capture_frac + const float spread_capture_frac, + const float q_gap_threshold ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; @@ -113,6 +122,13 @@ extern "C" __global__ void epsilon_greedy_routed( exposure_idx = a; } } + /* Q-gap conviction filter */ + if (q_gap_threshold > 0.0f && num_actions > 2) { + float q_flat = q[2]; + if (best_q - q_flat < q_gap_threshold) { + exposure_idx = 2; + } + } } /* Step 2: Order routing (spread/vol → order_type + urgency) */ @@ -161,14 +177,15 @@ extern "C" __global__ void branching_action_select( unsigned int* rng_states, /* [batch_size] */ unsigned int* actions_out, /* [batch_size] — factored index 0-44 */ const float epsilon, - const int batch_size + const int batch_size, + const float q_gap_threshold /* min Q-gap for trade entry (0.0 = disabled) */ ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; unsigned int rng = rng_states[idx]; - /* Head 1: Exposure (5 actions) */ + /* Head 1: Exposure (5 actions) — with Q-gap conviction filter */ int exposure; float r1 = gpu_random(&rng); if (r1 < epsilon) { @@ -181,6 +198,13 @@ extern "C" __global__ void branching_action_select( for (int a = 1; a < 5; a++) { if (qe[a] > best) { best = qe[a]; exposure = a; } } + /* Q-gap filter: force flat when conviction is low */ + if (q_gap_threshold > 0.0f) { + float q_flat = qe[2]; /* index 2 = Flat */ + if (best - q_flat < q_gap_threshold) { + exposure = 2; + } + } } /* Head 2: Order type (3 actions) */ diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index f8358b67a..1c73fbc5e 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -515,6 +515,8 @@ extern "C" __global__ void experience_env_step( float dsr_denom = (var_32 > 0.01f) ? var_32 : 0.01f; float dsr_t = (dsr_B * delta_A - 0.5f * dsr_A * delta_B) / dsr_denom; + /* NaN guard: if any DSR input was NaN (from prior accumulation), zero out */ + if (isnan(dsr_t) || isinf(dsr_t)) dsr_t = 0.0f; /* Clamp DSR individually to [-1, +1] BEFORE weighting */ dsr_t = fmaxf(-1.0f, fminf(1.0f, dsr_t)); @@ -527,6 +529,7 @@ extern "C" __global__ void experience_env_step( pnl_var = pnl_var + eta * (pnl_diff * pnl_diff - pnl_var); float pnl_std = sqrtf(pnl_var + 1e-8f); float normalized_pnl = pnl_diff / pnl_std; + if (isnan(normalized_pnl) || isinf(normalized_pnl)) normalized_pnl = 0.0f; /* Clamp z-score to [-3, +3] (99.7% of distribution). Old code * allowed ±10+ outliers that dominated after asymmetric scaling. */ normalized_pnl = fmaxf(-3.0f, fminf(3.0f, normalized_pnl)); diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 8f3e15b22..41f0989c5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -38,6 +38,7 @@ pub struct GpuActionSelector { fill_mask_buf: CudaSlice, max_batch_size: usize, stream: Arc, + q_gap_threshold: f32, } impl GpuActionSelector { @@ -60,21 +61,25 @@ impl GpuActionSelector { let mut rng_states = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?; stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| MLError::ModelError(format!("upload rng_states: {e}")))?; info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, kernel compiled"); - Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, actions_buf, fill_mask_buf, max_batch_size, stream }) + Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, actions_buf, fill_mask_buf, max_batch_size, stream, q_gap_threshold: 0.0 }) } pub fn stream(&self) -> &Arc { &self.stream } + /// Set the Q-gap conviction threshold (0.0 = disabled). + pub fn set_q_gap_threshold(&mut self, threshold: f32) { self.q_gap_threshold = threshold; } + pub fn select_actions(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize) -> Result, MLError> { if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); } let config = launch_config_1d(batch_size); let bs_i32 = batch_size as i32; let na_i32 = num_actions as i32; + let q_gap = self.q_gap_threshold; unsafe { self.stream.launch_builder(&self.kernel_func) .arg(q_values).arg(&mut self.rng_states).arg(&mut self.actions_buf) - .arg(&epsilon).arg(&bs_i32).arg(&na_i32) + .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&q_gap) .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy kernel launch: {e}")))?; } self.copy_actions_out(batch_size) @@ -97,7 +102,7 @@ impl GpuActionSelector { .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&step_offset) .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_vol) .arg(&spread_bps).arg(&ioc_fill_prob).arg(&limit_fill_min).arg(&limit_fill_max) - .arg(&spread_cost_frac).arg(&spread_capture_frac) + .arg(&spread_cost_frac).arg(&spread_capture_frac).arg(&self.q_gap_threshold) .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed kernel launch: {e}")))?; } self.copy_actions_out(batch_size) @@ -113,7 +118,7 @@ impl GpuActionSelector { unsafe { self.stream.launch_builder(&self.branching_kernel_func) .arg(q_exposure).arg(q_order).arg(q_urgency) - .arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&epsilon).arg(&bs_i32) + .arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&epsilon).arg(&bs_i32).arg(&self.q_gap_threshold) .launch(config).map_err(|e| MLError::ModelError(format!("branching kernel launch: {e}")))?; } self.copy_actions_out(batch_size) diff --git a/crates/ml/src/trainers/dqn/trainer/action.rs b/crates/ml/src/trainers/dqn/trainer/action.rs index 0d69057e3..434b10f8b 100644 --- a/crates/ml/src/trainers/dqn/trainer/action.rs +++ b/crates/ml/src/trainers/dqn/trainer/action.rs @@ -79,9 +79,10 @@ impl DQNTrainer { { if self.gpu_action_selector.is_none() && self.cuda_stream.is_some() { let stream = std::sync::Arc::clone(self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?); - let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(stream, self.hyperparams.batch_size.max(batch_size).max(8192), 0xDEAD_BEEF_CAFE_u64) + let mut selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(stream, self.hyperparams.batch_size.max(batch_size).max(8192), 0xDEAD_BEEF_CAFE_u64) .map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?; - info!("GPU action selector initialized for select_actions_batch"); + selector.set_q_gap_threshold(self.hyperparams.q_gap_threshold as f32); + info!("GPU action selector initialized for select_actions_batch (q_gap={:.3})", self.hyperparams.q_gap_threshold); self.gpu_action_selector = Some(selector); } let selector = self.gpu_action_selector.as_mut().ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?;