gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: zen precommit — all HIGH/MEDIUM/LOW issues resolved

Browse Source 
HIGH: c51_loss_kernel now uses atom_positions per-branch in shmem_support
(was ignoring adaptive positions → forward/loss atom mismatch).
MEDIUM: adaptive_gamma wired into C51 Bellman projection via
set_adaptive_gamma(). Config gamma replaced with adaptive_gamma in
both launch_c51_loss sites.
MEDIUM: c51_grad z_norm uses adaptive atom positions when available
(was assuming linear grid for spread gradient).
LOW: adaptive_gamma field added to GpuDqnTrainer, initialized from
config, updated via passthrough from training loop.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 22:03:39 +02:00
parent 62fda674e5
commit 00dc37f414
5 changed files with 44 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -85,10 +85,22 @@ 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. */
float z_norm = 2.0f * (float)j / fmaxf((float)(num_atoms - 1), 1.0f) - 1.0f;
/* 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];
/* z_norm: use adaptive atom positions if available, else linear grid */
float z_norm;
float delta_z;
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];
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)
? atom_positions[(long long)d * num_atoms + j + 1] - z_val
: 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];
}
float velocity_mod = liquid_mod[d];
float spread_scale = inv_batch * delta_z * velocity_mod;

7
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -344,6 +344,13 @@ 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) */
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();
}
/* ═══ STEP a: Current log-probs for taken action a_d ═══════ */
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)

14
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -735,6 +735,8 @@ pub struct GpuDqnTrainer {
regime_dropout_kernel: CudaFunction,
/// Epoch seed for regime dropout Philox PRNG (changes per epoch).
regime_dropout_epoch_seed: i32,
/// G4: Adaptive gamma for C51 Bellman projection. Updated from DQNTrainer.
adaptive_gamma: f32,
// ── G5: Epistemic-gated magnitude ──
/// Kernel: epistemic_gate_magnitude — sigmoid-gates magnitude Q-values by ensemble variance.
@@ -1901,6 +1903,11 @@ impl GpuDqnTrainer {
unsafe { *self.regime_util_pinned = util; }
}
/// G4: Set adaptive gamma for C51 Bellman projection.
pub fn set_adaptive_gamma(&mut self, gamma: f32) {
self.adaptive_gamma = gamma;
}
/// G9: Apply regime-conditioned dropout to h_s2 trunk output.
/// Must be called AFTER mamba2_step and BEFORE compute_expected_q.
pub(crate) fn apply_regime_dropout(&self, batch_size: usize, is_training: bool) -> Result<(), MLError> {
@@ -4676,6 +4683,8 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("alloc v_logits_blended: {e}")))?;
info!("GpuDqnTrainer: multi-horizon value heads allocated (5-bar + 20-bar)");
let initial_gamma = config.gamma;
Ok(Self {
config,
stream,
@@ -4782,6 +4791,7 @@ impl GpuDqnTrainer {
regime_util_dev_ptr,
regime_dropout_kernel,
regime_dropout_epoch_seed: 0,
adaptive_gamma: initial_gamma,
epistemic_gate_kernel,
var_ema_pinned,
var_ema_dev_ptr,
@@ -7515,7 +7525,7 @@ impl GpuDqnTrainer {
// N-step returns: use gamma^n for the Bellman projection.
// The experience collector pre-computes R_n = sum(gamma^i * r_i).
let gamma = self.config.gamma.powi(self.config.n_steps as i32);
let gamma = self.adaptive_gamma.powi(self.config.n_steps as i32);
let batch_i32 = b as i32;
let na_i32 = na as i32;
let b0_i32 = b0 as i32;
@@ -7723,7 +7733,7 @@ impl GpuDqnTrainer {
let on_next_b3_ptr = on_next_b2_ptr + (b * b2 * na * f32_sz) as u64;
// N-step returns: use gamma^n for the Bellman projection.
let gamma = self.config.gamma.powi(self.config.n_steps as i32);
let gamma = self.adaptive_gamma.powi(self.config.n_steps as i32);
let batch_i32 = b as i32;
let na_i32 = na as i32;
let b0_i32 = b0 as i32;

5
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -2307,6 +2307,11 @@ impl FusedTrainingCtx {
self.trainer.set_regime_dropout_seed(seed);
}
/// G4: Set adaptive gamma for C51 Bellman projection.
pub(crate) fn set_adaptive_gamma(&mut self, gamma: f32) {
self.trainer.set_adaptive_gamma(gamma);
}
/// G5: Set the EMA threshold for epistemic variance gating.
pub(crate) fn set_var_ema(&self, val: f32) {
self.trainer.set_var_ema(val);

4
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -557,6 +557,10 @@ impl DQNTrainer {
} else if util < 0.4 {
self.adaptive_gamma = (self.adaptive_gamma - 0.01).max(0.90);
}
// Wire adaptive gamma into GPU trainer for C51 Bellman projection
if let Some(ref mut fused) = self.fused_ctx {
fused.set_adaptive_gamma(self.adaptive_gamma);
}
info!("G4 gamma={:.3} util={:.2}", self.adaptive_gamma, util);
}