diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 03bbc3cf1..00f041628 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -162,6 +162,11 @@ pub struct GpuDqnTrainConfig { /// Number of epochs to use MSE loss on expected Q-values before switching to C51. /// MSE provides stronger gradient signal during early training. pub c51_warmup_epochs: usize, + /// Enable Conservative Q-Learning (CQL) penalty (Kumar et al. 2020). + /// Adds `alpha * (logsumexp(Q(s,·)) - Q(s, a_data))` to prevent Q-value overestimation. + pub use_cql: bool, + /// CQL regularization strength (0.0 = disabled, 0.1 = mild, 1.0 = full offline-RL). + pub cql_alpha: f32, } impl Default for GpuDqnTrainConfig { @@ -194,6 +199,8 @@ impl Default for GpuDqnTrainConfig { iqn_kappa: 1.0, entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning. c51_warmup_epochs: 5, + use_cql: true, + cql_alpha: 0.1, } } } @@ -529,6 +536,15 @@ pub struct GpuDqnTrainer { q_stats_kernel: CudaFunction, /// GPU buffer for Q-value statistics [5 floats] q_stats_buf: CudaSlice, + + // ── CQL (Conservative Q-Learning) penalty kernel ───────────────── + /// Computes CQL logit gradients: dCQL/d_value_logits and dCQL/d_adv_logits. + /// Only used when `config.use_cql == true && config.cql_alpha > 0`. + cql_logit_grad_kernel: Option, + /// CQL scratch: value logit gradients [B, NA] + cql_d_value_logits: CudaSlice, + /// CQL scratch: advantage logit gradients [B, (B0+B1+B2)*NA] + cql_d_adv_logits: CudaSlice, } impl Drop for GpuDqnTrainer { @@ -833,6 +849,146 @@ impl GpuDqnTrainer { Ok(()) } + /// Apply CQL (Conservative Q-Learning) gradient injection. + /// + /// Computes the CQL penalty gradient w.r.t. the C51 logits and runs a + /// second cuBLAS backward pass to produce parameter gradients, which are + /// accumulated into `grad_buf` (beta=1.0) on top of the C51 gradients. + /// + /// Called between `graph_forward` and `graph_adam` in `FusedTrainingCtx`. + /// Returns `true` if CQL was applied, `false` if skipped (disabled or no kernel). + pub fn apply_cql_gradient( + &mut self, + ) -> Result { + let cql_kernel = match &self.cql_logit_grad_kernel { + Some(k) => k.clone(), + None => return Ok(false), + }; + if self.config.cql_alpha <= 0.0 { + return Ok(false); + } + + let b = self.config.batch_size; + let na = self.config.num_atoms; + let b0 = self.config.branch_0_size; + let b1 = self.config.branch_1_size; + let b2 = self.config.branch_2_size; + let total_actions = b0 + b1 + b2; + let _total_branch_atoms = total_actions * na; + + let _evt_guard = EventTrackingGuard::new(self.stream.context()); + + // Step 1: Compute CQL logit gradients via kernel + let v_logits_ptr = raw_device_ptr(&self.on_v_logits_buf, &self.stream); + let adv_logits_ptr = raw_device_ptr(&self.on_b_logits_buf, &self.stream); + let actions_ptr = raw_device_ptr_i32(&self.actions_buf, &self.stream); + let d_v_ptr = raw_device_ptr(&self.cql_d_value_logits, &self.stream); + let d_adv_ptr = raw_device_ptr(&self.cql_d_adv_logits, &self.stream); + + let cql_alpha = self.config.cql_alpha; + let n_i32 = b as i32; + let na_i32 = na as i32; + let b0_i32 = b0 as i32; + let b1_i32 = b1 as i32; + let b2_i32 = b2 as i32; + let v_min = self.config.v_min; + let v_max = self.config.v_max; + + let blocks = ((b + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&cql_kernel) + .arg(&v_logits_ptr) + .arg(&adv_logits_ptr) + .arg(&actions_ptr) + .arg(&d_v_ptr) + .arg(&d_adv_ptr) + .arg(&cql_alpha) + .arg(&n_i32) + .arg(&na_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&v_min) + .arg(&v_max) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel: {e}")))?; + } + + // 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 + // combined C51+CQL logit gradients and produces combined parameter gradients. + // However, backward_full uses beta=1.0 accumulation, so calling it again + // would ADD to grad_buf (doubling the C51 contribution). + // + // Instead, we SAXPY the CQL logit gradients into the d_logits buffers and + // then rely on the fact that the NEXT graph_forward will use these updated + // logit gradients. But that's for the next step, not this one. + // + // For the current step, the correct approach is to add CQL gradients + // directly to grad_buf using a separate cuBLAS backward pass with + // cql_d_value_logits and cql_d_adv_logits as input. + // This is the same pattern as apply_iqn_trunk_gradient. + + // Extract weight pointers for cuBLAS backward + let param_sizes = compute_param_sizes(&self.config); + let w_ptrs = f32_weight_ptrs(&self.params_buf, ¶m_sizes, &self.stream); + + // Construct d_adv_logits pointers per branch + let f32_size = std::mem::size_of::(); + let d_adv_base = d_adv_ptr; + let d_adv_ptrs = [ + d_adv_base, + d_adv_base + (b0 * na * f32_size) as u64, + d_adv_base + ((b0 + b1) * na * f32_size) as u64, + ]; + + // Saved activations from the forward pass (still valid) + let states_ptr_fw = raw_device_ptr(&self.states_buf, &self.stream); + let h_s1_ptr = raw_device_ptr(&self.save_h_s1, &self.stream); + let h_s2_ptr = raw_device_ptr(&self.save_h_s2, &self.stream); + let h_v_ptr = raw_device_ptr(&self.save_h_v, &self.stream); + let h_b0_ptr = raw_device_ptr(&self.save_h_b0, &self.stream); + let h_b1_ptr = raw_device_ptr(&self.save_h_b1, &self.stream); + let h_b2_ptr = raw_device_ptr(&self.save_h_b2, &self.stream); + + // Scratch buffers for inter-layer gradients (can reuse bw_d_h_* buffers) + let scratch_d_h_s2 = raw_device_ptr(&self.bw_d_h_s2, &self.stream); + let scratch_d_h_s1 = raw_device_ptr(&self.bw_d_h_s1, &self.stream); + let scratch_d_h_v = raw_device_ptr(&self.bw_d_h_v, &self.stream); + let scratch_d_h_b0 = raw_device_ptr(&self.bw_d_h_b0, &self.stream); + let scratch_d_h_b1 = raw_device_ptr(&self.bw_d_h_b1, &self.stream); + let scratch_d_h_b2 = raw_device_ptr(&self.bw_d_h_b2, &self.stream); + + // Run full backward pass with CQL logit gradients. + // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM), + // adding CQL parameter gradients on top of C51's. + self.cublas_backward.backward_full( + &self.stream, + d_v_ptr, + &d_adv_ptrs, + states_ptr_fw, + h_s1_ptr, h_s2_ptr, h_v_ptr, + &[h_b0_ptr, h_b1_ptr, h_b2_ptr], + &w_ptrs, + raw_device_ptr(&self.grad_buf, &self.stream), + scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, + &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], + ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; + + Ok(true) + } + + /// Whether CQL is enabled and the kernel was compiled. + pub fn has_cql(&self) -> bool { + self.cql_logit_grad_kernel.is_some() && self.config.cql_alpha > 0.0 + } + /// Apply spectral normalization to trunk weight matrices W_s1 and W_s2. /// /// One step of power iteration per call (standard practice — single step @@ -1182,6 +1338,24 @@ impl GpuDqnTrainer { let q_stats_buf = alloc_f32(&stream, 5, "q_stats")?; info!("GpuDqnTrainer: expected_q + q_stats kernels compiled"); + // ── Compile CQL penalty kernel (if enabled) ────────────────────── + let cql_logit_grad_kernel = if config.use_cql && config.cql_alpha > 0.0 { + match compile_cql_logit_grad_kernel(&stream) { + Ok(k) => { + info!(cql_alpha = config.cql_alpha, "GpuDqnTrainer: CQL logit gradient kernel compiled"); + Some(k) + } + Err(e) => { + tracing::warn!("CQL kernel compilation failed (non-fatal): {e}"); + None + } + } + } else { + None + }; + let cql_d_value_logits = alloc_f32(&stream, b * config.num_atoms, "cql_d_value_logits")?; + let cql_d_adv_logits = alloc_f32(&stream, b * total_branch_atoms, "cql_d_adv_logits")?; + // ── Gradient output buffers for cuBLAS backward ────────────── let d_value_logits_buf = alloc_f32(&stream, b * config.num_atoms, "d_value_logits")?; let d_adv_logits_buf = alloc_f32(&stream, b * total_branch_atoms, "d_adv_logits")?; @@ -1399,6 +1573,9 @@ impl GpuDqnTrainer { expected_q_kernel, q_stats_kernel, q_stats_buf, + cql_logit_grad_kernel, + cql_d_value_logits, + cql_d_adv_logits, }) } @@ -4130,3 +4307,172 @@ pub(crate) fn launch_bf16_to_f32( } Ok(()) } + +/// Compile the CQL (Conservative Q-Learning) logit gradient kernel. +/// +/// Computes dCQL/d_value_logits and dCQL/d_adv_logits for the Branching Dueling +/// C51 architecture. The CQL penalty per branch is: +/// `penalty = alpha * (logsumexp_a(Q(a)) - Q(a_taken))` +/// +/// The gradient flows through the softmax expectation: +/// `dCQL/d_logit[a,j] = alpha * dCQL_dQ[a] * p[j] * (z[j] - Q[a])` +/// where `dCQL_dQ[a] = softmax_Q[a] - I(a == a_taken)` and `p[j] = softmax(combined_logits)[j]`. +/// +/// One thread per sample. Iterates over the 3 branches (exposure, order, urgency). +fn compile_cql_logit_grad_kernel( + stream: &Arc, +) -> Result { + let src = r#" +extern "C" __global__ void cql_logit_grad_kernel( + const float* __restrict__ v_logits, // [N, num_atoms] + const float* __restrict__ adv_logits, // [N, total_actions * num_atoms] + const int* __restrict__ actions, // [N] factored action indices (0-44) + float* __restrict__ d_v_logits, // [N, num_atoms] output + float* __restrict__ d_adv_logits, // [N, total_actions * num_atoms] output + float cql_alpha, + int N, int num_atoms, + int b0_size, int b1_size, int b2_size, + float v_min, float v_max) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + int total_actions = b0_size + b1_size + b2_size; + float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; + + // Decode factored action: action = a0 * b1 * b2 + a1 * b2 + a2 + int action = actions[i]; + int a2_taken = action % b2_size; + int a1_taken = (action / b2_size) % b1_size; + int a0_taken = action / (b1_size * b2_size); + + // Clamp to valid range (safety) + if (a0_taken >= b0_size) a0_taken = b0_size - 1; + if (a1_taken >= b1_size) a1_taken = b1_size - 1; + if (a2_taken >= b2_size) a2_taken = b2_size - 1; + + int branch_taken[3]; + branch_taken[0] = a0_taken; + branch_taken[1] = a1_taken; + branch_taken[2] = a2_taken; + + int branch_sizes[3]; + branch_sizes[0] = b0_size; + branch_sizes[1] = b1_size; + branch_sizes[2] = b2_size; + + const float* val = v_logits + (long long)i * num_atoms; + + // Zero value logit gradient accumulator + float d_val_accum[256]; // max num_atoms (C51: 51, generous buffer) + for (int j = 0; j < num_atoms && j < 256; j++) d_val_accum[j] = 0.0f; + + int adv_offset = 0; + for (int d = 0; d < 3; d++) { + int bd = branch_sizes[d]; + int a_taken_d = branch_taken[d]; + + // Step 1: Compute expected Q for each action in this branch + // using dueling: combined_logit[j] = val[j] + adv[a,j] - mean_adv + float q_values_branch[9]; // max branch size = 5 (exposure), 3+3=6 total per branch + + // Compute mean advantage per atom (for dueling centering) + float mean_adv_sum = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + for (int j = 0; j < num_atoms; j++) mean_adv_sum += adv_aa[j]; + } + float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); + + for (int a = 0; a < bd; a++) { + const float* adv = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + + // softmax over atoms + expected value + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float c = val[j] + adv[j] - mean_adv_per_atom; + if (c > max_logit) max_logit = c; + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + float eq = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float z = v_min + (float)j * dz; + eq += p * z; + } + q_values_branch[a] = eq; + } + + // Step 2: CQL penalty gradient w.r.t. Q-values + // dCQL/dQ[a] = softmax_Q[a] - I(a == a_taken) + float max_q = q_values_branch[0]; + for (int a = 1; a < bd; a++) + if (q_values_branch[a] > max_q) max_q = q_values_branch[a]; + float sum_exp_q = 0.0f; + for (int a = 0; a < bd; a++) + sum_exp_q += expf(q_values_branch[a] - max_q); + + float d_cql_dq[9]; + for (int a = 0; a < bd; a++) { + float softmax_q = expf(q_values_branch[a] - max_q) / (sum_exp_q + 1e-8f); + d_cql_dq[a] = cql_alpha * (softmax_q - ((a == a_taken_d) ? 1.0f : 0.0f)); + } + + // Step 3: Chain rule through softmax-expectation to get d_logits + for (int a = 0; a < bd; a++) { + const float* adv = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + + // Recompute p[j] for this action + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float c = val[j] + adv[j] - mean_adv_per_atom; + if (c > max_logit) max_logit = c; + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + float eq = q_values_branch[a]; + + for (int j = 0; j < num_atoms; j++) { + float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float z = v_min + (float)j * dz; + // d_combined_logit[j] = d_cql_dq[a] * p * (z - Q) + float d_combined = d_cql_dq[a] * p * (z - eq); + + // Split combined gradient to adv and val + d_adv[j] = d_combined; + if (j < 256) d_val_accum[j] += d_combined; + } + } + adv_offset += bd; + } + + // Write accumulated value logit gradient (summed across all branches and actions) + float* d_val = d_v_logits + (long long)i * num_atoms; + for (int j = 0; j < num_atoms && j < 256; j++) { + d_val[j] = d_val_accum[j]; + } +} +"#; + + let context = stream.context(); + let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) + .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel compilation: {e}")))?; + let module = context.load_module(ptx) + .map_err(|e| MLError::ModelError(format!("cql_logit_grad module load: {e}")))?; + module.load_function("cql_logit_grad_kernel") + .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel load: {e}"))) +} diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 67f86a7fb..c4f1d8e00 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -1277,6 +1277,18 @@ pub struct DQNHyperparameters { pub dt_num_layers: usize, /// DT target return-to-go for conditioning (in units of cumulative reward). pub dt_target_return: f64, + + // CVaR (Conditional Value at Risk) action selection for IQN + /// Enable CVaR-aware action selection when IQN is active. + /// When true, actions are selected by optimizing the worst-case quantile tail + /// (risk-averse) instead of the mean quantile (risk-neutral). + /// Default: true (enables risk-aware position scaling via IQN head). + pub use_cvar_action_selection: bool, + /// CVaR confidence level alpha (0.0-1.0). Lower alpha = more risk-averse. + /// 0.05 = optimize for worst 5% outcomes (extreme caution). + /// 0.25 = optimize for worst 25% outcomes (moderate risk aversion). + /// Default: 0.05. + pub cvar_alpha: f32, } impl Default for DQNHyperparameters { @@ -1532,6 +1544,10 @@ impl DQNHyperparameters { dt_embed_dim: 128, dt_num_layers: 3, dt_target_return: 2.0, + + // CVaR action selection: enabled by default (risk-aware IQN action scoring) + use_cvar_action_selection: true, + cvar_alpha: 0.05, // Optimize for worst 5% quantile tail } } } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 25708a0bc..5b3fdddde 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -191,6 +191,8 @@ impl FusedTrainingCtx { iqn_kappa: dqn.config.iqn_kappa, entropy_coefficient: dqn.config.entropy_coefficient as f32, c51_warmup_epochs: hyperparams.c51_warmup_epochs, + use_cql: hyperparams.use_cql, + cql_alpha: hyperparams.cql_alpha as f32, }; // Extract weight sets from VarMaps (online + target) @@ -740,9 +742,25 @@ impl FusedTrainingCtx { // Spectral norm moved to Step 1b (before graph_forward) — correct placement. + // ── Step 5c: CQL conservative penalty (if enabled) ─────────────── + // Computes CQL logit gradients from current Q-values and adds + // parameter gradients into grad_buf via a second cuBLAS backward. + // Same pattern as IQN trunk gradient injection. + if self.trainer.has_cql() { + match self.trainer.apply_cql_gradient() { + Ok(true) => { + tracing::trace!("CQL gradient injected into grad_buf"); + } + Ok(false) => {} // CQL disabled or alpha=0 + Err(e) => { + tracing::warn!("CQL gradient failed (non-fatal): {e}"); + } + } + } + // ── Step 5e: Replay graph_adam — Adam sees combined gradient ──── - // All auxiliary gradients (IQN, attention, ensemble) are now in grad_buf. - // The single Adam update sees: C51 + iqn_lambda*IQN + attention + diversity. + // All auxiliary gradients (IQN, attention, ensemble, CQL) are now in grad_buf. + // The single Adam update sees: C51 + iqn_lambda*IQN + attention + diversity + CQL. let fused_result = self.trainer.replay_adam_and_readback() .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?; diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index feca0e1b5..d85f00b78 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -310,8 +310,8 @@ impl DQNTrainer { use_branching: hyperparams.use_branching, branch_hidden_dim: hyperparams.branch_hidden_dim, use_regime_conditioning: true, // Always enable per-regime IS weights for branching loss - use_cvar_action_selection: false, - cvar_alpha: 0.05, + use_cvar_action_selection: hyperparams.use_cvar_action_selection, + cvar_alpha: hyperparams.cvar_alpha, #[allow(clippy::cast_possible_truncation)] minimum_profit_factor: hyperparams.minimum_profit_factor as f32,