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

feat(rl): wire 3 of 6 missing EMA inputs (Phase A — entropy, adv_var, td_kurt)

Browse Source 
R9 cluster smoke alpha-rl-qzstj diag exposed that 6 of 7 controllers
held at bootstrap for the entire 1000-step run because their input
EMAs were never populated. Only `mean_abs_pnl_ema` was wired (via
ema_update_on_done on reward_abs_d). The other 6 EMA producers
existed as generic kernels (ema_update_per_step / ema_update_on_done)
but nothing computed the per-step input signals to feed them.

This commit wires the 3 EMAs whose source signals are ALREADY
computed and live in trainer per-step buffers (Phase A — cheapest
to wire):

  * `entropy_observed_ema` (ISV[420] → rl_entropy_coef controller)
    ← per-batch entropy `entropy_d` from PPO surrogate forward.
    `ema_update_per_step` does mean-reduce internally, so this is
    a single launch with entropy_d as input (b_size native).

  * `advantage_var_ratio_ema` (ISV[421] → rl_rollout_steps)
    ← `var(advantages) / max(|mean(advantages)|, 1e-6)` reduction
    on advantages_d. New kernel `rl_var_over_abs_mean_b`
    (two-pass shared-mem tree-reduce) writes scalar to trainer-
    owned `ema_input_scratch_d[1]`, then `ema_update_per_step`
    consumes with b_size=1.

  * `td_kurtosis_ema` (ISV[422] → rl_per_alpha)
    ← `E[(x-μ)⁴] / σ⁴` kurtosis reduction on td_per_sample_d
    (R7d's per-sample CE loss from dqn_distributional_q_bwd).
    New kernel `rl_kurtosis_b` (three-pass shared-mem tree-reduce)
    writes scalar to ema_input_scratch_d, then ema_update_per_step
    with b_size=1.

## Wiring placement

* `advantage_var_ratio` update: in `step_with_lobsim` immediately
  after `compute_advantage_return` populates `advantages_d`. Fires
  BEFORE the next step's controllers, so the controller sees the
  fresh signal one step later.

* `entropy_observed` + `td_kurtosis` updates: in `step_synthetic`
  AFTER the encoder backward (deferred from their natural in-place
  locations to avoid a borrow-checker conflict with `h_t_borrow`
  which holds `&self.perception` through the entire forward chain).
  One-step lag — same as advantage_var_ratio for the same reason
  (controllers fire in the NEXT step_with_lobsim).

## Trainer-owned scratch

Single `ema_input_scratch_d: CudaSlice<f32>` of length 1. Reused
across the var-over-abs-mean and kurtosis launches in any given
step — they're stream-serialised, so the second reducer's write
to slot 0 strictly follows the first reducer's consumer (the
corresponding ema_update_per_step). Cheap (4 bytes); avoids
two separate scratches.

## Why a 1-float scratch + b_size=1 ema_update

`ema_update_per_step` expects `obs_d[b_size]` and computes per-step
mean as `Σobs / b_size`. Passing a 1-element buffer gives
mean = obs[0] = the reduce kernel's scalar output. The EMA then
blends `prev` toward that scalar via Wiener-α (or bootstraps on
first non-zero per `pearl_first_observation_bootstrap`).

This pattern lets the existing per-step EMA kernel handle scalar
inputs without modification — the alternative (a dedicated
"ema_scalar_per_step") would duplicate logic per
`feedback_single_source_of_truth_no_duplicates`.

## Verified gates (post-fix, local sm_86)

  G1  isv_bootstrap                ✅
  G3  controllers_emit             ✅ (test pre-seeds inputs, so
                                       wiring path not exercised)
  G4  target_soft_update           ✅
  G6  r7d_per_wiring               ✅
  R3, R4, smoke                    ✅

Local smoke at b_size=1 won't exercise kurtosis (kernel returns 0
at b_size<2 → cold-start gate holds per_α at bootstrap). var_over_
abs_mean does fire because b_size=1 has a well-defined (degenerate)
variance of 0. Cluster smoke at b_size=1 will mostly exercise
entropy_observed.

