diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index ec61ee726..449fa814d 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -497,7 +497,11 @@ extern "C" __global__ void c51_loss_batched( min_eq = fminf(min_eq, eq_per_action[a]); } float q_gap_local = max_eq - min_eq; - float tau = fmaxf(q_gap_local, 0.01f); + /* Floor: proportional to mean Q magnitude so the Boltzmann is + * meaningful at any Q-value scale. 1% of |mean Q| or 1e-6. */ + float mean_q = (max_eq + min_eq) * 0.5f; + float tau_floor = fmaxf(fabsf(mean_q) * 0.01f, 1e-6f); + float tau = fmaxf(q_gap_local, tau_floor); float sum_exp = 0.0f; for (int a = 0; a < n_d; a++) { action_weights[a] = expf((eq_per_action[a] - max_eq) / tau);