From 44ca6147a9ba7349f25afa60dc21049c855a7fcf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 7 Apr 2026 22:36:16 +0200 Subject: [PATCH] feat(reward-v7.1): track best_exposure_idx in experience kernel for auxiliary loss Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/src/cuda_pipeline/experience_kernels.cu | 14 ++++++++++++++ .../src/cuda_pipeline/gpu_experience_collector.rs | 12 ++++++++++++ 2 files changed, 26 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 2d0d0f7db..268e069e8 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -871,6 +871,7 @@ extern "C" __global__ void experience_env_step( float ofi_reward_weight, /* v7 gem: OFI-weighted reward confidence */ float kelly_sizing_weight, /* v7 gem: Kelly-optimal sizing signal */ float reward_noise_scale, /* v7 gem: reward label smoothing noise */ + int* out_best_exposure, /* v7.1: counterfactual best exposure idx (-1 = no exit) */ /* #33 Per-episode saboteur params [N, 3]: (spread_mult, fill_prob, slippage_mult). * NULL = disabled (use global scalars). When non-NULL, overrides * spread_cost, fill_ioc_fill_prob, tx_cost_multiplier per episode. */ @@ -1245,6 +1246,7 @@ extern "C" __global__ void experience_env_step( * ================================================================ */ float reward = 0.0f; + int best_exposure_k = -1; /* v7.1: counterfactual best exposure, -1 = no exit */ if (segment_complete && segment_hold_time > 0.0f) { float segment_pnl = __bfloat162float(ps[11]) + raw_pnl - trade_start_pnl; @@ -1282,6 +1284,7 @@ extern "C" __global__ void experience_env_step( reward = asymmetric_soft_clamp(base_reward); /* ── Layer 3: Counterfactual Exposure Advantage (CEA) ── */ + float best_exposure_reward = -1e30f; if (cea_weight > 0.0f) { float sum_alt_rewards = 0.0f; float price_delta = raw_next - raw_close; @@ -1292,6 +1295,10 @@ extern "C" __global__ void experience_env_step( float alt_vol_norm = alt_return / vol_norm; float alt_reward = asymmetric_soft_clamp(10.0f * alt_vol_norm); sum_alt_rewards += alt_reward; + if (alt_reward > best_exposure_reward) { + best_exposure_reward = alt_reward; + best_exposure_k = k; + } } float mean_alt_reward = sum_alt_rewards / (float)b0_size; float exposure_advantage = reward - mean_alt_reward; @@ -1486,6 +1493,13 @@ extern "C" __global__ void experience_env_step( /* Final NaN guard — if reward is NaN/Inf, write 0.0 instead of poisoning * the replay buffer. This prevents gradient explosion from propagating. */ if (isnan(reward) || isinf(reward)) reward = 0.0f; + + /* v7.1: Write counterfactual best exposure for auxiliary loss */ + if (out_best_exposure != NULL) { + out_best_exposure[out_off] = (segment_complete && segment_hold_time > 0.0f && cea_weight > 0.0f) + ? best_exposure_k : -1; + } + out_rewards[out_off] = reward; out_dones[out_off] = (float)done; diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 13ba79c76..97ea606cd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -528,6 +528,7 @@ pub struct GpuExperienceCollector { actions_out: CudaSlice, // [alloc_episodes * alloc_timesteps] rewards_out: CudaSlice, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN) done_out: CudaSlice, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN) + best_exposure_out: CudaSlice, // [alloc_episodes * alloc_timesteps] v7.1 aux target /// Raw per-bar portfolio returns (unshaped) for accurate Sharpe/MaxDD/Sortino. /// True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}. raw_returns_out: CudaSlice, // [alloc_episodes * alloc_timesteps] @@ -881,6 +882,9 @@ impl GpuExperienceCollector { let done_out = stream .alloc_zeros::(total_output * cf_mult) .map_err(|e| MLError::ModelError(format!("alloc done_out: {e}")))?; + let best_exposure_out = stream + .alloc_zeros::(alloc_episodes * alloc_timesteps) + .map_err(|e| MLError::ModelError(format!("alloc best_exposure_out: {e}")))?; let raw_returns_out = stream .alloc_zeros::(total_output) .map_err(|e| MLError::ModelError(format!("alloc raw_returns_out: {e}")))?; @@ -1113,6 +1117,7 @@ impl GpuExperienceCollector { actions_out, rewards_out, done_out, + best_exposure_out, raw_returns_out, epoch_state, reset_flags: 0, @@ -1856,6 +1861,7 @@ impl GpuExperienceCollector { .arg(&config.ofi_reward_weight) .arg(&config.kelly_sizing_weight) .arg(&config.reward_noise_scale) + .arg(&mut self.best_exposure_out) // v7.1: counterfactual best exposure // #33 Per-episode saboteur params (0 = NULL = disabled) .arg(&{ if self.saboteur_active { @@ -2131,6 +2137,12 @@ impl GpuExperienceCollector { &self.actions_out } + /// v7.1: Best counterfactual exposure index per step (for auxiliary loss). + /// Values: 0-8 at trade exits, -1 at non-exit steps. + pub fn best_exposure_buf(&self) -> &CudaSlice { + &self.best_exposure_out + } + /// Train curiosity forward model directly on GPU using experience data. pub fn train_curiosity_gpu( &mut self,