diff --git a/crates/ml/src/cuda_pipeline/nstep_kernel.cu b/crates/ml/src/cuda_pipeline/nstep_kernel.cu index 862473de1..f28204f6f 100644 --- a/crates/ml/src/cuda_pipeline/nstep_kernel.cu +++ b/crates/ml/src/cuda_pipeline/nstep_kernel.cu @@ -82,3 +82,74 @@ extern "C" __global__ void reward_normalize_kernel( rewards[i] = (rewards[i] - mean) * inv_std; } } + +/** + * TD(λ) truncated return kernel (v8 D1). + * + * Computes G^λ = (1-λ) Σ_{n=1}^{N-1} λ^{n-1} G^(n) + λ^{N-1} G^(N) + * where G^(n) = r_t + γr_{t+1} + ... + γ^{n-1}r_{t+n-1} + γ^n Q(s_{t+n}) + * + * λ=0: reduces to 1-step TD. λ=1: Monte Carlo (N-step, no intermediate bootstrap). + * Strict generalization of nstep_accumulate_kernel. + * + * Requires bootstrapped Q(s') values. When not available, pass raw_rewards + * as q_next for a self-bootstrap approximation. + * + * Launch: grid=(ceil(N*L/256)), block=(256). + */ +extern "C" __global__ void td_lambda_kernel( + const float* __restrict__ raw_rewards, + const float* __restrict__ raw_dones, + const float* __restrict__ q_next, + float* __restrict__ out_rewards, + float* __restrict__ out_dones, + float gamma, + float lambda, + int max_trace, + int L, + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * L) return; + + int ep = idx / L; + int t = idx % L; + int base = ep * L; + + float G_lambda = 0.0f; + float weight_sum = 0.0f; + float lambda_pow = 1.0f; + float any_done = 0.0f; + + for (int n = 1; n <= max_trace; n++) { + float G_n = 0.0f; + float gamma_pow = 1.0f; + int valid = 1; + + for (int k = 0; k < n && (t + k) < L; k++) { + int off = base + t + k; + float d_k = raw_dones[off]; + if (d_k > 0.5f && k > 0) { valid = 0; any_done = 1.0f; break; } + G_n += gamma_pow * raw_rewards[off]; + gamma_pow *= gamma; + if (d_k > 0.5f) { any_done = 1.0f; break; } + } + + if (!valid || (t + n) >= L) { + G_lambda += lambda_pow * G_n; + weight_sum += lambda_pow; + break; + } + + G_n += gamma_pow * q_next[base + t + n - 1]; + + float w = (n < max_trace) ? (1.0f - lambda) * lambda_pow : lambda_pow; + G_lambda += w * G_n; + weight_sum += w; + lambda_pow *= lambda; + } + + int out_idx = base + t; + out_rewards[out_idx] = (weight_sum > 1e-8f) ? G_lambda / weight_sum : raw_rewards[out_idx]; + out_dones[out_idx] = any_done; +}