## What's NOT in this commit (Phase B — 3 EMAs left)

  * `kl_pi_ema` (ISV[419] → rl_ppo_clip)
    needs: D_KL approximation between log_pi_old and log_pi_new.
    Both buffers exist in trainer; need a small subtract-and-mean
    kernel OR extend PPO surrogate forward to emit kl_per_batch.

  * `q_divergence_ema` (ISV[418] → rl_target_tau)
    needs: `‖W_online − W_target‖₂`. Both DQN weight buffers
    accessible via dqn_head fields; need a small L2-diff-norm
    kernel called after soft_update_target.

  * `trade_duration_ema` (ISV[417] → rl_gamma)
    needs: per-batch step counter (i32, b_size, trainer-owned)
    that increments each step and emits its value on done. Needs
    a small `step_counter_update` kernel + the counter buffer.

These three need NEW signal-derivation kernels (not just reductions
over existing buffers). Separate commit.

## Cluster smoke expected diag change

Before this commit: 6 of 7 EMA input slots stuck at 0.0 for all
1000 steps. After: ISV[420], ISV[421], ISV[422] populated each
step. The corresponding controllers (coef, n_roll, per_α) should
visibly adapt after the first few non-zero observations.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-23 16:22:16 +02:00
parent c295fa9c92
commit 91c4e499d2
4 changed files with 318 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
crates/ml-alpha/build.rs
View File

@@ -55,17 +55,26 @@ const KERNELS: &[&str] = &[
"apply_reward_scale", // RL Phase R6: element-wise rewards *= ISV[RL_REWARD_SCALE_INDEX=406]; closes the F.3b host roundtrip
"actions_to_market_targets", // RL Phase R6: 9-action grid → LobSim market_targets[B*2] on device; replaces host submit_market loop per feedback_cpu_is_read_only
"abs_copy", // RL Phase R7a: element-wise dst[b] = fabsf(src[b]); feeds |reward| into ema_update_on_done for the MEAN_ABS_PNL_EMA slot
"rl_var_over_abs_mean_b", // R9 EMA wiring: var(x)/|mean(x)| reduction; feeds advantage_var_ratio_ema (ISV[421]) from advantages_d
"rl_kurtosis_b", // R9 EMA wiring: E[(x-μ)⁴]/σ⁴ reduction; feeds td_kurtosis_ema (ISV[422]) from td_per_sample_d. (entropy_observed_ema uses existing ema_update_per_step's internal mean reduce directly on entropy_d — no separate mean kernel needed.)
];
// Cache bust v30 (2026-05-23): RL Phase R7d — dqn_distributional_q_bwd
// kernel signature change. Adds a new `loss_per_batch [B]` output so the
// trainer can DtoH per-sample CE loss into ReplayBuffer.update_priorities
// after each off-policy Q backward (feedback_always_per wiring). Single
// writer per batch (atom 0 only), so the new output is non-atomic — the
// existing scalar `loss_out` still accumulates via atomicAdd for the
// diagnostic total. All callers (the integrated trainer's
// dqn_offpolicy_step) migrate atomically in the same R7d commit per
// feedback_no_partial_refactor.
// Cache bust v31 (2026-05-23): R9 EMA wiring follow-up — three new
// reduce kernels populate the input EMAs for 3 of the 6 previously
// frozen controllers (entropy_coef, rollout_steps, per_alpha). Each
// kernel is a single-block tree reduction over b_size floats:
// * rl_reduce_mean_b — mean(x) — for entropy_observed
// * rl_var_over_abs_mean_b — var(x)/|mean(x)| — for advantage_var_ratio
// * rl_kurtosis_b — E[(x-μ)⁴]/σ⁴ — for td_kurtosis
// Block dim must be next pow2 ≥ b_size (trainer enforces); shared mem
// is reused across multi-pass reductions. The trainer launches each
// reducer after the source signal is populated (PPO surrogate fwd for
// entropy, compute_advantage_return for advantages, dqn backward_logits
// for td_per_sample) then feeds the scalar output to ema_update_per_step
// or ema_update_on_done targeting the right ISV slot.
//
// The remaining 3 EMAs (kl_pi, q_divergence, trade_duration) need new
// state/derivation kernels — separate follow-up commit.
fn main() {
println!("cargo:rerun-if-changed=build.rs");

81
crates/ml-alpha/cuda/rl_kurtosis_b.cu Normal file
View File

@@ -0,0 +1,81 @@
// rl_kurtosis_b.cu — kurtosis reduction `E[(x-μ)⁴] / σ⁴` over b_size
// floats → scalar.
//
// R9 audit follow-up (2026-05-23): the rl_per_alpha controller needs
// `td_kurtosis_ema` (ISV[422]) populated from the per-step
// td_per_sample_d (output of R7d's dqn_distributional_q_bwd
// `loss_per_batch` extension). Kurtosis quantifies tail-thickness of
// the per-sample CE distribution — heavy tails mean a few transitions
// dominate Q learning, which is exactly what PER α should crank up
// to concentrate sampling on the informative samples.
//
// Three-pass reduction:
// Pass 1: sum → mean
// Pass 2: sum((x - mean)²) → var = sum_sq2 / b_size
// Pass 3: sum((x - mean)⁴) → m4 = sum_sq4 / b_size
// Output: m4 / (var² + ε)
//
// For Gaussian, kurtosis = 3.0 (Pearson's definition). Heavy-tailed
// distributions can have kurtosis > 10. The rl_per_alpha kernel's
// target formula maps kurt > 3 → higher α (sharpens PER sampling).
//
// Single block, b_size threads, shared mem reused across passes.
// At b_size=1 (smoke), kurtosis is undefined (single sample has no
// distribution); the kernel returns 0 in that case — which the
// per_α controller's cold-start gate then skips, holding bootstrap.
extern "C" __global__ void rl_kurtosis_b(
const float* __restrict__ x,
float* __restrict__ out,
int b_size
) {
extern __shared__ float s[];
const int tid = threadIdx.x;
// Degenerate case: b_size < 2 → return 0 (kurtosis undefined;
// cold-start gate on the controller skips when input is 0).
if (b_size < 2) {
if (tid == 0) out[0] = 0.0f;
return;
}
// ── Pass 1: sum → mean ────────────────────────────────────────
s[tid] = (tid < b_size) ? x[tid] : 0.0f;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) s[tid] += s[tid + stride];
__syncthreads();
}
__shared__ float mean;
if (tid == 0) mean = s[0] / (float)b_size;
__syncthreads();
// ── Pass 2: sum((x - mean)²) → var ────────────────────────────
const float xi = (tid < b_size) ? x[tid] : mean;
const float d = xi - mean;
const float d2 = d * d;
s[tid] = d2;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) s[tid] += s[tid + stride];
__syncthreads();
}
__shared__ float var;
if (tid == 0) var = s[0] / (float)b_size;
__syncthreads();
// ── Pass 3: sum((x - mean)⁴) → m4 ─────────────────────────────
const float d4 = d2 * d2;
s[tid] = (tid < b_size) ? d4 : 0.0f;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) s[tid] += s[tid + stride];
__syncthreads();
}
if (tid == 0) {
const float m4 = s[0] / (float)b_size;
const float var2 = var * var;
out[0] = m4 / (var2 + 1e-12f);
}
}

