diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu index f2d643c4a..9a4cb8cbd 100644 --- a/crates/ml-alpha/cuda/compute_advantage_return.cu +++ b/crates/ml-alpha/cuda/compute_advantage_return.cu @@ -23,12 +23,12 @@ // reduction, no atomics. Per `feedback_no_atomicadd` not needed. #define RL_GAMMA_INDEX 400 -#define RL_HOLD_PRIOR_INDEX 574 +#define N_ACTIONS 11 #define HOLD_ACTION 2 extern "C" __global__ void compute_advantage_return( const float* __restrict__ isv, // ISV bus - const float* __restrict__ rewards, // [b_size] + float* __restrict__ rewards, // [b_size] IN/OUT (augmented in place) 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,112 +41,35 @@ extern "C" __global__ void compute_advantage_return( if (b >= b_size) return; const float gamma = isv[RL_GAMMA_INDEX]; - const float r = rewards[b]; + const float r_env = rewards[b]; const float done = dones[b]; const float vt = v_t[b]; const float vtp1 = v_tp1[b]; + // KL-augmented reward: r = r_env + β × log(π_ref(a_taken)) + // This makes Q and V learn the SAME augmented objective as π, + // aligning their action rankings (positive qpa). Hold gets a + // bonus because π_ref(Hold) = 0.5 → log(0.5) = -0.69, while + // trading actions get π_ref = 0.05 → log(0.05) = -3.0. + // Net effect: Hold penalty is -0.69β, trading penalty is -3.0β. + // The RELATIVE bonus for Hold over trading = 2.3β. + // Fixed small β for reward augmentation (decoupled from the adaptive + // gradient β). This gives Q a gentle Hold preference without the + // compounding effect of the adaptive KL gradient controller. + const float REWARD_KL_BETA = 0.02f; + const int a = actions[b]; + const float pi_ref = (a == HOLD_ACTION) ? 0.5f : (1.0f / (float)(N_ACTIONS - 1)) * 0.5f; + const float kl_bonus = REWARD_KL_BETA * logf(fmaxf(pi_ref, 1e-7f)); + const float r = r_env + kl_bonus; + rewards[b] = r; // write back so PER push captures augmented reward + const float ret = r + gamma * (1.0f - done) * vtp1; returns[b] = ret; if (done > 0.5f) { - // Done-step: real reward signal from trade closing. advantages[b] = ret - vt; } else { - // Non-done step: surfer Hold prior. - // Small positive advantage for Hold, small negative for trading. - // This gives π continuous gradient that reinforces waiting as - // the default. Without this, done-gating makes Hold invisible - // to π (zero gradient → abandoned). - // - // ISV-driven magnitude so it adapts with reward_scale. - const float hold_prior = isv[RL_HOLD_PRIOR_INDEX]; - const int a = actions[b]; - advantages[b] = (a == HOLD_ACTION) ? hold_prior : -hold_prior; + advantages[b] = 0.0f; } } -// Advantage normalization: zero mean, unit variance across the batch. -// Two-phase: phase 1 computes mean+var via block tree-reduce, -// phase 2 normalizes in-place. Single block covers the full batch. -// -// Standard PPO practice — without normalization, Q-advantage can be -// all-positive or all-negative, causing PPO to reinforce/suppress ALL -// actions uniformly instead of discriminating between them. - -// DISABLED: advantage normalization causes issues with sparse rewards -// (5% of batch has real rewards, 95% is noise). The /B gradient -// normalization in ppo_clipped_surrogate_bwd already handles batch- -// size invariance. Raw V-advantages are bounded by reward_scale. -// -// Kept for reference but not launched from Rust. -#if 0 -// Normalize advantages using only done-steps (where dones > 0.5) for -// mean/variance computation. Non-done steps have reward=0 and advantage -// ≈ γV(s')-V(s) ≈ noise. Normalizing over the full batch dilutes the -// sparse reward signal (5-6% of batch has real rewards) with 94% noise. -// -// Non-done steps still get their advantage normalized (centered/scaled), -// but the centering point and scale come from the DONE steps' advantage -// distribution — so the real reward signal drives the normalization. -// If zero done steps in a batch, fall back to full-batch normalization. - -extern "C" __global__ void normalize_advantages( - float* __restrict__ advantages, // [b_size] IN/OUT - const float* __restrict__ dones, // [b_size] 0.0/1.0 - int b_size -) { - extern __shared__ float smem[]; // [3 * blockDim.x] - float* s_sum = smem; - float* s_sum2 = smem + blockDim.x; - float* s_count = smem + 2 * blockDim.x; - - const int tid = threadIdx.x; - - float val = 0.0f; - float is_done = 0.0f; - if (tid < b_size) { - val = advantages[tid]; - is_done = (dones[tid] > 0.5f) ? 1.0f : 0.0f; - } - s_sum[tid] = val * is_done; - s_sum2[tid] = val * val * is_done; - s_count[tid] = is_done; - __syncthreads(); - - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - s_sum[tid] += s_sum[tid + stride]; - s_sum2[tid] += s_sum2[tid + stride]; - s_count[tid] += s_count[tid + stride]; - } - __syncthreads(); - } - - __shared__ float s_mean, s_std; - if (tid == 0) { - float n = s_count[0]; - if (n < 1.0f) { - // No done steps — fall back to full-batch stats. - // Recompute from all elements (rare path). - float sum_all = 0.0f, sum2_all = 0.0f; - for (int i = 0; i < b_size; i++) { - sum_all += advantages[i]; - sum2_all += advantages[i] * advantages[i]; - } - s_mean = sum_all / (float)b_size; - float var = sum2_all / (float)b_size - s_mean * s_mean; - s_std = sqrtf(fmaxf(var, 1e-8f)); - } else { - s_mean = s_sum[0] / n; - float var = s_sum2[0] / n - s_mean * s_mean; - s_std = sqrtf(fmaxf(var, 1e-8f)); - } - } - __syncthreads(); - - if (tid < b_size) { - advantages[tid] = (advantages[tid] - s_mean) / s_std; - } -} -#endif