diff --git a/crates/ml-dqn/src/branching_composition_tests.rs b/crates/ml-dqn/src/branching_composition_tests.rs new file mode 100644 index 000000000..35f9974c0 --- /dev/null +++ b/crates/ml-dqn/src/branching_composition_tests.rs @@ -0,0 +1,34 @@ +#[cfg(test)] +mod tests { + #[test] + fn test_factored_action_composition() { + // Verify: exposure * 9 + order * 3 + urgency covers 0-44 + let mut seen = std::collections::HashSet::new(); + for exposure in 0..5_u32 { + for order in 0..3_u32 { + for urgency in 0..3_u32 { + let factored = exposure * 9 + order * 3 + urgency; + assert!(factored < 45, "factored {} out of range", factored); + seen.insert(factored); + } + } + } + assert_eq!(seen.len(), 45, "Must cover all 45 factored actions"); + } + + #[test] + fn test_factored_action_decomposition() { + // Verify round-trip: compose then decompose + for original in 0..45_u32 { + let exposure = original / 9; + let order = (original % 9) / 3; + let urgency = original % 3; + + let recomposed = exposure * 9 + order * 3 + urgency; + assert_eq!(original, recomposed, "Round-trip failed for {}", original); + assert!(exposure < 5); + assert!(order < 3); + assert!(urgency < 3); + } + } +} diff --git a/crates/ml-dqn/src/lib.rs b/crates/ml-dqn/src/lib.rs index e7bd874f8..1f25a311e 100644 --- a/crates/ml-dqn/src/lib.rs +++ b/crates/ml-dqn/src/lib.rs @@ -43,6 +43,9 @@ pub mod multi_asset; // Branching DQN (Phase C+) pub mod branching; +#[cfg(test)] +mod branching_composition_tests; + // Rainbow DQN components pub mod distributional; pub mod distributional_dueling; diff --git a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu new file mode 100644 index 000000000..921caf9f3 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu @@ -0,0 +1,221 @@ +/** + * Fused epsilon-greedy action selection kernel with order routing + fill simulation. + * + * Given Q-values [batch_size, num_actions], selects actions entirely on GPU: + * - With probability epsilon: random action (GPU LCG RNG) + * - With probability 1-epsilon: argmax over Q-values + * - Routes action through OrderRouter (spread/vol-dependent order type + urgency) + * - Simulates fill (splitmix64 deterministic hash, order-type-dependent probability) + * - Unfilled orders override to Flat (exposure=2) + * + * Reuses gpu_random(), route_order(), simulate_fill_check() from + * common_device_functions.cuh (prepended via NVRTC). + * Launch config: grid=(ceil(batch_size/256),1,1), block=(256,1,1). + */ + +/** + * Basic epsilon-greedy: outputs exposure index (0-4) only. + * Used when fill simulation is not needed (e.g. inference, CPU routing fallback). + */ +extern "C" __global__ void epsilon_greedy_select( + const float* __restrict__ q_values, /* [batch_size, num_actions] */ + unsigned int* rng_states, /* [batch_size] — persistent LCG state */ + unsigned int* actions_out, /* [batch_size] — output action indices */ + const float epsilon, + const int batch_size, + const int num_actions +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch_size) return; + + unsigned int rng = rng_states[idx]; + float r = gpu_random(&rng); + + unsigned int action; + if (r < epsilon) { + /* Random exploration */ + action = (unsigned int)(gpu_random(&rng) * (float)num_actions); + if (action >= (unsigned int)num_actions) action = (unsigned int)(num_actions - 1); + } else { + /* Greedy: argmax over Q-values for this batch element */ + const float* q = q_values + idx * num_actions; + float best_q = q[0]; + action = 0; + for (int a = 1; a < num_actions; a++) { + if (q[a] > best_q) { + best_q = q[a]; + action = (unsigned int)a; + } + } + } + + rng_states[idx] = rng; + actions_out[idx] = action; +} + +/** + * Fused epsilon-greedy + order routing + fill simulation. + * + * Outputs post-fill action indices: if fill check fails, action is overridden + * to Flat (exposure_idx=2). This eliminates the CPU readback + route_action() + * + simulate_fill() roundtrip entirely. + * + * Parameters: + * step_offset — global training step offset for deterministic fill hash + * spread, median_spread — for order type selection + * volatility, median_vol — for urgency selection + * spread_bps — bid-ask spread in basis points (for fill cost) + * ioc_fill_prob, limit_fill_min, limit_fill_max — fill simulator config + * spread_cost_frac, spread_capture_frac — cost model config + * + * Outputs: + * actions_out — post-fill exposure indices (0-4), Flat=2 for unfilled + * fill_mask_out — 1 if filled, 0 if not (optional: NULL to skip) + */ +extern "C" __global__ void epsilon_greedy_routed( + const float* __restrict__ q_values, /* [batch_size, num_actions] */ + unsigned int* rng_states, /* [batch_size] */ + unsigned int* actions_out, /* [batch_size] — post-fill exposure */ + int* fill_mask_out, /* [batch_size] — 1=filled, 0=not (nullable) */ + const float epsilon, + const int batch_size, + const int num_actions, + const int step_offset, + const float spread, + const float median_spread, + const float volatility, + const float median_vol, + const float spread_bps, + const float ioc_fill_prob, + const float limit_fill_min, + const float limit_fill_max, + const float spread_cost_frac, + const float spread_capture_frac +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch_size) return; + + unsigned int rng = rng_states[idx]; + float r = gpu_random(&rng); + + /* Step 1: Epsilon-greedy action selection (exposure index 0-4) */ + int exposure_idx; + if (r < epsilon) { + exposure_idx = (int)(gpu_random(&rng) * (float)num_actions); + if (exposure_idx >= num_actions) exposure_idx = num_actions - 1; + } else { + const float* q = q_values + idx * num_actions; + float best_q = q[0]; + exposure_idx = 0; + for (int a = 1; a < num_actions; a++) { + if (q[a] > best_q) { + best_q = q[a]; + exposure_idx = a; + } + } + } + + /* Step 2: Order routing (spread/vol → order_type + urgency) */ + int order_type, urgency; + route_order(spread, median_spread, volatility, median_vol, + &order_type, &urgency); + + /* Step 3: Fill simulation (deterministic splitmix64 hash) */ + /* Normalized vol for fill probability calculation */ + float norm_vol = (median_vol > 0.0f) ? (volatility / median_vol) : 1.0f; + float cost_adj; + int filled = simulate_fill_check( + order_type, urgency, norm_vol, spread_bps, + step_offset + idx, exposure_idx, + ioc_fill_prob, limit_fill_min, limit_fill_max, + spread_cost_frac, spread_capture_frac, + &cost_adj + ); + + /* Step 4: Override to Flat if not filled */ + int final_action = filled ? exposure_idx : 2; /* 2 = Flat */ + + rng_states[idx] = rng; + actions_out[idx] = (unsigned int)final_action; + if (fill_mask_out != NULL) { + fill_mask_out[idx] = filled; + } +} + +/** + * Branching DQN action selection: 3 independent epsilon-greedy heads. + * + * Each head independently selects via epsilon-greedy: + * - Exposure head: 5 actions (Short100/Short50/Flat/Long50/Long100) + * - Order head: 3 actions (LimitMaker/IoC/Market) + * - Urgency head: 3 actions (Patient/Normal/Aggressive) + * + * Composes factored action index: exposure * 9 + order * 3 + urgency (0-44). + * + * Launch: grid=(ceil(batch_size/256),1,1), block=(256,1,1). + */ +extern "C" __global__ void branching_action_select( + const float* __restrict__ q_exposure, /* [batch_size, 5] */ + const float* __restrict__ q_order, /* [batch_size, 3] */ + const float* __restrict__ q_urgency, /* [batch_size, 3] */ + unsigned int* rng_states, /* [batch_size] */ + unsigned int* actions_out, /* [batch_size] — factored index 0-44 */ + const float epsilon, + const int batch_size +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch_size) return; + + unsigned int rng = rng_states[idx]; + + /* Head 1: Exposure (5 actions) */ + int exposure; + float r1 = gpu_random(&rng); + if (r1 < epsilon) { + exposure = (int)(gpu_random(&rng) * 5.0f); + if (exposure >= 5) exposure = 4; + } else { + const float* qe = q_exposure + idx * 5; + float best = qe[0]; + exposure = 0; + for (int a = 1; a < 5; a++) { + if (qe[a] > best) { best = qe[a]; exposure = a; } + } + } + + /* Head 2: Order type (3 actions) */ + int order; + float r2 = gpu_random(&rng); + if (r2 < epsilon) { + order = (int)(gpu_random(&rng) * 3.0f); + if (order >= 3) order = 2; + } else { + const float* qo = q_order + idx * 3; + float best = qo[0]; + order = 0; + for (int a = 1; a < 3; a++) { + if (qo[a] > best) { best = qo[a]; order = a; } + } + } + + /* Head 3: Urgency (3 actions) */ + int urgency; + float r3 = gpu_random(&rng); + if (r3 < epsilon) { + urgency = (int)(gpu_random(&rng) * 3.0f); + if (urgency >= 3) urgency = 2; + } else { + const float* qu = q_urgency + idx * 3; + float best = qu[0]; + urgency = 0; + for (int a = 1; a < 3; a++) { + if (qu[a] > best) { best = qu[a]; urgency = a; } + } + } + + /* Compose factored action: exposure * 9 + order * 3 + urgency */ + unsigned int factored = (unsigned int)(exposure * 9 + order * 3 + urgency); + + rng_states[idx] = rng; + actions_out[idx] = factored; +}