diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index e3d4dd7d1..b53f662c8 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -24,7 +24,8 @@ extern "C" __global__ void c51_grad_kernel( int total_branch_atoms, float entropy_coeff, const float* __restrict__ branch_scales, /* [B, 4] per-sample per-branch gradient scale */ - const float* __restrict__ per_sample_support) /* [B, 3] per-sample [v_min, v_max, delta_z] */ + const float* __restrict__ per_sample_support, /* [B, 3] per-sample [v_min, v_max, delta_z] */ + const float* __restrict__ spread_velocity) /* [1] pinned device-mapped modulator */ { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_elems = batch_size * num_atoms; @@ -86,7 +87,8 @@ extern "C" __global__ void c51_grad_kernel( /* Scale proportional to delta_z — tighter atoms = smaller spread needed. * No hardcoded constants: spread = inv_batch * delta_z (same order as CE grad). */ float delta_z = per_sample_support[b * 3 + 2]; - float spread_scale = inv_batch * delta_z; + float velocity_mod = spread_velocity[0]; + float spread_scale = inv_batch * delta_z * velocity_mod; for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index c1eb4e58c..7134e3d5f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -651,6 +651,13 @@ pub struct GpuDqnTrainer { /// EMA of Q-divergence for adaptive tau computation. q_div_ema: f32, + /// Q-gap EMA for momentum-modulated spread gradient. + q_gap_ema: f32, + /// Spread velocity scalar — written by CPU, read by GPU via pinned device-mapped memory. + spread_velocity_pinned: *mut f32, + /// Device pointer for spread_velocity_pinned (from cuMemHostGetDevicePointer). + spread_velocity_dev_ptr: u64, + // ── Training state ────────────────────────────────────────────── pub(crate) adam_step: i32, total_params: usize, @@ -997,6 +1004,9 @@ impl Drop for GpuDqnTrainer { if !self.adaptive_clip_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.adaptive_clip_pinned.cast()) }; } + if !self.spread_velocity_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.spread_velocity_pinned.cast()) }; + } if !self.t_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.t_pinned.cast()) }; } @@ -1035,6 +1045,20 @@ impl GpuDqnTrainer { self.eval_q_std_ema = (1.0 - alpha_std) * self.eval_q_std_ema + alpha_std * q_std; } + // Q-gap momentum: modulate spread gradient based on Q-gap velocity + let q_gap_val = q_gap; + if self.q_gap_ema < 1e-12 { + self.q_gap_ema = q_gap_val; + } else { + let err = (q_gap_val - self.q_gap_ema).abs(); + let alpha = (err / (err + self.q_gap_ema.max(0.001))).clamp(0.01, 0.3); + self.q_gap_ema = (1.0 - alpha) * self.q_gap_ema + alpha * q_gap_val; + } + let velocity = q_gap_val - self.q_gap_ema; + let delta_z_approx = self.eval_q_std_ema.max(0.01); + let modulator = (1.0 - velocity / delta_z_approx).clamp(0.1, 2.0); + unsafe { *self.spread_velocity_pinned = modulator; } + // Width from Q-gap (action differentiation range), not Q-std. // Floor = max(10 * q_gap, 3 * q_std) — enough atoms to resolve actions. // With 51 atoms, a width of 10*q_gap gives q_gap/delta_z ≈ 5 atoms of @@ -2469,6 +2493,23 @@ impl GpuDqnTrainer { dev_ptr }; + let spread_velocity_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned spread_velocity alloc: {e}")))? + as *mut f32 + }; + unsafe { *spread_velocity_pinned = 1.0; } // Full spread initially + let spread_velocity_dev_ptr = unsafe { + let mut dev_ptr: u64 = 0; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dev_ptr as *mut u64, + spread_velocity_pinned.cast(), + 0, + ); + dev_ptr + }; + // ── Allocate consolidated transfer buffers ───────────────── // Upload staging: states + next_states (f32) // Rewards/dones uploaded separately as f32 @@ -3056,6 +3097,9 @@ impl GpuDqnTrainer { adaptive_clip_dev_ptr, grad_norm_ema: 0.0, q_div_ema: 0.0, + q_gap_ema: 0.0, + spread_velocity_pinned, + spread_velocity_dev_ptr, adam_step: 0, total_params, params_initialized: false, @@ -5437,6 +5481,7 @@ impl GpuDqnTrainer { .arg(&entropy_coeff) .arg(&self.branch_scales_ptr) .arg(&self.per_sample_support_ptr) + .arg(&self.spread_velocity_dev_ptr) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1),