diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index cce082eb5..6877e0c10 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -860,6 +860,127 @@ extern "C" __global__ void barrier_gradient_direction( } } +/* ══════════════════════════════════════════════════════════════════════ + * F4/D5: INFORMATION BOTTLENECK VARIANCE GRADIENT KERNEL + * + * Per sample i, for the direction branch (b0): + * Q(a) = Σ_z p(z|a) * z_vals[z] via softmax over dueling logits. + * mean_q = (1/b0) * Σ_a Q(a) + * var_q = (1/b0) * Σ_a (Q(a) - mean_q)^2 (population variance) + * ib_weight = 0.05 * (1 - health) + * penalty = max(0, min_var - var_q) + * + * When var_q < min_var, adds to cql_d_adv_logits / cql_d_value_logits to + * spread Q values across actions (increase var_q → decrease ib_loss). + * + * Gradient derivation: + * ib_loss = ib_weight * max(0, min_var - var_q) + * d(ib_loss)/dQ(a) = -ib_weight * (2/b0) * (Q(a) - mean_q) when var_q < min_var + * + * Gradient descent: params -= lr * grad. + * So for Q(a) > mean_q: negative dib/dQ(a) means we subtract a negative value + * from Q(a)'s gradient, i.e. we INCREASE Q(a). Symmetrically for Q(a) < mean_q. + * Net effect: Q values are pushed AWAY from mean_q — spreading them. + * + * Then: d(ib_loss)/d_logit(a,z) = d(ib_loss)/dQ(a) * dQ(a)/d_logit(a,z) + * where dQ(a)/d_logit(a,z) = (z_vals[z] - Q(a)) * p(z|a) + * + * Piggybacked on the CQL d-logit accumulators (same as F5). No new cuBLAS pass. + * Direction branch only. Zero-op when var_q >= min_var or ib_weight < 1e-6. + * + * Launch: grid=(ceil(B/256)), block=(256). One thread per sample. + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void ib_gradient_direction( + const float* __restrict__ adv_logits, /* [B, b0_size, num_atoms] f32 */ + const float* __restrict__ v_logits, /* [B, num_atoms] f32 */ + const float* __restrict__ z_vals, /* [num_atoms] f32 */ + const float* __restrict__ isv_signals, /* [ISV_DIM=13] pinned; [12]=health */ + float* __restrict__ d_adv_logits, /* [B, b0_size, num_atoms] — atomicAdd */ + float* __restrict__ d_v_logits, /* [B, num_atoms] — atomicAdd */ + int batch_size, + int num_atoms, + int b0_size, + float min_var +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + /* Safety guard */ + if (b0_size <= 1 || b0_size > 16) return; + + float health = (isv_signals != NULL) ? isv_signals[12] : 0.5f; + float ib_weight = 0.05f * (1.0f - health); + if (ib_weight < 1e-6f) return; + + const float* v_row = v_logits + (long long)i * num_atoms; + const float* adv_row = adv_logits + (long long)i * b0_size * num_atoms; + + /* Compute Q(a) for each direction action via numerically-stable softmax. */ + float q_vals[16]; + for (int a = 0; a < b0_size; a++) { + const float* adv_a = adv_row + a * num_atoms; + float max_l = -INFINITY; + for (int z = 0; z < num_atoms; z++) { + float l = v_row[z] + adv_a[z]; + if (l > max_l) max_l = l; + } + float sum = 0.0f; + for (int z = 0; z < num_atoms; z++) + sum += expf(v_row[z] + adv_a[z] - max_l); + float q = 0.0f; + for (int z = 0; z < num_atoms; z++) { + float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f); + q += p * z_vals[z]; + } + q_vals[a] = q; + } + + /* Population variance across actions. */ + float mean_q = 0.0f; + for (int a = 0; a < b0_size; a++) mean_q += q_vals[a]; + mean_q /= (float)b0_size; + float var_q = 0.0f; + for (int a = 0; a < b0_size; a++) { + float d = q_vals[a] - mean_q; + var_q += d * d; + } + var_q /= (float)b0_size; + + /* Zero-op when variance already meets the minimum. */ + if (var_q >= min_var) return; + + /* grad_scale = -ib_weight * (2 / b0_size) + * Gradient descent subtracts this times (Q(a)-mean_q), spreading Q values. */ + float grad_scale = -ib_weight * (2.0f / (float)b0_size); + + float* d_v_row = d_v_logits + (long long)i * num_atoms; + + for (int a = 0; a < b0_size; a++) { + float dq_a = grad_scale * (q_vals[a] - mean_q); + if (fabsf(dq_a) < 1e-9f) continue; + + /* Recompute softmax probabilities for this action to get p(z|a). */ + const float* adv_a = adv_row + a * num_atoms; + float max_l = -INFINITY; + for (int z = 0; z < num_atoms; z++) { + float l = v_row[z] + adv_a[z]; + if (l > max_l) max_l = l; + } + float sum = 0.0f; + for (int z = 0; z < num_atoms; z++) + sum += expf(v_row[z] + adv_a[z] - max_l); + + float* d_adv_a = d_adv_logits + (long long)i * b0_size * num_atoms + a * num_atoms; + float q_a = q_vals[a]; + for (int z = 0; z < num_atoms; z++) { + float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f); + float contribution = dq_a * (z_vals[z] - q_a) * p; + atomicAdd(&d_adv_a[z], contribution); + atomicAdd(&d_v_row[z], contribution); + } + } +} + /* ══════════════════════════════════════════════════════════════════════ * C51 MANIFOLD MIXUP CE KERNEL (separate launch, no barrier) * diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 58a5e556b..0b64038ec 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1332,6 +1332,12 @@ pub struct GpuDqnTrainer { /// Loaded from c51_loss_kernel.cubin ("barrier_gradient_direction"). barrier_gradient_kernel: Option, + // ── F4/D5: Information Bottleneck variance gradient kernel ─────── + /// Spreads Q values across direction-branch actions when population + /// variance is below min_var (0.01). Piggybacks on the CQL d-logit + /// buffers. Loaded from c51_loss_kernel.cubin ("ib_gradient_direction"). + ib_gradient_kernel: Option, + /// v8: PopArt reward normalization kernel (running mean/variance). popart_normalize_kernel: CudaFunction, /// v8: Robust PopArt kernel (median/IQR normalization). @@ -4708,6 +4714,43 @@ impl GpuDqnTrainer { } } + // F4/D5: IB variance gradient — spreads Q values in the direction branch. + // Runs after F5 barrier (both atomicAdd into same buffers, races are fine). + // Zero-op when var_q >= 0.01 or health ≈ 1 (ib_weight ≈ 0). + if let Some(ref ib_kernel) = self.ib_gradient_kernel.clone() { + let adv_ptr_ib = self.on_b_logits_buf.raw_ptr(); + let v_ptr_ib = self.on_v_logits_buf.raw_ptr(); + let z_ptr_ib = self.atom_positions_buf.raw_ptr(); + let isv_dev_ib = self.isv_signals_dev_ptr; + let d_adv_ptr_ib = self.cql_d_adv_logits.raw_ptr(); + let d_v_ptr_ib = self.cql_d_value_logits.raw_ptr(); + let b_i32_ib = b as i32; + let na_i32_ib = na as i32; + let b0_i32_ib = b0 as i32; + let blocks_ib = blocks; + let min_var = 0.01_f32; + unsafe { + let _ = self.stream + .launch_builder(ib_kernel) + .arg(&adv_ptr_ib) + .arg(&v_ptr_ib) + .arg(&z_ptr_ib) + .arg(&isv_dev_ib) + .arg(&d_adv_ptr_ib) + .arg(&d_v_ptr_ib) + .arg(&b_i32_ib) + .arg(&na_i32_ib) + .arg(&b0_i32_ib) + .arg(&min_var) + .launch(LaunchConfig { + grid_dim: (blocks_ib, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }); + // Non-fatal: log failure but continue with CQL backward. + } + } + // Step 2: SAXPY — add CQL logit gradients to d_value_logits_buf and d_adv_logits_buf // These are the same buffers that already hold C51's logit gradients. // After this, when we run cuBLAS backward, the backward pass sees @@ -5593,6 +5636,8 @@ impl GpuDqnTrainer { let c51_grad_kernel = compile_c51_grad_kernel(&stream, &config)?; // F5/D2: Q-gap barrier gradient kernel (from same c51_loss cubin). let barrier_gradient_kernel = load_barrier_gradient_kernel(&stream)?; + // F4/D5: Information Bottleneck variance gradient kernel (from same c51_loss cubin). + let ib_gradient_kernel = load_ib_gradient_kernel(&stream)?; info!("GpuDqnTrainer: c51_loss + c51_grad kernels compiled"); // ── Compile MSE loss + gradient kernels (warmup before C51) ─ @@ -7353,6 +7398,7 @@ impl GpuDqnTrainer { cql_d_value_logits, cql_d_adv_logits, barrier_gradient_kernel, + ib_gradient_kernel, curiosity_error_buf, drawdown_depths_buf, asymmetric_dd_weight: dd_weight, @@ -13002,6 +13048,24 @@ fn load_barrier_gradient_kernel( } } +/// Load the F4/D5 Information Bottleneck variance gradient kernel from the c51_loss cubin. +/// +/// Returns Ok(None) gracefully if the symbol is not found (older cubin cached). +fn load_ib_gradient_kernel( + stream: &Arc, +) -> Result, MLError> { + let context = stream.context(); + let module = context.load_cubin(C51_LOSS_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("c51_loss cubin load (ib): {e}")))?; + match module.load_function("ib_gradient_direction") { + Ok(f) => Ok(Some(f)), + Err(e) => { + tracing::warn!("ib_gradient_direction not found in cubin (non-fatal): {e}"); + Ok(None) + } + } +} + /// Load the C51 loss gradient kernel from precompiled cubin. /// /// entropy_coeff is now passed as a runtime kernel parameter (not #define).