From 81538081bc2bd47aeecdb95a371f8e485caaf4c1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 11 Apr 2026 12:47:36 +0200 Subject: [PATCH] feat: hindsight magnitude relabeling via alternating counterfactual Even timesteps use the existing directional mirror (Short<->Long, negated reward). Odd timesteps substitute a different magnitude and linearly scale the reward by alt_mag_frac / actual_mag_frac, giving the magnitude branch direct gradient signal about what bigger or smaller positions would have earned. Falls back to directional mirror when reward is near-zero (scaling uninformative). No extra buffer needed -- reuses the existing N*L counterfactual slot. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 65 +++++++++++++++---- 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index e1022ef03..b898f3e93 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -1774,33 +1774,72 @@ extern "C" __global__ void experience_env_step( * because the bottleneck (#31) compresses regime information into the 2D * latent space automatically. */ - /* ── #7 Counterfactual augmentation: ALWAYS write mirror action + negated reward ── */ + /* ── #7 Counterfactual augmentation: alternating direction mirror / magnitude relabel ── + * + * Even timesteps: directional mirror (Short↔Long, negated reward). + * Odd timesteps: hindsight magnitude relabeling — substitute a different + * magnitude and linearly scale the reward by the ratio of magnitude + * fractions. This gives the magnitude branch direct gradient signal + * about what BIGGER or SMALLER positions would have earned, breaking + * the collapse toward Small (72%) caused by sparse magnitude reward. + * + * No extra buffer needed — reuses the existing N*L counterfactual slot. */ { - /* Write a second experience at offset N*L in the output buffers. - * Same state, mirror direction action, negated reward. - * This doubles effective data diversity with zero extra compute. */ long long cf_off = out_off + (long long)N * L; - /* Copy same state to counterfactual slot */ + /* Copy same state to counterfactual slot (shared by both paths) */ const float* cf_src = batch_states + (long long)i * state_dim; float* cf_dst = out_states + cf_off * state_dim; for (int k = 0; k < state_dim; k++) cf_dst[k] = cf_src[k]; - /* Mirror direction: Short(0)↔Long(2), Flat(1) stays */ - int cf_dir = (b0_size - 1) - dir_idx; - int orig_order = decode_order_4b(action_idx, b2_size, b3_size); + int orig_order = decode_order_4b(action_idx, b2_size, b3_size); int orig_urgency = decode_urgency_4b(action_idx, b3_size); - int cf_action = cf_dir * b1_size * b2_size * b3_size + + int cf_action; + float cf_reward; + + if ((current_t & 1) == 0) { + /* ── Even step: directional mirror (original logic) ── */ + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + mag_idx * b2_size * b3_size + orig_order * b3_size + orig_urgency; - out_actions[cf_off] = cf_action; + cf_reward = -reward; + } else { + /* ── Odd step: hindsight magnitude relabeling ── + * Pick the best alternative magnitude: + * If Full(2), try Half(1). Otherwise try Full(2). + * Scale reward linearly by (alt_mag_frac / actual_mag_frac). + * Only relabel when reward is non-trivial (|reward| > 0.001); + * for near-zero reward the linear scaling is meaningless, so + * fall back to directional mirror. */ + const float mag_fracs[3] = { 0.25f, 0.50f, 1.00f }; + float actual_mag_frac = mag_fracs[mag_idx < 3 ? mag_idx : 0]; + + if (fabsf(reward) > 0.001f) { + int best_alt_mag = (mag_idx == 2) ? 1 : 2; + float alt_mag_frac = mag_fracs[best_alt_mag]; + cf_action = dir_idx * b1_size * b2_size * b3_size + + best_alt_mag * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = reward * (alt_mag_frac / actual_mag_frac); + /* Clamp to prevent extreme rewards from small→full scaling */ + cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f); + } else { + /* Near-zero reward: magnitude scaling uninformative, use directional mirror */ + int cf_dir = (b0_size - 1) - dir_idx; + cf_action = cf_dir * b1_size * b2_size * b3_size + + mag_idx * b2_size * b3_size + + orig_order * b3_size + orig_urgency; + cf_reward = -reward; + } + } - /* Negated reward: if long earned +R, short would earn -R */ - float cf_reward = -reward; if (isnan(cf_reward) || isinf(cf_reward)) cf_reward = 0.0f; + out_actions[cf_off] = cf_action; out_rewards[cf_off] = cf_reward; - out_dones[cf_off] = (float)done; + out_dones[cf_off] = (float)done; } /* ---- Advance timestep counter ---- */