feat(reward-v7.1): track best_exposure_idx in experience kernel for auxiliary loss

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-07 22:36:16 +02:00
parent 28f8fc74b7
commit 44ca6147a9
2 changed files with 26 additions and 0 deletions

View File

@@ -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;

View File

@@ -528,6 +528,7 @@ pub struct GpuExperienceCollector {
actions_out: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
rewards_out: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN)
done_out: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN)
best_exposure_out: CudaSlice<i32>, // [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<half::bf16>, // [alloc_episodes * alloc_timesteps]
@@ -881,6 +882,9 @@ impl GpuExperienceCollector {
let done_out = stream
.alloc_zeros::<f32>(total_output * cf_mult)
.map_err(|e| MLError::ModelError(format!("alloc done_out: {e}")))?;
let best_exposure_out = stream
.alloc_zeros::<i32>(alloc_episodes * alloc_timesteps)
.map_err(|e| MLError::ModelError(format!("alloc best_exposure_out: {e}")))?;
let raw_returns_out = stream
.alloc_zeros::<half::bf16>(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<i32> {
&self.best_exposure_out
}
/// Train curiosity forward model directly on GPU using experience data.
pub fn train_curiosity_gpu(
&mut self,