From dd6143936e6c828e5c4818a5a433fd6c7423dfd1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 11 Apr 2026 22:40:47 +0200 Subject: [PATCH] fix: IQN 4-branch runtime + order counterfactual + per-branch entropy + histogram MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit IQN kernel rewrite (critical — was corrupting 75% of gradient budget): - Eliminate compile-time BRANCH_*_SIZE defines (BRANCH_0_SIZE was 5, should be 3) - All 9 IQN kernels now take runtime b0/b1/b2/b3 params - 4-branch forward, backward, decode (was 3-branch, missing magnitude) - branch_actions buffer B*3 → B*4 Action diversity infrastructure: - 3-cycle counterfactual: dir mirror / mag relabel / order-type cost delta (50× amplified) - Per-branch entropy: order (d==2) gets 3× boost in C51 + MSE grad kernels - Position histogram expanded 6→12 bins (all 4 branches get entropy bonus) - Boltzmann temperature floor raised to 0.5 for order + urgency branches - Smoke test epochs 3→10 for diversity verification Note: experience collector still needs bottleneck forward path — currently reads trainer's bottleneck-layout weights with state_dim K, producing misaligned Q-values. Next: eliminate weight copy, use direct pointer views into trainer's params_buf. Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-smoketest.toml | 2 +- .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 5 +- .../src/cuda_pipeline/experience_kernels.cu | 110 ++-- .../cuda_pipeline/gpu_experience_collector.rs | 4 +- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 64 +- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 314 ++++++---- .../ml/src/cuda_pipeline/mse_grad_kernel.cu | 9 +- .../2026-04-11-action-diversity-iqn4branch.md | 548 ++++++++++++++++++ 8 files changed, 882 insertions(+), 174 deletions(-) create mode 100644 docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index a3c9f77f7..792d549d3 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -10,7 +10,7 @@ # - q_clip ±50 (prevents logit explosion through bf16 forward) [training] -epochs = 3 +epochs = 10 batch_size = 64 learning_rate = 2e-6 gamma = 0.95 diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index c10e7a24a..091d08b79 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -51,7 +51,10 @@ extern "C" __global__ void c51_grad_kernel( * atom distribution collapse. Standardized advantages prevent Q-value scale * collapse; this prevents atom sharpness collapse. Two-pronged defense. */ if (entropy_coeff > 0.0f) { - float ent_scale = (d == 1) ? 5.0f : 1.0f; + float ent_scale; + if (d == 1) ent_scale = 5.0f; /* magnitude: existing */ + else if (d == 2) ent_scale = 3.0f; /* order: prevent order type collapse */ + else ent_scale = 1.0f; /* direction, urgency: baseline */ float lp_clamped = fmaxf(lp, -10.0f); d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped); } diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 276ad3290..f201ce18e 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -925,7 +925,12 @@ extern "C" __global__ void experience_action_select( q_max_ord = fmaxf(q_max_ord, q_b2[a]); q_min_ord = fminf(q_min_ord, q_b2[a]); } - float tau_ord = fmaxf(q_max_ord - q_min_ord, 0.01f); + /* Higher temperature floor (0.5 vs 0.01) for order: Q-values for different + * order types start nearly identical (same reward signal). The low 0.01 floor + * amplifies initialization noise 100×, locking in a single order type before + * learning can differentiate them. 0.5 gives meaningful Boltzmann stochasticity + * until the order branch develops genuine Q-gaps from counterfactual training. */ + float tau_ord = fmaxf(q_max_ord - q_min_ord, 0.5f); float sum_e = 0.0f; float exps_ord[3]; for (int a = 0; a < b2_size; a++) { @@ -953,7 +958,7 @@ extern "C" __global__ void experience_action_select( q_max_urg = fmaxf(q_max_urg, q_b3[a]); q_min_urg = fminf(q_min_urg, q_b3[a]); } - float tau_urg = fmaxf(q_max_urg - q_min_urg, 0.01f); + float tau_urg = fmaxf(q_max_urg - q_min_urg, 0.5f); float sum_e = 0.0f; float exps_urg[3]; for (int a = 0; a < b3_size; a++) { @@ -1661,48 +1666,44 @@ extern "C" __global__ void experience_env_step( reward = fmaxf(-10.0f, -5.0f - 5.0f * risk_taken); } - /* ── #19 Position entropy: track direction(3) + magnitude(3) diversity ── */ + /* ── #19 Position entropy: track all 4 branches (3 bins each = 12 total) ── */ if (position_histogram != NULL && position_entropy_weight > 0.0f) { - /* Histogram layout: [N, 6] — first 3 bins = direction, next 3 = magnitude */ - if (dir_idx >= 0 && dir_idx < 3) { - position_histogram[(long long)i * 6 + dir_idx] += 1.0f; - } - if (mag_idx >= 0 && mag_idx < 3) { - position_histogram[(long long)i * 6 + 3 + mag_idx] += 1.0f; - } + /* Histogram layout: [N, 12] — dir[0:3], mag[3:6], order[6:9], urgency[9:12] */ + int order_idx = decode_order_4b(action_idx, b2_size, b3_size); + int urgency_idx = decode_urgency_4b(action_idx, b3_size); + if (dir_idx >= 0 && dir_idx < 3) + position_histogram[(long long)i * 12 + dir_idx] += 1.0f; + if (mag_idx >= 0 && mag_idx < 3) + position_histogram[(long long)i * 12 + 3 + mag_idx] += 1.0f; + if (order_idx >= 0 && order_idx < 3) + position_histogram[(long long)i * 12 + 6 + order_idx] += 1.0f; + if (urgency_idx >= 0 && urgency_idx < 3) + position_histogram[(long long)i * 12 + 9 + urgency_idx] += 1.0f; } /* ---- Done detection ---- */ int next_bar = bar_idx + 1; int done = (next_bar >= total_bars || check_capital_floor(new_portfolio_value, peak_equity)) ? 1 : 0; - /* #19 Position entropy bonus at episode end — direction(3) + magnitude(3) */ + /* #19 Position entropy bonus at episode end — all 4 branches */ if (done && position_histogram != NULL && position_entropy_weight > 0.0f) { - float* hist = position_histogram + (long long)i * 6; - /* Compute entropy over direction bins [0:3] */ - float dir_total = 0.0f; - for (int b = 0; b < 3; b++) dir_total += hist[b]; - float dir_entropy = 0.0f; - if (dir_total > 1.0f) { - for (int b = 0; b < 3; b++) { - float p = hist[b] / dir_total; - if (p > 1e-6f) dir_entropy -= p * logf(p); + float* hist = position_histogram + (long long)i * 12; + float combined_entropy = 0.0f; + for (int branch = 0; branch < 4; branch++) { + float* bh = hist + branch * 3; + float total = 0.0f; + for (int b = 0; b < 3; b++) total += bh[b]; + float ent = 0.0f; + if (total > 1.0f) { + for (int b = 0; b < 3; b++) { + float p = bh[b] / total; + if (p > 1e-6f) ent -= p * logf(p); + } + ent /= 1.0986f; /* log(3) normalization → [0, 1] */ } - dir_entropy /= 1.0986f; /* Normalize by log(3) = 1.0986 → [0, 1] */ + combined_entropy += ent; } - /* Compute entropy over magnitude bins [3:6] */ - float mag_total = 0.0f; - for (int b = 3; b < 6; b++) mag_total += hist[b]; - float mag_entropy = 0.0f; - if (mag_total > 1.0f) { - for (int b = 3; b < 6; b++) { - float p = hist[b] / mag_total; - if (p > 1e-6f) mag_entropy -= p * logf(p); - } - mag_entropy /= 1.0986f; /* Normalize by log(3) → [0, 1] */ - } - /* Average of direction and magnitude entropy */ - float combined_entropy = 0.5f * (dir_entropy + mag_entropy); + combined_entropy *= 0.25f; /* average across 4 branches */ reward += position_entropy_weight * combined_entropy; } @@ -1796,15 +1797,16 @@ extern "C" __global__ void experience_env_step( int cf_action; float cf_reward; - if ((current_t & 1) == 0) { - /* ── Even step: directional mirror (original logic) ── */ + int cf_cycle = current_t % 3; + if (cf_cycle == 0) { + /* ── Step 0: directional mirror (original logic) ── */ int cf_dir = (b0_size - 1) - dir_idx; cf_action = cf_dir * b1_size * b2_size * b3_size + mag_idx * b2_size * b3_size + orig_order * b3_size + orig_urgency; cf_reward = -reward; - } else { - /* ── Odd step: hindsight magnitude relabeling ── + } else if (cf_cycle == 1) { + /* ── Step 1: hindsight magnitude relabeling ── * Pick the best alternative magnitude: * If Full(2), try Half(1). Otherwise try Full(2). * Scale reward linearly by (alt_mag_frac / actual_mag_frac). @@ -1818,7 +1820,7 @@ extern "C" __global__ void experience_env_step( /* Cycle: smaller, larger, opposite. All 3 magnitudes get gradient. * Was: only Half or Full — Small NEVER appeared as counterfactual. */ int alt_mag; - int cycle = (current_t / 2) % 3; + int cycle = (current_t / 3) % 3; if (cycle == 0) alt_mag = (mag_idx > 0) ? mag_idx - 1 : 2; else if (cycle == 1) alt_mag = (mag_idx < 2) ? mag_idx + 1 : 0; else alt_mag = 2 - mag_idx; @@ -1837,6 +1839,36 @@ extern "C" __global__ void experience_env_step( + orig_order * b3_size + orig_urgency; cf_reward = -reward; } + } else { + /* ── Step 2: Order-type counterfactual ── + * Cycle through alternative order types. Compute execution cost + * difference and amplify to match directional signal magnitude. */ + int alt_order = (orig_order + 1 + ((current_t / 3) & 1)) % b2_size; + + if (fabsf(position) > 0.001f) { + float taken_cost = compute_tx_cost( + position - pre_trade_position, raw_close, tx_cost_multiplier, + spread_cost, max_position, orig_order, spread_scale); + float alt_cost = compute_tx_cost( + position - pre_trade_position, raw_close, tx_cost_multiplier, + spread_cost, max_position, alt_order, spread_scale); + float cost_delta = (taken_cost - alt_cost) + / fmaxf(prev_equity, 1.0f); + + cf_action = dir_idx * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + alt_order * b3_size + orig_urgency; + /* Amplify 50×: raw cost delta ~0.0001-0.0005, directional ~0.01-0.1 */ + cf_reward = reward + 50.0f * cost_delta; + cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f); + } else { + /* Flat position: order irrelevant, directional mirror fallback */ + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = -reward; + } } if (isnan(cf_reward) || isinf(cf_reward)) cf_reward = 0.0f; diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 93fb35766..05699a168 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1045,8 +1045,8 @@ impl GpuExperienceCollector { "GPU experience collector initialized (cuBLAS + 3 focused kernels)" ); - // #19 Position entropy histogram [alloc_episodes * 6] — direction(3) + magnitude(3) - let position_histogram = stream.alloc_zeros::(alloc_episodes * 6) + // #19 Position entropy histogram [alloc_episodes * 12] — dir(3)+mag(3)+order(3)+urgency(3) + let position_histogram = stream.alloc_zeros::(alloc_episodes * 12) .map_err(|e| MLError::ModelError(format!("alloc position_histogram: {e}")))?; // #30 F32 forward pass buffers for full-precision experience collection diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 7d0dac7e0..8f77d7ff1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -178,7 +178,7 @@ pub struct GpuIqnHead { /// Precomputed cosine features [N, embed_dim] — cos(π·(d+1)·τ_i) /// Fixed at construction (τ are midpoints). Eliminates 16.7M cosf() per step. cos_features: CudaSlice, - /// Branch actions decoded from flat actions [B, 3] + /// Branch actions decoded from flat actions [B, 4] branch_actions: CudaSlice, /// Target h_s2 computed from next_states + target trunk weights [B, H] target_h_s2: CudaSlice, @@ -282,7 +282,7 @@ impl GpuIqnHead { MLError::ModelError(format!("IQN htod target_taus ({} f32): {e}", b * n)) })?; } - let branch_actions = stream.alloc_zeros::(b * 3).map_err(|e| { + let branch_actions = stream.alloc_zeros::(b * 4).map_err(|e| { MLError::ModelError(format!("IQN alloc branch_actions: {e}")) })?; let target_h_s2 = alloc_f32(&stream, b * h, "iqn_target_h_s2")?; @@ -306,9 +306,9 @@ impl GpuIqnHead { // IQR buffer: per-action interquartile range for exploration bonus let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?; - let vram_bytes = (total_params * 6 + b * n * 2 + b * 3 + let vram_bytes = (total_params * 6 + b * n * 2 + b * 4 + b * h + b * n * h * 2 + b * n * tba * 2 - + b * 3 + 2) * 4; + + b * 4 + 2) * 4; info!( hidden_dim = h, @@ -395,11 +395,15 @@ impl GpuIqnHead { let shared_h1_i32 = self.config.shared_h1 as i32; let hidden_dim_i32 = self.config.hidden_dim as i32; let embed_dim_i32 = self.config.embed_dim as i32; + let b0_i32 = self.config.branch_0_size as i32; + let b1_i32 = self.config.branch_1_size as i32; + let b2_i32 = self.config.branch_2_size as i32; + let b3_i32 = self.config.branch_3_size as i32; // Layer 3a: τ values are fixed midpoints, pre-computed in constructor. // No per-step sampling → deterministic → CUDA Graph compatible. - // 2. GPU decode: flat actions [B] → branch actions [B, 3] + // 2. GPU decode: flat actions [B] → branch actions [B, 4] let decode_blocks = (b + 255) / 256; let decode_config = LaunchConfig { grid_dim: (decode_blocks as u32, 1, 1), @@ -408,7 +412,7 @@ impl GpuIqnHead { }; let batch_size_i32 = b as i32; - // Safety: dqn_actions_buf[B] i32, branch_actions[B*3] i32, same context. + // Safety: dqn_actions_buf[B] i32, branch_actions[B*4] i32, same context. unsafe { self.stream .launch_builder(&self.decode_actions_kernel) @@ -418,6 +422,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(decode_config) .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?; } @@ -461,6 +469,10 @@ impl GpuIqnHead { let shared_h1_i32 = self.config.shared_h1 as i32; let hidden_dim_i32 = self.config.hidden_dim as i32; let embed_dim_i32 = self.config.embed_dim as i32; + let b0_i32 = self.config.branch_0_size as i32; + let b1_i32 = self.config.branch_1_size as i32; + let b2_i32 = self.config.branch_2_size as i32; + let b3_i32 = self.config.branch_3_size as i32; if let Some(cached_ptr) = self.cached_target_h_s2_ptr.take() { // Reuse pre-computed target h_s2 from DQN Pass 2 — skip trunk forward. @@ -496,6 +508,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(trunk_config) .map_err(|e| MLError::ModelError(format!("IQN trunk forward kernel: {e}")))?; } @@ -549,6 +565,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(fwd_config) .map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?; } @@ -580,6 +600,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(bwd_config) .map_err(|e| MLError::ModelError(format!("IQN backward kernel: {e}")))?; } @@ -603,6 +627,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(norm_config) .map_err(|e| MLError::ModelError(format!("IQN grad_norm kernel: {e}")))?; } @@ -652,6 +680,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(adam_config) .map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?; } @@ -687,6 +719,10 @@ impl GpuIqnHead { let shared_h1_i32 = self.config.shared_h1 as i32; let hidden_dim_i32 = self.config.hidden_dim as i32; let embed_dim_i32 = self.config.embed_dim as i32; + let b0_i32 = self.config.branch_0_size as i32; + let b1_i32 = self.config.branch_1_size as i32; + let b2_i32 = self.config.branch_2_size as i32; + let b3_i32 = self.config.branch_3_size as i32; let tau_ptr = self.tau_buf.raw_ptr(); // Safety: target_params and online_params have total_params elements each. @@ -700,6 +736,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(config) .map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?; } @@ -852,6 +892,10 @@ impl GpuIqnHead { block_dim: (256, 1, 1), shared_mem_bytes: 0, }; + let b0_i32 = self.config.branch_0_size as i32; + let b1_i32 = self.config.branch_1_size as i32; + let b2_i32 = self.config.branch_2_size as i32; + let b3_i32 = self.config.branch_3_size as i32; let rng_step = self.rng_step as u32; self.rng_step += 1; unsafe { @@ -863,6 +907,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(tau_config) .map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?; } @@ -887,6 +935,10 @@ impl GpuIqnHead { .arg(&shared_h1_i32) .arg(&hidden_dim_i32) .arg(&embed_dim_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) .launch(fwd_config) .map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?; } diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index d9ae33a0c..bd0e55d41 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -34,9 +34,7 @@ * IQN_EMBED_DIM -- cosine embedding dimension (default 64) * IQN_NUM_QUANTILES -- τ samples per forward pass (default 32) * IQN_KAPPA -- Huber threshold for quantile loss (default 1.0) - * BRANCH_0_SIZE -- exposure actions (default 5) - * BRANCH_1_SIZE -- order type actions (default 3) - * BRANCH_2_SIZE -- urgency actions (default 3) + * Branch sizes (b0, b1, b2, b3) are runtime kernel parameters. */ /* ── Compile-time defaults ──────────────────────────────────────────── */ @@ -59,50 +57,17 @@ #ifndef IQN_KAPPA #define IQN_KAPPA 1.0f #endif -#ifndef BRANCH_0_SIZE -#define BRANCH_0_SIZE 5 -#endif -#ifndef BRANCH_1_SIZE -#define BRANCH_1_SIZE 3 -#endif -#ifndef BRANCH_2_SIZE -#define BRANCH_2_SIZE 3 -#endif - -#define TOTAL_BRANCH_ACTIONS (BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE) +/* Branch sizes (b0_size, b1_size, b2_size, b3_size) are runtime kernel parameters. + * No compile-time BRANCH_*_SIZE defines — all offsets computed at runtime. */ /* ── Weight layout (flat f32 buffer) ─────────────────────────────────── */ -/* W_embed [hidden_dim, IQN_EMBED_DIM], b_embed [hidden_dim] - * W_b0 [BRANCH_0_SIZE, hidden_dim], b_b0 [BRANCH_0_SIZE] - * W_b1 [BRANCH_1_SIZE, hidden_dim], b_b1 [BRANCH_1_SIZE] - * W_b2 [BRANCH_2_SIZE, hidden_dim], b_b2 [BRANCH_2_SIZE] +/* W_embed [hidden_dim, embed_dim], b_embed [hidden_dim] + * W_b0 [b0_size, hidden_dim], b_b0 [b0_size] + * W_b1 [b1_size, hidden_dim], b_b1 [b1_size] + * W_b2 [b2_size, hidden_dim], b_b2 [b2_size] + * W_b3 [b3_size, hidden_dim], b_b3 [b3_size] * - * NOTE: IQN_OFF_* macros below use compile-time IQN_HIDDEN. They are kept - * only for documentation. Kernels compute offsets from runtime hidden_dim. */ - -#define IQN_W_EMBED_SIZE (IQN_HIDDEN * IQN_EMBED_DIM) -#define IQN_B_EMBED_SIZE (IQN_HIDDEN) -#define IQN_W_B0_SIZE (BRANCH_0_SIZE * IQN_HIDDEN) -#define IQN_B_B0_SIZE (BRANCH_0_SIZE) -#define IQN_W_B1_SIZE (BRANCH_1_SIZE * IQN_HIDDEN) -#define IQN_B_B1_SIZE (BRANCH_1_SIZE) -#define IQN_W_B2_SIZE (BRANCH_2_SIZE * IQN_HIDDEN) -#define IQN_B_B2_SIZE (BRANCH_2_SIZE) - -#define IQN_TOTAL_PARAMS (IQN_W_EMBED_SIZE + IQN_B_EMBED_SIZE \ - + IQN_W_B0_SIZE + IQN_B_B0_SIZE \ - + IQN_W_B1_SIZE + IQN_B_B1_SIZE \ - + IQN_W_B2_SIZE + IQN_B_B2_SIZE) - -/* Offsets into the flat parameter buffer (compile-time, for documentation only) */ -#define IQN_OFF_W_EMBED 0 -#define IQN_OFF_B_EMBED (IQN_OFF_W_EMBED + IQN_W_EMBED_SIZE) -#define IQN_OFF_W_B0 (IQN_OFF_B_EMBED + IQN_B_EMBED_SIZE) -#define IQN_OFF_B_B0 (IQN_OFF_W_B0 + IQN_W_B0_SIZE) -#define IQN_OFF_W_B1 (IQN_OFF_B_B0 + IQN_B_B0_SIZE) -#define IQN_OFF_B_B1 (IQN_OFF_W_B1 + IQN_W_B1_SIZE) -#define IQN_OFF_W_B2 (IQN_OFF_B_B1 + IQN_B_B1_SIZE) -#define IQN_OFF_B_B2 (IQN_OFF_W_B2 + IQN_W_B2_SIZE) + * Offsets computed at runtime by iqn_compute_offsets(). */ /* Block size for IQN forward/backward/trunk kernels (8 warps = 256 threads). */ #define IQN_BLOCK_SIZE 256 @@ -185,17 +150,25 @@ __device__ __forceinline__ float quantile_huber_grad(float tau, float delta, flo return -weight * huber_grad; } -/** Compute runtime weight offsets from hidden_dim and embed_dim. Returns offsets in an array: - * [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1, [6]=W_B2, [7]=B_B2 */ -__device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_dim, int offsets[8]) { +/** Compute runtime weight offsets from hidden_dim, embed_dim, and 4 branch sizes. + * Returns 10 offsets: + * [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1, + * [6]=W_B2, [7]=B_B2, [8]=W_B3, [9]=B_B3 */ +__device__ __forceinline__ void iqn_compute_offsets( + int hidden_dim, int embed_dim, + int b0, int b1, int b2, int b3, + int offsets[10]) +{ offsets[0] = 0; /* W_EMBED */ offsets[1] = offsets[0] + hidden_dim * embed_dim; /* B_EMBED */ offsets[2] = offsets[1] + hidden_dim; /* W_B0 */ - offsets[3] = offsets[2] + BRANCH_0_SIZE * hidden_dim; /* B_B0 */ - offsets[4] = offsets[3] + BRANCH_0_SIZE; /* W_B1 */ - offsets[5] = offsets[4] + BRANCH_1_SIZE * hidden_dim; /* B_B1 */ - offsets[6] = offsets[5] + BRANCH_1_SIZE; /* W_B2 */ - offsets[7] = offsets[6] + BRANCH_2_SIZE * hidden_dim; /* B_B2 */ + offsets[3] = offsets[2] + b0 * hidden_dim; /* B_B0 */ + offsets[4] = offsets[3] + b0; /* W_B1 */ + offsets[5] = offsets[4] + b1 * hidden_dim; /* B_B1 */ + offsets[6] = offsets[5] + b1; /* W_B2 */ + offsets[7] = offsets[6] + b2 * hidden_dim; /* B_B2 */ + offsets[8] = offsets[7] + b2; /* W_B3 */ + offsets[9] = offsets[8] + b3 * hidden_dim; /* B_B3 */ } /* ── Forward + Loss Kernel ──────────────────────────────────────────── */ @@ -208,7 +181,7 @@ __device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_di * τ values are pre-sampled and uploaded via the `taus` buffer (f32). * * One warp (32 threads) per sample. Quantiles are processed sequentially - * to conserve registers — each τ produces TOTAL_BRANCH_ACTIONS values. + * to conserve registers — each τ produces tba (= b0+b1+b2+b3) values. * * Saves per-quantile intermediate activations for backward pass. */ @@ -220,7 +193,7 @@ void iqn_forward_loss_kernel( /* Inputs (f32) */ const float* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */ const float* __restrict__ target_taus, /* [B, N] target τ samples */ - const int* __restrict__ actions, /* [B, 3] branch actions (exposure, order, urgency) */ + const int* __restrict__ actions, /* [B, 4] branch actions (direction, magnitude, order, urgency) */ const float* __restrict__ rewards, /* [B] */ const float* __restrict__ dones, /* [B] (0.0/1.0) */ float gamma, @@ -235,12 +208,16 @@ void iqn_forward_loss_kernel( /* Activation saves for backward (f32) */ float* __restrict__ save_embed, /* [B, N, hidden_dim] post-ReLU embeddings */ float* __restrict__ save_combined, /* [B, N, hidden_dim] h_s2 ⊙ embed */ - float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */ - float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */ + float* __restrict__ save_q_online, /* [B, N, tba] quantile Q-values */ + float* __restrict__ save_q_target, /* [B, N, tba] target Q-values */ int batch_size, int shared_h1, /* runtime: first hidden layer width (unused here, for consistency) */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ - int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int b0_size, /* runtime: branch 0 (direction) action count */ + int b1_size, /* runtime: branch 1 (magnitude) action count */ + int b2_size, /* runtime: branch 2 (order) action count */ + int b3_size /* runtime: branch 3 (urgency) action count */ ) { int sample = blockIdx.x; @@ -251,8 +228,9 @@ void iqn_forward_loss_kernel( __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ - int off[8]; - iqn_compute_offsets(hidden_dim, embed_dim, off); + int tba = b0_size + b1_size + b2_size + b3_size; + int off[10]; + iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off); /* ── Pointers into this sample's data ── */ const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim; @@ -260,9 +238,10 @@ void iqn_forward_loss_kernel( const float* my_taus = taus + sample * IQN_NUM_QUANTILES; const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES; - int a0 = actions[sample * 3 + 0]; /* exposure action */ - int a1 = actions[sample * 3 + 1]; /* order action */ - int a2 = actions[sample * 3 + 2]; /* urgency action */ + int a0 = actions[sample * 4 + 0]; /* direction action */ + int a1 = actions[sample * 4 + 1]; /* magnitude action */ + int a2 = actions[sample * 4 + 2]; /* order action */ + int a3 = actions[sample * 4 + 3]; /* urgency action */ float reward = rewards[sample]; float done = dones[sample]; @@ -275,6 +254,8 @@ void iqn_forward_loss_kernel( const float* b_b1 = online_params + off[5]; const float* w_b2 = online_params + off[6]; const float* b_b2 = online_params + off[7]; + const float* w_b3 = online_params + off[8]; + const float* b_b3 = online_params + off[9]; const float* tw_embed = target_params + off[0]; const float* tb_embed = target_params + off[1]; @@ -284,6 +265,8 @@ void iqn_forward_loss_kernel( const float* tb_b1 = target_params + off[5]; const float* tw_b2 = target_params + off[6]; const float* tb_b2 = target_params + off[7]; + const float* tw_b3 = target_params + off[8]; + const float* tb_b3 = target_params + off[9]; /* ── Load h_s2 into distributed registers (bf16 → f32 at boundary) ── */ float h_dist[IQN_DIST_MAX]; @@ -328,11 +311,11 @@ void iqn_forward_loss_kernel( } /* 4. Per-branch output: q_d = W_bd × combined + b_bd */ - int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba; - /* Branch 0 (exposure) */ + /* Branch 0 (direction) */ float q0_taken = 0.0f; - for (int a = 0; a < BRANCH_0_SIZE; a++) { + for (int a = 0; a < b0_size; a++) { float partial = 0.0f; const float* w_row = w_b0 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) @@ -348,16 +331,16 @@ void iqn_forward_loss_kernel( __syncthreads(); q0_taken = shmem_reduce[0]; - /* Branch 1 (order) */ + /* Branch 1 (magnitude) */ float q1_taken = 0.0f; - for (int a = 0; a < BRANCH_1_SIZE; a++) { + for (int a = 0; a < b1_size; a++) { float partial = 0.0f; const float* w_row = w_b1 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; float acc = b_b1[a] + iqn_block_sum(partial, shmem_reduce); if (tid == 0) { - save_q_online[q_save_offset + BRANCH_0_SIZE + a] = acc; + save_q_online[q_save_offset + b0_size + a] = acc; if (a == a1) q1_taken = acc; } } @@ -365,16 +348,16 @@ void iqn_forward_loss_kernel( __syncthreads(); q1_taken = shmem_reduce[0]; - /* Branch 2 (urgency) */ + /* Branch 2 (order) */ float q2_taken = 0.0f; - for (int a = 0; a < BRANCH_2_SIZE; a++) { + for (int a = 0; a < b2_size; a++) { float partial = 0.0f; const float* w_row = w_b2 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; float acc = b_b2[a] + iqn_block_sum(partial, shmem_reduce); if (tid == 0) { - save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a] = acc; + save_q_online[q_save_offset + b0_size + b1_size + a] = acc; if (a == a2) q2_taken = acc; } } @@ -382,8 +365,25 @@ void iqn_forward_loss_kernel( __syncthreads(); q2_taken = shmem_reduce[0]; + /* Branch 3 (urgency) */ + float q3_taken = 0.0f; + for (int a = 0; a < b3_size; a++) { + float partial = 0.0f; + const float* w_row = w_b3 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float acc = b_b3[a] + iqn_block_sum(partial, shmem_reduce); + if (tid == 0) { + save_q_online[q_save_offset + b0_size + b1_size + b2_size + a] = acc; + if (a == a3) q3_taken = acc; + } + } + if (tid == 0) shmem_reduce[0] = q3_taken; + __syncthreads(); + q3_taken = shmem_reduce[0]; + /* Sum of taken-action Q-values across branches (for loss computation) */ - online_q_a[t] = q0_taken + q1_taken + q2_taken; + online_q_a[t] = q0_taken + q1_taken + q2_taken + q3_taken; } /* ── Process TARGET quantiles ── */ @@ -410,7 +410,7 @@ void iqn_forward_loss_kernel( comb_dist[h / IQN_BLOCK_SIZE] = h_target_dist[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; /* Per-branch Q-values for taken actions */ - float q0, q1, q2; + float q0, q1, q2, q3; /* Branch 0 */ { @@ -436,17 +436,26 @@ void iqn_forward_loss_kernel( partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; q2 = tb_b2[a2] + iqn_block_sum(partial, shmem_reduce); } + /* Branch 3 */ + { + float partial = 0.0f; + const float* w_row = tw_b3 + a3 * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + q3 = tb_b3[a3] + iqn_block_sum(partial, shmem_reduce); + } /* Bellman target: r + γ(1-done) × Q_target */ - float raw_target = q0 + q1 + q2; + float raw_target = q0 + q1 + q2 + q3; target_q_a[t] = reward + gamma * (1.0f - done) * raw_target; /* Save target Q for backward */ - int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba; if (tid == 0) { save_q_target[q_save_offset + a0] = target_q_a[t]; - save_q_target[q_save_offset + BRANCH_0_SIZE + a1] = target_q_a[t]; - save_q_target[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a2] = target_q_a[t]; + save_q_target[q_save_offset + b0_size + a1] = target_q_a[t]; + save_q_target[q_save_offset + b0_size + b1_size + a2] = target_q_a[t]; + save_q_target[q_save_offset + b0_size + b1_size + b2_size + a3] = target_q_a[t]; } } @@ -491,10 +500,10 @@ void iqn_backward_kernel( const float* __restrict__ h_s2, /* [B, hidden_dim] (bf16 from DQN trunk) */ const float* __restrict__ save_embed, /* [B, N, hidden_dim] */ const float* __restrict__ save_combined, /* [B, N, hidden_dim] */ - const float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] */ - const float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] */ + const float* __restrict__ save_q_online, /* [B, N, tba] */ + const float* __restrict__ save_q_target, /* [B, N, tba] */ const float* __restrict__ taus, /* [B, N] */ - const int* __restrict__ actions, /* [B, 3] */ + const int* __restrict__ actions, /* [B, 4] */ /* Weights (f32, for backward through linear layers) */ const float* __restrict__ online_params, /* Precomputed cosine features (Optimization C) */ @@ -505,7 +514,11 @@ void iqn_backward_kernel( int batch_size, int shared_h1, /* runtime: first hidden layer width (unused here) */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ - int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int b0_size, /* runtime: branch 0 (direction) action count */ + int b1_size, /* runtime: branch 1 (magnitude) action count */ + int b2_size, /* runtime: branch 2 (order) action count */ + int b3_size /* runtime: branch 3 (urgency) action count */ ) { int sample = blockIdx.x; @@ -516,19 +529,22 @@ void iqn_backward_kernel( __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ - int off[8]; - iqn_compute_offsets(hidden_dim, embed_dim, off); + int tba = b0_size + b1_size + b2_size + b3_size; + int off[10]; + iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off); const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim; const float* my_taus = taus + sample * IQN_NUM_QUANTILES; - int a0 = actions[sample * 3 + 0]; - int a1 = actions[sample * 3 + 1]; - int a2 = actions[sample * 3 + 2]; + int a0 = actions[sample * 4 + 0]; + int a1 = actions[sample * 4 + 1]; + int a2 = actions[sample * 4 + 2]; + int a3 = actions[sample * 4 + 3]; const float* w_embed = online_params + off[0]; const float* w_b0 = online_params + off[2]; const float* w_b1 = online_params + off[4]; const float* w_b2 = online_params + off[6]; + const float* w_b3 = online_params + off[8]; /* Load h_s2 into distributed registers (bf16 → f32 at boundary) */ float h_dist[IQN_DIST_MAX]; @@ -540,11 +556,12 @@ void iqn_backward_kernel( float online_q[IQN_NUM_QUANTILES]; float target_q[IQN_NUM_QUANTILES]; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - int qoff = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + int qoff = (sample * IQN_NUM_QUANTILES + t) * tba; if (tid == 0) { online_q[t] = save_q_online[qoff + a0] - + save_q_online[qoff + BRANCH_0_SIZE + a1] - + save_q_online[qoff + BRANCH_0_SIZE + BRANCH_1_SIZE + a2]; + + save_q_online[qoff + b0_size + a1] + + save_q_online[qoff + b0_size + b1_size + a2] + + save_q_online[qoff + b0_size + b1_size + b2_size + a3]; target_q[t] = save_q_target[qoff + a0]; /* target_q already contains Bellman target */ } } @@ -623,6 +640,15 @@ void iqn_backward_kernel( if (tid == 0) atomicAdd(&grad_buf[off[7] + a2], inv_batch * dL_dq); + /* Branch 3 */ + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dcomb[h / IQN_BLOCK_SIZE] += w_b3[a3 * hidden_dim + h] * dL_dq; + atomicAdd(&grad_buf[off[8] + a3 * hidden_dim + h], + inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); + } + if (tid == 0) + atomicAdd(&grad_buf[off[9] + a3], inv_batch * dL_dq); + /* ── Gradient through element-wise product ── */ /* combined = h_s2 ⊙ embed * dL/d(embed) = dL/d(combined) ⊙ h_s2 @@ -671,7 +697,11 @@ void iqn_grad_norm_kernel( int total_params, int shared_h1, /* runtime: unused, for consistent interface */ int hidden_dim, /* runtime: unused, for consistent interface */ - int embed_dim /* runtime: unused, for consistent interface */ + int embed_dim, /* runtime: unused, for consistent interface */ + int b0_size, /* runtime: unused, for consistent interface */ + int b1_size, /* runtime: unused, for consistent interface */ + int b2_size, /* runtime: unused, for consistent interface */ + int b3_size /* runtime: unused, for consistent interface */ ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -713,7 +743,11 @@ void iqn_adam_kernel( int total_params, int shared_h1, /* runtime: unused, for consistent interface */ int hidden_dim, /* runtime: unused, for consistent interface */ - int embed_dim /* runtime: unused, for consistent interface */ + int embed_dim, /* runtime: unused, for consistent interface */ + int b0_size, /* runtime: unused, for consistent interface */ + int b1_size, /* runtime: unused, for consistent interface */ + int b2_size, /* runtime: unused, for consistent interface */ + int b3_size /* runtime: unused, for consistent interface */ ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -760,7 +794,7 @@ void iqn_adam_kernel( * Computes expected Q-values per action per branch by averaging over * quantile samples. Used for action selection blending with C51. * - * Outputs: expected_q [B, TOTAL_BRANCH_ACTIONS] — mean Q per action. + * Outputs: expected_q [B, tba] — mean Q per action. */ extern "C" __global__ void iqn_forward_kernel( @@ -768,11 +802,15 @@ void iqn_forward_kernel( const float* __restrict__ taus, /* [B, N] (f32) */ const float* __restrict__ params, /* IQN weights (f32) */ const float* __restrict__ cos_features, /* [N, embed_dim] precomputed cosine features */ - float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] (f32) */ + float* __restrict__ expected_q, /* [B, tba] (f32) */ int batch_size, int shared_h1, /* runtime: unused here, for consistency */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ - int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */ + int b0_size, /* runtime: branch 0 (direction) action count */ + int b1_size, /* runtime: branch 1 (magnitude) action count */ + int b2_size, /* runtime: branch 2 (order) action count */ + int b3_size /* runtime: branch 3 (urgency) action count */ ) { int sample = blockIdx.x; @@ -783,8 +821,9 @@ void iqn_forward_kernel( __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ - int off[8]; - iqn_compute_offsets(hidden_dim, embed_dim, off); + int tba = b0_size + b1_size + b2_size + b3_size; + int off[10]; + iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off); const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim; const float* my_taus = taus + sample * IQN_NUM_QUANTILES; @@ -796,15 +835,18 @@ void iqn_forward_kernel( const float* b_b1 = params + off[5]; const float* w_b2 = params + off[6]; const float* b_b2 = params + off[7]; + const float* w_b3 = params + off[8]; + const float* b_b3 = params + off[9]; /* Load h_s2 (bf16 → f32 at boundary) */ float h_dist[IQN_DIST_MAX]; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h]; - /* Accumulate Q-values across quantiles (for mean) */ - float q_acc[TOTAL_BRANCH_ACTIONS]; - for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++) + /* Accumulate Q-values across quantiles (for mean). + * Max tba = 3+3+3+3 = 12 — fits easily in registers. */ + float q_acc[16]; /* sized >= max tba */ + for (int a = 0; a < tba; a++) q_acc[a] = 0.0f; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { @@ -829,7 +871,7 @@ void iqn_forward_kernel( /* Branch outputs */ /* Branch 0 */ - for (int a = 0; a < BRANCH_0_SIZE; a++) { + for (int a = 0; a < b0_size; a++) { float partial = 0.0f; const float* w_row = w_b0 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) @@ -838,30 +880,39 @@ void iqn_forward_kernel( q_acc[a] += q_val; } /* Branch 1 */ - for (int a = 0; a < BRANCH_1_SIZE; a++) { + for (int a = 0; a < b1_size; a++) { float partial = 0.0f; const float* w_row = w_b1 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; float q_val = iqn_block_sum(partial, shmem_reduce) + b_b1[a]; - q_acc[BRANCH_0_SIZE + a] += q_val; + q_acc[b0_size + a] += q_val; } /* Branch 2 */ - for (int a = 0; a < BRANCH_2_SIZE; a++) { + for (int a = 0; a < b2_size; a++) { float partial = 0.0f; const float* w_row = w_b2 + a * hidden_dim; for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; float q_val = iqn_block_sum(partial, shmem_reduce) + b_b2[a]; - q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val; + q_acc[b0_size + b1_size + a] += q_val; + } + /* Branch 3 */ + for (int a = 0; a < b3_size; a++) { + float partial = 0.0f; + const float* w_row = w_b3 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float q_val = iqn_block_sum(partial, shmem_reduce) + b_b3[a]; + q_acc[b0_size + b1_size + b2_size + a] += q_val; } } /* Write mean Q-values (averaged over quantiles) */ float inv_n = 1.0f / (float)IQN_NUM_QUANTILES; if (tid == 0) { - for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++) - expected_q[sample * TOTAL_BRANCH_ACTIONS + a] = q_acc[a] * inv_n; + for (int a = 0; a < tba; a++) + expected_q[sample * tba + a] = q_acc[a] * inv_n; } } @@ -891,7 +942,11 @@ void iqn_trunk_forward_kernel( int state_dim, /* runtime state dimension (replaces IQN_STATE_DIM) */ int shared_h1, /* runtime: first hidden layer width */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ - int embed_dim /* runtime: unused here, for consistent interface */ + int embed_dim, /* runtime: unused here, for consistent interface */ + int b0_size, /* runtime: unused here, for consistent interface */ + int b1_size, /* runtime: unused here, for consistent interface */ + int b2_size, /* runtime: unused here, for consistent interface */ + int b3_size /* runtime: unused here, for consistent interface */ ) { int sample = blockIdx.x; @@ -941,7 +996,11 @@ void iqn_ema_kernel( int n, int shared_h1, int hidden_dim, - int embed_dim + int embed_dim, + int b0_size, + int b1_size, + int b2_size, + int b3_size ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -953,11 +1012,12 @@ void iqn_ema_kernel( /* ── Flat action → branch action decode kernel ───────────────────────── * - * Decodes flat action index ∈ [0, B0×B1×B2) into 3 branch indices. + * Decodes flat action index ∈ [0, b0×b1×b2×b3) into 4 branch indices. + * action = dir*b1*b2*b3 + mag*b2*b3 + order*b3 + urgency * Grid: (ceil(batch_size/256), 1, 1), Block: (256, 1, 1) * * Input: flat_actions[B] -- i32 flat action indices - * Output: branch_actions[B*3] -- i32 [exposure, order, urgency] per sample + * Output: branch_actions[B*4] -- i32 [direction, magnitude, order, urgency] per sample */ extern "C" __global__ void iqn_decode_actions_kernel( @@ -966,7 +1026,11 @@ void iqn_decode_actions_kernel( int batch_size, int shared_h1, /* runtime: unused, for consistent interface */ int hidden_dim, /* runtime: unused, for consistent interface */ - int embed_dim /* runtime: unused, for consistent interface */ + int embed_dim, /* runtime: unused, for consistent interface */ + int b0_size, /* runtime: branch 0 (direction) action count */ + int b1_size, /* runtime: branch 1 (magnitude) action count */ + int b2_size, /* runtime: branch 2 (order) action count */ + int b3_size /* runtime: branch 3 (urgency) action count */ ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -974,17 +1038,21 @@ void iqn_decode_actions_kernel( int flat = flat_actions[i]; /* Bounds check: clamp to valid action range [0, total_actions) */ - int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE; + int max_action = b0_size * b1_size * b2_size * b3_size; if (flat < 0 || flat >= max_action) flat = 0; - int stride = BRANCH_1_SIZE * BRANCH_2_SIZE; - int exposure = flat / stride; - int remainder = flat % stride; - int order = remainder / BRANCH_2_SIZE; - int urgency = remainder % BRANCH_2_SIZE; + int s1 = b1_size * b2_size * b3_size; + int s2 = b2_size * b3_size; + int direction = flat / s1; + int remainder = flat % s1; + int magnitude = remainder / s2; + remainder = remainder % s2; + int order = remainder / b3_size; + int urgency = remainder % b3_size; - branch_actions[i * 3 + 0] = exposure; - branch_actions[i * 3 + 1] = order; - branch_actions[i * 3 + 2] = urgency; + branch_actions[i * 4 + 0] = direction; + branch_actions[i * 4 + 1] = magnitude; + branch_actions[i * 4 + 2] = order; + branch_actions[i * 4 + 3] = urgency; } /* ── GPU τ sampling kernel (Philox-based PRNG) ───────────────────────── @@ -1022,7 +1090,11 @@ void iqn_sample_taus_kernel( int total, int shared_h1, /* runtime: unused, for consistent interface */ int hidden_dim, /* runtime: unused, for consistent interface */ - int embed_dim /* runtime: unused, for consistent interface */ + int embed_dim, /* runtime: unused, for consistent interface */ + int b0_size, /* runtime: unused, for consistent interface */ + int b1_size, /* runtime: unused, for consistent interface */ + int b2_size, /* runtime: unused, for consistent interface */ + int b3_size /* runtime: unused, for consistent interface */ ) { int i = blockIdx.x * blockDim.x + threadIdx.x; diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index fab9891e2..7cbfd8a43 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -50,11 +50,12 @@ extern "C" __global__ void mse_grad_kernel( float inv_batch = 1.0f / (float)batch_size; float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q); - /* Magnitude entropy boost: prevent atom distribution collapse for d==1. - * MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */ - if (d == 1) { + /* Per-branch entropy boost: prevent distribution collapse. + * Magnitude (d==1): 0.005 = 5× base. Order (d==2): 0.003 = 3× base. */ + if (d == 1 || d == 2) { + float ent_weight = (d == 1) ? 0.005f : 0.003f; float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f); - d_combined += 0.005f * (1.0f + lp_approx); + d_combined += ent_weight * (1.0f + lp_approx); } /* Route through dueling: d_value[b,j] += d_combined. diff --git a/docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md b/docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md new file mode 100644 index 000000000..e3bc68704 --- /dev/null +++ b/docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md @@ -0,0 +1,548 @@ +# Action Diversity + IQN 4-Branch Fix Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Fix IQN 3→4 branch kernel mismatch (corrupting 75% of gradient budget), add order-type counterfactual relabeling, per-branch entropy regularization, and position histogram for all 4 branches — raising order diversity from 1/3 to 3/3. + +**Architecture:** Five compounding causes of order branch collapse, fixed bottom-up: (1) IQN kernel uses pre-4-branch 3-branch defines with BRANCH_0_SIZE=5, corrupting 2/3 of gradient signal; (2) no counterfactual experience for order types; (3) no per-branch entropy boost for order; (4) position histogram only tracks dir+mag; (5) Flat-state trains order Q-values on null transitions. All changes are in CUDA kernels + build.rs + Rust config — no CUDA graph structure changes. + +**Tech Stack:** CUDA 12.4 (nvcc cubin), Rust 1.85, cudarc + +--- + +## File Map + +| File | Action | Responsibility | +|------|--------|---------------| +| `crates/ml/build.rs` | Modify | Add `-D` defines for IQN branch sizes | +| `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` | Modify | 3→4 branch: defines, offsets, forward, backward, decode | +| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Modify | Update buffer sizes for 4 branches, fix `total_branch_actions` in save buffers | +| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | 3-cycle counterfactual + order/urgency histogram | +| `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` | Modify | Per-branch entropy boost for order (d==2) | +| `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` | Modify | Per-branch entropy boost for order (d==2) | + +--- + +### Task 1: IQN kernel — runtime branch sizes, eliminate compile-time defines + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:60-99,188-199` + +No build.rs changes needed — branch sizes become runtime kernel parameters like `hidden_dim` and `embed_dim` already are. This prevents the entire class of stale-define bugs (the compile-time BRANCH_0_SIZE=5 was never updated when 4-branch landed). + +- [ ] **Step 1: Remove compile-time BRANCH_*_SIZE defines, add runtime parameters** + +In `iqn_dual_head_kernel.cu`, delete the compile-time branch defines (lines 62-72) and the `IQN_W_B*_SIZE` / `IQN_TOTAL_PARAMS` macros (lines 83-95). Replace with a comment: + +```c +/* Branch sizes are RUNTIME parameters — passed to every kernel. + * Eliminates stale compile-time define bugs (BRANCH_0_SIZE was 5 when 4-branch uses 3). + * Weight layout: W_embed[H,D], b_embed[H], then per-branch W_bd[Bd,H], b_bd[Bd]. */ +``` + +- [ ] **Step 2: Update offset function to use runtime branch sizes** + +Replace `iqn_compute_offsets` (lines 188-199) to accept runtime branch sizes: + +```c +/** Compute runtime weight offsets from hidden_dim, embed_dim, and 4 branch sizes. + * Returns offsets[10]: W_EMBED, B_EMBED, W_B0, B_B0, W_B1, B_B1, W_B2, B_B2, W_B3, B_B3 */ +__device__ __forceinline__ void iqn_compute_offsets( + int hidden_dim, int embed_dim, + int b0, int b1, int b2, int b3, + int offsets[10] +) { + offsets[0] = 0; /* W_EMBED */ + offsets[1] = offsets[0] + hidden_dim * embed_dim; /* B_EMBED */ + offsets[2] = offsets[1] + hidden_dim; /* W_B0 */ + offsets[3] = offsets[2] + b0 * hidden_dim; /* B_B0 */ + offsets[4] = offsets[3] + b0; /* W_B1 */ + offsets[5] = offsets[4] + b1 * hidden_dim; /* B_B1 */ + offsets[6] = offsets[5] + b1; /* W_B2 */ + offsets[7] = offsets[6] + b2 * hidden_dim; /* B_B2 */ + offsets[8] = offsets[7] + b2; /* W_B3 */ + offsets[9] = offsets[8] + b3 * hidden_dim; /* B_B3 */ +} +``` + +- [ ] **Step 3: Add branch size parameters to all 9 kernel signatures** + +Every IQN kernel signature gets 4 new `int` parameters: `int b0_size, int b1_size, int b2_size, int b3_size`. Replace all internal uses of `BRANCH_0_SIZE` etc. with these runtime params. Also compute `int tba = b0_size + b1_size + b2_size + b3_size;` at the top of each kernel (replaces `TOTAL_BRANCH_ACTIONS`). + +For `iqn_forward_loss_kernel`, `iqn_backward_kernel`, `iqn_forward_kernel`: +```c +/* Add after existing params (hidden_dim, embed_dim): */ + int b0_size, int b1_size, int b2_size, int b3_size +``` + +For `iqn_decode_actions_kernel`, `iqn_grad_norm_kernel`, `iqn_adam_kernel`, `iqn_ema_kernel`, `iqn_trunk_forward_kernel`, `iqn_sample_taus_kernel`: add the same 4 params even if unused (consistent interface, like `shared_h1` is already passed to some kernels unused). + +- [ ] **Step 4: Update Rust kernel launch calls in `gpu_iqn_head.rs`** + +Every `.launch_builder()` call for IQN kernels needs 4 additional `.arg()` calls: +```rust +.arg(&(self.config.branch_0_size as i32)) +.arg(&(self.config.branch_1_size as i32)) +.arg(&(self.config.branch_2_size as i32)) +.arg(&(self.config.branch_3_size as i32)) +``` + +- [ ] **Step 5: Verify kernel compiles and all tests pass** + +```bash +SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "iqn_dual_head|error|FAILED" +SQLX_OFFLINE=true cargo check -p ml --lib +``` +Expected: `Compiled iqn_dual_head_kernel.cu -> iqn_dual_head_kernel.cubin` with no errors. + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +git commit -m "fix(iqn): runtime branch sizes — eliminate stale compile-time defines" +``` + +--- + +### Task 2: IQN kernel — 4-branch forward, backward, and decode + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:215-475,488-664,766-988` + +- [ ] **Step 1: Update decode kernel to 4-branch** + +Replace `iqn_decode_actions_kernel` (lines 962-988). Output becomes `[B, 4]` (was `[B, 3]`): + +```c +/* Output: branch_actions[B*4] -- i32 [dir, order, urgency, magnitude] per sample */ +extern "C" __global__ +void iqn_decode_actions_kernel( + const int* __restrict__ flat_actions, + int* __restrict__ branch_actions, + int batch_size, + int shared_h1, + int hidden_dim, + int embed_dim +) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + int flat = flat_actions[i]; + int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE * BRANCH_3_SIZE; + if (flat < 0 || flat >= max_action) flat = 0; + + /* 4-branch decode: action = dir*b1*b2*b3 + mag*b2*b3 + order*b3 + urgency */ + int stride_0 = BRANCH_1_SIZE * BRANCH_2_SIZE * BRANCH_3_SIZE; + int stride_1 = BRANCH_2_SIZE * BRANCH_3_SIZE; + int stride_2 = BRANCH_3_SIZE; + + int dir = flat / stride_0; + int rem = flat % stride_0; + int mag = rem / stride_1; + rem = rem % stride_1; + int order = rem / stride_2; + int urgency = rem % stride_2; + + branch_actions[i * 4 + 0] = dir; + branch_actions[i * 4 + 1] = mag; + branch_actions[i * 4 + 2] = order; + branch_actions[i * 4 + 3] = urgency; +} +``` + +- [ ] **Step 2: Update forward+loss kernel for 4 branches** + +In `iqn_forward_loss_kernel`, add branch 3 to the online forward (after branch 2 block at line ~380), target forward (after line ~437), and sum of taken Q-values. The pattern is identical to branches 0-2 — just add: + +```c +/* Branch 3 (magnitude) — online */ +float q3_taken = 0.0f; +for (int a = 0; a < BRANCH_3_SIZE; a++) { + float partial = 0.0f; + const float* w_row = w_b3 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float acc = b_b3[a] + iqn_block_sum(partial, shmem_reduce); + if (tid == 0) { + save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE + a] = acc; + if (a == a3) q3_taken = acc; + } +} +if (tid == 0) shmem_reduce[0] = q3_taken; +__syncthreads(); +q3_taken = shmem_reduce[0]; +``` + +Update the sum line: `online_q_a[t] = q0_taken + q1_taken + q2_taken + q3_taken;` + +Similarly for the target branch and `target_q_a` save. + +Also update action decode at the top of the kernel: +```c +int a0 = actions[sample * 4 + 0]; // was sample * 3 +int a1 = actions[sample * 4 + 1]; +int a2 = actions[sample * 4 + 2]; +int a3 = actions[sample * 4 + 3]; +``` + +And add weight pointers: +```c +const float* w_b3 = online_params + off[8]; +const float* b_b3 = online_params + off[9]; +const float* tw_b3 = target_params + off[8]; +const float* tb_b3 = target_params + off[9]; +``` + +- [ ] **Step 3: Update backward kernel for 4 branches** + +In `iqn_backward_kernel`, add branch 3 gradient computation after the branch 2 block (after line ~624). Pattern identical: + +```c +/* Branch 3 (magnitude) */ +for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dcomb[h / IQN_BLOCK_SIZE] += w_b3[a3 * hidden_dim + h] * dL_dq; + atomicAdd(&grad_buf[off[8] + a3 * hidden_dim + h], + inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); +} +if (tid == 0) + atomicAdd(&grad_buf[off[9] + a3], inv_batch * dL_dq); +``` + +Update action reads: `int a3 = actions[sample * 4 + 3];` and add `const float* w_b3 = online_params + off[8];` + +Update offset array: `int off[10]; iqn_compute_offsets(hidden_dim, embed_dim, off);` + +- [ ] **Step 4: Update forward-only (inference) kernel for 4 branches** + +The `iqn_forward_kernel` (line 766) also needs branch 3. Same pattern as forward+loss but without loss computation. Add branch 3 Q-value computation and update `q_total += q3`. + +- [ ] **Step 5: Verify kernel compiles** + +```bash +SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "iqn_dual_head|error" +``` + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +git commit -m "fix(iqn): 4-branch forward, backward, decode — was 3-branch corrupting 75% gradient" +``` + +--- + +### Task 3: IQN Rust side — update buffer sizes and decode for 4 branches + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs:113-127,233-293,339-345` + +- [ ] **Step 1: Update `branch_actions` allocation from `B*3` to `B*4`** + +In `GpuIqnHead::new()`, change: +```rust +// Was: alloc_zeros::(b * 3) +let branch_actions = stream.alloc_zeros::(b * 4).map_err(|e| { + MLError::ModelError(format!("IQN alloc branch_actions: {e}")) +})?; +``` + +- [ ] **Step 2: Update `save_q_online` and `save_q_target` buffer sizes** + +These use `tba = config.total_branch_actions()` which already returns `b0+b1+b2+b3=12`. Verify `tba` is used (not hardcoded). Check lines 292-293: +```rust +let save_q_online = alloc_f32(&stream, b * n * tba, "iqn_save_q_online")?; +let save_q_target = alloc_f32(&stream, b * n * tba, "iqn_save_q_target")?; +``` +These should already be correct since `total_branch_actions()` includes b3. + +- [ ] **Step 3: Update IQR buffer size** + +Line 306: `let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;` — already uses `tba=12`. OK. + +- [ ] **Step 4: Verify all IQN-related tests pass** + +```bash +SQLX_OFFLINE=true cargo check -p ml --lib +SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget --nocapture +``` + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +git commit -m "fix(iqn): update Rust buffer sizes for 4-branch decode (B*4 actions)" +``` + +--- + +### Task 4: Order-type counterfactual relabeling (3-cycle) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu:1799-1846` + +- [ ] **Step 1: Replace 2-cycle counterfactual with 3-cycle** + +Replace the `if ((current_t & 1) == 0)` block (lines 1799-1846) with a 3-cycle that adds order-type counterfactuals: + +```c +int cf_cycle = current_t % 3; + +if (cf_cycle == 0) { + /* ── Directional mirror (existing) ── */ + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = -reward; +} else if (cf_cycle == 1) { + /* ── Hindsight magnitude relabeling (existing) ── */ + const float mag_fracs[3] = { 0.50f, 0.707f, 1.00f }; + float actual_mag_frac = mag_fracs[mag_idx < 3 ? mag_idx : 0]; + + if (fabsf(reward) > 0.001f) { + int alt_mag; + int cycle = (current_t / 3) % 3; + if (cycle == 0) alt_mag = (mag_idx > 0) ? mag_idx - 1 : 2; + else if (cycle == 1) alt_mag = (mag_idx < 2) ? mag_idx + 1 : 0; + else alt_mag = 2 - mag_idx; + + float alt_mag_frac = mag_fracs[alt_mag]; + cf_action = dir_idx * b1_size * b2_size * b3_size + + alt_mag * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = reward * (alt_mag_frac / actual_mag_frac); + cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f); + } else { + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = -reward; + } +} else { + /* ── NEW: Order-type counterfactual ── + * Cycle through alternative order types. Compute execution cost + * difference and amplify it to be comparable to directional signal. + * Only meaningful when position is non-zero (order type affects cost). */ + int alt_order = (orig_order + 1 + ((current_t / 3) & 1)) % b2_size; + + if (fabsf(position) > 0.001f) { + /* Compute cost delta between taken and alternative order type */ + float taken_cost = compute_tx_cost( + position - pre_trade_position, raw_close, tx_cost_multiplier, + spread_cost, max_position, orig_order, spread_scale); + float alt_cost = compute_tx_cost( + position - pre_trade_position, raw_close, tx_cost_multiplier, + spread_cost, max_position, alt_order, spread_scale); + float cost_delta = (taken_cost - alt_cost) + / fmaxf(prev_equity, 1.0f); + + cf_action = dir_idx * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + alt_order * b3_size + orig_urgency; + /* Amplify 50×: raw cost delta ~0.0001-0.0005, directional ~0.01-0.1 */ + cf_reward = reward + 50.0f * cost_delta; + cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f); + } else { + /* Flat position: order type irrelevant, fall back to directional mirror */ + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = -reward; + } +} +``` + +Note: `position`, `pre_trade_position`, `raw_close`, `tx_cost_multiplier`, `spread_cost`, `max_position`, `spread_scale` are all already in scope from the experience_env_step kernel's local variables. `compute_tx_cost` is declared in `trade_physics.cuh` which is already included. + +- [ ] **Step 2: Verify kernel compiles** + +```bash +SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "experience_kernels|error" +``` + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/experience_kernels.cu +git commit -m "feat: order-type counterfactual relabeling (3-cycle: dir/mag/order)" +``` + +--- + +### Task 5: Per-branch entropy regularization for order + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu:53-57` +- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu:53-58` + +- [ ] **Step 1: Add order entropy boost to C51 grad kernel** + +Replace lines 53-57 in `c51_grad_kernel.cu`: + +```c +/* Entropy regularization — per-branch coefficients. + * Magnitude (d==1): 5× boost for atom distribution collapse prevention. + * Order (d==2): 3× boost to prevent order type collapse. + * Direction (d==0), urgency (d==3): baseline 1×. */ +if (entropy_coeff > 0.0f) { + float ent_scale; + if (d == 1) ent_scale = 5.0f; + else if (d == 2) ent_scale = 3.0f; + else ent_scale = 1.0f; + float lp_clamped = fmaxf(lp, -10.0f); + d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped); +} +``` + +- [ ] **Step 2: Add order entropy boost to MSE grad kernel** + +Replace lines 53-58 in `mse_grad_kernel.cu`: + +```c +/* Per-branch entropy boost: prevent distribution collapse. + * Magnitude (d==1): 0.005 = 5× base (0.001). Order (d==2): 0.003 = 3× base. */ +if (d == 1 || d == 2) { + float ent_weight = (d == 1) ? 0.005f : 0.003f; + float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f); + d_combined += ent_weight * (1.0f + lp_approx); +} +``` + +- [ ] **Step 3: Verify kernel compiles** + +```bash +SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "c51_grad\|mse_grad|error" +``` + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +git commit -m "feat: per-branch entropy regularization — order gets 3× boost" +``` + +--- + +### Task 6: Position histogram for all 4 branches + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu:1664-1707` + +- [ ] **Step 1: Expand histogram from 6 bins to 12 bins (3×4 branches)** + +Replace the histogram tracking block (lines 1664-1707): + +```c +/* ── #19 Position entropy: track all 4 branches (3 bins each = 12 total) ── */ +if (position_histogram != NULL && position_entropy_weight > 0.0f) { + /* Histogram layout: [N, 12] — 3 bins per branch: dir[0:3], mag[3:6], order[6:9], urgency[9:12] */ + int order_idx = decode_order_4b(action_idx, b2_size, b3_size); + int urgency_idx = decode_urgency_4b(action_idx, b3_size); + if (dir_idx >= 0 && dir_idx < 3) + position_histogram[(long long)i * 12 + dir_idx] += 1.0f; + if (mag_idx >= 0 && mag_idx < 3) + position_histogram[(long long)i * 12 + 3 + mag_idx] += 1.0f; + if (order_idx >= 0 && order_idx < 3) + position_histogram[(long long)i * 12 + 6 + order_idx] += 1.0f; + if (urgency_idx >= 0 && urgency_idx < 3) + position_histogram[(long long)i * 12 + 9 + urgency_idx] += 1.0f; +} +``` + +- [ ] **Step 2: Update entropy bonus computation to cover all 4 branches** + +Replace the done-block entropy bonus (lines 1679-1707): + +```c +/* #19 Position entropy bonus at episode end — all 4 branches */ +if (done && position_histogram != NULL && position_entropy_weight > 0.0f) { + float* hist = position_histogram + (long long)i * 12; + float combined_entropy = 0.0f; + /* Compute entropy over each branch's 3 bins */ + for (int branch = 0; branch < 4; branch++) { + float* branch_hist = hist + branch * 3; + float total = 0.0f; + for (int b = 0; b < 3; b++) total += branch_hist[b]; + float ent = 0.0f; + if (total > 1.0f) { + for (int b = 0; b < 3; b++) { + float p = branch_hist[b] / total; + if (p > 1e-6f) ent -= p * logf(p); + } + ent /= 1.0986f; /* Normalize by log(3) → [0, 1] */ + } + combined_entropy += ent; + } + /* Average across all 4 branches (was 2) */ + combined_entropy *= 0.25f; + reward += position_entropy_weight * combined_entropy; +} +``` + +- [ ] **Step 3: Update position_histogram allocation in Rust** + +Search for `position_histogram` allocation in `gpu_experience_collector.rs` and change from `alloc_episodes * 6` to `alloc_episodes * 12`: + +```rust +// Was: alloc_episodes * 6 +let position_histogram = stream.alloc_zeros::(alloc_episodes * 12) + .map_err(|e| MLError::ModelError(format!("alloc position_histogram: {e}")))?; +``` + +- [ ] **Step 4: Verify compilation and run tests** + +```bash +SQLX_OFFLINE=true cargo check -p ml --lib +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget config monitoring +``` + +- [ ] **Step 5: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/experience_kernels.cu crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +git commit -m "feat: position histogram for all 4 branches — order+urgency get entropy bonus" +``` + +--- + +### Task 7: Full integration test — smoke test + compute-sanitizer + +**Files:** (no changes — verification only) + +- [ ] **Step 1: Run full smoke test** + +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::feature_coverage --ignored --test-threads=1 --nocapture 2>&1 | grep -E "Action diversity|order|entropy|LOW ORDER" +``` + +Expected: `Action diversity` should show order > 1/3. `LOW ORDER DIVERSITY` warnings should disappear or reduce. + +- [ ] **Step 2: Run compute-sanitizer** + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib --no-run 2>&1 | tail -1 +FOXHUNT_TEST_DATA=test_data/futures-baseline compute-sanitizer --tool memcheck target/debug/deps/ml-*.test "smoke_tests::feature_coverage" --ignored --test-threads=1 +``` + +Expected: `ERROR SUMMARY: 0 errors` + +- [ ] **Step 3: Run full test suite** + +```bash +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +SQLX_OFFLINE=true cargo test -p ml --lib +FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests --ignored --test-threads=1 +``` + +Expected: All pass (286 + 903 + 19). + +- [ ] **Step 4: Final commit with all changes** + +```bash +git add -A +git commit -m "fix: IQN 4-branch + order counterfactual + per-branch entropy — action diversity restored" +```