diff --git a/crates/ml-alpha/cuda/bellman_target_projection.cu b/crates/ml-alpha/cuda/bellman_target_projection.cu index a1ab9c9f7..aba1f4896 100644 --- a/crates/ml-alpha/cuda/bellman_target_projection.cu +++ b/crates/ml-alpha/cuda/bellman_target_projection.cu @@ -90,12 +90,13 @@ extern "C" __global__ void dqn_select_action_atoms( extern "C" __global__ void bellman_target_projection( - const float* __restrict__ target_logits, // [B × Q_N_ATOMS] - const float* __restrict__ rewards, // [B] - const float* __restrict__ dones, // [B] - const float* __restrict__ isv, // ≥ 401 + const float* __restrict__ target_logits, // [B × Q_N_ATOMS] + const float* __restrict__ rewards, // [B] (n-step discounted R_n) + const float* __restrict__ dones, // [B] + const float* __restrict__ n_step_gammas, // [B] (γⁿ per transition, 0 if done) + const float* __restrict__ isv, // ≥ 401 int B, - float* __restrict__ target_dist // [B × Q_N_ATOMS] + float* __restrict__ target_dist // [B × Q_N_ATOMS] ) { const int batch = blockIdx.x; const int atom = threadIdx.x; @@ -133,10 +134,11 @@ extern "C" __global__ void bellman_target_projection( const float p = s_softmax[atom] / s_sumexp; - // ── γ from ISV + done masking ──────────────────────────────────── - const float gamma = fmaxf(0.0f, fminf(1.0f, isv[RL_GAMMA_INDEX])); + // ── n-step γ: per-transition γⁿ replaces the ISV γ for the + // bootstrap discount. Already zero when any done occurred in the + // n-step window (set at push time). For n=1, n_step_gamma = γ. const float r = rewards[batch]; - const float gamma_eff = gamma * (1.0f - dones[batch]); + const float gamma_eff = n_step_gammas[batch]; // ── Adaptive atom span from ISV (audit follow-up). DELTA_Z is // recomputed every kernel launch — ratchet semantics in the diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index 6b24a35cd..e94ea0306 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -506,6 +506,7 @@ impl DqnHead { target_logits_d: &CudaSlice, rewards_d: &CudaSlice, dones_d: &CudaSlice, + n_step_gammas_d: &CudaSlice, isv_d: &CudaSlice, b_size: usize, target_dist_d: &mut CudaSlice, @@ -513,6 +514,7 @@ impl DqnHead { debug_assert_eq!(target_logits_d.len(), b_size * Q_N_ATOMS); debug_assert_eq!(rewards_d.len(), b_size); debug_assert_eq!(dones_d.len(), b_size); + debug_assert_eq!(n_step_gammas_d.len(), b_size); debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS); let b_i = b_size as i32; @@ -526,6 +528,7 @@ impl DqnHead { .arg(target_logits_d) .arg(rewards_d) .arg(dones_d) + .arg(n_step_gammas_d) .arg(isv_d) .arg(&b_i) .arg(target_dist_d); diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 3e5684372..c061ce359 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -660,6 +660,7 @@ pub struct IntegratedTrainer { /// Per-step gather buffer for PER-sampled done flags. Length `B`. pub sampled_dones_d: CudaSlice, + pub sampled_n_step_gammas_d: CudaSlice, /// Per-step gather buffer for PER-sampled next-state argmax actions /// (Double-DQN online-Q argmax on `sampled_h_tp1`). Computed @@ -1341,6 +1342,9 @@ impl IntegratedTrainer { let sampled_dones_d = stream .alloc_zeros::(b_size) .context("alloc sampled_dones_d")?; + let sampled_n_step_gammas_d = stream + .alloc_zeros::(b_size) + .context("alloc sampled_n_step_gammas_d")?; let sampled_next_actions_d = stream .alloc_zeros::(b_size) .context("alloc sampled_next_actions_d")?; @@ -1537,6 +1541,7 @@ impl IntegratedTrainer { sampled_actions_d, sampled_rewards_d, sampled_dones_d, + sampled_n_step_gammas_d, sampled_next_actions_d, td_per_sample_d, grad_h_t_combined_d, @@ -2894,6 +2899,7 @@ impl IntegratedTrainer { &q_target_action_d, &self.sampled_rewards_d, &self.sampled_dones_d, + &self.sampled_n_step_gammas_d, &self.isv_d, b_size, &mut target_dist_d, @@ -3495,6 +3501,7 @@ impl IntegratedTrainer { &q_target_action_d, &self.sampled_rewards_d, &self.sampled_dones_d, + &self.sampled_n_step_gammas_d, &self.isv_d, b_size, &mut target_dist_d, @@ -4773,12 +4780,6 @@ impl IntegratedTrainer { fn push_to_replay(&mut self, b_size: usize) -> Result<()> { use crate::rl::common::{N_ACTIONS, Transition}; - // Per-batch metadata via mapped-pinned staging (the only - // permitted CPU↔GPU path per - // `feedback_no_htod_htoh_only_mapped_pinned`). Each helper - // allocates a mapped-pinned buffer, issues a DtoD copy from - // the device source, syncs, and returns a host Vec read from - // the page-mapped host_ptr. let actions_host = read_slice_i32_d(&self.stream, &self.actions_d, b_size) .context("push_to_replay: read actions_d")?; let rewards_host = read_slice_d(&self.stream, &self.rewards_d, b_size) @@ -4790,26 +4791,17 @@ impl IntegratedTrainer { let log_pi_old_host = read_slice_d(&self.stream, &self.log_pi_old_d, b_size) .context("push_to_replay: read log_pi_old_d")?; - // For each batch index, alloc per-transition device buffers - // and DtoD-copy the per-batch slice of h_t / h_tp1 in. + let n_step = self.isv_host[crate::rl::isv_slots::RL_N_STEP_INDEX] as usize; + let n_step = n_step.max(1); + let gamma = self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX]; + let hidden_nbytes = HIDDEN_DIM * std::mem::size_of::(); for b in 0..b_size { + // Alloc h_t for this step's n-step entry. let mut h_t_per_b = self .stream .alloc_zeros::(HIDDEN_DIM) .context("push_to_replay: alloc h_t per-batch")?; - let mut h_tp1_per_b = self - .stream - .alloc_zeros::(HIDDEN_DIM) - .context("push_to_replay: alloc h_tp1 per-batch")?; - - // Two separate inner scopes for the per-buffer DtoD so - // each `SyncOnDrop` guard drops before the next - // `device_ptr_mut` borrow / before the move into - // `Transition`. Without this, the borrow checker rejects - // the moves into `Transition { h_t, next_h_t, .. }` as - // overlapping the guard lifetimes (canonical borrow-checker - // dance for the cudarc raw-pointer API). let off = (b * HIDDEN_DIM * std::mem::size_of::()) as u64; unsafe { let s = self.stream.cu_stream(); @@ -4820,6 +4812,50 @@ impl IntegratedTrainer { ) .context("push_to_replay: DtoD per-batch h_t slice")?; } + + let is_done = dones_host[b] > 0.5; + self.n_step_buffer[b].push(NStepEntry { + h_t: h_t_per_b, + action: actions_host[b] as u32, + raw_reward: raw_rewards_host[b], + scaled_reward: rewards_host[b], + done: is_done, + log_pi_old: log_pi_old_host[b], + }); + + // Flush when buffer reaches n OR a done event truncates. + let should_flush = self.n_step_buffer[b].len() >= n_step || is_done; + if !should_flush { + continue; + } + + let buf = &self.n_step_buffer[b]; + let buf_len = buf.len(); + + // Compute n-step discounted return: R_n = Σ γᵏ rₖ + let mut r_n_scaled = 0.0_f32; + let mut r_n_raw = 0.0_f32; + let mut gamma_acc = 1.0_f32; + let mut any_done = false; + for entry in buf.iter() { + r_n_scaled += gamma_acc * entry.scaled_reward; + r_n_raw += gamma_acc * entry.raw_reward; + if entry.done { + any_done = true; + } + gamma_acc *= gamma; + } + + // γⁿ for the bootstrap term. Zero if any done in window + // (terminal state zeroes the bootstrap). + let n_step_gamma_val = if any_done { 0.0 } else { gamma_acc }; + + // h_t from oldest entry, h_tp1 from current step. + let oldest = &buf[0]; + let mut h_tp1_per_b = self + .stream + .alloc_zeros::(HIDDEN_DIM) + .context("push_to_replay: alloc h_tp1 per-batch")?; unsafe { let s = self.stream.cu_stream(); let (h_tp1_base, _g2) = self.h_tp1_d.device_ptr(&self.stream); @@ -4829,15 +4865,36 @@ impl IntegratedTrainer { ) .context("push_to_replay: DtoD per-batch h_tp1 slice")?; } + + // Clone oldest h_t (it's about to be consumed by the drain). + let mut h_t_oldest = self + .stream + .alloc_zeros::(HIDDEN_DIM) + .context("push_to_replay: alloc h_t_oldest")?; + unsafe { + let s = self.stream.cu_stream(); + let (src, _g1) = oldest.h_t.device_ptr(&self.stream); + let (dst, _gd) = h_t_oldest.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, src, hidden_nbytes, s, + ) + .context("push_to_replay: DtoD clone h_t_oldest")?; + } + let oldest_action = oldest.action; + let oldest_log_pi = oldest.log_pi_old; + + // Drain the buffer. + self.n_step_buffer[b].clear(); + self.replay.push(Transition { - h_t: h_t_per_b, - action: actions_host[b] as u32, - reward: rewards_host[b], - raw_reward: raw_rewards_host[b], + h_t: h_t_oldest, + action: oldest_action, + reward: r_n_scaled, + raw_reward: r_n_raw, next_h_t: h_tp1_per_b, - done: dones_host[b] > 0.5, - n_step_gamma: self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX], - log_pi_old: log_pi_old_host[b], + done: any_done, + n_step_gamma: n_step_gamma_val, + log_pi_old: oldest_log_pi, q_value_old: [0.0; N_ACTIONS], }); } @@ -4874,21 +4931,23 @@ impl IntegratedTrainer { let mut actions_host = vec![0i32; b_size]; let mut rewards_host = vec![0.0f32; b_size]; let mut dones_host = vec![0.0f32; b_size]; + let mut n_step_gammas_host = vec![0.0f32; b_size]; let current_scale = self.isv_host[crate::rl::isv_slots::RL_REWARD_SCALE_INDEX]; for (i, &idx) in indices.iter().enumerate() { let t = &self.replay.transitions[idx]; actions_host[i] = t.action as i32; rewards_host[i] = t.raw_reward * current_scale; dones_host[i] = if t.done { 1.0 } else { 0.0 }; + n_step_gammas_host[i] = t.n_step_gamma; } - // Mapped-pinned staging — the only permitted CPU→GPU path - // per `feedback_no_htod_htoh_only_mapped_pinned`. write_slice_i32_d(&self.stream, &actions_host, &mut self.sampled_actions_d) .context("sample_and_gather: write sampled_actions_d")?; write_slice_f32_d(&self.stream, &rewards_host, &mut self.sampled_rewards_d) .context("sample_and_gather: write sampled_rewards_d")?; write_slice_f32_d(&self.stream, &dones_host, &mut self.sampled_dones_d) .context("sample_and_gather: write sampled_dones_d")?; + write_slice_f32_d(&self.stream, &n_step_gammas_host, &mut self.sampled_n_step_gammas_d) + .context("sample_and_gather: write sampled_n_step_gammas_d")?; // Per-batch DtoD: per-transition h_t / h_tp1 device buffers → // contiguous `sampled_h_t_d` / `sampled_h_tp1_d` at slot i.