From 278eb152c0bbbed724a3368cf05ddb1b93a91f96 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 09:26:26 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20TLOB=20microstructure=20injection=20?= =?UTF-8?q?=E2=80=94=20OFI=20features=20to=20order/urgency=20branches=20(3?= =?UTF-8?q?,072=20params)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Order branch gets [bid_ask_spread, depth_imbalance, queue_pressure]. Urgency branch gets [spread_velocity, depth_change_rate, trade_arrival_rate]. Same concat pattern as Layer 3 direction→magnitude conditioning. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/batched_backward.rs | 66 +++++++- .../ml/src/cuda_pipeline/batched_forward.rs | 82 +++++++-- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 155 +++++++++++++++++- 3 files changed, 281 insertions(+), 22 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 94f2ca3f8..c74c9f87c 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -800,6 +800,12 @@ impl CublasBackwardSet { // Magnitude branch conditioning: saved forward concat and dX output mag_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW) d_mag_concat_ptr: u64, // [B, SH2+3] dX output for branch 1 (caller accumulates) + // Order branch OFI conditioning: saved forward concat and dX output + ord_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW) + d_ord_concat_ptr: u64, // [B, SH2+3] dX output for branch 2 (caller accumulates) + // Urgency branch OFI conditioning: saved forward concat and dX output + urg_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW) + d_urg_concat_ptr: u64, // [B, SH2+3] dX output for branch 3 (caller accumulates) // GLU saved forward activations for backward pass glu_gate_pre: &[u64; 4], // [B, AH] each — saved pre-sigmoid gate activations glu_value: &[u64; 4], // [B, AH] each — saved value path outputs @@ -859,11 +865,11 @@ impl CublasBackwardSet { let goff_w_b1out: u64 = goff_b_b1fc + ah * f32; let goff_b_b1out: u64 = goff_w_b1out + b1 * na64 * ah * f32; let goff_w_b2fc: u64 = goff_b_b1out + b1 * na64 * f32; - let goff_b_b2fc: u64 = goff_w_b2fc + ah * sh2 * f32; + let goff_b_b2fc: u64 = goff_w_b2fc + ah * (sh2 + 3) * f32; // order uses wider input: SH2+3 let goff_w_b2out: u64 = goff_b_b2fc + ah * f32; let goff_b_b2out: u64 = goff_w_b2out + b2 * na64 * ah * f32; let goff_w_b3fc: u64 = goff_b_b2out + b2 * na64 * f32; - let goff_b_b3fc: u64 = goff_w_b3fc + ah * sh2 * f32; + let goff_b_b3fc: u64 = goff_w_b3fc + ah * (sh2 + 3) * f32; // urgency uses wider input: SH2+3 let goff_w_b3out: u64 = goff_b_b3fc + ah * f32; let goff_b_b3out: u64 = goff_w_b3out + b3 * na64 * ah * f32; @@ -974,9 +980,13 @@ impl CublasBackwardSet { // ── Step 3: Value path weight gradients ──────────────────── // dW_bdf += d_value^T @ vsn_input, db_bdf += sum(d_value) - // Branch 1 (magnitude) uses wider concat input [B, SH2+3] for dW. + // Branches 1,2,3 use wider concat input [B, SH2+3] for dW. let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 { (mag_concat_ptr, self.shared_h2 + 3) + } else if d == 2 && ord_concat_ptr != 0 { + (ord_concat_ptr, self.shared_h2 + 3) + } else if d == 3 && urg_concat_ptr != 0 { + (urg_concat_ptr, self.shared_h2 + 3) } else { (save_h_s2, self.shared_h2) }; @@ -1033,6 +1043,56 @@ impl CublasBackwardSet { )?; // Caller will call accumulate_d_h_s2_from_concat(d_mag_concat, d_h_s2, B, 1.0) // after backward_full returns. + } else if d == 2 && ord_concat_ptr != 0 { + // Order: dX writes to d_ord_concat [B, SH2+3]. + // Value path: beta=0 (overwrite) + self.launch_dx_only( + stream, + scratch_d_glu_value, + w_fc, + d_ord_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 0.0_f32, + )?; + // Gate path: beta=1 (accumulate) + self.launch_dx_only( + stream, + scratch_d_glu_gate, + w_gate, + d_ord_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, + )?; + // Caller will call accumulate_d_h_s2_from_concat(d_ord_concat, d_h_s2, B, 1.0) + } else if d == 3 && urg_concat_ptr != 0 { + // Urgency: dX writes to d_urg_concat [B, SH2+3]. + // Value path: beta=0 (overwrite) + self.launch_dx_only( + stream, + scratch_d_glu_value, + w_fc, + d_urg_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 0.0_f32, + )?; + // Gate path: beta=1 (accumulate) + self.launch_dx_only( + stream, + scratch_d_glu_gate, + w_gate, + d_urg_concat_ptr, + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, + )?; + // Caller will call accumulate_d_h_s2_from_concat(d_urg_concat, d_h_s2, B, 1.0) } else { // Other branches: accumulate into d_h_s2 directly. let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 }; diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 719de730b..216af12ec 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -189,6 +189,10 @@ pub struct CublasGemmSet { glu_value_ptrs: [u64; 4], /// VSN bottleneck rank (R=16). vsn_rank: usize, + /// Order branch (d=2) OFI concat buffer ptr [B, SH2+3] (0 if not wired). + ord_concat_ptr: u64, + /// Urgency branch (d=3) OFI concat buffer ptr [B, SH2+3] (0 if not wired). + urg_concat_ptr: u64, } /// Backward-compatible type alias — existing callers (`GpuDqnTrainer`, @@ -349,6 +353,8 @@ impl CublasGemmSet { glu_gate_pre_ptrs: [0; 4], glu_value_ptrs: [0; 4], vsn_rank: 0, + ord_concat_ptr: 0, + urg_concat_ptr: 0, }) } @@ -381,6 +387,15 @@ impl CublasGemmSet { self.vsn_rank = vsn_rank; } + /// Wire OFI concat buffer pointers for order (d=2) and urgency (d=3) branches. + /// + /// Called once after construction from `GpuDqnTrainer`, after `wire_vsn_glu`. + /// These pointers point to [B, SH2+3] buffers pre-built by `launch_concat_ofi`. + pub fn wire_ofi_concat(&mut self, ord_concat_ptr: u64, urg_concat_ptr: u64) { + self.ord_concat_ptr = ord_concat_ptr; + self.urg_concat_ptr = urg_concat_ptr; + } + /// Rebind the cuBLAS handle to a (potentially new) stream. /// // Accessors for diagnostic test @@ -985,22 +1000,23 @@ impl CublasGemmSet { )))?; } - // ── 2. Rebuild mag_concat from VSN-masked output for d==1 ── - // The caller's launch_mag_concat used raw h_s2 (pre-VSN). Now that VSN - // has produced vsn_masked, re-run mag_concat_qdir with vsn_masked as - // source so magnitude sees feature-selected [vsn_masked; Q_dir]. - // The Q_dir values (last 3 cols) are recomputed from the same logits. + // ── 2. Rebuild concat from VSN-masked output ── + // + // d==1 (magnitude): Scatter vsn_masked → first SH2 cols of mag_concat [B, SH2+3]. + // Q_dir (last 3 cols) were written by the pre-VSN launch_mag_concat_from call. + // + // d==2 (order) / d==3 (urgency): Scatter vsn_masked → first SH2 cols of + // ord_concat / urg_concat [B, SH2+3]. + // OFI (last 3 cols) were written by the pre-VSN launch_concat_ofi call. + // // This is one kernel launch (B threads, negligible cost). + let scatter_k = self.strided_scatter_kernel.as_ref().unwrap(); + let total_scatter = (b * sh2) as i32; + let src_stride_scatter = sh2 as i32; + let dst_stride_scatter = self.mag_concat_dim as i32; + let scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1); + let (vsn_input, fc_k) = if d == 1 && mag_concat_ptr != 0 { - // Re-launch mag_concat with vsn_masked as source instead of raw h_s2 - // The mag_concat_kernel reads source[B, SH2] + computes Q_dir from logits - // Scatter vsn_masked [B, SH2] → first SH2 cols of mag_concat [B, SH2+3]. - // Q_dir (last 3 cols) preserved from the pre-VSN launch_mag_concat_from call. - let scatter_k = self.strided_scatter_kernel.as_ref().unwrap(); - let total_scatter = (b * sh2) as i32; - let src_stride_scatter = sh2 as i32; - let dst_stride_scatter = self.mag_concat_dim as i32; - let scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1); unsafe { stream .launch_builder(scatter_k) @@ -1019,6 +1035,44 @@ impl CublasGemmSet { )))?; } (mag_concat_ptr, self.mag_concat_dim) + } else if d == 2 && self.ord_concat_ptr != 0 { + unsafe { + stream + .launch_builder(scatter_k) + .arg(&self.vsn_masked_ptr) + .arg(&self.ord_concat_ptr) + .arg(&src_stride_scatter) + .arg(&dst_stride_scatter) + .arg(&total_scatter) + .launch(LaunchConfig { + grid_dim: (scatter_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "strided_scatter vsn→ord_concat: {e}" + )))?; + } + (self.ord_concat_ptr, self.mag_concat_dim) + } else if d == 3 && self.urg_concat_ptr != 0 { + unsafe { + stream + .launch_builder(scatter_k) + .arg(&self.vsn_masked_ptr) + .arg(&self.urg_concat_ptr) + .arg(&src_stride_scatter) + .arg(&dst_stride_scatter) + .arg(&total_scatter) + .launch(LaunchConfig { + grid_dim: (scatter_blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "strided_scatter vsn→urg_concat: {e}" + )))?; + } + (self.urg_concat_ptr, self.mag_concat_dim) } else { (self.vsn_masked_ptr, sh2) }; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ab139043e..43c84b63a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -630,6 +630,16 @@ pub struct GpuDqnTrainer { strided_accumulate_kernel: CudaFunction, /// Kernel: strided_scatter — writes vsn_masked into first SH2 cols of mag_concat. strided_scatter_kernel: CudaFunction, + /// Order branch concat input [B, SH2+3]: [vsn_masked; OFI_order] for microstructure conditioning. + ord_concat_buf: CudaSlice, + /// Backward scratch for order concat dX [B, SH2+3]. + d_ord_concat_buf: CudaSlice, + /// Urgency branch concat input [B, SH2+3]: [vsn_masked; OFI_urgency] for microstructure conditioning. + urg_concat_buf: CudaSlice, + /// Backward scratch for urgency concat dX [B, SH2+3]. + d_urg_concat_buf: CudaSlice, + /// Kernel: concat_ofi_features — builds [vsn_masked; OFI] for order/urgency branches. + concat_ofi_kernel: CudaFunction, // ── Cross-Branch Q Attention ── q_attn_params: CudaSlice, // [624] @@ -1324,6 +1334,44 @@ impl GpuDqnTrainer { Ok(()) } + /// Build OFI concat for order (d=2) or urgency (d=3) branch. + /// + /// Reads vsn_masked[B, SH2] and raw feature vector states[B, SD], appends 3 OFI + /// features at ofi_base offset, producing [B, SH2+3] in ord_concat_buf / urg_concat_buf. + /// + /// For order (branch_idx=2): OFI features at indices 42, 43, 44 + /// For urgency (branch_idx=3): OFI features at indices 45, 46, 47 + pub(crate) fn launch_concat_ofi(&self, branch_idx: usize, states_ptr: u64) -> Result<(), MLError> { + let (concat_ptr, ofi_base) = match branch_idx { + 2 => (self.ptrs.ord_concat_buf, 42_i32), + 3 => (self.ptrs.urg_concat_buf, 45_i32), + _ => return Ok(()), + }; + let b = self.config.batch_size as i32; + let sh2 = self.config.shared_h2 as i32; + let sd = self.config.state_dim as i32; + let total = b * (sh2 + 3); + let blocks = ((total as u32 + 255) / 256).max(1); + let vsn_masked_ptr = self.vsn_masked_buf.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.concat_ofi_kernel) + .arg(&concat_ptr) + .arg(&vsn_masked_ptr) + .arg(&states_ptr) + .arg(&b) + .arg(&sh2) + .arg(&sd) + .arg(&ofi_base) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("concat_ofi branch {branch_idx}: {e}")))?; + } + Ok(()) + } /// Run 2-head cross-branch Q attention over the 12 raw Q-values in q_out_buf. /// Writes coordinated Q-values to q_coord_buf [batch_size, 12]. @@ -2349,6 +2397,10 @@ impl GpuDqnTrainer { 0, // CQL backward: no bottleneck dX needed (separate gradient budget) self.ptrs.mag_concat_buf, self.ptrs.d_mag_concat_buf, + self.ptrs.ord_concat_buf, + self.d_ord_concat_buf.raw_ptr(), + self.ptrs.urg_concat_buf, + self.d_urg_concat_buf.raw_ptr(), &glu_gate_pre_ptrs, &glu_value_ptrs, self.ptrs.bw_d_glu_value, @@ -2364,6 +2416,22 @@ impl GpuDqnTrainer { 1.0, // beta=1: d==0 already wrote to scratch_d_h_s2 )?; } + if self.ptrs.ord_concat_buf != 0 { + self.accumulate_d_h_s2_from_concat( + self.d_ord_concat_buf.raw_ptr(), + scratch_d_h_s2, + self.config.batch_size, + 1.0, + )?; + } + if self.ptrs.urg_concat_buf != 0 { + self.accumulate_d_h_s2_from_concat( + self.d_urg_concat_buf.raw_ptr(), + scratch_d_h_s2, + self.config.batch_size, + 1.0, + )?; + } Ok(true) } @@ -2743,6 +2811,10 @@ impl GpuDqnTrainer { let mag_concat_dim = config.shared_h2 + 3; let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?; let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?; + let ord_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "ord_concat_buf")?; + let d_ord_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_ord_concat_buf")?; + let urg_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "urg_concat_buf")?; + let d_urg_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_urg_concat_buf")?; let save_current_lp = alloc_f32( &stream, b * num_branches * config.num_atoms, @@ -2957,7 +3029,9 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("strided_accumulate load: {e}")))?; let strided_scatter_kernel = exp_module_for_mag.load_function("strided_scatter") .map_err(|e| MLError::ModelError(format!("strided_scatter load: {e}")))?; - info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter kernels loaded"); + let concat_ofi_kernel = exp_module_for_mag.load_function("concat_ofi_features") + .map_err(|e| MLError::ModelError(format!("concat_ofi_features load: {e}")))?; + info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi kernels loaded"); // ── Cross-Branch Q Attention buffers ────────────────────────── let q_attn_params = stream.alloc_zeros::(624) @@ -3281,7 +3355,11 @@ impl GpuDqnTrainer { ], 16, // VSN bottleneck rank R=16 ); - info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+KAN wired)"); + cublas_forward.wire_ofi_concat( + ord_concat_buf.raw_ptr(), + urg_concat_buf.raw_ptr(), + ); + info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+KAN+OFI wired)"); // DDQN forward shares the same cuBLAS handle — runs sequentially // on the main stream for NVIDIA bit-wise reproducibility. @@ -3317,7 +3395,11 @@ impl GpuDqnTrainer { ], 16, ); - info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+KAN wired)"); + cublas_forward_ddqn.wire_ofi_concat( + ord_concat_buf.raw_ptr(), + urg_concat_buf.raw_ptr(), + ); + info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+KAN+OFI wired)"); // ── Initialize cuBLAS backward context (required) ────────── let cublas_backward = CublasBackwardSet::new( @@ -3414,6 +3496,8 @@ impl GpuDqnTrainer { save_h_b3: save_h_b3.raw_ptr(), mag_concat_buf: mag_concat_buf.raw_ptr(), d_mag_concat_buf: d_mag_concat_buf.raw_ptr(), + ord_concat_buf: ord_concat_buf.raw_ptr(), + urg_concat_buf: urg_concat_buf.raw_ptr(), bw_d_h_s1: bw_d_h_s1.raw_ptr(), bw_d_h_s2: bw_d_h_s2.raw_ptr(), bw_d_h_v: bw_d_h_v.raw_ptr(), @@ -3653,6 +3737,11 @@ impl GpuDqnTrainer { mag_concat_kernel, strided_accumulate_kernel, strided_scatter_kernel, + ord_concat_buf, + d_ord_concat_buf, + urg_concat_buf, + d_urg_concat_buf, + concat_ofi_kernel, q_attn_params, q_attn_adam_m, q_attn_adam_v, @@ -4883,6 +4972,8 @@ impl GpuDqnTrainer { let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); // Build mag_concat from previous step's logits (one-step lag) self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?; + self.launch_concat_ofi(2, self.ptrs.states_buf)?; + self.launch_concat_ofi(3, self.ptrs.states_buf)?; self.cublas_forward.forward_online_raw( &self.stream, self.ptrs.states_buf, &on_w_ptrs, self.ptrs.save_h_s1, self.ptrs.save_h_s2, self.ptrs.save_h_v, @@ -4946,6 +5037,8 @@ impl GpuDqnTrainer { let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); // Build mag_concat from previous step's logits (zeros at step 0) self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?; + self.launch_concat_ofi(2, self.ptrs.states_buf)?; + self.launch_concat_ofi(3, self.ptrs.states_buf)?; self.cublas_forward.forward_online_raw( &self.stream, self.ptrs.states_buf, &on_w_ptrs, self.ptrs.save_h_s1, self.ptrs.save_h_s2, self.ptrs.save_h_v, @@ -5846,6 +5939,11 @@ impl GpuDqnTrainer { // This call is captured in the CUDA graph — on replay it reads the logits // from the previous graph execution (still in the buffer). self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?; + // ── Build OFI concat for order (d=2) and urgency (d=3) branches (one-step lag). + // Uses prev-step vsn_masked (one-step lag, same as mag_concat). The vsn_masked + // portion is updated inside launch_vsn_glu_branch after VSN runs. + self.launch_concat_ofi(2, self.ptrs.states_buf)?; + self.launch_concat_ofi(3, self.ptrs.states_buf)?; // ── Pass 1: Online forward on STATES (graph-safe: CachedPtrs, no device_ptr) // When bottleneck is active, s1_input_ptr = bn_concat (compressed features). @@ -5899,6 +5997,10 @@ impl GpuDqnTrainer { // Build target mag_concat from previous step's target logits (one-step lag). // Reuses the same mag_concat_buf — target forward runs after online forward. self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.tg_v_logits_buf, self.ptrs.tg_b_logits_buf)?; + // OFI concat for target forward (branches 2 and 3) — reuses ord/urg concat buffers + // with next_states features (target forward sees next_states). + self.launch_concat_ofi(2, self.ptrs.next_states_buf)?; + self.launch_concat_ofi(3, self.ptrs.next_states_buf)?; cublas.forward_target_raw( &self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, self.ptrs.tg_h_s1_scratch, self.ptrs.tg_h_s2_buf, @@ -6485,6 +6587,10 @@ impl GpuDqnTrainer { s1_dx_output, self.ptrs.mag_concat_buf, self.ptrs.d_mag_concat_buf, + self.ptrs.ord_concat_buf, + self.d_ord_concat_buf.raw_ptr(), + self.ptrs.urg_concat_buf, + self.d_urg_concat_buf.raw_ptr(), &glu_gate_pre_ptrs, &glu_value_ptrs, self.ptrs.bw_d_glu_value, @@ -6500,6 +6606,24 @@ impl GpuDqnTrainer { 1.0, // beta=1: d==0 already wrote to d_h_s2 in backward_full )?; } + // Accumulate order branch dX (first SH2 columns) into d_h_s2 + if self.ptrs.ord_concat_buf != 0 { + self.accumulate_d_h_s2_from_concat( + self.d_ord_concat_buf.raw_ptr(), + d_h_s2_ptr, + self.config.batch_size, + 1.0, + )?; + } + // Accumulate urgency branch dX (first SH2 columns) into d_h_s2 + if self.ptrs.urg_concat_buf != 0 { + self.accumulate_d_h_s2_from_concat( + self.d_urg_concat_buf.raw_ptr(), + d_h_s2_ptr, + self.config.batch_size, + 1.0, + )?; + } // ── #31 Bottleneck backward: chain rule through tanh + GEMM ── // After backward_full, d_bn_concat_buf has dL/d(bn_concat) [B, concat_dim] f32. @@ -6845,11 +6969,11 @@ impl GpuDqnTrainer { (0, 0), // [13] b_b1fc (cfg.branch_1_size * cfg.num_atoms, cfg.adv_h), // [14] w_b1out (0, 0), // [15] b_b1out - (cfg.adv_h, cfg.shared_h2), // [16] w_b2fc + (cfg.adv_h, cfg.shared_h2 + 3), // [16] w_b2fc (OFI-conditioned: SH2+3 input) (0, 0), // [17] b_b2fc (cfg.branch_2_size * cfg.num_atoms, cfg.adv_h), // [18] w_b2out (0, 0), // [19] b_b2out - (cfg.adv_h, cfg.shared_h2), // [20] w_b3fc + (cfg.adv_h, cfg.shared_h2 + 3), // [20] w_b3fc (OFI-conditioned: SH2+3 input) (0, 0), // [21] b_b3fc (cfg.branch_3_size * cfg.num_atoms, cfg.adv_h), // [22] w_b3out (0, 0), // [23] b_b3out @@ -6906,6 +7030,27 @@ impl GpuDqnTrainer { offset += align4(n); } + // ── Zero the last 3 OFI columns of w_b2fc [16] and w_b3fc [20] ── + // xavier_init fills the full tensor (adv_h × (shared_h2+3)) but the OFI + // columns should start at zero so the network learns to use them gradually. + // (gate weights w_gate_2/w_gate_3 are already zero via (0,0) fan_dims.) + for &fc_idx in &[16_usize, 20_usize] { + let base: usize = sizes[..fc_idx].iter().map(|&s| align4(s)).sum(); + let sh2 = cfg.shared_h2; + let sh2p3 = sh2 + 3; + let ah = cfg.adv_h; + // w_b2fc/w_b3fc layout: [AH, SH2+3] row-major. + // Zero columns [SH2, SH2+1, SH2+2] for all AH rows. + for row in 0..ah { + for col in sh2..sh2p3 { + let idx = base + row * sh2p3 + col; + if idx < host_buf.len() { + host_buf[idx] = 0.0_f32; + } + } + } + } + // ── KAN spline coefficient init: uniform 0.125 (flat 0.5 gate ≈ sigmoid(0)) ── // Indices 42, 44, 46, 48 are kan_coeff_d [AH, 8]. Fill with 1/8 so the // spline sums to 0.5 regardless of input (mimics sigmoid(0) = 0.5).