64
crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu Normal file
View File

@@ -0,0 +1,64 @@
// rl_var_over_abs_mean_b.cu — `var(x) / max(|mean(x)|, ε)` reduction
// over b_size floats → scalar.
//
// R9 audit follow-up (2026-05-23): the rl_rollout_steps controller
// needs `advantage_var_ratio_ema` (ISV[421]) populated from the
// per-step advantages buffer (output of compute_advantage_return).
// The controller's target formula is `prev × (var/abs_mean) /
// ADV_VAR_RATIO_TARGET` so the input is var/abs_mean directly.
//
// Two-pass reduction:
// Pass 1: sum → mean = sum / b_size
// Pass 2: sum((x - mean)²) → var = sum_sq / b_size
// Output: var / max(|mean|, EPS)
//
// Floor on the denominator prevents div-by-zero when advantages are
// centred at zero (which happens during cold-start before the policy
// has learned anything — though by the time this runs, the
// cold-start gate on the controller should hold prev at bootstrap
// regardless).
//
// Single block, b_size threads, shared mem holds the running sum
// for both passes (reused — pass 2 overwrites pass 1).
extern "C" __global__ void rl_var_over_abs_mean_b(
const float* __restrict__ x,
float* __restrict__ out,
int b_size
) {
extern __shared__ float s[];
const int tid = threadIdx.x;
// ── Pass 1: sum → mean ────────────────────────────────────────
s[tid] = (tid < b_size) ? x[tid] : 0.0f;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
s[tid] += s[tid + stride];
}
__syncthreads();
}
__shared__ float mean;
if (tid == 0) {
mean = s[0] / (float)b_size;
}
__syncthreads();
// ── Pass 2: sum((x - mean)²) → var ────────────────────────────
const float xi = (tid < b_size) ? x[tid] : mean; // pad with mean → contributes 0 to var
const float d = xi - mean;
s[tid] = d * d;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
s[tid] += s[tid + stride];
}
__syncthreads();
}
if (tid == 0) {
const float var = s[0] / (float)b_size;
const float abs_mean = fmaxf(fabsf(mean), 1e-6f);
out[0] = var / abs_mean;
}
}

