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

feat(isv): drift-conditioned gamma — C51 scalar→per-sample buffer

Browse Source 
c51_loss_kernel: float gamma → const float* gamma_buf [B].
fill_gamma_buf kernel: broadcasts base_gamma * gamma_mod[0] to [B].
gamma_mod from isv_forward ∈ [0.5, 1.0] → effective_gamma ∈
[base_gamma*0.5, base_gamma]. Model learns its own planning horizon.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-16 23:21:07 +02:00
parent 2acaaeaacd
commit ad402deec7
3 changed files with 51 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -199,7 +199,7 @@ extern "C" __global__ void c51_loss_batched(
const float* __restrict__ curiosity_errors, /* [B] per-sample curiosity prediction error (0 = no penalty) */
float curiosity_q_penalty_lambda, /* scaling factor: gamma *= 1/(1 + lambda * error) */
float gamma,
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] */
@@ -331,10 +331,10 @@ extern "C" __global__ void c51_loss_batched(
* gamma down so the Bellman target approaches the immediate reward,
* preventing overconfident extrapolation into novel states.
* gamma_eff = gamma / (1 + lambda * prediction_error) */
float gamma_eff = gamma;
float gamma_eff = gamma_buf[sample_id];
if (curiosity_q_penalty_lambda > 0.0f) {
float cur_err = curiosity_errors[sample_id];
gamma_eff = gamma / (1.0f + curiosity_q_penalty_lambda * cur_err);
gamma_eff = gamma_buf[sample_id] / (1.0f + curiosity_q_penalty_lambda * cur_err);
}
float total_ce = 0.0f;

14
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -4930,3 +4930,17 @@ extern "C" __global__ void isv_forward(
gamma_raw += w_gamma[k] * embedding[k];
gamma_mod_out[0] = 0.5f + 0.5f / (1.0f + expf(-gamma_raw));
}
/* ================================================================== */
/* Kernel: fill_gamma_buf — broadcast scalar gamma to per-sample buf */
/* ================================================================== */
extern "C" __global__ void fill_gamma_buf(
float* __restrict__ gamma_buf, /* [B] output */
float base_gamma,
const float* __restrict__ gamma_mod, /* [1] from isv_forward */
int B
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= B) return;
gamma_buf[i] = base_gamma * gamma_mod[0];
}

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

@@ -1338,6 +1338,7 @@ pub struct GpuDqnTrainer {
// ── ISV forward outputs ──
isv_forward_kernel: CudaFunction,
fill_gamma_buf_kernel: CudaFunction,
isv_embedding_buf: CudaSlice<f32>, // [ISV_DIM] = [8]
branch_gate_buf: CudaSlice<f32>, // [4]
gamma_mod_buf: CudaSlice<f32>, // [1]
@@ -2067,6 +2068,28 @@ impl GpuDqnTrainer {
Ok(())
}
/// Broadcast base_gamma * gamma_mod[0] into gamma_buf [B] for per-sample C51 loss.
pub(crate) fn fill_gamma_buf(&self) -> Result<(), MLError> {
let base_gamma = self.adaptive_gamma.powi(self.config.n_steps as i32);
let blocks = ((self.config.batch_size as u32 + 255) / 256).max(1);
let b = self.config.batch_size as i32;
let gamma_mod_ptr = self.gamma_mod_buf.raw_ptr();
unsafe {
self.stream.launch_builder(&self.fill_gamma_buf_kernel)
.arg(&self.gamma_buf)
.arg(&base_gamma)
.arg(&gamma_mod_ptr)
.arg(&b)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("fill_gamma_buf: {e}")))?;
}
Ok(())
}
/// LR scale factor for new components: min(1.0, step / 500).
fn new_component_lr_scale(&self) -> f32 {
(self.new_component_warmup_step as f32 / NEW_COMPONENT_WARMUP_STEPS as f32).min(1.0)
@@ -4160,7 +4183,9 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("isv_signal_update load: {e}")))?;
let isv_forward_kernel = exp_module_for_mag.load_function("isv_forward")
.map_err(|e| MLError::ModelError(format!("isv_forward load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward kernels loaded");
let fill_gamma_buf_kernel = exp_module_for_mag.load_function("fill_gamma_buf")
.map_err(|e| MLError::ModelError(format!("fill_gamma_buf load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward + fill_gamma_buf kernels loaded");
// ── G5: Epistemic-gated magnitude — pinned var_ema threshold ─
let (var_ema_pinned, var_ema_dev_ptr) = {
@@ -5654,6 +5679,7 @@ impl GpuDqnTrainer {
lagged_td_error_dev_ptr,
isv_signal_update_kernel,
isv_forward_kernel,
fill_gamma_buf_kernel,
isv_embedding_buf,
branch_gate_buf,
gamma_mod_buf,
@@ -6292,6 +6318,7 @@ impl GpuDqnTrainer {
self.launch_mse_loss()?;
self.launch_loss_reduce(self.mse_loss_dev_ptr)?;
self.launch_mse_grad_to_scratch()?;
self.fill_gamma_buf()?;
self.launch_c51_loss()?;
self.launch_loss_reduce(self.total_loss_dev_ptr)?;
self.launch_c51_grad()?;
@@ -7545,6 +7572,7 @@ impl GpuDqnTrainer {
self.q_divergence_dev_ptr, 0, std::mem::size_of::<f32>(), self.stream.cu_stream(),
);
}
self.fill_gamma_buf()?;
self.launch_c51_loss()?;
self.launch_loss_reduce(self.total_loss_dev_ptr)?;
self.launch_c51_grad()?;
@@ -7697,6 +7725,7 @@ impl GpuDqnTrainer {
self.q_divergence_dev_ptr, 0, std::mem::size_of::<f32>(), self.stream.cu_stream(),
);
}
self.fill_gamma_buf()?;
self.launch_c51_loss()?;
self.launch_loss_reduce(self.total_loss_dev_ptr)?;
self.launch_c51_grad()?;
@@ -8169,9 +8198,9 @@ impl GpuDqnTrainer {
let on_next_b2_ptr = on_next_b1_ptr + (b * b1 * na * f32_sz) as u64;
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.
// The experience collector pre-computes R_n = sum(gamma^i * r_i).
let gamma = self.adaptive_gamma.powi(self.config.n_steps as i32);
// N-step returns: gamma_buf [B] is pre-filled by fill_gamma_buf kernel
// (base_gamma^n * gamma_mod[0] per sample).
let gamma_buf_ptr = self.gamma_buf.raw_ptr();
let batch_i32 = b as i32;
let na_i32 = na as i32;
let b0_i32 = b0 as i32;
@@ -8222,7 +8251,7 @@ impl GpuDqnTrainer {
.arg(&self.curiosity_error_buf)
.arg(&self.config.curiosity_q_penalty_lambda)
// ── Config (8 — per_sample_support replaces v_min+v_max) ──
.arg(&gamma)
.arg(&gamma_buf_ptr)
.arg(&batch_i32)
.arg(&na_i32)
.arg(&self.per_sample_support_ptr)