feat(v8): TD(λ) truncated return kernel — strict generalization of n-step

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;
+}