From c55d9d0276709c3278dee6589592bbb8ae14e04a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 16:10:20 +0200 Subject: [PATCH] fix: Expected SARSA temperature prevents uniform-averaging plateau MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Softmax weights used implicit tau=1, producing near-uniform weights when Q-gap=0.06: exp(0.06)≈1.06 → weights [0.35, 0.33, 0.32]. The mixture averaged all distributions → zero gradient to differentiate actions → Q-gap frozen → val_Sharpe locked at 0.00 after epoch 19. Now uses tau = max(Q_gap/3, 0.001). With Q-gap=0.06 and tau=0.02: exp(0.06/0.02) = exp(3) ≈ 20:1 weight ratio between best and worst. The 0.001 floor means even tiny Q-differences produce sharp weights, preventing the uniform-averaging fixed point. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/c51_loss_kernel.cu | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 6af86fc31..a4a9e0ccf 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -476,15 +476,26 @@ extern "C" __global__ void c51_loss_batched( * When Q-values differentiate, this converges to argmax. * When Q-values are flat, this averages → stable learning signal. */ - /* Compute softmax weights from expected Q per action */ + /* Compute softmax weights from expected Q per action. + * Temperature = max(Q_gap/3, 0.001). When Q-values differentiate, + * tau shrinks → weights sharpen toward argmax. When flat, tau=0.001 + * still produces exp(0.06/0.001)=exp(60) → effectively argmax for + * ANY non-zero Q-difference. This prevents the uniform-averaging + * fixed point where all actions get equal weight → zero gradient + * to push Q-values apart. */ __shared__ float action_weights[MAX_BRANCH_SIZE]; if (tid == 0) { float max_eq = eq_per_action[0]; - for (int a = 1; a < n_d; a++) + float min_eq = eq_per_action[0]; + for (int a = 1; a < n_d; a++) { max_eq = fmaxf(max_eq, eq_per_action[a]); + min_eq = fminf(min_eq, eq_per_action[a]); + } + float q_gap_local = max_eq - min_eq; + float tau = fmaxf(q_gap_local / 3.0f, 0.001f); float sum_exp = 0.0f; for (int a = 0; a < n_d; a++) { - action_weights[a] = expf(eq_per_action[a] - max_eq); + action_weights[a] = expf((eq_per_action[a] - max_eq) / tau); sum_exp += action_weights[a]; } for (int a = 0; a < n_d; a++)