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

feat(F4/D5): Information Bottleneck variance gradient — spreads Q within direction branch

Browse Source 
Adds ib_gradient_direction CUDA kernel that fires when population variance
of Q(a) across b0 actions falls below min_var=0.01, pushing each Q(a) away
from mean_q. Piggybacked on cql_d_adv_logits / cql_d_value_logits (same
atomicAdd path as F5 barrier). ib_weight = 0.05 × (1−health) → zero-op when
network is healthy (health≈1). Wired inside apply_cql_gradient immediately
after the F5 barrier launch, flows through the existing CQL SAXPY + backward.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 21:40:54 +02:00
parent 5bc712478e
commit 85b7c10421
2 changed files with 185 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -860,6 +860,127 @@ extern "C" __global__ void barrier_gradient_direction(
}
}
/* ══════════════════════════════════════════════════════════════════════
* F4/D5: INFORMATION BOTTLENECK VARIANCE GRADIENT KERNEL
*
* Per sample i, for the direction branch (b0):
* Q(a) = Σ_z p(z|a) * z_vals[z] via softmax over dueling logits.
* mean_q = (1/b0) * Σ_a Q(a)
* var_q = (1/b0) * Σ_a (Q(a) - mean_q)^2 (population variance)
* ib_weight = 0.05 * (1 - health)
* penalty = max(0, min_var - var_q)
*
* When var_q < min_var, adds to cql_d_adv_logits / cql_d_value_logits to
* spread Q values across actions (increase var_q → decrease ib_loss).
*
* Gradient derivation:
* ib_loss = ib_weight * max(0, min_var - var_q)
* d(ib_loss)/dQ(a) = -ib_weight * (2/b0) * (Q(a) - mean_q) when var_q < min_var
*
* Gradient descent: params -= lr * grad.
* So for Q(a) > mean_q: negative dib/dQ(a) means we subtract a negative value
* from Q(a)'s gradient, i.e. we INCREASE Q(a). Symmetrically for Q(a) < mean_q.
* Net effect: Q values are pushed AWAY from mean_q — spreading them.
*
* Then: d(ib_loss)/d_logit(a,z) = d(ib_loss)/dQ(a) * dQ(a)/d_logit(a,z)
* where dQ(a)/d_logit(a,z) = (z_vals[z] - Q(a)) * p(z|a)
*
* Piggybacked on the CQL d-logit accumulators (same as F5). No new cuBLAS pass.
* Direction branch only. Zero-op when var_q >= min_var or ib_weight < 1e-6.
*
* Launch: grid=(ceil(B/256)), block=(256). One thread per sample.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void ib_gradient_direction(
const float* __restrict__ adv_logits, /* [B, b0_size, num_atoms] f32 */
const float* __restrict__ v_logits, /* [B, num_atoms] f32 */
const float* __restrict__ z_vals, /* [num_atoms] f32 */
const float* __restrict__ isv_signals, /* [ISV_DIM=13] pinned; [12]=health */
float* __restrict__ d_adv_logits, /* [B, b0_size, num_atoms] — atomicAdd */
float* __restrict__ d_v_logits, /* [B, num_atoms] — atomicAdd */
int batch_size,
int num_atoms,
int b0_size,
float min_var
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= batch_size) return;
/* Safety guard */
if (b0_size <= 1 || b0_size > 16) return;
float health = (isv_signals != NULL) ? isv_signals[12] : 0.5f;
float ib_weight = 0.05f * (1.0f - health);
if (ib_weight < 1e-6f) return;
const float* v_row = v_logits + (long long)i * num_atoms;
const float* adv_row = adv_logits + (long long)i * b0_size * num_atoms;
/* Compute Q(a) for each direction action via numerically-stable softmax. */
float q_vals[16];
for (int a = 0; a < b0_size; a++) {
const float* adv_a = adv_row + a * num_atoms;
float max_l = -INFINITY;
for (int z = 0; z < num_atoms; z++) {
float l = v_row[z] + adv_a[z];
if (l > max_l) max_l = l;
}
float sum = 0.0f;
for (int z = 0; z < num_atoms; z++)
sum += expf(v_row[z] + adv_a[z] - max_l);
float q = 0.0f;
for (int z = 0; z < num_atoms; z++) {
float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f);
q += p * z_vals[z];
}
q_vals[a] = q;
}
/* Population variance across actions. */
float mean_q = 0.0f;
for (int a = 0; a < b0_size; a++) mean_q += q_vals[a];
mean_q /= (float)b0_size;
float var_q = 0.0f;
for (int a = 0; a < b0_size; a++) {
float d = q_vals[a] - mean_q;
var_q += d * d;
}
var_q /= (float)b0_size;
/* Zero-op when variance already meets the minimum. */
if (var_q >= min_var) return;
/* grad_scale = -ib_weight * (2 / b0_size)
* Gradient descent subtracts this times (Q(a)-mean_q), spreading Q values. */
float grad_scale = -ib_weight * (2.0f / (float)b0_size);
float* d_v_row = d_v_logits + (long long)i * num_atoms;
for (int a = 0; a < b0_size; a++) {
float dq_a = grad_scale * (q_vals[a] - mean_q);
if (fabsf(dq_a) < 1e-9f) continue;
/* Recompute softmax probabilities for this action to get p(z|a). */
const float* adv_a = adv_row + a * num_atoms;
float max_l = -INFINITY;
for (int z = 0; z < num_atoms; z++) {
float l = v_row[z] + adv_a[z];
if (l > max_l) max_l = l;
}
float sum = 0.0f;
for (int z = 0; z < num_atoms; z++)
sum += expf(v_row[z] + adv_a[z] - max_l);
float* d_adv_a = d_adv_logits + (long long)i * b0_size * num_atoms + a * num_atoms;
float q_a = q_vals[a];
for (int z = 0; z < num_atoms; z++) {
float p = expf(v_row[z] + adv_a[z] - max_l) / (sum + 1e-8f);
float contribution = dq_a * (z_vals[z] - q_a) * p;
atomicAdd(&d_adv_a[z], contribution);
atomicAdd(&d_v_row[z], contribution);
}
}
}
/* ══════════════════════════════════════════════════════════════════════
* C51 MANIFOLD MIXUP CE KERNEL (separate launch, no barrier)
*

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

@@ -1332,6 +1332,12 @@ pub struct GpuDqnTrainer {
/// Loaded from c51_loss_kernel.cubin ("barrier_gradient_direction").
barrier_gradient_kernel: Option<CudaFunction>,
// ── F4/D5: Information Bottleneck variance gradient kernel ───────
/// Spreads Q values across direction-branch actions when population
/// variance is below min_var (0.01). Piggybacks on the CQL d-logit
/// buffers. Loaded from c51_loss_kernel.cubin ("ib_gradient_direction").
ib_gradient_kernel: Option<CudaFunction>,
/// v8: PopArt reward normalization kernel (running mean/variance).
popart_normalize_kernel: CudaFunction,
/// v8: Robust PopArt kernel (median/IQR normalization).
@@ -4708,6 +4714,43 @@ impl GpuDqnTrainer {
}
}
// F4/D5: IB variance gradient — spreads Q values in the direction branch.
// Runs after F5 barrier (both atomicAdd into same buffers, races are fine).
// Zero-op when var_q >= 0.01 or health ≈ 1 (ib_weight ≈ 0).
if let Some(ref ib_kernel) = self.ib_gradient_kernel.clone() {
let adv_ptr_ib = self.on_b_logits_buf.raw_ptr();
let v_ptr_ib = self.on_v_logits_buf.raw_ptr();
let z_ptr_ib = self.atom_positions_buf.raw_ptr();
let isv_dev_ib = self.isv_signals_dev_ptr;
let d_adv_ptr_ib = self.cql_d_adv_logits.raw_ptr();
let d_v_ptr_ib = self.cql_d_value_logits.raw_ptr();
let b_i32_ib = b as i32;
let na_i32_ib = na as i32;
let b0_i32_ib = b0 as i32;
let blocks_ib = blocks;
let min_var = 0.01_f32;
unsafe {
let _ = self.stream
.launch_builder(ib_kernel)
.arg(&adv_ptr_ib)
.arg(&v_ptr_ib)
.arg(&z_ptr_ib)
.arg(&isv_dev_ib)
.arg(&d_adv_ptr_ib)
.arg(&d_v_ptr_ib)
.arg(&b_i32_ib)
.arg(&na_i32_ib)
.arg(&b0_i32_ib)
.arg(&min_var)
.launch(LaunchConfig {
grid_dim: (blocks_ib, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
});
// Non-fatal: log failure but continue with CQL backward.
}
}
// 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
@@ -5593,6 +5636,8 @@ impl GpuDqnTrainer {
let c51_grad_kernel = compile_c51_grad_kernel(&stream, &config)?;
// F5/D2: Q-gap barrier gradient kernel (from same c51_loss cubin).
let barrier_gradient_kernel = load_barrier_gradient_kernel(&stream)?;
// F4/D5: Information Bottleneck variance gradient kernel (from same c51_loss cubin).
let ib_gradient_kernel = load_ib_gradient_kernel(&stream)?;
info!("GpuDqnTrainer: c51_loss + c51_grad kernels compiled");
// ── Compile MSE loss + gradient kernels (warmup before C51) ─
@@ -7353,6 +7398,7 @@ impl GpuDqnTrainer {
cql_d_value_logits,
cql_d_adv_logits,
barrier_gradient_kernel,
ib_gradient_kernel,
curiosity_error_buf,
drawdown_depths_buf,
asymmetric_dd_weight: dd_weight,
@@ -13002,6 +13048,24 @@ fn load_barrier_gradient_kernel(
}
}
/// Load the F4/D5 Information Bottleneck variance gradient kernel from the c51_loss cubin.
///
/// Returns Ok(None) gracefully if the symbol is not found (older cubin cached).
fn load_ib_gradient_kernel(
stream: &Arc<CudaStream>,
) -> Result<Option<CudaFunction>, MLError> {
let context = stream.context();
let module = context.load_cubin(C51_LOSS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("c51_loss cubin load (ib): {e}")))?;
match module.load_function("ib_gradient_direction") {
Ok(f) => Ok(Some(f)),
Err(e) => {
tracing::warn!("ib_gradient_direction not found in cubin (non-fatal): {e}");
Ok(None)
}
}
}
/// Load the C51 loss gradient kernel from precompiled cubin.
///
/// entropy_coeff is now passed as a runtime kernel parameter (not #define).