diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 2566baec8..0637f171e 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -120,6 +120,7 @@ const KERNELS: &[&str] = &[ "rl_outcome_fwd", // Outcome aux: linear forward h_t → 3-class logits "rl_outcome_ce", // Outcome aux: softmax CE loss + gradient (masked by sentinel -1) "rl_outcome_label", // Outcome aux: assign labels from reward/done (Profit/Timeout/Loss) + "rl_outcome_bwd", // Outcome aux: single linear layer backward — dW (per-batch), db (per-batch), dh_t (per-batch); 1 block per batch, 128 threads ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_outcome_bwd.cu b/crates/ml-alpha/cuda/rl_outcome_bwd.cu new file mode 100644 index 000000000..67df8a812 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_outcome_bwd.cu @@ -0,0 +1,70 @@ +// rl_outcome_bwd.cu — Outcome aux head backward (single linear layer). +// +// Given dL/dlogits [B, K_CLASSES=3] (produced by rl_outcome_ce), +// compute the linear layer gradients via the chain rule: +// +// logits[b, k] = Σ_i h_t[b, i] × W[i, k] + b[k] +// +// dL/dW_per_batch[b, i, k] = h_t[b, i] × dL/dlogits[b, k] +// dL/db_per_batch[b, k] = dL/dlogits[b, k] +// dL/dh_t[b, i] = Σ_k W[i, k] × dL/dlogits[b, k] +// +// Block layout: 1 block per batch, HIDDEN_DIM=128 threads. +// * Thread i (i < 128) computes dL/dh_t[b, i] (inner product +// across K_CLASSES=3 outputs — trivially unrolled). +// * Thread i also writes grad_W_per_batch[b, i, 0..2] (3 writes). +// * Threads 0..2 write grad_b_per_batch[b, tid] (thread 3..127 idle +// for that single-cycle step). +// +// Per `feedback_no_atomicadd`: sole-writer per (b, i, k). +// Per `feedback_cpu_is_read_only`: pure device-side. + +#include + +#define HIDDEN_DIM 128 +#define K_CLASSES 3 + +extern "C" __global__ void rl_outcome_bwd( + const float* __restrict__ h_t, // [B, HIDDEN_DIM] + const float* __restrict__ grad_logits, // [B, K_CLASSES] + const float* __restrict__ w, // [HIDDEN_DIM, K_CLASSES] + int b_size, + float* __restrict__ grad_w_per_batch, // [B, HIDDEN_DIM, K_CLASSES] + float* __restrict__ grad_b_per_batch, // [B, K_CLASSES] + float* __restrict__ grad_h_t // [B, HIDDEN_DIM] +) { + const int b = blockIdx.x; + const int i = threadIdx.x; + if (b >= b_size || i >= HIDDEN_DIM) return; + + // Stage grad_logits into shared — K_CLASSES=3 so only threads 0..2 + // participate, but all 128 threads wait on the sync. + __shared__ float s_grad_logits[K_CLASSES]; + if (i < K_CLASSES) { + s_grad_logits[i] = grad_logits[b * K_CLASSES + i]; + } + __syncthreads(); + + const float h_bi = h_t[b * HIDDEN_DIM + i]; + + // grad_W_per_batch[b, i, k] = h_bi × s_grad_logits[k] + const int row_off = (b * HIDDEN_DIM + i) * K_CLASSES; + #pragma unroll + for (int k = 0; k < K_CLASSES; ++k) { + grad_w_per_batch[row_off + k] = h_bi * s_grad_logits[k]; + } + + // grad_h_t[b, i] = Σ_k W[i, k] × s_grad_logits[k] + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < K_CLASSES; ++k) { + acc += w[i * K_CLASSES + k] * s_grad_logits[k]; + } + grad_h_t[b * HIDDEN_DIM + i] = acc; + + // grad_b_per_batch[b, k] = s_grad_logits[k] — sole writer per + // (b, k); thread i writes slot k == i (covers k < 3 only). + if (i < K_CLASSES) { + grad_b_per_batch[b * K_CLASSES + i] = s_grad_logits[i]; + } +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 60fb07b02..b9b6662b5 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -86,6 +86,12 @@ use ml_alpha::rl::isv_slots::{ RL_REWARD_SCALE_INDEX, RL_TARGET_TAU_INDEX, RL_TD_KURTOSIS_CLAMP_INDEX, RL_TD_KURTOSIS_EMA_INDEX, RL_TD_KURT_STREAM_M2_INDEX, RL_TD_KURT_STREAM_M4_INDEX, RL_TD_KURT_STREAM_MEAN_INDEX, RL_V_GRAD_NORM_EMA_INDEX, + // P1+P2 feature metrics + RL_POPART_MEAN_INDEX, RL_POPART_SIGMA_INDEX, RL_POPART_VAR_INDEX, + RL_SPECTRAL_NORM_MAX_INDEX, RL_SPECTRAL_DECOUPLE_LAMBDA_INDEX, + RL_Q_BIAS_EMA_INDEX, RL_Q_BIAS_CORRECTION_INDEX, + RL_LR_SCALE_Q_INDEX, RL_LR_SCALE_PI_INDEX, RL_LR_SCALE_V_INDEX, RL_LR_SCALE_IQN_INDEX, + RL_OUTCOME_AUX_LAMBDA_INDEX, }; use ml_alpha::trainer::diag_staging::DiagStaging; use cudarc::driver::DevicePtrMut; @@ -1113,6 +1119,31 @@ fn main() -> Result<()> { // away from the uniform baseline is the supervised-learning signal // for the FRD head (CE loss + bwd land in F.3). "frd": frd_diag, + // P1+P2 feature metrics — PopArt reward normalization, + // spectral norm/decoupling, Q-bias correction, per-branch + // LR scaling, and outcome aux head lambda. + "popart": { + "mean": isv[RL_POPART_MEAN_INDEX], + "sigma": isv[RL_POPART_SIGMA_INDEX], + "var": isv[RL_POPART_VAR_INDEX], + }, + "spectral": { + "norm_max_config": isv[RL_SPECTRAL_NORM_MAX_INDEX], + "decouple_lambda": isv[RL_SPECTRAL_DECOUPLE_LAMBDA_INDEX], + }, + "q_bias": { + "ema": isv[RL_Q_BIAS_EMA_INDEX], + "correction": isv[RL_Q_BIAS_CORRECTION_INDEX], + }, + "per_branch_lr": { + "scale_q": isv[RL_LR_SCALE_Q_INDEX], + "scale_pi": isv[RL_LR_SCALE_PI_INDEX], + "scale_v": isv[RL_LR_SCALE_V_INDEX], + "scale_iqn": isv[RL_LR_SCALE_IQN_INDEX], + }, + "outcome_aux": { + "lambda": isv[RL_OUTCOME_AUX_LAMBDA_INDEX], + }, }); writeln!(diag, "{}", record).context("diag: writeln jsonl")?; diff --git a/crates/ml-alpha/src/rl/outcome_head.rs b/crates/ml-alpha/src/rl/outcome_head.rs index abe1f290c..bb2dddc08 100644 --- a/crates/ml-alpha/src/rl/outcome_head.rs +++ b/crates/ml-alpha/src/rl/outcome_head.rs @@ -37,10 +37,12 @@ pub struct OutcomeHead { _module_fwd: Arc, _module_ce: Arc, _module_label: Arc, + _module_bwd: Arc, kernel_fwd: CudaFunction, kernel_ce: CudaFunction, kernel_label: CudaFunction, kernel_fill_sentinel: CudaFunction, + kernel_bwd: CudaFunction, } impl OutcomeHead { @@ -60,6 +62,8 @@ impl OutcomeHead { include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_ce.cubin")); static CUBIN_LABEL: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_label.cubin")); + static CUBIN_BWD: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_bwd.cubin")); let module_fwd = stream .context() @@ -88,6 +92,14 @@ impl OutcomeHead { .load_function("rl_outcome_fill_sentinel") .map_err(|e| anyhow!("OutcomeHead: fill_sentinel kernel resolve: {e}"))?; + let module_bwd = stream + .context() + .load_cubin(CUBIN_BWD.to_vec()) + .map_err(|e| anyhow!("OutcomeHead: bwd cubin load: {e}"))?; + let kernel_bwd = module_bwd + .load_function("rl_outcome_bwd") + .map_err(|e| anyhow!("OutcomeHead: bwd kernel resolve: {e}"))?; + // Xavier init W at 0.01x scale: near-zero for near-uniform softmax at init. let w_d = ml_core::cuda_autograd::init::near_zero_xavier( HIDDEN_DIM, K_CLASSES, &stream, @@ -143,10 +155,12 @@ impl OutcomeHead { _module_fwd: module_fwd, _module_ce: module_ce, _module_label: module_label, + _module_bwd: module_bwd, kernel_fwd, kernel_ce, kernel_label, kernel_fill_sentinel, + kernel_bwd, }) } @@ -212,6 +226,58 @@ impl OutcomeHead { Ok(()) } + /// Backward pass: grad_logits (from `compute_loss`) -> grad_W, grad_b, grad_h_t. + /// + /// Requires `forward()` and `compute_loss()` to have been called first + /// (populates `grad_logits_d`). Computes per-batch weight/bias gradients + /// and the encoder-upstream gradient `grad_h_t = grad_logits × W^T`. + /// + /// The caller is responsible for: + /// 1. Zeroing the output buffers before the call. + /// 2. `reduce_axis0` on `grad_w_per_batch_d` and `grad_b_per_batch_d`. + /// 3. Adam step on `self.w_d` / `self.b_d` with the reduced grads. + /// 4. `grad_h_accumulate` to fold `grad_h_t_d` into the encoder grad. + #[allow(clippy::too_many_arguments)] + pub fn backward( + &self, + h_t: &CudaSlice, + grad_w_per_batch_d: &mut CudaSlice, + grad_b_per_batch_d: &mut CudaSlice, + grad_h_t_d: &mut CudaSlice, + actual_b: usize, + ) -> Result<()> { + if actual_b > self.b_size { + return Err(anyhow!( + "OutcomeHead::backward: actual_b ({actual_b}) > allocated b_size ({})", + self.b_size + )); + } + debug_assert_eq!(h_t.len(), actual_b * HIDDEN_DIM); + debug_assert_eq!(grad_w_per_batch_d.len(), actual_b * HIDDEN_DIM * K_CLASSES); + debug_assert_eq!(grad_b_per_batch_d.len(), actual_b * K_CLASSES); + debug_assert_eq!(grad_h_t_d.len(), actual_b * HIDDEN_DIM); + let b_i32 = actual_b as i32; + let cfg = LaunchConfig { + grid_dim: (actual_b as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + self.stream + .launch_builder(&self.kernel_bwd) + .arg(h_t) + .arg(&self.grad_logits_d) + .arg(&self.w_d) + .arg(&b_i32) + .arg(grad_w_per_batch_d) + .arg(grad_b_per_batch_d) + .arg(grad_h_t_d) + .launch(cfg) + .map_err(|e| anyhow!("rl_outcome_bwd launch: {e}"))?; + } + Ok(()) + } + /// Assign labels from reward/done signals. /// /// `rewards` and `dones` must be device slices of length `actual_b`. diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 1df82bb1d..f51aa5fec 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -856,6 +856,8 @@ pub struct IntegratedTrainer { // ── P1.3 Outcome auxiliary head ────────────────────────────────── pub outcome_head: crate::rl::outcome_head::OutcomeHead, + pub outcome_w_adam: AdamW, + pub outcome_b_adam: AdamW, // ── P2.1 Q-bias correction ─────────────────────────────────────── _rl_q_bias_correction_module: Arc, @@ -1044,6 +1046,13 @@ pub struct IntegratedTrainer { // FRD encoder-upstream gradient. pub ss_frd_grad_h_t_d: CudaSlice, + // Outcome head backward scratch buffers. + pub ss_outcome_grad_w_pb_d: CudaSlice, // [B, HIDDEN_DIM, K_CLASSES] + pub ss_outcome_grad_b_pb_d: CudaSlice, // [B, K_CLASSES] + pub ss_outcome_grad_h_t_d: CudaSlice, // [B, HIDDEN_DIM] + pub ss_outcome_grad_w_d: CudaSlice, // [HIDDEN_DIM, K_CLASSES] + pub ss_outcome_grad_b_d: CudaSlice, // [K_CLASSES] + // Bellman projection scratch. pub ss_q_target_full_d: CudaSlice, pub ss_q_target_action_d: CudaSlice, @@ -1928,6 +1937,10 @@ impl IntegratedTrainer { let outcome_head = crate::rl::outcome_head::OutcomeHead::new( b_size, stream.clone(), ).context("OutcomeHead::new")?; + let outcome_w_adam = AdamW::new(dev, outcome_head.w_d.len(), lr_placeholder) + .context("outcome_w_adam")?; + let outcome_b_adam = AdamW::new(dev, outcome_head.b_d.len(), lr_placeholder) + .context("outcome_b_adam")?; // ── P2.1 Q-bias correction cubin load ──────────────────────── let rl_q_bias_correction_module = ctx .load_cubin(RL_Q_BIAS_CORRECTION_CUBIN.to_vec()) @@ -2095,6 +2108,24 @@ impl IntegratedTrainer { let ss_frd_grad_h_t_d = stream.alloc_zeros::(b_size * HIDDEN_DIM) .context("alloc ss_frd_grad_h_t_d")?; + // Outcome head backward scratch buffers. + let k_classes = crate::rl::outcome_head::K_CLASSES; + let ss_outcome_grad_w_pb_d = stream + .alloc_zeros::(b_size * HIDDEN_DIM * k_classes) + .context("alloc ss_outcome_grad_w_pb_d")?; + let ss_outcome_grad_b_pb_d = stream + .alloc_zeros::(b_size * k_classes) + .context("alloc ss_outcome_grad_b_pb_d")?; + let ss_outcome_grad_h_t_d = stream + .alloc_zeros::(b_size * HIDDEN_DIM) + .context("alloc ss_outcome_grad_h_t_d")?; + let ss_outcome_grad_w_d = stream + .alloc_zeros::(HIDDEN_DIM * k_classes) + .context("alloc ss_outcome_grad_w_d")?; + let ss_outcome_grad_b_d = stream + .alloc_zeros::(k_classes) + .context("alloc ss_outcome_grad_b_d")?; + // Bellman projection scratch. let ss_q_target_full_d = stream.alloc_zeros::(b_size * k_dqn_ss) .context("alloc ss_q_target_full_d")?; @@ -2402,6 +2433,8 @@ impl IntegratedTrainer { spectral_v_buffer_pi, // P1.3 Outcome outcome_head, + outcome_w_adam, + outcome_b_adam, // P2.1 Q-bias _rl_q_bias_correction_module: rl_q_bias_correction_module, rl_q_bias_correction_fn, @@ -2476,6 +2509,11 @@ impl IntegratedTrainer { ss_v_grad_w_d, ss_v_grad_b_d, ss_frd_grad_h_t_d, + ss_outcome_grad_w_pb_d, + ss_outcome_grad_b_pb_d, + ss_outcome_grad_h_t_d, + ss_outcome_grad_w_d, + ss_outcome_grad_b_d, ss_q_target_full_d, ss_q_target_action_d, ss_target_dist_d, @@ -3719,6 +3757,9 @@ impl IntegratedTrainer { self.frd_b1_adam.lr = lr_frd; self.frd_w2_adam.lr = lr_frd; self.frd_b2_adam.lr = lr_frd; + // Outcome head shares the FRD LR slot (both are lightweight aux heads). + self.outcome_w_adam.lr = lr_frd; + self.outcome_b_adam.lr = lr_frd; // Borrow encoder hidden state for forward kernels. let h_t_borrow: &CudaSlice = self.perception.h_t_view(); @@ -3813,6 +3854,18 @@ impl IntegratedTrainer { self.stream.memset_zeros(&mut self.ss_frd_grad_h_t_d) .map_err(|e| anyhow::anyhow!("zero ss_frd_grad_h_t_d: {e}"))?; + // Outcome head backward scratch. + self.stream.memset_zeros(&mut self.ss_outcome_grad_w_pb_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_pb_d: {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_b_pb_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_pb_d: {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_h_t_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_h_t_d: {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_w_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_d: {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_b_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_d: {e}"))?; + // Bellman projection scratch. self.stream.memset_zeros(&mut self.ss_q_target_full_d) .map_err(|e| anyhow::anyhow!("zero ss_q_target_full_d: {e}"))?; @@ -4110,15 +4163,48 @@ impl IntegratedTrainer { } } - // ── Step 9c: Outcome head forward + CE loss ────────────────── + // ── Step 9c: Outcome head forward + CE loss + backward ──────── // K=3 outcome classifier trained from reward/done labels assigned // in step_with_lobsim. Forward and loss run on sampled_h_t_d. + // Backward: grad_logits → grad_W/b (per-batch), grad_h_t. + // reduce_axis0 → batch-summed grad_W/b. Adam on outcome W/b. + // Lambda from ISV[RL_OUTCOME_AUX_LAMBDA_INDEX] scales the + // encoder-upstream grad_h_t contribution in Step 10. { let h_t_borrow_for_outcome: &CudaSlice = &self.sampled_h_t_d; self.outcome_head.forward(h_t_borrow_for_outcome, b_size) .context("step_synthetic: outcome_head.forward")?; self.outcome_head.compute_loss(b_size) .context("step_synthetic: outcome_head.compute_loss")?; + // Backward: grad_logits → per-batch grad_W, grad_b, grad_h_t. + self.outcome_head.backward( + h_t_borrow_for_outcome, + &mut self.ss_outcome_grad_w_pb_d, + &mut self.ss_outcome_grad_b_pb_d, + &mut self.ss_outcome_grad_h_t_d, + b_size, + ).context("step_synthetic: outcome_head.backward")?; + // Reduce-axis-0: batch-summed weight and bias gradients. + let k_classes = crate::rl::outcome_head::K_CLASSES; + reduce_axis0_free( + &self.stream, &self.reduce_axis0_fn, + &self.ss_outcome_grad_w_pb_d, b_size, + HIDDEN_DIM * k_classes, + &mut self.ss_outcome_grad_w_d, + )?; + reduce_axis0_free( + &self.stream, &self.reduce_axis0_fn, + &self.ss_outcome_grad_b_pb_d, b_size, + k_classes, + &mut self.ss_outcome_grad_b_d, + )?; + // Adam step on outcome W and b. + self.outcome_w_adam + .step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d) + .context("outcome_w_adam.step")?; + self.outcome_b_adam + .step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d) + .context("outcome_b_adam.step")?; } // ── SP20 P3 FRD backward chain (F.4) ───────────────────────── @@ -4209,9 +4295,10 @@ impl IntegratedTrainer { // shared-encoder pattern (Q is trained on PER-sampled past h_t // values, whose gradient should NOT flow into the encoder // — accumulating stale-state grad would poison the encoder). - // Encoder learns from π + V (on-policy current-step) + BCE + - // aux (supervised via perception.step_batched). Standard - // off-policy AC pattern (SAC / R2D2 / IMPALA do the same). + // Encoder learns from π + V (on-policy current-step) + FRD + + // outcome (aux supervised) + BCE + aux (supervised via + // perception.step_batched). Standard off-policy AC pattern + // (SAC / R2D2 / IMPALA do the same for the Q stop-grad). // // The combined grad slot feeds `backward_encoder_with_grad_h_t` // below, which runs the shared encoder backward chain that @@ -4246,6 +4333,23 @@ impl IntegratedTrainer { lambdas.frd, &mut self.grad_h_t_combined_d, )?; + // Outcome head encoder-upstream gradient — λ from + // ISV[RL_OUTCOME_AUX_LAMBDA_INDEX]. With sentinel -1 labels + // the CE grad is zero → grad_h_t is all-zeros → this + // accumulate is a weighted zero-add (no signal). Once labels + // are populated the outcome head injects its supervised + // gradient into the encoder at the configured λ strength. + { + let outcome_lambda = + self.isv_host[crate::rl::isv_slots::RL_OUTCOME_AUX_LAMBDA_INDEX]; + accumulate_grad_h( + &self.stream, + &self.grad_h_accumulate_fn, + &self.ss_outcome_grad_h_t_d, + outcome_lambda, + &mut self.grad_h_t_combined_d, + )?; + } // ── Step 11: encoder backward (Phase E.3a) ─────────────────── // Seed the perception trainer's per-K hidden-state grad buffer @@ -4697,6 +4801,17 @@ impl IntegratedTrainer { .map_err(|e| anyhow::anyhow!("zero ss_q_grad_w_d: {e}"))?; self.stream.memset_zeros(&mut self.ss_q_grad_b_d) .map_err(|e| anyhow::anyhow!("zero ss_q_grad_b_d: {e}"))?; + // Outcome head backward scratch (replay steps also train the head). + self.stream.memset_zeros(&mut self.ss_outcome_grad_w_pb_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_pb_d (replay): {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_b_pb_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_pb_d (replay): {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_h_t_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_h_t_d (replay): {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_w_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_d (replay): {e}"))?; + self.stream.memset_zeros(&mut self.ss_outcome_grad_b_d) + .map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_d (replay): {e}"))?; let b_size_i = b_size as i32; @@ -5069,25 +5184,37 @@ impl IntegratedTrainer { // CE loss + grad w.r.t. logits (masked by sentinel -1 labels). self.outcome_head.compute_loss(b_size) .context("dqn_replay_step: outcome_head.compute_loss")?; - - // Backward: grad_logits → grad_w, grad_b, grad_h_t. - // Manual outer-product per batch: grad_w[b] = h_t[b]^T × grad_logits[b]. - // For now, we use the reduce_axis0 pattern on per-batch outer products. - // The outcome head is small (128×3), so we compute gradients inline. - // grad_w = sum_b(h_t[b] outer grad_logits[b]) is done by the same - // backward_to_w_b_h pattern as other heads. - // We compute grad_h_t[b] = grad_logits[b] × W^T in-place. - // Since OutcomeHead doesn't have a backward_to_w_b_h method, we - // skip the gradient for now — the outcome head's gradient flows - // only through its own Adam, not to the encoder. The CE loss - // provides the training signal. Encoder gradient is additive from - // Q+pi+V heads which already flow through grad_h_t_combined_d. - // - // Adam step on outcome W and b with reduced gradients. - // The outcome head is lightweight (128×3 + 3 params) — Adam on - // the per-batch mean gradient is sufficient. - // For the outcome head we skip full per-batch gradient reduce and - // just step on the CE grad directly through the logits. + // Backward: grad_logits → per-batch grad_W, grad_b, grad_h_t. + // grad_h_t is computed but discarded (no encoder backward in + // replay steps — same stop-grad rationale as Q's grad_h_t). + self.outcome_head.backward( + &self.sampled_h_t_d, + &mut self.ss_outcome_grad_w_pb_d, + &mut self.ss_outcome_grad_b_pb_d, + &mut self.ss_outcome_grad_h_t_d, + b_size, + ).context("dqn_replay_step: outcome_head.backward")?; + { + let k_classes = crate::rl::outcome_head::K_CLASSES; + reduce_axis0_free( + &self.stream, &self.reduce_axis0_fn, + &self.ss_outcome_grad_w_pb_d, b_size, + HIDDEN_DIM * k_classes, + &mut self.ss_outcome_grad_w_d, + )?; + reduce_axis0_free( + &self.stream, &self.reduce_axis0_fn, + &self.ss_outcome_grad_b_pb_d, b_size, + k_classes, + &mut self.ss_outcome_grad_b_d, + )?; + self.outcome_w_adam + .step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d) + .context("dqn_replay_step: outcome_w_adam.step")?; + self.outcome_b_adam + .step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d) + .context("dqn_replay_step: outcome_b_adam.step")?; + } // ── 9. Per-branch LR controller ────────────────────────────── // Single-thread kernel adjusts LR scale factors from loss deltas.