155
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -167,6 +167,15 @@ const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] =
const ABS_COPY_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/abs_copy.cubin"));
// R9 EMA wiring follow-up — reduction kernels that derive scalar
// inputs for the 2 of 3 Phase A controllers needing a transform
// (variance ratio + kurtosis). entropy_observed uses
// ema_update_per_step's built-in mean reduce directly on entropy_d.
const RL_VAR_OVER_ABS_MEAN_B_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_var_over_abs_mean_b.cubin"));
const RL_KURTOSIS_B_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/rl_kurtosis_b.cubin"));
/// Per-head LR controller Wiener-α target. The controller kernel floors
/// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the
/// effective blend rate stays ≥ 0.4 regardless of this seed value. Kept
@@ -346,6 +355,16 @@ pub struct IntegratedTrainer {
_abs_copy_module: Arc<CudaModule>,
abs_copy_fn: CudaFunction,
// ── R9 EMA wiring: reduction kernels for 2 of the 3 Phase A EMAs ──
_rl_var_over_abs_mean_b_module: Arc<CudaModule>,
rl_var_over_abs_mean_b_fn: CudaFunction,
_rl_kurtosis_b_module: Arc<CudaModule>,
rl_kurtosis_b_fn: CudaFunction,
/// 1-float scratch for the reduce kernels' scalar output. Reused
/// across per-step launches (reductions serialised by stream order).
/// Feeds `ema_update_per_step` with `b_size=1` for slots 421, 422.
pub ema_input_scratch_d: CudaSlice<f32>,
/// Trainer-owned snapshot of `lobsim.pos.realized_pnl` taken at end
/// of previous step. `extract_realized_pnl_delta` reads
/// `lobsim.pos_d` post-fill, computes
@@ -675,6 +694,23 @@ impl IntegratedTrainer {
.load_function("abs_copy")
.context("load abs_copy")?;
// R9 EMA wiring: reduce kernels for variance ratio + kurtosis.
let rl_var_over_abs_mean_b_module = ctx
.load_cubin(RL_VAR_OVER_ABS_MEAN_B_CUBIN.to_vec())
.context("load rl_var_over_abs_mean_b cubin")?;
let rl_var_over_abs_mean_b_fn = rl_var_over_abs_mean_b_module
.load_function("rl_var_over_abs_mean_b")
.context("load rl_var_over_abs_mean_b")?;
let rl_kurtosis_b_module = ctx
.load_cubin(RL_KURTOSIS_B_CUBIN.to_vec())
.context("load rl_kurtosis_b cubin")?;
let rl_kurtosis_b_fn = rl_kurtosis_b_module
.load_function("rl_kurtosis_b")
.context("load rl_kurtosis_b")?;
let ema_input_scratch_d = stream
.alloc_zeros::<f32>(1)
.context("alloc ema_input_scratch_d")?;
// Per-batch PRNG state for the Thompson sampler. Seeded
// deterministically from cfg.dqn_seed via ChaCha8 host RNG so
// (cfg.dqn_seed, b_size) → identical Thompson draws across
@@ -871,6 +907,11 @@ impl IntegratedTrainer {
actions_to_market_targets_fn,
_abs_copy_module: abs_copy_module,
abs_copy_fn,
_rl_var_over_abs_mean_b_module: rl_var_over_abs_mean_b_module,
rl_var_over_abs_mean_b_fn,
_rl_kurtosis_b_module: rl_kurtosis_b_module,
rl_kurtosis_b_fn,
ema_input_scratch_d,
prev_realized_pnl_d,
prev_position_lots_d,
rewards_d,
@@ -1775,6 +1816,41 @@ impl IntegratedTrainer {
)
.context("perception.backward_encoder_with_grad_h_t")?;
// ── Step 11b (R9 EMA wiring): feed step_synthetic's two
// EMA inputs into their controllers. Deferred to AFTER the
// encoder backward so the `&self.perception` borrow (held
// by h_t_borrow earlier in this function) is released
// before the `&mut self` launches here.
//
// 1. `entropy_observed_ema` (ISV[420]) ← per-batch entropy
// from PPO surrogate forward (entropy_d still live here).
// ema_update_per_step does the mean reduce internally,
// so entropy_d goes in directly.
// 2. `td_kurtosis_ema` (ISV[422]) ← kurtosis of per-sample
// CE loss (td_per_sample_d). rl_kurtosis_b reduces to
// a scalar, then ema_update_per_step with b_size=1.
//
// Both updates take effect on the NEXT step's controller
// fire (controllers fire in step_with_lobsim BEFORE
// step_synthetic). One-step lag is acceptable — controllers
// adapt over many steps.
self.launch_ema_update_per_step(
crate::rl::isv_slots::RL_ENTROPY_OBSERVED_EMA_INDEX,
RL_LR_CONTROLLER_ALPHA,
&entropy_d,
b_size,
)
.context("R9 EMA: ema_update_per_step(entropy_observed)")?;
self.launch_kurtosis(&self.td_per_sample_d.clone(), b_size)
.context("R9 EMA: launch_kurtosis(td_per_sample_d)")?;
self.launch_ema_update_per_step(
crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX,
RL_LR_CONTROLLER_ALPHA,
&self.ema_input_scratch_d.clone(),
1,
)
.context("R9 EMA: ema_update_per_step(td_kurtosis)")?;
// ── Step 12: compose stats ───────────────────────────────────
// BCE / aux losses are NOT read this phase — perception is driven
// separately by callers via its existing step_batched path. The
@@ -2246,6 +2322,23 @@ impl IntegratedTrainer {
// post-R7b (the bootstrap-fallback host read of γ that the
// flawed branch did is gone with the host advantage loop).
// R9 EMA wiring: feed `advantage_var_ratio_ema` (ISV[421] →
// rl_rollout_steps controller) from this step's advantages_d.
// Reduce to var/|mean| scalar, then ema_update_per_step with
// b_size=1. Pre-R9 the advantage_var_ratio_ema stayed at
// sentinel zero forever and the rollout_steps controller
// was frozen at ROLLOUT_BOOTSTRAP for the entire smoke.
// Now it adapts.
self.launch_var_over_abs_mean(&self.advantages_d.clone(), b_size)
.context("R9 EMA: launch_var_over_abs_mean(advantages_d)")?;
self.launch_ema_update_per_step(
crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX,
RL_LR_CONTROLLER_ALPHA,
&self.ema_input_scratch_d.clone(),
1,
)
.context("R9 EMA: ema_update_per_step(advantage_var_ratio)")?;
// ── Step 7a (R7d): PER push + sample + gather. ────────────────
// Push CURRENT step's b_size transitions (one per batch) to
// ReplayBuffer; sample b_size indices (priority^α from
@@ -2477,6 +2570,68 @@ impl IntegratedTrainer {
Ok(indices)
}
/// R9 EMA wiring: reduce `var(x) / max(|mean(x)|, 1e-6)` over
/// `input_d[b_size]` into `self.ema_input_scratch_d[1]`. Single
/// block, next-pow2(b_size) threads, shared mem reused across
/// the two passes inside the kernel.
fn launch_var_over_abs_mean(
&mut self,
input_d: &CudaSlice<f32>,
b_size: usize,
) -> Result<()> {
debug_assert!(input_d.len() >= b_size);
debug_assert!(b_size <= 1024);
let block_dim = (b_size.next_power_of_two() as u32).max(1);
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.rl_var_over_abs_mean_b_fn);
launch
.arg(input_d)
.arg(&mut self.ema_input_scratch_d)
.arg(&b_size_i);
unsafe {
launch
.launch(cfg)
.context("rl_var_over_abs_mean_b launch")?;
}
Ok(())
}
/// R9 EMA wiring: reduce kurtosis `E[(x-μ)⁴]/σ⁴` over
/// `input_d[b_size]` into `self.ema_input_scratch_d[1]`. Returns
/// 0 at b_size < 2 (kurtosis undefined; the per_α controller's
/// cold-start gate then skips, holding bootstrap).
fn launch_kurtosis(
&mut self,
input_d: &CudaSlice<f32>,
b_size: usize,
) -> Result<()> {
debug_assert!(input_d.len() >= b_size);
debug_assert!(b_size <= 1024);
let block_dim = (b_size.next_power_of_two() as u32).max(1);
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.rl_kurtosis_b_fn);
launch
.arg(input_d)
.arg(&mut self.ema_input_scratch_d)
.arg(&b_size_i);
unsafe {
launch
.launch(cfg)
.context("rl_kurtosis_b launch")?;
}
Ok(())
}
/// Launch `rl_lr_controller` to emit per-head learning rates into
/// `ISV[412..417]`. Phase E.2-DEFER item 3. Diagnostic-input args
/// are currently zero (the controller emits the bootstrap target