gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: Q-gap filter in all action selection kernels + C51 boundary + NaN guards

Browse Source 
Bug hunt findings:
- epsilon_greedy_kernel.cu: all 3 kernels (select, routed, branching) lacked
  the Q-gap conviction filter. Training learned with q_gap=0.1 but inference
  paths bypassed it entirely. Now all action selection paths are consistent.
- c51_loss_kernel.cu: Bellman projection boundary fix — b_pos clamped away
  from exact NUM_ATOMS-1 to prevent phantom upper==lower atom collapse.
- experience_kernels.cu: NaN/Inf guards on DSR and normalized PnL outputs.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-23 08:17:35 +01:00
parent 35d417f08e
commit 4f5daaa755
5 changed files with 48 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -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);

32
crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu
View File

@@ -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) */

3
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -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));

13
crates/ml/src/cuda_pipeline/gpu_action_selector.rs
View File

@@ -38,6 +38,7 @@ pub struct GpuActionSelector {
fill_mask_buf: CudaSlice<i32>,
max_batch_size: usize,
stream: Arc<CudaStream>,
q_gap_threshold: f32,
}
impl GpuActionSelector {
@@ -60,21 +61,25 @@ impl GpuActionSelector {
let mut rng_states = stream.alloc_zeros::<u32>(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<CudaStream> { &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<f32>, epsilon: f32, batch_size: usize, num_actions: usize) -> Result<CudaSlice<u32>, MLError> {
if batch_size == 0 { return self.stream.alloc_zeros::<u32>(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)

5
crates/ml/src/trainers/dqn/trainer/action.rs
View File

@@ -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"))?;