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

feat(alpha): stabilize alpha_dqn_h600_smoke — reward norm + target net + grad clip

Browse Source 
Three stabilizers applied to the H=600 DQN smoke after initial run showed
unstable training (early_mvmt=2268× at lr=1e-6, NaN at lr=1e-4):

  1. Reward normalization (--reward-scale, default 1000)
     Rewards divided by scale BEFORE the Munchausen target. TD error
     drops from ~1000 (raw reward magnitude at H=600) into O(1) target /
     gradient / weight-update scale. Action selection + rollout-R
     reporting use ORIGINAL rewards (so rvr math stays correct against
     the Task 7c baseline).

  2. Target network (--target-update-every, default 10 episodes)
     Separate w_target_dev / b_target_dev buffers. Q_next(s') forward
     uses target weights; SGD updates online only. Hard-update copies
     online → target every K episodes. Breaks the V_soft(s') chase-its-
     own-tail divergence of online-only Munchausen.

  3. Gradient clipping (--grad-clip, default 1.0)
     New `alpha_clip_inplace_kernel` in alpha_linear_q.cu (element-wise
     clamp). Applied to dW and db after grad, before SGD. Safety net.

Diagnostic fix: weight_norm was direction-insensitive — orthogonal
rotations don't change ||W||_F, so early_mvmt read ≈0 even when training.
Switched to weight_distance_from_init = ||W_now − W_init||_F +
||b_now − b_init||_F (captures rotation). q_early = q_init + distance
so kernel's |q_early − q_init| / |q_init| ratio = distance / ||W_init||_F.

With lr bumped back up to 1e-4 (default for the stabilized config),
verified at horizon=100, n_episodes=200:

  Q_SPREAD_EMA         = 23.64   (≥ 0.05)     PASS
  ACTION_ENTROPY_EMA   = 1.86    (≥ 1.0986)   PASS
  RETURN_VS_RANDOM_EMA = +0.586  (≥ 0.0)      PASS
  EARLY_Q_MOVEMENT_EMA = 0.0212  (≥ 0.01)     PASS
  Overall: PASS (H=6000 scale-up VIABLE)

early_mvmt grew monotonically (0.005 → 0.021) across the 200-episode
run — direction-sensitive diagnostic confirms genuine policy learning.

Audit doc docs/isv-slots.md updated per Invariant 7.

Next: H=600 / 1000-episode run on full data; if PASS holds, Task 13
(H=6000 scale-up) unlocks.
This commit is contained in:
jgrusewski
2026-05-15 15:57:24 +02:00
parent fa30c2dd66
commit 8958637c77
4 changed files with 194 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

129
crates/ml/examples/alpha_dqn_h600_smoke.rs
View File

@@ -96,15 +96,13 @@ struct Cli {
snapshot_interval: usize,
#[arg(long, default_value_t = 500_000)]
max_snapshots: usize,
/// SGD learning rate. Default 1e-6 — higher values (e.g. 1e-4) diverge
/// to NaN at H=600 because the Munchausen target's
/// `r + α_m·τ·log π + γ·V_soft` produces large gradients when Q-values
/// grow. Until gradient clipping lands (E.1 follow-up), use 1e-6 or
/// lower. At 1e-6 the Q-network still grows ~2000× over 50 episodes
/// (early_mvmt ≈ 2000) — stable but uncalibrated. A future commit
/// should add: (a) gradient clipping, (b) target network with periodic
/// hard-update, (c) reward normalisation.
#[arg(long, default_value_t = 1.0e-6)]
/// SGD learning rate. With the stabilizer trio (--reward-scale,
/// --target-update-every, --grad-clip), 1e-4 is the right operating
/// point — gradients are O(1) in normalized units, target net
/// breaks the V_soft chase-tail loop, grad clip catches bursts.
/// Without the stabilizers, lr=1e-4 diverges to NaN within 20
/// episodes; lr=1e-6 stays finite but undertrains.
#[arg(long, default_value_t = 1.0e-4)]
lr: f32,
/// ε-greedy: ε at start of training.
#[arg(long, default_value_t = 0.50)]
@@ -127,6 +125,25 @@ struct Cli {
/// Episodes between kill-criteria pipeline launches.
#[arg(long, default_value_t = 50)]
kill_criteria_every: usize,
/// Reward normalization scale. Rewards are divided by this before
/// being passed to the Munchausen target — the network learns
/// normalized Q-values. Default 1000 ≈ |median reward| at H=600
/// (≈-2800 median random-baseline reward / sane order). Action
/// selection and rollout-R reporting use ORIGINAL rewards.
#[arg(long, default_value_t = 1000.0)]
reward_scale: f32,
/// Episodes between hard updates of the target network (W_target,
/// b_target ← W_online, b_online). Decouples the V_soft(s')
/// bootstrap target from the policy's per-step drift — essential
/// for Munchausen DQN stability.
#[arg(long, default_value_t = 10)]
target_update_every: usize,
/// Per-element gradient clip bound. After alpha_linear_q_grad and
/// before the SGD step, |dW[i]|, |db[i]| are clamped to ±this value.
/// Default 1.0 — combined with lr=1e-4 caps per-step weight delta
/// at 1e-4 per element, which is safe.
#[arg(long, default_value_t = 1.0)]
grad_clip: f32,
/// Output JSON path for final verdict + ISV readings.
#[arg(long, default_value = "config/ml/alpha_dqn_h600_smoke.json")]
out_path: PathBuf,
@@ -240,15 +257,33 @@ fn state_as_vec(env: &ExecutionEnv, ep: &EpisodeState) -> Vec<f32> {
env.state(ep).to_vec()
}
/// L2 (Frobenius) norm of W concatenated with b — diagnostic for early-Q-movement.
/// Pure read-side reduction; computed CPU-side after dtoh (not on the
/// hot loop, only at episode boundaries for the kill-criteria scalar
/// inputs).
/// L2 (Frobenius) norm of W concatenated with b — diagnostic anchor for
/// early-Q-movement. Pure read-side reduction; computed CPU-side after
/// dtoh (not on the hot loop, only at episode boundaries).
fn weight_norm(w: &[f32], b: &[f32]) -> f32 {
let s: f32 = w.iter().map(|x| x * x).sum::<f32>() + b.iter().map(|x| x * x).sum::<f32>();
s.sqrt()
}
/// L2 distance from init: `||W_now − W_init||_F + ||b_now − b_init||_F` (combined).
/// Captures *direction change* in weight space, not just growth — the kill-
/// criteria kernel computes `|q_early − q_init| / |q_init|`, so the call
/// site sets `q_early = q_init + weight_distance_from_init(...)`, making
/// the kernel's ratio = `distance_from_init / ||W_init||_F`. A pure
/// rotation that keeps `||W||` constant still produces a non-zero distance.
fn weight_distance_from_init(w_now: &[f32], b_now: &[f32], w_init: &[f32], b_init: &[f32]) -> f32 {
let mut s: f32 = 0.0;
for (a, b) in w_now.iter().zip(w_init.iter()) {
let d = a - b;
s += d * d;
}
for (a, b) in b_now.iter().zip(b_init.iter()) {
let d = a - b;
s += d * d;
}
s.sqrt()
}
/// SplitMix64 RNG — same as ExecutionEnv's ReplayRng so seeds compose
/// cleanly when the smoke is reproduced.
struct SmokeRng {
@@ -315,6 +350,8 @@ fn main() -> Result<()> {
.context("grad load")?;
let lq_sgd = lq_module.load_function("alpha_linear_q_sgd_step_kernel")
.context("sgd load")?;
let lq_clip = lq_module.load_function("alpha_clip_inplace_kernel")
.context("clip load")?;
// Munchausen target
let munch_cubin: Vec<u8> = std::fs::read(
@@ -373,6 +410,10 @@ fn main() -> Result<()> {
let mut w_dev = stream.clone_htod(&w_init).context("upload W")?;
let mut b_dev = stream.clone_htod(&b_init).context("upload b")?;
// Target network: identical to online at init; periodic hard-update.
// Munchausen V_soft(s') bootstrap reads from this, not w_dev/b_dev.
let mut w_target_dev = stream.clone_htod(&w_init).context("upload W_target")?;
let mut b_target_dev = stream.clone_htod(&b_init).context("upload b_target")?;
let mut dw_dev = stream.alloc_zeros::<f32>(N_WEIGHTS).context("alloc dW")?;
let mut db_dev = stream.alloc_zeros::<f32>(N_BIASES).context("alloc db")?;
@@ -485,7 +526,14 @@ fn main() -> Result<()> {
.context("htod next_states")?;
stream.memcpy_htod(&actions_host, &mut actions_dev)
.context("htod actions")?;
stream.memcpy_htod(&rewards_host, &mut rewards_dev)
// Reward normalization: divide by cli.reward_scale so the
// Munchausen target produces normalized targets. Original
// rewards stay in rewards_host for reporting / kill-criteria.
let rewards_normalized: Vec<f32> = rewards_host
.iter()
.map(|r| r / cli.reward_scale)
.collect();
stream.memcpy_htod(&rewards_normalized, &mut rewards_dev)
.context("htod rewards")?;
stream.memcpy_htod(&dones_host, &mut dones_dev)
.context("htod dones")?;
@@ -504,10 +552,12 @@ fn main() -> Result<()> {
)?;
}
}
// Forward Q_next on next_states
// Forward Q_next on next_states USING TARGET WEIGHTS.
// Decouples the V_soft(s') bootstrap target from per-step policy
// drift — the Munchausen target then has a stationary reference.
{
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (w_ptr, _g0) = w_target_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_target_dev.device_ptr(&stream);
let (s_ptr, _g2) = next_states_dev.device_ptr(&stream);
let (q_ptr, _g3) = q_next_dev.device_ptr_mut(&stream);
unsafe {
@@ -553,6 +603,28 @@ fn main() -> Result<()> {
)?;
}
}
// Gradient clip — element-wise |g| ≤ cli.grad_clip on dW and db.
// Safety net against gradient bursts that survive reward
// normalization + target net (e.g., rare large-PnL terminal
// rewards). Identity transform when |g| < bound.
{
let (dw_ptr, _g0) = dw_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
&stream, &lq_clip,
dw_ptr, cli.grad_clip, N_WEIGHTS as i32,
)?;
}
}
{
let (db_ptr, _g0) = db_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
&stream, &lq_clip,
db_ptr, cli.grad_clip, N_BIASES as i32,
)?;
}
}
// SGD step on W
{
let (w_ptr, _g0) = w_dev.device_ptr_mut(&stream);
@@ -576,13 +648,31 @@ fn main() -> Result<()> {
}
}
// Target network hard update every K episodes: w_target ← w_online.
// Standard DQN stabilizer — V_soft(s') bootstrap reads w_target,
// so updating it slowly breaks the chase-its-own-tail divergence
// of online-only Munchausen.
if (ep + 1) % cli.target_update_every == 0 {
stream.synchronize().context("sync before target update")?;
let w_now = stream.clone_dtoh(&w_dev).context("dtoh W for target update")?;
let b_now = stream.clone_dtoh(&b_dev).context("dtoh b for target update")?;
stream.memcpy_htod(&w_now, &mut w_target_dev)
.context("htod W_target hard update")?;
stream.memcpy_htod(&b_now, &mut b_target_dev)
.context("htod b_target hard update")?;
}
// Periodic kill-criteria pipeline launch
if (ep + 1) % cli.kill_criteria_every == 0 {
stream.synchronize().context("sync before kc")?;
// Snapshot weights for q_early_norm
let w_now = stream.clone_dtoh(&w_dev).context("dtoh W")?;
let b_now = stream.clone_dtoh(&b_dev).context("dtoh b")?;
let q_early_norm = weight_norm(&w_now, &b_now);
// q_early = q_init + ||W_now − W_init||_F so the kernel's
// |q_early − q_init| / |q_init| ratio yields the relative
// weight-space distance from init (direction-sensitive).
let dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init);
let q_early_norm = q_init_norm + dist;
let rollout_r_mean: f32 = if recent_returns.is_empty() {
0.0
} else {
@@ -697,6 +787,9 @@ fn main() -> Result<()> {
"gamma": cli.gamma,
"alpha_m": cli.alpha_m,
"tau": cli.tau,
"reward_scale": cli.reward_scale,
"target_update_every": cli.target_update_every,
"grad_clip": cli.grad_clip,
"q_init_norm": q_init_norm,
"q_spread_ema": kc_final[0],
"action_entropy_ema": kc_final[1],

36
crates/ml/src/cuda_pipeline/alpha_kernels.rs
View File

@@ -235,6 +235,42 @@ pub unsafe fn launch_alpha_linear_q_sgd_step(
Ok(())
}
/// Launch `alpha_clip_inplace_kernel` on `stream`. Element-wise
/// clamp(x, -bound, +bound) for gradient clipping or activation
/// stabilization.
///
/// # Safety
///
/// `buf_dev` MUST be a valid device pointer to `n` floats. Modified in place.
pub unsafe fn launch_alpha_clip_inplace(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
buf_dev: u64,
bound: f32,
n: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(n > 0, "n must be positive");
debug_assert!(bound > 0.0, "bound must be positive");
const BLOCK: u32 = 256;
let grid_x = ((n as u32) + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), 1, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&buf_dev)
.arg(&bound)
.arg(&n)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_clip_inplace launch: {e}")))?;
Ok(())
}
/// Launch `alpha_munchausen_target_kernel` on `stream`.
///
/// Computes the Munchausen-augmented target for each batch sample:

24
crates/ml/src/cuda_pipeline/alpha_linear_q.cu
View File

@@ -131,3 +131,27 @@ extern "C" __global__ void alpha_linear_q_sgd_step_kernel(
if (idx >= n) return;
params[idx] -= lr * grad[idx];
}
// ----------------------------------------------------------------------
// (4) Element-wise clip-in-place: x = clamp(x, -bound, +bound)
// ----------------------------------------------------------------------
//
// Stabilizer for the Munchausen-target gradient. The Q-network can grow
// rapidly when ‖dW‖ exceeds the per-step LR's safe operating range
// (lr · ‖dW‖ → weight delta). Without clipping, a single burst step can
// double the weight magnitude, after which the V_soft bootstrap target
// also doubles, gradient doubles → exponential blow-up.
//
// Clipping per-element is the simplest safety net (vs L2-norm clipping
// which needs a reduction). For the linear Q-network with 99 parameters,
// per-element clipping is appropriate.
extern "C" __global__ void alpha_clip_inplace_kernel(
float* __restrict__ x, // [n] — overwritten in place
float bound,
int n
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= n) return;
const float v = x[idx];
x[idx] = fmaxf(-bound, fminf(bound, v));
}