merge: SP6 Pearl 3 — NoisyLinear per-branch σ array

Brings in worktree-agent-a4d8a879 (commit ed3fa066b): per-branch σ via
[4]-element mapped-pinned device buffer. add_advantage_noise kernel
indexes σ by branch derived from action_idx % total_actions; Q-value
layout is branch-major contiguous so per-branch σ derivation requires
no forward-pass restructuring.

3 ExperienceCollectorConfig constructors updated.

Resolves Pearl 3 averaging from SP5 Layer B which collapsed 4 per-branch
σ values into a single scalar via training_loop.rs:1747.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-02 09:10:51 +02:00
4 changed files with 78 additions and 30 deletions

View File

@@ -3907,31 +3907,50 @@ extern "C" __global__ void step_counter_advance(int* counter) {
}
/**
* Add per-action Gaussian noise to Q-values for exploration (replaces dead NoisyLinear).
* Uses Philox PRNG seeded from (thread_id, step). Only fires during experience collection.
* Breaks the dueling symmetry trap: identical Q-values → uniform Boltzmann → zero advantage gradient.
* SP6 Pearl 3: Add per-branch Gaussian noise to Q-values for exploration.
* Replaces dead NoisyLinear. Uses Philox PRNG seeded from (thread_id, step).
* Each action slot receives noise scaled by its branch's sigma.
* branch_sizes {b0, b1, b2, b3}: cumulative offsets derive branch index.
* Sigma=0 for a branch disables noise for that branch entirely.
*
* Grid: ceil(N * total_actions / 256), Block: 256.
*/
extern "C" __global__ void add_advantage_noise(
float* __restrict__ q_values, /* [N, total_actions] — modified in-place */
float noise_sigma, /* noise std dev (e.g. 0.1) */
int N, /* number of episodes */
int total_actions, /* 13 = 4+3+3+3 */
unsigned int seed /* monotonic step counter for decorrelation */
float* __restrict__ q_values, /* [N, total_actions] — modified in-place */
const float* __restrict__ noise_sigma, /* [4] per-branch std dev */
int N, /* number of episodes */
int total_actions, /* 13 = 4+3+3+3 */
unsigned int seed, /* monotonic step counter for decorrelation */
int b0_size, /* branch 0 action count (direction = 4) */
int b1_size, /* branch 1 action count (magnitude = 3) */
int b2_size, /* branch 2 action count (order = 3) */
int b3_size /* branch 3 action count (urgency = 3) */
) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total = N * total_actions;
if (tid >= total) return;
int action = tid % total_actions;
/* Derive branch index from action offset */
int branch_idx;
int b01 = b0_size + b1_size;
int b012 = b01 + b2_size;
if (action < b0_size) branch_idx = 0;
else if (action < b01) branch_idx = 1;
else if (action < b012) branch_idx = 2;
else branch_idx = 3;
float sigma = noise_sigma[branch_idx];
if (sigma <= 0.0f) return;
/* Stateless Philox hash — deterministic per (tid, seed) */
int episode = tid / total_actions;
int action = tid % total_actions;
float u1 = fmaxf(philox_uniform(episode, (int)seed, 4000 + action * 2), 1e-6f);
float u2 = philox_uniform(episode, (int)seed, 4000 + action * 2 + 1);
float gauss = sqrtf(-2.0f * logf(u1)) * cosf(6.2831853f * u2);
q_values[tid] += noise_sigma * gauss;
q_values[tid] += sigma * gauss;
}
/* ══════════════════════════════════════════════════════════════════════

View File

@@ -292,10 +292,10 @@ pub struct ExperienceCollectorConfig {
pub hindsight_fraction: f32,
/// v8: Lookahead bars for hindsight optimal exit.
pub hindsight_lookahead: i32,
/// Advantage noise std for exploration (replaces dead NoisyLinear).
/// Per-action Gaussian noise breaks the dueling symmetry trap.
/// 0.0 = disabled. Default: 0.1.
pub noise_sigma: f32,
/// SP6 Pearl 3: per-branch advantage noise std for exploration.
/// Branch order: [direction(0), magnitude(1), order(2), urgency(3)].
/// 0.0 per-branch disables noise for that branch. Default: [0.1; 4].
pub noise_sigma_per_branch: [f32; 4],
}
impl Default for ExperienceCollectorConfig {
@@ -366,7 +366,7 @@ impl Default for ExperienceCollectorConfig {
max_trace_length: 7,
hindsight_fraction: 0.0,
hindsight_lookahead: 10,
noise_sigma: 0.1,
noise_sigma_per_branch: [0.1; 4],
}
}
}
@@ -757,6 +757,9 @@ pub struct GpuExperienceCollector {
/// Advantage noise kernel for dueling symmetry breaking.
/// Per-action Gaussian noise replaces dead NoisyLinear (cuBLAS bypasses Candle).
noise_kernel: CudaFunction,
/// SP6 Pearl 3: [4] mapped-pinned device buffer for per-branch noise sigma.
/// CPU writes via host_ptr before each noise kernel launch; kernel reads via dev_ptr.
noise_sigma_dev: MappedF32Buffer,
// Pre-allocated pinned host buffers for DtoH transfers.
pinned_states: PinnedHostBuf<f32>,
@@ -1565,6 +1568,14 @@ impl GpuExperienceCollector {
.map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf (mapped pinned): {e}")))?
};
// SP6 Pearl 3: [4] mapped-pinned buffer for per-branch noise sigma.
// CPU writes the 4 sigma values before each noise kernel launch;
// kernel reads via dev_ptr. Zero memcpy per feedback_no_htod.
let noise_sigma_dev = unsafe {
MappedF32Buffer::new(4)
.map_err(|e| MLError::ModelError(format!("alloc noise_sigma_dev (mapped pinned): {e}")))?
};
// HOT path: hindsight relabel bar_indices buffer.
// Capacity = alloc_episodes * alloc_timesteps * 2 (cf doubles total).
let bar_indices_pinned = unsafe {
@@ -1729,6 +1740,7 @@ impl GpuExperienceCollector {
step_counter_gpu,
step_counter_kernel,
noise_kernel,
noise_sigma_dev,
pinned_states,
pinned_actions,
pinned_rewards,
@@ -3617,11 +3629,17 @@ impl GpuExperienceCollector {
}
}
// ── 3d. Add advantage noise for dueling symmetry breaking ──
// Replaces dead NoisyLinear (cuBLAS bypasses Candle).
// ── 3d. Add per-branch advantage noise for dueling symmetry breaking ──
// SP6 Pearl 3: noise_sigma_per_branch[4] replaces scalar noise_sigma.
// Per-action Gaussian noise creates non-uniform Q-values → non-uniform
// Boltzmann → advantage heads get differentiated gradient.
if config.noise_sigma > 0.0 {
let any_branch_noise = config.noise_sigma_per_branch.iter().any(|&s| s > 0.0);
if any_branch_noise {
// Write per-branch sigma values to mapped-pinned buffer.
// Kernel reads via dev_ptr — no HtoD copy per feedback_no_htod.
self.noise_sigma_dev.write_from_slice(&config.noise_sigma_per_branch);
let sigma_dev_ptr = self.noise_sigma_dev.dev_ptr;
let total_q = n_episodes * (self.branch_sizes[0] + self.branch_sizes[1]
+ self.branch_sizes[2] + self.branch_sizes[3]);
let noise_blocks = ((total_q + 255) / 256) as u32;
@@ -3632,15 +3650,23 @@ impl GpuExperienceCollector {
};
let total_actions_i32 = (self.branch_sizes[0] + self.branch_sizes[1]
+ self.branch_sizes[2] + self.branch_sizes[3]) as i32;
let b0 = self.branch_sizes[0] as i32;
let b1 = self.branch_sizes[1] as i32;
let b2 = self.branch_sizes[2] as i32;
let b3 = self.branch_sizes[3] as i32;
let seed = t as u32; // monotonic step counter
unsafe {
self.stream
.launch_builder(&self.noise_kernel)
.arg(&mut self.q_values)
.arg(&config.noise_sigma)
.arg(&sigma_dev_ptr)
.arg(&n_i32)
.arg(&total_actions_i32)
.arg(&seed)
.arg(&b0)
.arg(&b1)
.arg(&b2)
.arg(&b3)
.launch(noise_cfg)
.map_err(|e| MLError::ModelError(format!(
"add_advantage_noise t={t}: {e}"

View File

@@ -1741,16 +1741,19 @@ impl DQNTrainer {
hindsight_fraction: self.hyperparams.hindsight_fraction as f32,
hindsight_lookahead: self.hyperparams.hindsight_lookahead as i32,
// SP5 Layer B (Pearl 3): read per-branch NoisyNet σ from ISV[210..214)
// and collapse to a single effective scale by averaging across 4 branches.
// Falls back to the static hyperparam when fused_ctx is unavailable
// (pre-initialization path). 0.01f floor = Invariant 1 carve-out.
noise_sigma: self.fused_ctx.as_ref().map(|ctx| {
let mean = (0..4_usize)
.map(|b| ctx.read_isv_signal_at(NOISY_SIGMA_BASE + b))
.sum::<f32>()
/ 4.0_f32;
mean.max(0.01_f32)
}).unwrap_or(self.hyperparams.noise_sigma as f32),
// SP6 Pearl 3: read per-branch sigma directly from ISV[210..214).
// No averaging. Cold-start floor 0.01 = Invariant 1 carve-out.
// Falls back to static hyperparam repeated 4 times when fused_ctx
// is unavailable (pre-initialization path).
noise_sigma_per_branch: {
let mut arr = [self.hyperparams.noise_sigma as f32; 4];
if let Some(ctx) = self.fused_ctx.as_ref() {
for b in 0..4_usize {
arr[b] = ctx.read_isv_signal_at(NOISY_SIGMA_BASE + b).max(0.01_f32);
}
}
arr
},
..Default::default()
};

View File

@@ -138,7 +138,7 @@ Constants live in `crates/ml/src/cuda_pipeline/sp5_isv_slots.rs`.
| [198..202) | `BUDGET_CQL_BASE` | per-branch [4] | Pearl 2 | CQL loss budget weight. SP6 Pearl 2: `compute_adaptive_budgets()` reads individually, applies correction-factor sub-launches via `apply_cql_saxpy_branch`. |
| [202..206) | `BUDGET_ENS_BASE` | per-branch [4] | Pearl 2 | Ensemble loss budget weight. SP6 Pearl 2: used as trunk-mean only. |
| [206..210) | `FLATNESS_BASE` | per-branch [4] | Pearl 2 | Loss flatness diagnostic |
| [210..214) | `NOISY_SIGMA_BASE` | per-branch [4] | Pearl 3 | NoisyNet σ level |
| [210..214) | `NOISY_SIGMA_BASE` | per-branch [4] | Pearl 3 | NoisyNet σ level — SP6 Pearl 3 consumer wired: `add_advantage_noise` kernel reads per-branch σ via mapped-pinned dev_ptr; `training_loop.rs` reads slots directly (no averaging) |
| [214..218) | `SIGMA_FRACTION_BASE` | per-branch [4] | Pearl 3 | NoisyNet σ fraction |
| [218..222) | `BRANCH_ENTROPY_BASE` | per-branch [4] | Pearl 3 | Branch action entropy |
| [222..226) | `Q_VAR_PER_BRANCH_BASE` | per-branch [4] | shared | Q-value variance |