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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-11 12:47:36 +02:00
parent 8d568158a7
commit 81538081bc

View File

@@ -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 ---- */