diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 477344aa7..4a239f6b2 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -68,6 +68,7 @@ const KERNELS: &[&str] = &[ "rl_isv_write", // RL R9: generic single-slot ISV seeder for tunable design constants — no HtoD "rl_q_pi_agree_b", // RL R9 audit: per-batch fraction (argmax Q == argmax π) → ISV[407] EMA; wires previously-dead slot "rl_pi_action_kernel", // audit Option B: π drives action selection via multinomial sampling from softmax(pi_logits); Q becomes pure critic + "rl_reward_clamp_controller", // audit 2026-05-24: adaptive [-LOSS, +WIN] clamp from positive-tail EMA; replaces static [-3, +1] that crushed winning-trade signal in rmgm5 ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/apply_reward_scale.cu b/crates/ml-alpha/cuda/apply_reward_scale.cu index 4156c0204..72634e680 100644 --- a/crates/ml-alpha/cuda/apply_reward_scale.cu +++ b/crates/ml-alpha/cuda/apply_reward_scale.cu @@ -54,19 +54,37 @@ // Win cap default 1.0 matches the C51 atom span (V_MAX=1.0); loss cap // default 3.0 preserves loss-aversion asymmetry (HFT P&L losses are // typically 2-3× larger than wins per close). Seeded at trainer init -// by `rl_isv_write`; tunable at runtime by re-seeding. +// by `rl_isv_write`; tunable at runtime by re-seeding. Also adaptive +// per-step via `rl_reward_clamp_controller` reading the per-step +// positive-scaled max published below. #define RL_REWARD_CLAMP_WIN_INDEX 452 #define RL_REWARD_CLAMP_LOSS_INDEX 453 +// audit 2026-05-24: per-step max(positive scaled reward, 0) published +// here so `rl_reward_clamp_controller` can adapt WIN/LOSS bounds from +// the actual positive-tail distribution rather than the static 1.0/3.0 +// defaults. alpha-rl-rmgm5 evidence: clamp fired on 85% of steps with +// p95 pre_clamp = 15.5× the WIN bound, crushing the gradient signal +// of profitable trades down to 1.0 while losses (avg -3.84) routinely +// exceeded the LOSS bound 3.0. Adaptive clamp lets the winning tail +// reach the C51 atom span properly. +#define RL_POS_SCALED_REWARD_MAX_INDEX 478 + // Single block; grid-stride for b_size > BLOCK_DIM. Caller launches // with block_dim = min(b_size, 256), grid_dim = 1, shared bytes = // block_dim * sizeof(float). +// Shared-mem layout: caller passes `shared_mem_bytes = 2 * block_dim * +// sizeof(float)`. Lower half [0..block_dim) tracks max(|scaled|) for +// the diag slot; upper half [block_dim..2*block_dim) tracks max(positive +// scaled, 0) for the adaptive clamp controller's input signal. extern "C" __global__ void apply_reward_scale( float* __restrict__ isv, // ISV bus IN/OUT float* __restrict__ rewards, // [b_size] IN/OUT (scaled+clamped) int b_size ) { - extern __shared__ float s_max[]; + extern __shared__ float smem[]; + float* s_abs = smem; + float* s_pos = smem + blockDim.x; const int tid = threadIdx.x; const float scale = isv[RL_REWARD_SCALE_INDEX]; // ISV-driven clamp bounds (was hardcoded 1.0 / 3.0). Reading per @@ -75,8 +93,11 @@ extern "C" __global__ void apply_reward_scale( const float clamp_win = isv[RL_REWARD_CLAMP_WIN_INDEX]; const float clamp_loss = isv[RL_REWARD_CLAMP_LOSS_INDEX]; - // Pass 1 — scale, clamp, accumulate per-thread max |scaled_pre_clamp|. - float local_max = 0.0f; + // Pass 1 — scale, clamp, accumulate per-thread max |scaled| AND max + // max(positive scaled, 0). Two independent local maxes tracked in + // a single grid-stride loop so the kernel's HBM read pass is shared. + float local_abs = 0.0f; + float local_pos = 0.0f; for (int b = tid; b < b_size; b += blockDim.x) { const float raw = rewards[b]; const float scaled = raw * scale; @@ -84,7 +105,17 @@ extern "C" __global__ void apply_reward_scale( // Track magnitude BEFORE clamping so the diag sees what the // controller failed to bound on its own. const float a = fabsf(scaled); - if (a > local_max) local_max = a; + if (a > local_abs) local_abs = a; + + // Track positive tail BEFORE clamping for the adaptive clamp + // controller — winning-trade signal that the current static + // WIN=1.0 bound clips. Negative-tail (losses) tracked + // implicitly via max(|scaled|) at the diag slot since losses + // dominate the abs-distribution; the controller assumes the + // documented loss-aversion ratio (LOSS = ratio × WIN) so we + // don't need a separate negative-tail tracker. + const float p = fmaxf(scaled, 0.0f); + if (p > local_pos) local_pos = p; // Asymmetric clamp: [-clamp_loss, +clamp_win] from ISV. // Defaults preserve loss aversion (loss > win) while bounding @@ -93,24 +124,29 @@ extern "C" __global__ void apply_reward_scale( fminf(scaled, clamp_win)); rewards[b] = clamped; } - s_max[tid] = local_max; + s_abs[tid] = local_abs; + s_pos[tid] = local_pos; __syncthreads(); // Pass 2 — block tree reduce (no atomics per pearl_no_atomicadd). + // Both reductions advance together to share the __syncthreads barrier. for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { if (tid < stride) { - const float a = s_max[tid]; - const float b = s_max[tid + stride]; - s_max[tid] = a > b ? a : b; + const float a_a = s_abs[tid]; + const float a_b = s_abs[tid + stride]; + s_abs[tid] = a_a > a_b ? a_a : a_b; + const float p_a = s_pos[tid]; + const float p_b = s_pos[tid + stride]; + s_pos[tid] = p_a > p_b ? p_a : p_b; } __syncthreads(); } - // Thread 0 publishes the per-step pre-clamp max to the diagnostic - // ISV slot. This is a POINT measurement (not an EMA) — each - // training step overwrites with the current batch's max so the - // diag emitter sees the value for THIS step's reward batch. + // Thread 0 publishes per-step maxes to ISV. Both are POINT + // measurements; the controller maintains the EMA over the positive + // signal in its own slot. if (tid == 0) { - isv[RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX] = s_max[0]; + isv[RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX] = s_abs[0]; + isv[RL_POS_SCALED_REWARD_MAX_INDEX] = s_pos[0]; } } diff --git a/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu b/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu new file mode 100644 index 000000000..c4232cbeb --- /dev/null +++ b/crates/ml-alpha/cuda/rl_reward_clamp_controller.cu @@ -0,0 +1,95 @@ +// rl_reward_clamp_controller.cu — emits adaptive reward-clamp bounds to +// ISV[RL_REWARD_CLAMP_WIN_INDEX=452] and ISV[RL_REWARD_CLAMP_LOSS_INDEX=453]. +// +// Audit 2026-05-24: the static clamp `[-3.0, +1.0]` set up by +// `rl_isv_write` at trainer init was load-bearing on 85% of steps in +// alpha-rl-rmgm5 — pre-clamp p95 = 15.5× WIN, max = 2830× WIN. The +// adaptive bound below tracks the EMA of the per-step positive-tail +// max (published by `apply_reward_scale` into slot 478) so the C51 +// Bellman target receives the actual winning-trade magnitude +// distribution rather than a near-constant +1.0 clipped signal. +// +// Reads: +// isv[RL_POS_SCALED_REWARD_MAX_INDEX = 478] — per-step max(scaled, 0) +// isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX = 479] — controller's EMA state +// isv[RL_REWARD_CLAMP_MARGIN_INDEX = 480] — k multiplier (default 1.5) +// isv[RL_REWARD_CLAMP_RATIO_INDEX = 481] — loss/win ratio (default 3.0) +// +// Writes: +// isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX = 479] — updated EMA +// isv[RL_REWARD_CLAMP_WIN_INDEX = 452] — adaptive WIN bound +// isv[RL_REWARD_CLAMP_LOSS_INDEX = 453] — adaptive LOSS = ratio × WIN +// +// EMA discipline: Wiener-α blend with caller-supplied α (host-floored +// at WIENER_ALPHA_FLOOR per `pearl_wiener_alpha_floor_for_nonstationary`). +// Bootstrap on sentinel 0 per `pearl_first_observation_bootstrap` — +// first non-zero raw observation replaces the EMA slot directly so the +// blend always has a real prev value to mix with. +// +// Bounds: WIN ∈ [MIN_WIN=1.0, MAX_WIN=20.0]. MIN_WIN preserves the +// floor from the static design — clamp never tightens below the C51 +// V_MAX. MAX_WIN caps runaway during early training before the EMA +// stabilises (atom support × 20 still well within float headroom). +// +// LOSS = RATIO × WIN preserves the loss-aversion asymmetry from the +// `pearl_audit_unboundedness_for_implicit_asymmetry` design. Driven +// from a separate ISV slot so the asymmetry ratio is itself tunable +// without recompile. +// +// Per `feedback_no_atomicadd`: single-thread kernel. Per +// `feedback_cpu_is_read_only`: all state in ISV. Per +// `feedback_isv_for_adaptive_bounds`: hardcoded MIN/MAX retained per +// the user-stated "floors and clamp bounds" exemption (2026-05-24). + +#define RL_POS_SCALED_REWARD_MAX_INDEX 478 +#define RL_POS_SCALED_REWARD_MAX_EMA_INDEX 479 +#define RL_REWARD_CLAMP_MARGIN_INDEX 480 +#define RL_REWARD_CLAMP_RATIO_INDEX 481 +#define RL_REWARD_CLAMP_WIN_INDEX 452 +#define RL_REWARD_CLAMP_LOSS_INDEX 453 + +#define MIN_WIN 1.0f +#define MAX_WIN 20.0f +#define WIENER_ALPHA_FLOOR 0.4f + +extern "C" __global__ void rl_reward_clamp_controller( + float* __restrict__ isv, + float alpha +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + const float pos_max = isv[RL_POS_SCALED_REWARD_MAX_INDEX]; + const float ema_prev = isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX]; + + // EMA maintenance — bootstrap on sentinel 0 with first non-zero + // raw observation; otherwise Wiener-α blend with floor. + float ema_new; + if (ema_prev == 0.0f) { + if (pos_max == 0.0f) { + // Both prev and obs sentinels — defer adaptation, leave the + // statically-seeded WIN(1.0) / LOSS(3.0) bounds untouched. + // Next call where the kernel observes a positive reward + // will bootstrap the EMA and start adapting. + return; + } + ema_new = pos_max; + } else { + const float a = fmaxf(alpha, WIENER_ALPHA_FLOOR); + // If pos_max==0 this step (no positive reward sample), still + // blend toward 0 — captures regimes where wins disappear. The + // EMA decays gracefully toward 0 over ~1/alpha steps. + ema_new = (1.0f - a) * ema_prev + a * pos_max; + } + isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX] = ema_new; + + // Compute adaptive WIN bound: MARGIN × EMA, clamped to [MIN, MAX]. + const float margin = isv[RL_REWARD_CLAMP_MARGIN_INDEX]; + const float win_eff = fmaxf(MIN_WIN, fminf(MAX_WIN, margin * ema_new)); + + // LOSS = RATIO × WIN preserves asymmetry. RATIO seeded to 3.0. + const float ratio = isv[RL_REWARD_CLAMP_RATIO_INDEX]; + const float loss_eff = ratio * win_eff; + + isv[RL_REWARD_CLAMP_WIN_INDEX] = win_eff; + isv[RL_REWARD_CLAMP_LOSS_INDEX] = loss_eff; +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 96d1a859f..9b2744dc0 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -1,3 +1,8 @@ +// The per-step diag `json!{...}` block now spans ≥ 130 fields after the +// adaptive-clamp additions, exceeding serde_json's default macro recursion +// budget of 128. Bumping at the crate root keeps the expansion within limits. +#![recursion_limit = "256"] + //! Integrated RL trainer CLI (DQN + PPO on shared Mamba2 -> CfC encoder). //! //! Drives `IntegratedTrainer::step_with_lobsim` against a real `LobSimCuda` @@ -59,6 +64,8 @@ use ml_alpha::rl::isv_slots::{ RL_ROLLOUT_BOOTSTRAP_INDEX, RL_SCHULMAN_ADJUST_RATE_INDEX, RL_SCHULMAN_TOLERANCE_INDEX, RL_STREAM_ALPHA_INDEX, RL_TAU_BOOTSTRAP_INDEX, RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX, RL_MEAN_ABS_PNL_EMA_INDEX, + RL_POS_SCALED_REWARD_MAX_INDEX, RL_POS_SCALED_REWARD_MAX_EMA_INDEX, + RL_REWARD_CLAMP_MARGIN_INDEX, RL_REWARD_CLAMP_RATIO_INDEX, RL_MEAN_TRADE_DURATION_EMA_INDEX, RL_N_ROLLOUT_STEPS_INDEX, RL_PER_ALPHA_INDEX, RL_PI_GRAD_NORM_EMA_INDEX, RL_PPO_CLIP_INDEX, RL_PPO_LOG_RATIO_ABS_MAX_INDEX, RL_PPO_RATIO_CLAMP_MAX_INDEX, RL_Q_DIVERGENCE_EMA_INDEX, RL_Q_GRAD_NORM_EMA_INDEX, @@ -656,6 +663,13 @@ fn main() -> Result<()> { "abs_max": reward_abs_max, "scaled_pre_clamp_max": isv[RL_MAX_ABS_SCALED_REWARD_PRE_CLAMP_INDEX], + // Adaptive clamp signals — surface the controller's + // POINT input + EMA + emitted bounds so post-hoc + // analysis can verify the adaptation chain is moving. + "pos_scaled_max": + isv[RL_POS_SCALED_REWARD_MAX_INDEX], + "pos_scaled_max_ema": + isv[RL_POS_SCALED_REWARD_MAX_EMA_INDEX], }, // audit — PPO importance-ratio clamp diagnostics. // @@ -764,6 +778,8 @@ fn main() -> Result<()> { "rollout_bootstrap": isv[RL_ROLLOUT_BOOTSTRAP_INDEX], "reward_scale_bootstrap":isv[RL_REWARD_SCALE_BOOTSTRAP_INDEX], "ppo_ratio_clamp_bootstrap": isv[RL_PPO_RATIO_CLAMP_BOOTSTRAP_INDEX], + "reward_clamp_margin": isv[RL_REWARD_CLAMP_MARGIN_INDEX], + "reward_clamp_ratio": isv[RL_REWARD_CLAMP_RATIO_INDEX], }, // audit — Q vs π action agreement EMA at slot 407 // (previously dead). 1.0 = perfect ranking consistency, diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index c3aa77618..43f9237aa 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -38,8 +38,9 @@ //! | 452-461 | 10 ISV-driven design constants (first batch) | rl_isv_write (seed loop) | //! | 462-467 | 6 LR-controller numerics + aux λ | rl_isv_write (seed loop) | //! | 468-477 | 10 design constants (Schulman + bootstraps + streaming α + kurt refs) | rl_isv_write | +//! | 478-481 | 4 adaptive reward-clamp slots (raw + EMA + margin + ratio) | apply_reward_scale + rl_reward_clamp_controller | //! -//! Total: 78 slots; `RL_SLOTS_END = 478`. +//! Total: 82 slots; `RL_SLOTS_END = 482`. /// Discount factor γ. Controller input: mean trade duration / anchor. /// Bootstrap 0.99. @@ -548,6 +549,52 @@ pub const RL_REWARD_SCALE_BOOTSTRAP_INDEX: usize = 476; /// PPO ratio-clamp bootstrap. Default 10.0. pub const RL_PPO_RATIO_CLAMP_BOOTSTRAP_INDEX: usize = 477; +// ─── Adaptive reward-clamp (audit 2026-05-24) ─────────────────────── +// +// alpha-rl-rmgm5 (commit a776fab31) diag showed the static +// `[-RL_REWARD_CLAMP_LOSS, +RL_REWARD_CLAMP_WIN] = [-3.0, +1.0]` +// clamp firing on 85% of steps with p95 pre-clamp = 15.5× the WIN +// bound. Winning trade signal (avg +$2.06, max +$11 in scaled units) +// was squished to +1.0, while losses (avg -$3.84) routinely exceeded +// the LOSS bound. This kills the gradient differential between +// profitable and unprofitable actions in the C51 Bellman target. +// +// Fix: `rl_reward_clamp_controller` reads the per-step positive-tail +// max published by `apply_reward_scale` into slot 478, EMAs it into +// slot 479, and writes +// +// WIN_eff = clamp(MARGIN(480) × EMA(479), MIN_WIN=1.0, MAX_WIN=20.0) +// LOSS_eff = RATIO(481) × WIN_eff +// +// back into slots 452/453. The 3:1 ratio preserves the loss-aversion +// asymmetry per `pearl_audit_unboundedness_for_implicit_asymmetry`. + +/// Per-step max(positive scaled reward, 0) — POINT measurement +/// published by `apply_reward_scale` for the adaptive clamp +/// controller. Overwritten every step. +pub const RL_POS_SCALED_REWARD_MAX_INDEX: usize = 478; + +/// EMA of `RL_POS_SCALED_REWARD_MAX_INDEX` — maintained by +/// `rl_reward_clamp_controller` via Wiener-α blend (floor 0.4 per +/// `pearl_wiener_alpha_floor_for_nonstationary`; the reward +/// distribution drifts as the policy co-adapts). Sentinel 0 → +/// bootstrap to first non-zero raw observation per +/// `pearl_first_observation_bootstrap`. +pub const RL_POS_SCALED_REWARD_MAX_EMA_INDEX: usize = 479; + +/// Adaptive reward-clamp margin multiplier — `WIN_eff = MARGIN × ema`. +/// Default 1.5 — gives 50% headroom above the EMA so occasional +/// fat-tail wins still register without overshooting the C51 atom +/// support. Seeded by `rl_isv_write` at trainer init. +pub const RL_REWARD_CLAMP_MARGIN_INDEX: usize = 480; + +/// Adaptive reward-clamp loss/win ratio — `LOSS_eff = RATIO × WIN_eff`. +/// Default 3.0 — preserves the loss-aversion asymmetry from the +/// static `[-3, +1]` design per +/// `pearl_audit_unboundedness_for_implicit_asymmetry`. Seeded by +/// `rl_isv_write`. +pub const RL_REWARD_CLAMP_RATIO_INDEX: usize = 481; + /// Last RL-allocated slot index (exclusive). The integrated trainer /// extends `ISV_TOTAL_DIM` to at least this value at trainer init time. -pub const RL_SLOTS_END: usize = 478; +pub const RL_SLOTS_END: usize = 482; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index efcc88b89..871892df9 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -159,6 +159,12 @@ const EXTRACT_REALIZED_PNL_DELTA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/extract_realized_pnl_delta.cubin")); const APPLY_REWARD_SCALE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/apply_reward_scale.cubin")); +// Adaptive reward-clamp controller — audit fix 2026-05-24. Consumes the +// per-step positive-tail max written by apply_reward_scale into slot +// 478 and emits adaptive WIN/LOSS bounds into slots 452/453. See the +// kernel header for the design rationale (rmgm5 clamp-firing audit). +const RL_REWARD_CLAMP_CONTROLLER_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_reward_clamp_controller.cubin")); const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/actions_to_market_targets.cubin")); @@ -374,6 +380,9 @@ pub struct IntegratedTrainer { extract_realized_pnl_delta_fn: CudaFunction, _apply_reward_scale_module: Arc, apply_reward_scale_fn: CudaFunction, + // Adaptive reward-clamp controller (audit 2026-05-24). + _rl_reward_clamp_controller_module: Arc, + rl_reward_clamp_controller_fn: CudaFunction, _actions_to_market_targets_module: Arc, actions_to_market_targets_fn: CudaFunction, @@ -759,6 +768,12 @@ impl IntegratedTrainer { let apply_reward_scale_fn = apply_reward_scale_module .load_function("apply_reward_scale") .context("load apply_reward_scale")?; + let rl_reward_clamp_controller_module = ctx + .load_cubin(RL_REWARD_CLAMP_CONTROLLER_CUBIN.to_vec()) + .context("load rl_reward_clamp_controller cubin")?; + let rl_reward_clamp_controller_fn = rl_reward_clamp_controller_module + .load_function("rl_reward_clamp_controller") + .context("load rl_reward_clamp_controller")?; let actions_to_market_targets_module = ctx .load_cubin(ACTIONS_TO_MARKET_TARGETS_CUBIN.to_vec()) .context("load actions_to_market_targets cubin")?; @@ -1065,6 +1080,8 @@ impl IntegratedTrainer { extract_realized_pnl_delta_fn, _apply_reward_scale_module: apply_reward_scale_module, apply_reward_scale_fn, + _rl_reward_clamp_controller_module: rl_reward_clamp_controller_module, + rl_reward_clamp_controller_fn, _actions_to_market_targets_module: actions_to_market_targets_module, actions_to_market_targets_fn, _abs_copy_module: abs_copy_module, @@ -1182,9 +1199,17 @@ impl IntegratedTrainer { // (slot, value) pair — pure device write, no HtoD per // `feedback_no_htod_htoh_only_mapped_pinned`. { - let isv_constants: [(usize, f32); 26] = [ + let isv_constants: [(usize, f32); 28] = [ + // Static seeds for the adaptive reward-clamp controller — + // these are the initial values that + // `rl_reward_clamp_controller` will replace once it + // observes the first positive reward. Until then, kernels + // reading slots 452/453 see the documented [-3, +1] + // defaults. (crate::rl::isv_slots::RL_REWARD_CLAMP_WIN_INDEX, 1.0), (crate::rl::isv_slots::RL_REWARD_CLAMP_LOSS_INDEX, 3.0), + (crate::rl::isv_slots::RL_REWARD_CLAMP_MARGIN_INDEX, 1.5), + (crate::rl::isv_slots::RL_REWARD_CLAMP_RATIO_INDEX, 3.0), (crate::rl::isv_slots::RL_KL_TARGET_INDEX, 0.01), (crate::rl::isv_slots::RL_IMPROVEMENT_THRESHOLD_INDEX, 0.99), (crate::rl::isv_slots::RL_PLATEAU_PATIENCE_INDEX, 1000.0), @@ -1672,10 +1697,13 @@ impl IntegratedTrainer { ) -> Result<()> { debug_assert_eq!(rewards_d.len(), b_size); let block_x: u32 = (b_size as u32).min(256).max(1); + // Shared mem: 2 × block × f32 — kernel runs two parallel + // reductions (max|scaled| + max(positive scaled)) per audit + // 2026-05-24, both stored in `extern __shared__ float smem[]`. let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (block_x, 1, 1), - shared_mem_bytes: block_x * (std::mem::size_of::() as u32), + shared_mem_bytes: 2 * block_x * (std::mem::size_of::() as u32), }; let b_size_i = b_size as i32; let mut launch = self.stream.launch_builder(&self.apply_reward_scale_fn); @@ -1683,6 +1711,30 @@ impl IntegratedTrainer { unsafe { launch.launch(cfg).context("apply_reward_scale launch")?; } + + // Adaptive reward-clamp controller — reads the per-step + // positive-tail max just published into slot 478, maintains + // an EMA in slot 479, writes adaptive WIN/LOSS bounds into + // slots 452/453 for the NEXT step's apply_reward_scale to + // consume. Single-thread kernel; α driven by the same + // controller-cohort floor as the other R5 EMAs. + { + let cfg_ctrl = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let alpha = RL_LR_CONTROLLER_ALPHA; + let mut launch = self + .stream + .launch_builder(&self.rl_reward_clamp_controller_fn); + launch.arg(&self.isv_d).arg(&alpha); + unsafe { + launch + .launch(cfg_ctrl) + .context("rl_reward_clamp_controller launch")?; + } + } Ok(()) } @@ -2962,20 +3014,25 @@ impl IntegratedTrainer { .context("launch_rl_controllers_per_step")?; // Apply the reward scale on device (in place on rewards_d) + - // asymmetric clamp [-3, +1] + per-step pre-clamp max diag. - // Reads ISV[406] which the controller just updated; writes - // ISV[439] for the diag emitter. + // adaptive clamp + per-step pre-clamp diag + per-step + // positive-tail max for the adaptive clamp controller. + // Reads ISV[406] (scale) + ISV[452]/ISV[453] (current clamp + // bounds, refreshed by rl_reward_clamp_controller via the + // helper below); writes ISV[439] (diag) + ISV[478] (positive + // tail input for the controller). // // Single-block layout: tree reduction needs every thread in // the same block to participate in __syncthreads, so we cap // block_x at min(b_size, 256). For b_size > 256 the kernel's - // grid-stride loop walks the rest. + // grid-stride loop walks the rest. Shared mem now 2× block × + // f32 to fit the parallel max(|scaled|) + max(positive) + // reductions. { let block_x: u32 = (b_size as u32).min(256).max(1); let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (block_x, 1, 1), - shared_mem_bytes: block_x * (std::mem::size_of::() as u32), + shared_mem_bytes: 2 * block_x * (std::mem::size_of::() as u32), }; let mut launch = self.stream.launch_builder(&self.apply_reward_scale_fn); launch @@ -2987,6 +3044,27 @@ impl IntegratedTrainer { } } + // Adaptive reward-clamp controller — refresh slots 452/453 from + // the just-published positive-tail max in slot 478. The new + // bounds take effect on the NEXT step's apply_reward_scale. + { + let cfg_ctrl = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let alpha = RL_LR_CONTROLLER_ALPHA; + let mut launch = self + .stream + .launch_builder(&self.rl_reward_clamp_controller_fn); + launch.arg(&self.isv_d).arg(&alpha); + unsafe { + launch + .launch(cfg_ctrl) + .context("rl_reward_clamp_controller launch (step_with_lobsim path)")?; + } + } + // GPU compute_advantage_return: returns_d, advantages_d // populated for step_synthetic to consume. {