diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index f9fe2f683..b39fc694b 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -24,7 +24,7 @@ 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, 4, 3] stride-12 per-sample per-branch [v_min, v_max, delta_z] (Phase 2d) */ const float* __restrict__ liquid_mod, /* [4] pinned device-mapped per-branch modulators */ const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = linear. */ const float* __restrict__ q_mean_ema_ptr, /* [1] pinned device-mapped Q-mean EMA */ @@ -203,13 +203,17 @@ extern "C" __global__ void c51_grad_kernel( * The gradient is orthogonal to Bellman — it depends on atom * position, not on the target match. Active on ALL samples, * providing perpetual differentiation pressure. */ - /* z_norm: use adaptive atom positions if available, else linear grid */ + /* z_norm: use adaptive atom positions if available, else linear grid. + * Phase 2d: per_sample_support is [B, 4, 3] stride-12; index with + * (b, d, {0,1,2}) so each branch's normalisation uses its own + * v_min/v_max/delta_z. */ float z_norm; float delta_z; + int support_base = b * 12 + d * 3; if (atom_positions != NULL) { float z_val = atom_positions[(long long)d * num_atoms + j]; - float z_min = per_sample_support[b * 3 + 0]; - float z_max = per_sample_support[b * 3 + 1]; + float z_min = per_sample_support[support_base + 0]; + float z_max = per_sample_support[support_base + 1]; z_norm = 2.0f * (z_val - z_min) / fmaxf(z_max - z_min, 1e-6f) - 1.0f; /* Approximate delta_z from adjacent atom spacing */ delta_z = (j < num_atoms - 1) @@ -217,7 +221,7 @@ extern "C" __global__ void c51_grad_kernel( : z_val - atom_positions[(long long)d * num_atoms + j - 1]; } else { z_norm = 2.0f * (float)j / fmaxf((float)(num_atoms - 1), 1.0f) - 1.0f; - delta_z = per_sample_support[b * 3 + 2]; + delta_z = per_sample_support[support_base + 2]; } float velocity_mod = liquid_mod[d]; float spread_scale = inv_batch * delta_z * velocity_mod; diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index f08fec235..0f35e66d5 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -554,7 +554,7 @@ extern "C" __global__ void c51_loss_batched( const float* __restrict__ gamma_buf, /* [B] per-sample effective gamma */ int batch_size, int num_atoms, - const float* __restrict__ per_sample_support, /* [B, 3] per-sample: [v_min, v_max, delta_z] */ + const float* __restrict__ per_sample_support, /* [B, 4, 3] stride-12 per-sample per-branch: [v_min, v_max, delta_z] (Phase 2d) */ int b0_size, int b1_size, int b2_size, @@ -612,28 +612,16 @@ extern "C" __global__ void c51_loss_batched( int sample_id = blockIdx.x; if (sample_id >= batch_size) return; - /* Per-sample C51 support centered on V(s) */ - float v_min = per_sample_support[sample_id * 3]; - float v_max = per_sample_support[sample_id * 3 + 1]; - float delta_z = per_sample_support[sample_id * 3 + 2]; + /* Phase 2d: per_sample_support is [B, 4, 3] stride-12. Per-branch + * (v_min, v_max, delta_z) is read INSIDE the branch loop below; the + * degenerate-support skip also becomes per-branch. No single-range + * read here — the tile no longer carries a shared per-sample range. */ - /* Skip sample if support is degenerate (all Q-values identical) */ - if (delta_z < 1e-7f) { - if (tid == 0) { - per_sample_loss[sample_id] = 0.0f; - td_errors[sample_id] = 0.0f; - } - return; - } int b0_atoms = b0_size * num_atoms; int b1_atoms = b1_size * num_atoms; int b2_atoms = b2_size * num_atoms; int b3_atoms = b3_size * num_atoms; - for (int j = tid; j < num_atoms; j += BLOCK_THREADS) - shmem_support[j] = v_min + (float)j * delta_z; - __syncthreads(); - int factored_action = actions[sample_id]; int max_action = b0_size * b1_size * b2_size * b3_size; if (factored_action < 0 || factored_action >= max_action) @@ -702,12 +690,34 @@ extern "C" __global__ void c51_loss_batched( int n_atoms = n_d * num_atoms; long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * num_atoms; - /* Update atom support for this branch (adaptive positions if available) */ + /* Phase 2d: per-branch (v_min, v_max, delta_z) from the [B, 4, 3] + * tile. Direction/magnitude/order/urgency own independent ranges + * — matches the per-branch atom_positions layout that landed in + * Phase 2a/2b. */ + int support_base = sample_id * 12 + d * 3; + float v_min = per_sample_support[support_base + 0]; + float v_max = per_sample_support[support_base + 1]; + float delta_z = per_sample_support[support_base + 2]; + + /* Degenerate-support skip is now per-branch: when a branch's + * Q-spread collapses to zero, skip just that branch's loss + * contribution instead of the whole sample. Unskipped branches + * still contribute to the sample's total_ce. */ + bool branch_degenerate = (delta_z < 1e-7f); + + /* Update atom support for this branch: prefer adaptive positions + * when provided, otherwise fall back to the per-branch linear + * grid derived from v_min/delta_z. */ if (atom_positions != NULL) { for (int j = tid; j < num_atoms; j += BLOCK_THREADS) shmem_support[j] = atom_positions[(long long)d * num_atoms + j]; - __syncthreads(); + } else { + for (int j = tid; j < num_atoms; j += BLOCK_THREADS) + shmem_support[j] = v_min + (float)j * delta_z; } + __syncthreads(); + + if (branch_degenerate) continue; /* ═══ STEP a: Current log-probs for taken action a_d ═══════ */ diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index c3f0a657a..d191507d7 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -2721,7 +2721,7 @@ extern "C" __global__ void compute_expected_q( int b1_size, int b2_size, int b3_size, - const float* __restrict__ per_sample_support, /* [N*3]: per-sample [v_min, v_max, delta_z] from IQL */ + const float* __restrict__ per_sample_support, /* [N, 4, 3] stride-12: per-sample per-branch [v_min, v_max, delta_z] from IQL (Phase 2d) */ float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2]: per-block partial sums (phase 1). * NULL to skip atom-stat collection entirely. */ float* __restrict__ q_variance, /* [N, total_actions] Var[Q] per action. NULL to skip. */ @@ -2730,12 +2730,10 @@ extern "C" __global__ void compute_expected_q( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; - /* Per-sample support from IQL — breaks circular eval_v_range feedback */ - float v_min = per_sample_support[i * 3]; - float v_max = per_sample_support[i * 3 + 1]; + /* Phase 2d: per-branch (v_min, v_max, delta_z) read INSIDE the d-loop + * below — direction/magnitude/order/urgency own independent ranges. */ int total_actions = b0_size + b1_size + b2_size + b3_size; - float dz = per_sample_support[i * 3 + 2]; float util_threshold = 0.5f / (float)num_atoms; /* half of uniform */ /* Per-sample value logits — sample-major [N, num_atoms] */ @@ -2757,6 +2755,11 @@ extern "C" __global__ void compute_expected_q( const float* branch_base = b_logits + (long long)branch_logit_offset + (long long)i * n_d * num_atoms; + /* Phase 2d: per-branch support triple from the [N, 4, 3] tile. */ + int support_base = i * 12 + d * 3; + float v_min = per_sample_support[support_base + 0]; + float dz = per_sample_support[support_base + 2]; + for (int a = 0; a < n_d; a++) { const float* adv_a = branch_base + (long long)a * num_atoms; @@ -3163,7 +3166,7 @@ extern "C" __global__ void mag_concat_qdir( int SH2, int NA, /* num_atoms */ int b0_size, /* direction branch size (4) */ - const float* __restrict__ per_sample_support /* [B, 3]: v_min, v_max, delta_z */ + const float* __restrict__ per_sample_support /* [B, 4, 3] stride-12: v_min, v_max, delta_z per-branch (Phase 2d) */ ) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; @@ -3174,9 +3177,11 @@ extern "C" __global__ void mag_concat_qdir( concat_out[b * out_stride + i] = h_s2[b * SH2 + i]; } - /* Step 2: compute E[Q_dir] for each direction action */ - float v_min = per_sample_support[b * 3 + 0]; - float dz = per_sample_support[b * 3 + 2]; + /* Step 2: compute E[Q_dir] for each direction action. + * Phase 2d: read direction-branch (d=0) slots from the [B, 4, 3] tile. */ + int support_base = b * 12 + 0 * 3; /* direction = branch 0 */ + float v_min = per_sample_support[support_base + 0]; + float dz = per_sample_support[support_base + 2]; for (int a = 0; a < b0_size; a++) { /* Combined value + advantage logits for direction action a. @@ -3676,15 +3681,14 @@ extern "C" __global__ void quantile_q_select( int b1_size, int b2_size, int b3_size, - const float* __restrict__ per_sample_support, + const float* __restrict__ per_sample_support, /* [N, 4, 3] stride-12 per-branch (Phase 2d) */ float iqn_readiness, float util_ema) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; - float v_min = per_sample_support[i * 3]; - float dz = per_sample_support[i * 3 + 2]; + /* Phase 2d: per-branch (v_min, delta_z) read INSIDE the d-loop below. */ int total_actions = b0_size + b1_size + b2_size + b3_size; const float* v_row = v_logits + (long long)i * num_atoms; @@ -3703,6 +3707,11 @@ extern "C" __global__ void quantile_q_select( const float* branch_base = b_logits + (long long)branch_logit_offset + (long long)i * n_d * num_atoms; + /* Phase 2d: per-branch (v_min, dz) from the [N, 4, 3] tile. */ + int support_base = i * 12 + d * 3; + float v_min = per_sample_support[support_base + 0]; + float dz = per_sample_support[support_base + 2]; + for (int a = 0; a < n_d; a++) { const float* adv_a = branch_base + (long long)a * num_atoms; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 2d157a879..7abb2bc55 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1249,7 +1249,9 @@ pub struct GpuDqnTrainer { tau_pinned: *mut f32, tau_dev_ptr: u64, - /// Per-sample C51 support [B, 3] — pointer set by IQL trainer. Used by c51_loss_batched only. + /// Per-sample, per-branch C51 support [B, 4, 3] stride-12 — pointer set by IQL trainer + /// (Phase 2d). Consumed by c51_loss_batched / c51_grad_kernel / compute_expected_q / + /// mag_concat_qdir / quantile_q_select, which index `sample_id * 12 + d * 3 + {0,1,2}`. per_sample_support_ptr: u64, /// Per-branch gradient scales [B, 4] — pointer set by IQL trainer. branch_scales_ptr: u64, diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 437fc4877..b692fe406 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -439,8 +439,10 @@ pub struct GpuExperienceCollector { network_dims: (usize, usize, usize, usize), /// Number of C51 atoms (1 for non-distributional). num_atoms: usize, - /// Per-sample C51 support tiled [alloc_episodes, 3]: [v_min, v_max, delta_z] per episode. - /// Filled once per epoch via `update_per_sample_support()` — matches kernel expectation. + /// Per-sample, per-branch C51 support tiled [alloc_episodes, 4, 3]: + /// [v_min, v_max, delta_z] per (episode, branch). Filled once per epoch via + /// `update_per_sample_support()` — matches the Phase 2d stride-12 contract + /// consumed by compute_expected_q / quantile_q_select / mag_concat_qdir. per_sample_support_buf: CudaSlice, /// Per-sample epsilon from IQL expectile gap (0 = use cosine schedule). per_sample_epsilon_ptr: u64, @@ -1274,10 +1276,11 @@ impl GpuExperienceCollector { let exp_bn_concat = stream.alloc_zeros::(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128) .map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?; - // Per-sample C51 support buffer [alloc_episodes, 3] — tiled v_min/v_max/delta_z. - // Filled once per epoch in update_per_sample_support() — replaces the 2-float - // eval_v_range pointer that caused OOB reads in compute_expected_q/quantile_q_select. - let per_sample_support_buf = stream.alloc_zeros::(alloc_episodes * 3) + // Phase 2d: per-sample, per-branch C51 support buffer [alloc_episodes, 4, 3] + // stride-12 — tiled v_min/v_max/delta_z per branch. Filled once per epoch + // in update_per_sample_support(); consumers compute_expected_q / + // quantile_q_select / mag_concat_qdir index with (sample, branch, {0,1,2}). + let per_sample_support_buf = stream.alloc_zeros::(alloc_episodes * 4 * 3) .map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf: {e}")))?; // Task 8: GPU-resident step counter for experience collection loop @@ -1323,7 +1326,7 @@ impl GpuExperienceCollector { market_dim, network_dims, num_atoms, - per_sample_support_buf, // Tiled [alloc_episodes, 3]; filled via update_per_sample_support() + per_sample_support_buf, // Phase 2d: [alloc_episodes, 4, 3] stride-12; update_per_sample_support() tiles per_sample_epsilon_ptr: 0, // Set via set_per_sample_epsilon_ptr() after IQL init branch_sizes, alloc_episodes, @@ -2715,7 +2718,7 @@ impl GpuExperienceCollector { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&self.per_sample_support_buf) // [N, 3] per-sample support + .arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support .arg(&null_atom_stats) .arg(&mut self.q_var_buf) .arg(&null_atom_positions) // atom_positions (NULL = linear) @@ -2835,7 +2838,7 @@ impl GpuExperienceCollector { .arg(&b1) .arg(&b2) .arg(&b3) - .arg(&self.per_sample_support_buf) // [N, 3] per-sample support + .arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support .arg(&iqn_r) .arg(&util_e) .launch(launch_cfg) @@ -3089,17 +3092,23 @@ impl GpuExperienceCollector { info!(trainer_params_ptr = ptr, "Experience collector: zero-copy trainer params pointer set"); } - /// Fill per_sample_support_buf [alloc_episodes, 3] with uniform v_min/v_max/delta_z. + /// Fill per_sample_support_buf [alloc_episodes, 4, 3] with uniform v_min/v_max/delta_z + /// broadcast across all 4 branches. /// - /// Called once per epoch (not hot-path) when the eval v_range changes. - /// Replaces the old `set_eval_v_range_ptr` which passed a 2-float pointer - /// to kernels that expected [N, 3], causing OOB reads for every sample > 0. + /// Called once per epoch (not hot-path) when the eval v_range changes. The + /// experience collector does not own per-branch Q-stats, so all 4 branches + /// receive the same (v_min, v_max, delta_z) triple here; this preserves + /// existing training-path semantics while making the buffer layout match the + /// [N, 4, 3] stride-12 contract that compute_expected_q / quantile_q_select / + /// mag_concat_qdir read (Phase 2d). pub fn update_per_sample_support(&mut self, v_min: f32, v_max: f32) -> Result<(), MLError> { let n = self.alloc_episodes; let delta_z = (v_max - v_min) / (self.num_atoms as f32 - 1.0).max(1.0); - let data: Vec = (0..n).flat_map(|_| [v_min, v_max, delta_z]).collect(); + let data: Vec = (0..n) + .flat_map(|_| (0..4).flat_map(|_| [v_min, v_max, delta_z])) + .collect(); super::htod_f32(&self.stream, &data, &mut self.per_sample_support_buf)?; - debug!(v_min, v_max, delta_z, n, "per_sample_support tiled for experience collector"); + debug!(v_min, v_max, delta_z, n, "per_sample_support tiled [N, 4, 3] for experience collector"); Ok(()) } pub fn set_per_sample_epsilon_ptr(&mut self, ptr: u64) { diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index ebf825389..a70b215f2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -375,20 +375,30 @@ impl GpuIqlTrainer { let adv_sigma_ema_buf = alloc_f32(&stream, 1, "iql_adv_sigma_ema")?; let readiness_buf = alloc_f32(&stream, 4, "iql_readiness")?; let p5_state_buf = alloc_f32(&stream, 2, "iql_p5_state")?; - let mut per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?; + // Phase 2d: per-sample, per-branch support tile [B, 4, 3] stride-12. + // Each sample owns 4 branch triples (v_min, v_max, delta_z) — direction, + // magnitude, order, urgency — read by C51 loss/grad + experience kernels. + let mut per_sample_support_buf = alloc_f32(&stream, b * 4 * 3, "iql_per_sample_support")?; let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?; let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?; let gap_mean_buf = alloc_f32(&stream, 1, "iql_gap_mean")?; let per_sample_epsilon_buf = alloc_f32(&stream, b, "iql_per_sample_epsilon")?; - // Seed per_sample_support with defaults for step 0 + // Seed per_sample_support with defaults for step 0. + // Phase 2d: broadcast identical (v_min=-1, v_max=1, delta_z) to all + // 4 branches per sample. Bootstrap identity — the per-branch kernel + // will overwrite with ISV-driven centres + V(s)-derived half once + // the first training step lands. let na = config.num_atoms.max(2) as f32; let default_delta_z = 2.0 / (na - 1.0); - let mut default_support = vec![0.0_f32; b * 3]; + let mut default_support = vec![0.0_f32; b * 4 * 3]; for i in 0..b { - default_support[i * 3] = -1.0; - default_support[i * 3 + 1] = 1.0; - default_support[i * 3 + 2] = default_delta_z; + for d in 0..4 { + let base = i * 12 + d * 3; + default_support[base] = -1.0; + default_support[base + 1] = 1.0; + default_support[base + 2] = default_delta_z; + } } super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?; @@ -1208,10 +1218,19 @@ impl GpuIqlTrainer { Ok(()) } - /// Compute per-sample C51 atom support centered on V(s). + /// Compute per-sample, per-branch C51 atom support. + /// + /// Phase 2d: fills the [B, 4, 3] tile with per-branch (v_min, v_max, + /// delta_z) where the centre comes from ISV slots `V_CENTER_DIR + 2*d` + /// and the half-width from the per-sample V(s)-derived Q spread. + /// + /// `isv_signals_dev_ptr`: raw device pointer to the [ISV_TOTAL_DIM=31] + /// broadcast bus. Pass `0` to skip ISV (null in the kernel) — centres + /// collapse to 0 and the readiness blend still provides a safe default. pub fn compute_per_sample_support( &mut self, q_out_buf: &CudaSlice, + isv_signals_dev_ptr: u64, ) -> Result<(), MLError> { let b = self.config.batch_size; let batch_i32 = b as i32; @@ -1227,6 +1246,7 @@ impl GpuIqlTrainer { .arg(q_out_buf) .arg(&mut self.per_sample_support_buf) .arg(&self.readiness_buf) + .arg(&isv_signals_dev_ptr) .arg(&gamma) .arg(&batch_i32) .arg(&total_actions_i32) @@ -1349,7 +1369,9 @@ impl GpuIqlTrainer { Ok(()) } - /// Raw pointer to per-sample support buffer [B, 3]. + /// Raw pointer to per-sample, per-branch support buffer [B, 4, 3] stride-12 + /// (Phase 2d). Consumers read per-branch triples at offset + /// `sample_id * 12 + branch_idx * 3 + {0,1,2}`. pub fn per_sample_support_ptr(&self) -> u64 { self.per_sample_support_buf.raw_ptr() } diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index a0ce74e5b..677f3b27d 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -539,25 +539,43 @@ void iql_adv_variance_reduce( } /* ------------------------------------------------------------------ */ -/* Per-Sample C51 Atom Support Kernel */ +/* Per-Sample C51 Atom Support Kernel — per-branch [B, 4, 3] tile */ /* ------------------------------------------------------------------ */ /** - * Compute per-sample [v_min, v_max, delta_z] for C51 atom projection: + * Compute per-sample, per-branch [v_min, v_max, delta_z] for C51 atom + * projection. Spec: docs/superpowers/specs/2026-04-23-isv-v-range-unification.md + * (Phase 2d). + * * spread = max_{a} |Q(s,a) - V(s)| - * half_w = spread * (1 + gamma) - * v_min[b] = V(s) - half_w - * v_max[b] = V(s) + half_w + * half_w = spread * (1 + gamma) (shared across branches) + * centre_b = isv_signals[V_CENTER_DIR + 2*b] (per-branch from ISV) + * v_min_b = centre_b - half_w + * v_max_b = centre_b + half_w * delta_z = (v_max - v_min) / (num_atoms - 1) * - * Output: per_sample_support[b*3+0]=v_min, [b*3+1]=v_max, [b*3+2]=delta_z + * Rationale: the V(s)-derived half_w captures per-sample Q magnitude + * (how wide the atoms need to span around the centre). The centre + * differs per branch because direction/magnitude/order/urgency have + * genuinely different Q scales — the ISV bus broadcasts each branch's + * centre EMA. + * + * Bootstrap safety: when ISV centres are 0 and readiness=0 blend falls + * back to [-1, 1], all 4 branch triples become [-1, 1, 2/(NA-1)] — + * byte-identical to the pre-spec single-range behaviour at epoch 1. + * + * Output layout stride-12: per_sample_support[b*12 + d*3 + {0,1,2}] per + * branch d in 0..4. The 4-branch tile is read by the C51 loss / grad + * kernels and downstream experience kernels that migrated in lockstep. + * * Launch: grid=ceil(B/256), block=256. */ extern "C" __global__ void iql_compute_per_sample_support( const float* __restrict__ v_out, /* [B] */ const float* __restrict__ q_out, /* [B, total_actions] */ - float* __restrict__ per_sample_support, /* [B*3] */ + float* __restrict__ per_sample_support, /* [B, 4, 3] */ const float* __restrict__ readiness_buf, /* [1] CV-based readiness */ + const float* __restrict__ isv_signals, /* [ISV_TOTAL_DIM]; NULL → 0 centres */ float gamma, int batch_size, int total_actions, @@ -578,18 +596,34 @@ void iql_compute_per_sample_support( } float half_w = spread * (1.0f + gamma); - float iql_vmin = v - half_w; - float iql_vmax = v + half_w; - float iql_dz = (iql_vmax - iql_vmin) / (float)(num_atoms - 1); + float default_dz = 2.0f / (float)(num_atoms - 1); - /* Blend: readiness=0 → default [-1,1], readiness=1 → V(s)-centered */ - float v_min = r * iql_vmin + (1.0f - r) * (-1.0f); - float v_max = r * iql_vmax + (1.0f - r) * (1.0f); - float delta_z = r * iql_dz + (1.0f - r) * (2.0f / (float)(num_atoms - 1)); + /* Per-branch centre from ISV slots 23 + 2*d. When ISV is null (smoke + * paths / pre-init) the centre collapses to 0; in that regime the + * readiness blend below falls back to the uniform [-1, 1] default + * while r ramps from 0 → 1. */ + for (int d = 0; d < 4; d++) { + float isv_centre = (isv_signals != nullptr) + ? isv_signals[23 + 2 * d] + : 0.0f; - per_sample_support[b * 3 + 0] = v_min; - per_sample_support[b * 3 + 1] = v_max; - per_sample_support[b * 3 + 2] = delta_z; + /* Per-branch IQL range: centre from ISV, half from per-sample + * V(s)-derived Q spread. Direction/magnitude/order/urgency + * share the half but own independent centres. */ + float iql_vmin = isv_centre - half_w; + float iql_vmax = isv_centre + half_w; + float iql_dz = (iql_vmax - iql_vmin) / (float)(num_atoms - 1); + + /* Blend: readiness=0 → default [-1,1], readiness=1 → adaptive. */ + float v_min = r * iql_vmin + (1.0f - r) * (-1.0f); + float v_max = r * iql_vmax + (1.0f - r) * ( 1.0f); + float delta_z = r * iql_dz + (1.0f - r) * default_dz; + + int base = b * 12 + d * 3; + per_sample_support[base + 0] = v_min; + per_sample_support[base + 1] = v_max; + per_sample_support[base + 2] = delta_z; + } } /* ------------------------------------------------------------------ */ @@ -600,11 +634,18 @@ void iql_compute_per_sample_support( * p5 quantile of half-widths. Uses Frugal-1U streaming quantile estimator * (one GPU scalar, O(B) per step, scale-free). * + * Phase 2d layout: per_sample_support is now [B, 4, 3] stride-12; the + * p5 estimator aggregates half-widths across ALL (sample, branch) pairs + * (one scalar p5 for the entire tile) and the floor is applied per + * (sample, branch) independently. Sharing a single p5 across branches + * keeps the estimator statistically stable and preserves the anti- + * collapse guarantee for every branch's smallest half-width. + * * Launch: grid=1, block=1 (after iql_compute_per_sample_support). */ extern "C" __global__ void iql_support_floor( - float* __restrict__ per_sample_support, /* [B*3] in-place */ + float* __restrict__ per_sample_support, /* [B, 4, 3] in-place */ float* __restrict__ p5_state, /* [2]: [0]=p5_estimate, [1]=step_count */ int batch_size, int num_atoms @@ -612,16 +653,19 @@ void iql_support_floor( { float est = p5_state[0]; float step_count = p5_state[1]; + int total_branches = batch_size * 4; - /* First call: init from batch median spread (not hardcoded 1.0) */ + /* First call: init from aggregated mean spread across all (sample, branch). */ if (step_count < 0.5f) { - /* Find median via single pass: compute mean as proxy (O(B), no sort) */ float hw_sum = 0.0f; for (int b = 0; b < batch_size; b++) { - float hw = (per_sample_support[b * 3 + 1] - per_sample_support[b * 3 + 0]) * 0.5f; - hw_sum += hw; + for (int d = 0; d < 4; d++) { + int base = b * 12 + d * 3; + float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f; + hw_sum += hw; + } } - est = hw_sum / fmaxf((float)batch_size, 1.0f); + est = hw_sum / fmaxf((float)total_branches, 1.0f); est = fmaxf(est, 1e-6f); } @@ -629,30 +673,37 @@ void iql_support_floor( * Asymmetric ratio 0.05/0.95 is mathematically correct for p5 quantile. */ float adaptive_rate = 1.0f / fmaxf(sqrtf(step_count + 1.0f), 1.0f); for (int b = 0; b < batch_size; b++) { - float hw = (per_sample_support[b * 3 + 1] - per_sample_support[b * 3 + 0]) * 0.5f; - float step = fmaxf(est * adaptive_rate, 1e-8f); - if (hw < est) { - est -= step; - } else { - est += step * (0.05f / 0.95f); + for (int d = 0; d < 4; d++) { + int base = b * 12 + d * 3; + float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f; + float step = fmaxf(est * adaptive_rate, 1e-8f); + if (hw < est) { + est -= step; + } else { + est += step * (0.05f / 0.95f); + } + est = fmaxf(est, 1e-8f); } - est = fmaxf(est, 1e-8f); } p5_state[0] = est; p5_state[1] = step_count + 1.0f; /* Floor: 1/num_atoms of p5 — guarantees at least 1 atom of resolution. - * Adaptive to atom count: 51 atoms → 2% of p5, 11 atoms → 9% of p5. */ + * Adaptive to atom count: 51 atoms → 2% of p5, 11 atoms → 9% of p5. + * Applied per (sample, branch) so each branch's floor is independent. */ float floor_fraction = 1.0f / fmaxf((float)num_atoms, 1.0f); float floor_hw = est * floor_fraction; float floor_dz = (2.0f * floor_hw) / fmaxf((float)(num_atoms - 1), 1.0f); for (int b = 0; b < batch_size; b++) { - if (per_sample_support[b * 3 + 2] < floor_dz) { - float center = (per_sample_support[b * 3 + 0] + per_sample_support[b * 3 + 1]) * 0.5f; - per_sample_support[b * 3 + 0] = center - floor_hw; - per_sample_support[b * 3 + 1] = center + floor_hw; - per_sample_support[b * 3 + 2] = floor_dz; + for (int d = 0; d < 4; d++) { + int base = b * 12 + d * 3; + if (per_sample_support[base + 2] < floor_dz) { + float center = (per_sample_support[base + 0] + per_sample_support[base + 1]) * 0.5f; + per_sample_support[base + 0] = center - floor_hw; + per_sample_support[base + 1] = center + floor_hw; + per_sample_support[base + 2] = floor_dz; + } } } } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 184ee6846..0c33a8509 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1430,9 +1430,12 @@ impl FusedTrainingCtx { self.gpu_iql.update_adv_sigma() .map_err(|e| anyhow::anyhow!("IQL adv_sigma: {e}"))?; - // 6. Compute per-sample C51 support centered on V(s). + // 6. Compute per-sample, per-branch C51 support. + // Phase 2d: tile is [B, 4, 3]; centres from ISV slots 23..30, + // half-width from V(s)-derived Q spread. let q_out = self.trainer.q_out_buf(); - self.gpu_iql.compute_per_sample_support(q_out) + let isv_ptr = self.trainer.isv_signals_dev_ptr(); + self.gpu_iql.compute_per_sample_support(q_out, isv_ptr) .map_err(|e| anyhow::anyhow!("IQL per_sample_support: {e}"))?; // Floor support using Frugal p5 quantile — no silent zero-loss samples. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 7889ee534..aaac8ceeb 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1292,9 +1292,11 @@ impl DQNTrainer { return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}")); } - // Tile per-sample support [N, 3] from current eval v_range before experience collection. - // Reads v_min/v_max from fused trainer (which may have been updated last epoch), - // falls back to config defaults for epoch 0 / pre-fused-init. + // Tile per-sample, per-branch support [N, 4, 3] from current eval v_range + // before experience collection. Reads v_min/v_max from fused trainer (which + // may have been updated last epoch), falls back to config defaults for + // epoch 0 / pre-fused-init. All 4 branches receive the same triple here; + // see update_per_sample_support doc (Phase 2d). { let (v_min, v_max) = if let Some(ref fused_ctx) = self.fused_ctx { let vr = fused_ctx.eval_v_range();