diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu index 9369556f7..fbc58ca26 100644 --- a/crates/ml-alpha/cuda/compute_advantage_return.cu +++ b/crates/ml-alpha/cuda/compute_advantage_return.cu @@ -1,19 +1,26 @@ -// compute_advantage_return.cu — KL-regularized advantage computation. +// compute_advantage_return.cu — KL-regularized advantage with split +// static/dynamic components for on-policy/off-policy consistency. // -// Combines environment reward with an RLHF-style KL penalty that feeds -// through the STANDARD advantage pipeline (no separate gradient term). +// The KL penalty has two parts: +// -β × (log π_θ(a|s) - log π_ref(a)) +// = -β × log π_θ(a|s) + β × log π_ref(a) // -// r_modified = r_env - β_kl × (log π_θ(a|s) - log π_ref(a)) +// Static part: β × log π_ref(a) — depends only on which action was +// taken, not on the current policy. Written to rewards_d so PER +// captures it. Q's Bellman targets see this consistently across +// replays, giving Q a stable Hold preference. // -// On non-done steps (r_env=0): the KL penalty is the ONLY reward signal. -// It rewards actions that stay near the reference (Hold-heavy) and -// penalizes actions that drift away. On done-steps: the trade PnL -// dominates the KL penalty. +// Dynamic part: -β × log π_θ(a|s) — depends on the CURRENT policy. +// Only added to the on-policy advantage (not stored in PER). +// Provides the adaptive signal that pushes π back toward the +// reference when it drifts. // -// This avoids the structural imbalance of a separate KL gradient -// (which fired 100% of steps vs PPO's 6%, overwhelming PPO at any β). -// As a reward signal, KL feeds through PPO's own importance ratio -// mechanism with matched per-step coverage. +// This split ensures Q and π optimize the SAME objective: +// Q learns: E[r_env + β·log π_ref(a) + γV(s')] (stable across replays) +// π maximizes: E[A] where A = (r + β·log π_ref - β·log π_θ) + γV(s') - V(s) +// The β·log π_ref terms cancel between reward and KL, leaving: +// π effectively maximizes: E[r_env - β·log π_θ + γV(s') - V(s)] +// = standard entropy-regularized objective (SAC-like) #define RL_GAMMA_INDEX 400 #define RL_KL_REF_BETA_INDEX 578 @@ -22,7 +29,7 @@ extern "C" __global__ void compute_advantage_return( const float* __restrict__ isv, // ISV bus - float* __restrict__ rewards, // [b_size] IN/OUT (augmented in place) + float* __restrict__ rewards, // [b_size] IN/OUT const float* __restrict__ dones, // [b_size] 0.0 / 1.0 const float* __restrict__ v_t, // [b_size] V(s_t) const float* __restrict__ v_tp1, // [b_size] V(s_{t+1}) @@ -41,32 +48,26 @@ extern "C" __global__ void compute_advantage_return( const float done = dones[b]; const float vt = v_t[b]; const float vtp1 = v_tp1[b]; + const int a = actions[b]; - // RLHF-style KL penalty: -β × (log π_θ(a|s) - log π_ref(a)) - // = -β × log(π_θ / π_ref) - // - // When π_θ(a) > π_ref(a) (over-represented action): penalty < 0 - // When π_θ(a) < π_ref(a) (under-represented action): penalty > 0 (bonus) - // - // For Hold: π_ref = 0.5. If π_θ(Hold) dropped to 0.1: - // penalty = -β × (log(0.1) - log(0.5)) = -β × (-1.61) = +1.61β (bonus!) - // For BuyL1: π_ref = 0.05. If π_θ(Buy) rose to 0.3: - // penalty = -β × (log(0.3) - log(0.05)) = -β × (1.79) = -1.79β (penalty) - const int a = actions[b]; + // Static log π_ref — consistent across replays. const float log_pi_ref = (a == HOLD_ACTION) ? -0.6931472f // log(0.5) : -2.9957323f; // log(0.05) - const float kl_penalty = -beta * (log_pi_old[b] - log_pi_ref); - const float r = r_env + kl_penalty; - rewards[b] = r; // write back so PER push captures augmented reward + // Reward with static KL component only (stored in PER for Q). + const float r = r_env + beta * log_pi_ref; + rewards[b] = r; const float ret = r + gamma * (1.0f - done) * vtp1; returns[b] = ret; - // All steps get advantages — KL penalty provides signal on non-done - // steps, trade PnL provides signal on done steps. No done-gating - // needed because the KL reward prevents the noise problem (non-done - // advantages are small but directional, not zero/noise). - advantages[b] = ret - vt; + // Advantage with BOTH static (in ret) and dynamic KL components. + // Dynamic: -β × log π_θ(a|s) adds entropy-regularized signal. + // When π_θ(a) is low (unlikely action taken): -β × negative = positive + // → reinforce exploring unlikely actions + // When π_θ(a) is high (likely action taken): -β × ~0 = small + // → don't reinforce further concentration + const float dynamic_kl = -beta * log_pi_old[b]; + advantages[b] = ret - vt + dynamic_kl; }